diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 3db49d579e3..a064e53ebed 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -91,7 +91,7 @@ steps: - tests/quantization/test_cpu_wna16.py commands: - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs pytest -x -v -s tests/quantization/test_cpu_wna16.py" diff --git a/.buildkite/scripts/ci-clean-log.sh b/.buildkite/scripts/ci-clean-log.sh index 69d8a3a2883..e2e21483d54 100644 --- a/.buildkite/scripts/ci-clean-log.sh +++ b/.buildkite/scripts/ci-clean-log.sh @@ -13,5 +13,8 @@ INPUT_FILE="$1" # Strip timestamps sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE" +# Strip Buildkite inline timestamp markers (ESC _bk;t= 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" diff --git a/.buildkite/scripts/ci-fetch-log.sh b/.buildkite/scripts/ci-fetch-log.sh index 3f99bc50a57..4830135a112 100755 --- a/.buildkite/scripts/ci-fetch-log.sh +++ b/.buildkite/scripts/ci-fetch-log.sh @@ -1,74 +1,178 @@ #!/bin/bash -# Usage: ./ci-fetch-log.sh [output_file] -# ./ci-fetch-log.sh [output_file] +# Fetch vLLM Buildkite CI logs (public; no login required). # -# Downloads the raw log for a Buildkite job from the public, unauthenticated -# /organizations//pipelines//builds//jobs//download -# endpoint, then strips ANSI/timestamps via ci-clean-log.sh. +# Usage: +# ci-fetch-log.sh [--soft|--all] --pr [] failed jobs in the PR's latest +# build (current branch if omitted) +# ci-fetch-log.sh [--soft|--all] failed jobs in that build +# ci-fetch-log.sh [output] one job; both # and +# ?sid= URL forms work +# ci-fetch-log.sh [output] # -# Find and via: -# gh pr checks --repo vllm-project/vllm -# Each failing row's URL is .../builds/#. -# -# Default output path: ci--.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--.log (ANSI/timestamps stripped) and +# prints "\t" 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 [output_file]" - echo " $0 [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:-}" + 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//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() { # + 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 diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 186f7222539..148aea73c7f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -398,7 +398,7 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion==1.0.0 + - pip install helion==1.1.0 - pytest -v -s kernels/helion/ - label: Kernels Mamba Test # TBD diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 10b5b7527b8..159f940530e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -75,6 +75,19 @@ steps: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 +- label: Kernels Attention DiffKV Test (H100) + key: kernels-attention-diffkv-test-h100 + timeout_in_minutes: 20 + device: h100 + num_devices: 1 + source_file_dependencies: + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py + - vllm/v1/attention/backends/triton_attn_diffkv.py + - vllm/v1/attention/backends/flash_attn_diffkv.py + - tests/kernels/attention/test_triton_unified_attention_diffkv.py + commands: + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py + - label: Kernels Quantization Test %N key: kernels-quantization-test timeout_in_minutes: 90 @@ -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/ diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cda2bb4dafe..67fecf06df3 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -138,11 +138,26 @@ steps: - vllm/v1/spec_decode/extract_hidden_states.py - vllm/model_executor/models/extract_hidden_states.py - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py - tests/v1/kv_connector/extract_hidden_states_integration commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s v1/kv_connector/extract_hidden_states_integration +- label: Extract Hidden States Integration (2 GPUs) + key: extract-hidden-states-integration-2-gpus + timeout_in_minutes: 20 + num_devices: 2 + source_file_dependencies: + - vllm/v1/spec_decode/extract_hidden_states.py + - vllm/model_executor/models/extract_hidden_states.py + - vllm/transformers_utils/configs/extract_hidden_states.py + - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py + - tests/v1/kv_connector/extract_hidden_states_integration + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration + - label: Regression key: regression timeout_in_minutes: 20 diff --git a/AGENTS.md b/AGENTS.md index 441b8d9fb73..1f3a083f80c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,6 +105,26 @@ The line length limit for Python code is 88 characters. If you are not sure, use Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`). +### Coding style guidelines + +Follow these rules for all code changes in this repository: + +- Try to match existing code style. +- Code should be self-documenting and self-explanatory. +- Keep comments and docstrings minimal and concise. +- Assume the reader is familiar with vLLM. + +### Diagnosing CI failures + +Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md). + +```bash +# All failed-job logs for a PR's latest build (current branch's PR if omitted): +.buildkite/scripts/ci-fetch-log.sh --pr +# Any Buildkite build or job URL also works: +.buildkite/scripts/ci-fetch-log.sh "" +``` + ### Commit messages Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: diff --git a/CMakeLists.txt b/CMakeLists.txt index 0a48ddca68a..6f60759550b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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}") diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index f885b1e0952..5d0876f9125 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -792,6 +792,12 @@ def get_model_params(config): topk = text_config.num_experts_per_tok intermediate_size = text_config.moe_intermediate_size hidden_size = text_config.hidden_size + elif architecture == "DiffusionGemmaForBlockDiffusion": + text_config = config.get_text_config() + E = text_config.num_experts + topk = text_config.top_k_experts + intermediate_size = text_config.moe_intermediate_size + hidden_size = text_config.hidden_size elif architecture == "HunYuanMoEV1ForCausalLM": E = config.num_experts topk = config.moe_topk[0] diff --git a/benchmarks/kv_cache_watermark.sh b/benchmarks/kv_cache_watermark.sh new file mode 100755 index 00000000000..258afa9fce1 --- /dev/null +++ b/benchmarks/kv_cache_watermark.sh @@ -0,0 +1,248 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Reproducible demonstration of the KV cache watermark (`--watermark`) for +# reducing preemption thrashing. +# +# The watermark is the fraction of total KV cache blocks the scheduler keeps +# free when admitting a waiting/preempted request into the running queue. +# +# Why this workload triggers thrashing: +# Requests are admitted based on the KV cache they need *at admission time*. +# With `--scheduler-reserve-full-isl` (default) the input length is reserved up +# front, but the *output* length is unknown and unreserved. A decode-heavy +# workload (output >> input) at high concurrency therefore over-admits while +# requests are short, then runs out of KV cache as they all grow during decode +# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills +# them later, and repeats. The watermark keeps a block of KV cache free so +# running requests can grow into it instead of triggering this churn. +# +# This script launches `vllm serve` under a deliberately KV-constrained config +# and a decode-heavy workload, sweeping the watermark across several values, and +# reports the preemption count (scraped from /metrics), throughput, and latency +# percentiles for each. It then plots the results. +# +# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens +# (+/- 20% variance), sized to run each config for ~5 minutes. +# +# Usage: +# benchmarks/kv_cache_watermark.sh +# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh +# +# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it). +set -euo pipefail + +# ---- Config (override via environment) ------------------------------------- +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} +TP=${TP:-1} +PORT=${PORT:-8000} +URL="http://127.0.0.1:${PORT}" +# Constrain the KV cache to a *near-critical* size: large enough that the engine +# can run stably, but small enough that greedy over-admission tips it into +# preemption thrashing. (Independent of GPU size, so the demo is reproducible.) +# At the default workload this fits ~1.5x the mean concurrent KV demand. +KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} +# Optional weight loader (e.g. fastsafetensors on the GCP cluster). +LOAD_FORMAT=${LOAD_FORMAT:-auto} +# Decode-heavy workload: moderate input, long output, with length variance. The +# long output means preempted requests have generated a lot before eviction, so +# resuming them re-prefills a long sequence (high recomputation cost). +INPUT_LEN=${INPUT_LEN:-1000} +OUTPUT_LEN=${OUTPUT_LEN:-5000} +RANGE_RATIO=${RANGE_RATIO:-0.2} +CONCURRENCY=${CONCURRENCY:-128} +# Enough prompts to keep each config saturated for ~5+ minutes. +NUM_PROMPTS=${NUM_PROMPTS:-450} +OUTDIR=${OUTDIR:-./watermark_bench_results} +# Watermark fractions compared. "label value" per line; value=0 disables it. +CONFIGS=${CONFIGS:-"off 0 +w0.02 0.02 +w0.05 0.05 +w0.10 0.10 +w0.15 0.15"} + +KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024)) +mkdir -p "$OUTDIR" + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +scrape_preemptions() { + # Sum the vllm:num_preemptions_total counter across engines. + python - "${URL}/metrics" <<'PY' +import sys, urllib.request +total = 0.0 +try: + body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace") + for line in body.splitlines(): + if line.startswith("vllm:num_preemptions_total"): + total += float(line.rsplit(" ", 1)[-1]) +except Exception as e: # noqa: BLE001 + print(f"scrape error: {e}", file=sys.stderr) +print(int(total)) +PY +} + +wait_for_server() { + for _ in $(seq 1 300); do + if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: server process exited during startup" >&2; return 1 + fi + sleep 5 + done + echo "ERROR: server did not become ready" >&2; return 1 +} + +run_one() { + local label=$1 watermark=$2 + echo + echo "==================== watermark: ${label} (${watermark}) ====================" + vllm serve "$MODEL" \ + --tensor-parallel-size "$TP" \ + --load-format "$LOAD_FORMAT" \ + --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \ + --max-model-len "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --no-enable-prefix-caching \ + --watermark "$watermark" \ + --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 & + SERVER_PID=$! + wait_for_server + sleep 5 + + local pre post + pre=$(scrape_preemptions) + vllm bench serve \ + --backend vllm \ + --base-url "$URL" \ + --model "$MODEL" \ + --dataset-name random \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --ignore-eos \ + --num-prompts "$NUM_PROMPTS" \ + --max-concurrency "$CONCURRENCY" \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --metric-percentiles "50,90,99" \ + --save-result \ + --result-dir "$OUTDIR" \ + --result-filename "bench_${label}.json" + post=$(scrape_preemptions) + echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt" + + kill "$SERVER_PID" 2>/dev/null || true + for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done + SERVER_PID="" + sleep 10 +} + +: >"${OUTDIR}/preemptions.txt" +while read -r label watermark; do + [[ -z "${label:-}" ]] && continue + run_one "$label" "$watermark" +done <<<"$CONFIGS" + +echo +echo "==================== summary ====================" +python - "$OUTDIR" <<'PY' +import json, os, sys +outdir = sys.argv[1] +pre = {} +order = [] +for line in open(os.path.join(outdir, "preemptions.txt")): + label, watermark, n = line.split() + pre[label] = (float(watermark), int(n)) + order.append(label) + +def g(d, *names): + for n in names: + if d.get(n) is not None: + return d[n] + return float("nan") + +cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s", + "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"] +print(" ".join(f"{c:>10}" for c in cols)) +rows = [] +for label in order: + watermark, n = pre[label] + d = json.load(open(os.path.join(outdir, f"bench_{label}.json"))) + rows.append(dict( + label=label, watermark=watermark, preempt=n, + out_tok_s=g(d, "output_throughput"), + req_s=g(d, "request_throughput"), + ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"), + ttft_p99=g(d, "p99_ttft_ms"), + itl_p99=g(d, "p99_itl_ms"), + e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"), + )) + print(" ".join(f"{str(v):>10}" for v in [ + label, watermark, n, + f"{rows[-1]['out_tok_s']:.0f}", + f"{rows[-1]['req_s']:.3f}", + f"{rows[-1]['ttft_p50']/1000:.2f}", + f"{rows[-1]['ttft_p99']/1000:.2f}", + f"{rows[-1]['itl_p99']:.2f}", + f"{rows[-1]['e2el_p50']/1000:.1f}", + ])) +print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)") + +# ---- Plot ------------------------------------------------------------------- +try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception as e: # noqa: BLE001 + print(f"\n(skip plot: matplotlib unavailable: {e})") + sys.exit(0) + +x = [r["watermark"] for r in rows] +xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows] +idx = list(range(len(rows))) + +fig, axes = plt.subplots(2, 2, figsize=(12, 8)) +fig.suptitle( + f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}", + fontsize=12, +) + +ax = axes[0][0] +ax.bar(idx, [r["preempt"] for r in rows], color="tab:red") +ax.set_title("Preemptions (lower is better)") +ax.set_ylabel("preemptions") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[0][1] +ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green") +ax.set_title("Output throughput (higher is better)") +ax.set_ylabel("tokens/s") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][0] +ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue") +ax.set_title("Inter-token latency p99 (lower is better)") +ax.set_ylabel("ITL p99 (ms)") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][1] +ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50") +ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99") +ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50") +ax.set_title("Latency (lower is better)") +ax.set_ylabel("seconds") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) +ax.legend() + +fig.tight_layout(rect=(0, 0, 1, 0.95)) +out_png = os.path.join(outdir, "watermark_results.png") +fig.savefig(out_png, dpi=120) +print(f"\nWrote plot: {out_png}") +PY diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 1e4feb0ff9e..ea7ac544b9d 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25 + GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py index 34fb64c413d..d692502f3ff 100644 --- a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py +++ b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py @@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = { } VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = { - DataType.u8: "at::ScalarType::Byte", - DataType.s8: "at::ScalarType::Char", - DataType.e4m3: "at::ScalarType::Float8_e4m3fn", - DataType.s32: "at::ScalarType::Int", - DataType.f16: "at::ScalarType::Half", - DataType.bf16: "at::ScalarType::BFloat16", - DataType.f32: "at::ScalarType::Float", + DataType.u8: "torch::headeronly::ScalarType::Byte", + DataType.s8: "torch::headeronly::ScalarType::Char", + DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn", + DataType.s32: "torch::headeronly::ScalarType::Int", + DataType.f16: "torch::headeronly::ScalarType::Half", + DataType.bf16: "torch::headeronly::ScalarType::BFloat16", + DataType.f32: "torch::headeronly::ScalarType::Float", } VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = { diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h index 09ed1a470bd..783736ab509 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h index 9858df94573..04f90101be4 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" -#include "quantization/marlin/dequant.h" -#include "quantization/marlin/marlin_mma.h" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/dequant.h" +#include "libtorch_stable/quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh index ce96c2d11fe..ac33d5f2ce6 100644 --- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh +++ b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh @@ -6,7 +6,7 @@ #include -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" using marlin::MarlinScalarType2; namespace allspark { diff --git a/csrc/quantization/machete/Readme.md b/csrc/libtorch_stable/quantization/machete/Readme.md similarity index 100% rename from csrc/quantization/machete/Readme.md rename to csrc/libtorch_stable/quantization/machete/Readme.md diff --git a/csrc/quantization/machete/generate.py b/csrc/libtorch_stable/quantization/machete/generate.py similarity index 95% rename from csrc/quantization/machete/generate.py rename to csrc/libtorch_stable/quantization/machete/generate.py index e12601e9e97..11a5bbdd13c 100644 --- a/csrc/quantization/machete/generate.py +++ b/csrc/libtorch_stable/quantization/machete/generate.py @@ -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 maybe_scalartype( - std::optional const& t) { +static inline std::optional maybe_scalartype( + std::optional const& t) { if (!t) { return std::nullopt; } else { @@ -74,7 +74,7 @@ static inline std::optional 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>(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 diff --git a/csrc/quantization/machete/machete_collective_builder.cuh b/csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh similarity index 100% rename from csrc/quantization/machete/machete_collective_builder.cuh rename to csrc/libtorch_stable/quantization/machete/machete_collective_builder.cuh diff --git a/csrc/quantization/machete/machete_interleaving_utils.cuh b/csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh similarity index 100% rename from csrc/quantization/machete/machete_interleaving_utils.cuh rename to csrc/libtorch_stable/quantization/machete/machete_interleaving_utils.cuh diff --git a/csrc/quantization/machete/machete_mainloop.cuh b/csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh similarity index 100% rename from csrc/quantization/machete/machete_mainloop.cuh rename to csrc/libtorch_stable/quantization/machete/machete_mainloop.cuh diff --git a/csrc/quantization/machete/machete_mm_kernel.cuh b/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh similarity index 87% rename from csrc/quantization/machete/machete_mm_kernel.cuh rename to csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh index cc50e68b058..db3321a39db 100644 --- a/csrc/quantization/machete/machete_mm_kernel.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_mm_kernel.cuh @@ -1,8 +1,6 @@ #pragma once -#include -#include -#include +#include // 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 const& maybe_g_scales, // scale_KxN matrix - std::optional 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 const& + maybe_g_scales, // scale_KxN matrix + std::optional const& + maybe_g_zeros, // scale_KxN matrix std::optional maybe_group_size, - std::optional const& maybe_ch_scales, // len N vector - std::optional const& maybe_tok_scales) // len M vector + std::optional const& + maybe_ch_scales, // len N vector + std::optional 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(A, "A"); auto layout_D = make_cute_layout(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"); } }; diff --git a/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh b/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh new file mode 100644 index 00000000000..fcf7f18aac2 --- /dev/null +++ b/csrc/libtorch_stable/quantization/machete/machete_mm_launcher.cuh @@ -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 +#include + +#include +#include +#include + +namespace machete { + +struct MMArgs { + torch::stable::Tensor const& A; + torch::stable::Tensor const& B; + vllm::ScalarType const& b_type; + std::optional const& maybe_out_type; + std::optional const& maybe_group_scales; + std::optional const& maybe_group_zeros; + std::optional maybe_group_size; + std::optional const& maybe_channel_scales; + std::optional const& maybe_token_scales; + std::optional maybe_schedule; +}; + +struct SupportedSchedulesArgs { + torch::headeronly::ScalarType a_type; + vllm::ScalarType b_type; + std::optional maybe_group_scales_type; + std::optional maybe_group_zeros_type; + std::optional maybe_channel_scales_type; + std::optional maybe_token_scales_type; + std::optional maybe_out_type; +}; + +torch::stable::Tensor mm_dispatch(MMArgs args); + +std::vector supported_schedules_dispatch( + SupportedSchedulesArgs args); + +template +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, + 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 \ No newline at end of file diff --git a/csrc/quantization/machete/machete_prepack_kernel.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh similarity index 94% rename from csrc/quantization/machete/machete_prepack_kernel.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh index d002355ca49..e1e054e5a00 100644 --- a/csrc/quantization/machete/machete_prepack_kernel.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepack_kernel.cuh @@ -3,6 +3,7 @@ #include "machete_mm_kernel.cuh" #include "cutlass_extensions/cute_utils.cuh" #include "cutlass_extensions/torch_utils.hpp" +#include 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{}); diff --git a/csrc/quantization/machete/machete_prepack_launcher.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh similarity index 65% rename from csrc/quantization/machete/machete_prepack_launcher.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh index 634b651a4d1..94f6f684bc0 100644 --- a/csrc/quantization/machete/machete_prepack_launcher.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepack_launcher.cuh @@ -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 +#include + +#include 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 maybe_group_scales_type; + std::optional maybe_group_scales_type; }; template -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(B.const_data_ptr()); // elements per storage item for B auto eles_per_storage = - (B.dtype().itemsize() * 8) / cute::sizeof_bits_v; + (B.element_size() * 8) / cute::sizeof_bits_v; // 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; auto const l_Bt_packed = make_cute_layout(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( stream, B_ptr, layout_Bt, static_cast(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 \ No newline at end of file diff --git a/csrc/quantization/machete/machete_prepacked_layout.cuh b/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh similarity index 99% rename from csrc/quantization/machete/machete_prepacked_layout.cuh rename to csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh index 4a7d6341e6c..c16a2ab8a33 100644 --- a/csrc/quantization/machete/machete_prepacked_layout.cuh +++ b/csrc/libtorch_stable/quantization/machete/machete_prepacked_layout.cuh @@ -1,9 +1,5 @@ #pragma once -#include -#include -#include - // clang-format off // The cutlass include order matters (annoyingly) diff --git a/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu b/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu new file mode 100644 index 00000000000..7736d5b3ece --- /dev/null +++ b/csrc/libtorch_stable/quantization/machete/machete_pytorch.cu @@ -0,0 +1,77 @@ +#include "machete_mm_launcher.cuh" +#include "machete_prepack_launcher.cuh" +#include "core/scalar_type.hpp" + +#include +#include +#include + +namespace machete { + +using namespace vllm; + +std::vector supported_schedules( + torch::headeronly::ScalarType a_type, int64_t b_type_id, + std::optional maybe_group_scales_type, + std::optional maybe_group_zeros_type, + std::optional maybe_channel_scales_type, + std::optional maybe_token_scales_type, + std::optional 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 const& maybe_out_type, + std::optional const& maybe_group_scales, + std::optional const& maybe_group_zeros, + std::optional maybe_group_size, + std::optional const& maybe_channel_scales, + std::optional const& maybe_token_scales, + std::optional 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 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 diff --git a/csrc/quantization/marlin/.gitignore b/csrc/libtorch_stable/quantization/marlin/.gitignore similarity index 100% rename from csrc/quantization/marlin/.gitignore rename to csrc/libtorch_stable/quantization/marlin/.gitignore diff --git a/csrc/quantization/marlin/awq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/awq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu index 307bae6738e..55ce5b4e732 100644 --- a/csrc/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#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(b_q_weight.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(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)); } diff --git a/csrc/quantization/marlin/dequant.h b/csrc/libtorch_stable/quantization/marlin/dequant.h similarity index 100% rename from csrc/quantization/marlin/dequant.h rename to csrc/libtorch_stable/quantization/marlin/dequant.h diff --git a/csrc/quantization/marlin/generate_kernels.py b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py similarity index 99% rename from csrc/quantization/marlin/generate_kernels.py rename to csrc/libtorch_stable/quantization/marlin/generate_kernels.py index 7b316037ec6..2a038479893 100644 --- a/csrc/quantization/marlin/generate_kernels.py +++ b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py @@ -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.");' ) diff --git a/csrc/quantization/marlin/gptq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/gptq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu index 796e6c5359d..cafa212bccb 100644 --- a/csrc/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#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(b_q_weight.data_ptr()); - uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t const* perm_ptr = + reinterpret_cast(perm.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(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)); } diff --git a/csrc/quantization/marlin/kernel.h b/csrc/libtorch_stable/quantization/marlin/kernel.h similarity index 100% rename from csrc/quantization/marlin/kernel.h rename to csrc/libtorch_stable/quantization/marlin/kernel.h diff --git a/csrc/quantization/marlin/marlin.cu b/csrc/libtorch_stable/quantization/marlin/marlin.cu similarity index 61% rename from csrc/quantization/marlin/marlin.cu rename to csrc/libtorch_stable/quantization/marlin/marlin.cu index 721c206c33f..63fea239e4a 100644 --- a/csrc/quantization/marlin/marlin.cu +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cu @@ -24,7 +24,15 @@ #endif #include "kernel.h" -#include "core/registration.h" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::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 c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional 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 c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional 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 c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& a_scales_or_none, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional 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 c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional 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)); } diff --git a/csrc/quantization/marlin/marlin.cuh b/csrc/libtorch_stable/quantization/marlin/marlin.cuh similarity index 93% rename from csrc/quantization/marlin/marlin.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin.cuh index d3a91568349..bfb65e874b3 100644 --- a/csrc/quantization/marlin/marlin.cuh +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cuh @@ -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 - #include - #include - #endif #include #include #include diff --git a/csrc/quantization/marlin/marlin_dtypes.cuh b/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/quantization/marlin/marlin_dtypes.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu new file mode 100644 index 00000000000..f8ef6b12a01 --- /dev/null +++ b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -0,0 +1,118 @@ + +#include "marlin.cuh" + +#include +#include +#include +#include +#include +#include + +#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 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<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(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<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(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)); +} diff --git a/csrc/quantization/marlin/marlin_mma.h b/csrc/libtorch_stable/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/quantization/marlin/marlin_mma.h rename to csrc/libtorch_stable/quantization/marlin/marlin_mma.h diff --git a/csrc/quantization/marlin/marlin_template.h b/csrc/libtorch_stable/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/quantization/marlin/marlin_template.h rename to csrc/libtorch_stable/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 816f2665048..c805ecba1ba 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -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 diff --git a/csrc/quantization/machete/machete_mm_launcher.cuh b/csrc/quantization/machete/machete_mm_launcher.cuh deleted file mode 100644 index cabe0af46f0..00000000000 --- a/csrc/quantization/machete/machete_mm_launcher.cuh +++ /dev/null @@ -1,75 +0,0 @@ -#pragma once - -#include -#include - -#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 const& maybe_out_type; - std::optional const& maybe_group_scales; - std::optional const& maybe_group_zeros; - std::optional maybe_group_size; - std::optional const& maybe_channel_scales; - std::optional const& maybe_token_scales; - std::optional maybe_schedule; -}; - -struct SupportedSchedulesArgs { - at::ScalarType a_type; - vllm::ScalarType b_type; - std::optional maybe_group_scales_type; - std::optional maybe_group_zeros_type; - std::optional maybe_channel_scales_type; - std::optional maybe_token_scales_type; - std::optional maybe_out_type; -}; - -torch::Tensor mm_dispatch(MMArgs args); - -std::vector supported_schedules_dispatch( - SupportedSchedulesArgs args); - -template -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) - .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 \ No newline at end of file diff --git a/csrc/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu deleted file mode 100644 index 05a51ee21dd..00000000000 --- a/csrc/quantization/machete/machete_pytorch.cu +++ /dev/null @@ -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 supported_schedules( - at::ScalarType a_type, int64_t b_type_id, - std::optional maybe_group_scales_type, - std::optional maybe_group_zeros_type, - std::optional maybe_channel_scales_type, - std::optional maybe_token_scales_type, - std::optional 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 const& maybe_out_type, - std::optional const& maybe_group_scales, - std::optional const& maybe_group_zeros, - std::optional maybe_group_size, - std::optional const& maybe_channel_scales, - std::optional const& maybe_token_scales, - std::optional 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 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 diff --git a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu deleted file mode 100644 index 7d4c97fb57e..00000000000 --- a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ /dev/null @@ -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 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<<>>( - (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<<>>( - (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); -} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 58524c4c5db..cfd185394a4 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -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 } diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 0c7a8ab246e..90aaf01a0b7 100644 Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 3d8fda95a34..22406f2eaa2 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -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 \ diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index a0038f461a0..c57c430478f 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -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--.log`, stripped of timestamps and +ANSI codes: ```bash -# Find the failing job. Each row's URL is .../builds/#: -gh pr checks --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 -# 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/" + +# One job — `gh pr checks` URLs (#) 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/#" ``` diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 5d366253ef7..a585cd77ffb 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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`. diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 8cbbedf9d0b..dd0e47a1950 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -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 diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md new file mode 100644 index 00000000000..b99ba6659f7 --- /dev/null +++ b/docs/design/nixl_kv_push_connector.md @@ -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:`` (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--`` 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: +``` + +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 +`:` 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:`). + +Behavior on the engine main thread is otherwise unchanged. The writer +thread is event-driven and idle when there is no push work. diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 8ab29b43888..03b05751c14 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -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": }}'`. + ## Example Scripts/Code Refer to these example scripts in the vLLM repository: diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index d1a56e83cd4..43010c406f5 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -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`) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6f7cc6dab4b..31a550b95fa 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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 + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ | | `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + IE+ + VE+ | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ | | `InternS2PreviewForConditionalGeneration` | Intern-S2-Preview | T + IE+ + VE+ | `internlm/Intern-S2-Preview`, etc. | ✅︎ | ✅︎ | -| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `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 + IE+ + (VE+) | `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 + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | | `KananaVForConditionalGeneration` | Kanana-V | T + I+ | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ | | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ | diff --git a/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py new file mode 100644 index 00000000000..9f1a0a7f413 --- /dev/null +++ b/examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py @@ -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() diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 40a4b8ae6d1..a7df5b00c3b 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -2532,6 +2532,7 @@ MODELS_NEED_VIDEO_METADATA = [ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ + "llama4", "internvl_chat", "qwen2_5_vl", "qwen3_vl", diff --git a/requirements/common.txt b/requirements/common.txt index d6e2031f534..e42b8600412 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -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 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index ce18ce456cc..a6fc7242174 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -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 diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 12a85421bd3..b49d100da67 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -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, ) } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index e351e7e1c8d..c6bd7c2b12d 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -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(); diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index d70870dc32a..b6619b7a49c 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -72,6 +72,8 @@ impl ManagedEngineArgs { model: String, max_model_len: Option, 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()); diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index cc425ca076f..ce716bb65f7 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -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:")); + } +} diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index ffbf28048da..c11e4c79ca5 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -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(); } }; diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index e93c049b2d1..a8c70d273d0 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -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; } diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2b3e3ddb360..aa430db76cc 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -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( diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index fb64428e4b2..a623925e649 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -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(()) } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index b6e4383c7d1..3dc3bbff6fe 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -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( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 1dd73a4f530..2f6c760a990 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -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, @@ -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!({ diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index 2af8c8add11..2af41877bfd 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -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(()) } diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 039df87f9dd..7ec1251ddf3 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -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; diff --git a/rust/src/server/src/routes/openai/utils/usage.rs b/rust/src/server/src/routes/openai/utils/usage.rs new file mode 100644 index 00000000000..c8c9d1e7262 --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/usage.rs @@ -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) + } +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 68ffe04a3b7..c6de4034026 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -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 { + 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 + Send + 'a>>; fn boxed_test_future<'a>(future: impl Future + 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() { diff --git a/setup.py b/setup.py index 0a820587958..657a65161e7 100644 --- a/setup.py +++ b/setup.py @@ -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 diff --git a/tests/benchmarks/test_audio_dataset.py b/tests/benchmarks/test_audio_dataset.py new file mode 100644 index 00000000000..5957011c484 --- /dev/null +++ b/tests/benchmarks/test_audio_dataset.py @@ -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" diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 93f3abfc088..85307403200 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -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(), diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index ad9fed1d355..21d5154c675 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -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" diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 839793fde85..a3e05027b38 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -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 diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 22077bd4a31..e523cc2d4a3 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -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: diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py b/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py deleted file mode 100644 index 1c058adaf0a..00000000000 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py +++ /dev/null @@ -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 diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 8ca0d1604b1..a16fa83fe32 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -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( diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index d2985264e0c..0027c2763fa 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -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"}) diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py index ac5b8bcd915..bd52863342d 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -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 diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index 9095f80e20f..8e6fdb70452 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -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( diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index fedbd74795b..af61ebc5264 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -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) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 9e4e7c2ec9a..d92dabe9d3e 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -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) diff --git a/tests/kernels/attention/test_mixed_causal_attn.py b/tests/kernels/attention/test_mixed_causal_attn.py new file mode 100644 index 00000000000..5343f701f28 --- /dev/null +++ b/tests/kernels/attention/test_mixed_causal_attn.py @@ -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) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index d4fa9697cb7..f328f339332 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -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) diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py new file mode 100644 index 00000000000..1a19cf34379 --- /dev/null +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -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))}", + ) diff --git a/tests/kernels/helion/test_per_token_group_fp8_quant.py b/tests/kernels/helion/test_per_token_group_fp8_quant.py new file mode 100644 index 00000000000..304734c77e5 --- /dev/null +++ b/tests/kernels/helion/test_per_token_group_fp8_quant.py @@ -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 diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index c82c3c8358e..9876135056b 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -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: diff --git a/tests/kernels/helion/utils.py b/tests/kernels/helion/utils.py new file mode 100644 index 00000000000..38893fc8fec --- /dev/null +++ b/tests/kernels/helion/utils.py @@ -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") diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index f8967b19922..d8c1b9f2cb6 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -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"])) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 83cd2f09d1e..4080ca18459 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -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] diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index c7158dae537..82c0f8813e2 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -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 diff --git a/tests/model_executor/test_mistral_large_3_eagle.py b/tests/model_executor/test_mistral_large_3_eagle.py new file mode 100644 index 00000000000..d8ef109af98 --- /dev/null +++ b/tests/model_executor/test_mistral_large_3_eagle.py @@ -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 diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index e2dd0d9de76..a9afe73cad6 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -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 diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index f781caf492b..a1dc4e5bdd8 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -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, diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index d6f2b86c7af..03c19b0bf62 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -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 diff --git a/tests/models/registry.py b/tests/models/registry.py index d2d2794962f..86641c9b155 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -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", diff --git a/tests/models/utils.py b/tests/models/utils.py index a5d1844a307..8a629552131 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -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, diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py new file mode 100644 index 00000000000..2740ccbca04 --- /dev/null +++ b/tests/parser/test_harmony.py @@ -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"), + ] diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py index ba8bc1427f2..39c5c2e3d5a 100644 --- a/tests/parser/test_parse.py +++ b/tests/parser/test_parse.py @@ -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): diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index 5414d7571a5..db8783c9211 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -16,6 +16,9 @@ MODELS = [ "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear "Qwen/Qwen3-30B-A3B-FP8", # FP8 W8A16 block-quantized MoE "openai/gpt-oss-20b", # MXFP4 W4A16 + "QuixiAI/Qwen3-30B-A3B-AWQ", # AWQ W4A16 MoE + "Qwen/Qwen3-30B-A3B-GPTQ-Int4", # GPTQ W4A16 MoE + "RedHatAI/Qwen3-30B-A3B-quantized.w4a16", # compressed-tensors W4A16 MoE ] DTYPE = ["bfloat16"] diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index e17e6d8ae39..51044696637 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -5,11 +5,16 @@ Run `pytest tests/samplers/test_beam_search.py`. """ +import json + +import jsonschema import pytest from transformers import AutoModelForSeq2SeqLM from vllm.assets.audio import AudioAsset +from vllm.entrypoints.llm import LLM from vllm.platforms import current_platform +from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams # Extra engine kwargs needed for numerically deterministic beam search. # On ROCm, floating-point reductions in attention and GEMM kernels are @@ -223,3 +228,61 @@ def test_beam_search_passes_multimodal_data( # NOTE: encoder/decoder tests are currently located under # tests/models/multimodal/generation/test_whisper.py + + +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("dtype", ["half"]) +@pytest.mark.parametrize("beam_width", BEAM_WIDTHS) +def test_beam_search_structured_output( + model: str, + dtype: str, + beam_width: int, +) -> None: + """Ensure beam search with structured output produces valid JSON.""" + json_schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name", "age"], + "additionalProperties": False, + } + + llm = LLM( + model=model, + dtype=dtype, + max_model_len=512, + structured_outputs_config=dict( + backend="xgrammar", + disable_any_whitespace=True, + ), + **(dict(enforce_eager=True) | EXTRA_ENGINE_KWARGS), + ) + + params = BeamSearchParams( + beam_width=beam_width, + max_tokens=64, + structured_outputs=StructuredOutputsParams(json=json_schema), + ) + + prompts = [ + "Generate a JSON object for a person with name and age:", + ] + + outputs = llm.beam_search(prompts, params) + + assert len(outputs) == len(prompts) + for output in outputs: + assert len(output.sequences) > 0 + for seq in output.sequences: + assert seq.text is not None + print(f"Full text: {seq.text!r}") + # seq.text includes the prompt, extract generated JSON. + gen_start = seq.text.find("{") + assert gen_start != -1, f"No JSON found in output: {seq.text!r}" + generated = seq.text[gen_start:] + generated = generated.replace("", "").strip() + print(f"Generated JSON: {generated!r}") + parsed = json.loads(generated) + jsonschema.validate(instance=parsed, schema=json_schema) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 6f3709e19a4..eea084a2bb4 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -702,6 +702,88 @@ class TestStreamingExtraction: ' \n' ) + def _collect_tool_calls_by_index(self, results): + """Group streamed tool-call fragments by their ``index``. + + Returns ``{index: {"name": str | None, "arguments": str}}`` where + ``arguments`` is the concatenation of every streamed argument + fragment for that index (which should form valid JSON once complete). + """ + by_index: dict[int, dict[str, Any]] = {} + for delta, _ in results: + if not (delta and delta.tool_calls): + continue + for tc in delta.tool_calls: + entry = by_index.setdefault(tc.index, {"name": None, "arguments": ""}) + func = tc.function + if isinstance(func, dict): + name = func.get("name") + arg = func.get("arguments", "") + else: + name = getattr(func, "name", None) + arg = getattr(func, "arguments", "") or "" + if name: + entry["name"] = name + if arg: + entry["arguments"] += arg + return by_index + + def test_streaming_single_chunk_complete_tool_call(self, parser, mock_request): + """A backend may deliver a whole tool call in one streaming delta. + + The start token, ``call:name{...}`` payload and the end token all + arrive in a single chunk. The parser must still emit one + ``DeltaToolCall`` with the correct name + complete arguments JSON + (rather than swallowing it and finishing with finish_reason="stop"). + """ + chunks = [ + '<|tool_call>call:name_a_color{color_hex:<|"|>00ff11<|"|>}', + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + # Exactly one delta should carry tool_calls, and it must not be + # emitted as plain content (which would yield finish_reason="stop"). + tool_call_deltas = [ + delta for delta, _ in results if delta is not None and delta.tool_calls + ] + assert len(tool_call_deltas) == 1, ( + "Expected exactly one delta carrying the batched tool call" + ) + assert all( + delta.content is None for delta, _ in results if delta is not None + ), "Complete tool call must not leak as content" + + by_index = self._collect_tool_calls_by_index(results) + assert set(by_index) == {0} + assert by_index[0]["name"] == "name_a_color" + assert json.loads(by_index[0]["arguments"]) == {"color_hex": "00ff11"} + + def test_streaming_multi_chunk_batched_tool_calls(self, parser, mock_request): + """A single delta may batch MULTIPLE complete tool calls. + + ``<|tool_call>...<|tool_call>...`` arriving in + one chunk must emit BOTH calls (one DeltaToolCall each, with distinct + indices), not just the first. + """ + chunks = [ + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + '<|tool_call>call:get_time{timezone:<|"|>GMT<|"|>}', + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + by_index = self._collect_tool_calls_by_index(results) + assert set(by_index) == {0, 1}, ( + f"Expected two tool calls (indices 0 and 1), got {sorted(by_index)}" + ) + + assert by_index[0]["name"] == "get_weather" + assert json.loads(by_index[0]["arguments"]) == {"location": "London"} + + assert by_index[1]["name"] == "get_time" + assert json.loads(by_index[1]["arguments"]) == {"timezone": "GMT"} + def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request): """Trailing bare boolean must not be streamed twice.""" chunks = [ diff --git a/tests/tool_parsers/test_openai_tool_parser.py b/tests/tool_parsers/test_openai_tool_parser.py deleted file mode 100644 index 843fbca621f..00000000000 --- a/tests/tool_parsers/test_openai_tool_parser.py +++ /dev/null @@ -1,415 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json - -import pytest -from openai_harmony import ( - Conversation, - DeveloperContent, - HarmonyEncodingName, - Message, - Role, - SystemContent, - load_harmony_encoding, -) - -from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall -from vllm.tokenizers import get_tokenizer -from vllm.tool_parsers.openai_tool_parser import OpenAIToolParser - -MODEL = "gpt2" - - -@pytest.fixture(scope="module") -def openai_tokenizer(): - # The parser does not use the tokenizer, but the constructor requires it. - return get_tokenizer(MODEL) - - -@pytest.fixture -def openai_tool_parser(openai_tokenizer): - return OpenAIToolParser(openai_tokenizer) - - -@pytest.fixture(scope="module") -def harmony_encoding(): - return load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) - - -def assert_tool_calls( - actual_tool_calls: list[ToolCall], - expected_tool_calls: list[ToolCall], -): - assert len(actual_tool_calls) == len(expected_tool_calls) - - for actual_tool_call, expected_tool_call in zip( - actual_tool_calls, expected_tool_calls - ): - assert isinstance(actual_tool_call.id, str) - assert len(actual_tool_call.id) > 16 # Default from protocol.py - assert actual_tool_call.type == "function" - assert actual_tool_call.function == expected_tool_call.function - - -def test_extract_tool_calls_no_tools(openai_tool_parser, harmony_encoding): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.SYSTEM, - SystemContent.new(), - ), - Message.from_role_and_content( - Role.DEVELOPER, - DeveloperContent.new().with_instructions("Talk like a pirate!"), - ), - Message.from_role_and_content(Role.USER, "Arrr, how be you?"), - Message.from_role_and_content( - Role.ASSISTANT, "This is a test" - ).with_channel("final"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert not extracted_info.tools_called - assert extracted_info.tool_calls == [] - assert extracted_info.content == "This is a test" - - -@pytest.mark.parametrize( - "tool_args", - [ - '{"location": "Tokyo"}', - '{\n"location": "Tokyo"\n}', - ], -) -def test_extract_tool_calls_single_tool( - openai_tool_parser, harmony_encoding, tool_args -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?"), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" We need to use get_current_weather tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, tool_args) - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_multiple_tools( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" based on their location. We need to use get_current_weather tool and get_user_location tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_user_location") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.no_content_type"), - Message.from_role_and_content(Role.ASSISTANT, "foo") - .with_channel("commentary") - .with_recipient("functions.not_json_no_content_type"), - Message.from_role_and_content(Role.ASSISTANT, "{}") - .with_channel("commentary") - .with_recipient("functions.empty_args") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, "") - .with_channel("commentary") - .with_recipient("functions.no_args") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, - Role.ASSISTANT, - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="get_user_location", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="no_content_type", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="not_json_no_content_type", - arguments="foo", - ) - ), - ToolCall( - function=FunctionCall( - name="empty_args", - arguments=json.dumps({}), - ) - ), - ToolCall( - function=FunctionCall( - name="no_args", - arguments="", - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_bare_function_name( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?"), - Message.from_role_and_content( - Role.ASSISTANT, - "We need to use get_current_weather tool.", - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("get_current_weather") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_bare_function_name_multiple( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - "We need to use both tools.", - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, "{}") - .with_channel("commentary") - .with_recipient("get_user_location") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="get_user_location", - arguments=json.dumps({}), - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - - -def test_extract_tool_calls_assistant_recipient_ignored( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "Hello"), - Message.from_role_and_content(Role.ASSISTANT, "Some tool response") - .with_channel("commentary") - .with_recipient("assistant"), - Message.from_role_and_content( - Role.ASSISTANT, "Here is the answer" - ).with_channel("final"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert not extracted_info.tools_called - assert extracted_info.tool_calls == [] - assert extracted_info.content == "Here is the answer" - - -def test_extract_tool_calls_dotted_function_name( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "Compute 2+3"), - Message.from_role_and_content(Role.ASSISTANT, '{"a": 2, "b": 3}') - .with_channel("commentary") - .with_recipient("math.sum") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="math.sum", - arguments=json.dumps({"a": 2, "b": 3}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - - -def test_extract_tool_calls_with_content( - openai_tool_parser, - harmony_encoding, -): - final_content = "This tool call will get the weather." - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" based on their location. We need to use get_current_weather tool and get_user_location tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, final_content).with_channel( - "final" - ), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, - Role.ASSISTANT, - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content == final_content diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index cec531ca07f..300bae5c52b 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -3,6 +3,7 @@ import json from collections.abc import Generator +from unittest.mock import MagicMock import pytest from openai.types.responses.function_tool import FunctionTool @@ -19,15 +20,12 @@ from vllm.entrypoints.openai.engine.protocol import ( FunctionCall, ToolCall, ) +from vllm.parser.abstract_parser import DelegatingParser from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.detokenizer_utils import detokenize_incrementally from vllm.tool_parsers.qwen3coder_tool_parser import ( Qwen3CoderToolParser, ) -from vllm.tool_parsers.qwen3xml_tool_parser import ( - Qwen3XMLToolParser, - StreamingXMLToolCallParser, -) MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" @@ -43,17 +41,8 @@ def qwen3_tool_parser(qwen3_tokenizer, sample_tools): @pytest.fixture -def qwen3_xml_tool_parser(qwen3_tokenizer, sample_tools): - return Qwen3XMLToolParser(qwen3_tokenizer, tools=sample_tools) - - -@pytest.fixture(params=["xml"]) -def qwen3_tool_parser_parametrized(qwen3_tool_parser, qwen3_xml_tool_parser, request): - """Parameterized fixture that provides both parser types for testing""" - if request.param == "original": - return qwen3_tool_parser - else: - return qwen3_xml_tool_parser +def qwen3_tool_parser_parametrized(qwen3_tool_parser): + return qwen3_tool_parser WEATHER_PARAMS = { @@ -168,47 +157,6 @@ def assert_tool_calls( ) -def test_qwen3xml_deferred_array_parses_json_literals(): - parser = StreamingXMLToolCallParser() - parser.set_tools( - [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "AskUserQuestion", - "parameters": QUESTION_PARAMS, - }, - ) - ] - ) - - delta = parser.parse_single_streaming_chunks( - """ - - -[{"question": "Pick a color", "multiSelect": false, "answer": null}] - - -""" - ) - - arguments = "".join( - tool_call.function.arguments or "" - for tool_call in delta.tool_calls or [] - if tool_call.function and tool_call.function.arguments is not None - ) - - assert json.loads(arguments) == { - "questions": [ - { - "question": "Pick a color", - "multiSelect": False, - "answer": None, - } - ] - } - - def stream_delta_message_generator( qwen3_tool_parser, qwen3_tokenizer: TokenizerLike, @@ -523,7 +471,7 @@ hello world """ - parser = Qwen3XMLToolParser(qwen3_tokenizer, tools=tools) + parser = Qwen3CoderToolParser(qwen3_tokenizer, tools=tools) request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) extracted_tool_calls = parser.extract_tool_calls(model_output, request=request) @@ -1146,125 +1094,6 @@ TX assert parsed_args["state"] == "TX" -def test_extract_tool_calls_complex_type_with_single_quote( - qwen3_tokenizer, -): - """Test parameter type conversion based on tool schema""" - tools = [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "test_types", - "parameters": { - "type": "object", - "properties": { - "int_param": {"type": "integer"}, - "float_param": {"type": "float"}, - "bool_param": {"type": "boolean"}, - "str_param": {"type": "string"}, - "obj_param": {"type": "object"}, - }, - }, - }, - ) - ] - - model_output = """ - - -{'key': 'value'} - - -""" - - parser = Qwen3XMLToolParser(qwen3_tokenizer, tools=tools) - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - extracted_tool_calls = parser.extract_tool_calls(model_output, request=request) - - args = json.loads(extracted_tool_calls.tool_calls[0].function.arguments) - assert args["obj_param"] == {"key": "value"} - - -def test_extract_tool_calls_streaming_missing_opening_tag( - qwen3_tool_parser_parametrized, qwen3_tokenizer -): - """Test streaming with missing opening tag - - This tests that the streaming parser correctly handles - tool calls that start directly with - """ - model_output = """I'll check the weather for you. - - - -Dallas - - -TX - - -fahrenheit - - -""" - - request = ChatCompletionRequest(model=MODEL, messages=[]) - - other_content = "" - tool_states = {} - - for delta_message in stream_delta_message_generator( - qwen3_tool_parser_parametrized, qwen3_tokenizer, model_output, request - ): - if delta_message.content: - other_content += delta_message.content - - if delta_message.tool_calls: - for tool_call in delta_message.tool_calls: - idx = tool_call.index - - if idx not in tool_states: - tool_states[idx] = { - "id": None, - "name": None, - "arguments": "", - "type": None, - } - - if tool_call.id: - tool_states[idx]["id"] = tool_call.id - - if tool_call.type: - assert tool_call.type == "function" - tool_states[idx]["type"] = tool_call.type - - if tool_call.function: - if tool_call.function.name: - tool_states[idx]["name"] = tool_call.function.name - - if tool_call.function.arguments is not None: - tool_states[idx]["arguments"] += tool_call.function.arguments - - # Verify content was streamed - assert "I'll check the weather for you." in other_content - - # Verify we got the tool call - assert len(tool_states) == 1 - assert len(qwen3_tool_parser_parametrized.prev_tool_call_arr) == 1 - - state = tool_states[0] - assert state["id"] is not None - assert state["type"] == "function" - assert state["name"] == "get_current_weather" - - # Verify arguments were parsed correctly despite missing opening tag - assert state["arguments"] is not None - args = json.loads(state["arguments"]) - assert args["city"] == "Dallas" - assert args["state"] == "TX" - assert args["unit"] == "fahrenheit" - - def test_malformed_xml_no_gt_delimiter(qwen3_tool_parser): """Regression: malformed XML without '>' must not crash (PR #36774).""" model_output = ( @@ -1456,15 +1285,12 @@ def test_get_vllm_registry_structural_tag_returns_structural_tag( @pytest.mark.parametrize("include_reasoning", [True, False]) def test_adjust_request_auto_uses_vllm_registry_structural_tag( - monkeypatch: pytest.MonkeyPatch, - qwen3_tool_parser: Qwen3CoderToolParser, sample_tools: list[ChatCompletionToolsParam], include_reasoning: bool, ) -> None: - monkeypatch.setattr( - "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", - True, - ) + class TestParser(DelegatingParser): + tool_parser_cls = Qwen3CoderToolParser + request_tools = _as_chat_completion_tools(sample_tools) req = ChatCompletionRequest( messages=[], @@ -1473,7 +1299,7 @@ def test_adjust_request_auto_uses_vllm_registry_structural_tag( tool_choice="auto", include_reasoning=include_reasoning, ) - out = qwen3_tool_parser.adjust_request(req) + out = TestParser(MagicMock(), tools=sample_tools).adjust_request(req) assert out.structured_outputs is not None assert out.structured_outputs.structural_tag is not None assert isinstance(out.structured_outputs.structural_tag, str) @@ -1482,14 +1308,11 @@ def test_adjust_request_auto_uses_vllm_registry_structural_tag( def test_adjust_request_required_prefers_structural_tag( - monkeypatch: pytest.MonkeyPatch, - qwen3_tool_parser: Qwen3CoderToolParser, sample_tools: list[ChatCompletionToolsParam], ) -> None: - monkeypatch.setattr( - "vllm.tool_parsers.abstract_tool_parser.VLLM_ENFORCE_STRICT_TOOL_CALLING", - True, - ) + class TestParser(DelegatingParser): + tool_parser_cls = Qwen3CoderToolParser + request_tools = _as_chat_completion_tools(sample_tools) req = ChatCompletionRequest( messages=[], @@ -1497,6 +1320,6 @@ def test_adjust_request_required_prefers_structural_tag( tools=request_tools, tool_choice="required", ) - out = qwen3_tool_parser.adjust_request(req) + out = TestParser(MagicMock(), tools=sample_tools).adjust_request(req) assert out.structured_outputs is not None assert out.structured_outputs.structural_tag is not None diff --git a/tests/tool_parsers/test_qwen3xml_tool_parser.py b/tests/tool_parsers/test_qwen3xml_tool_parser.py deleted file mode 100644 index 1ea9a1d65c0..00000000000 --- a/tests/tool_parsers/test_qwen3xml_tool_parser.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import pytest - -from tests.tool_parsers.common_tests import ( - ToolParserTestConfig, - ToolParserTests, -) - - -class TestQwen3xmlToolParser(ToolParserTests): - @pytest.fixture - def test_config(self) -> ToolParserTestConfig: - return ToolParserTestConfig( - parser_name="qwen3_xml", - # Test data - no_tool_calls_output="This is a regular response without any tool calls.", - single_tool_call_output="\n\nTokyo\n\n", - parallel_tool_calls_output="\n\nTokyo\n\n\n\nAsia/Tokyo\n\n", - various_data_types_output=( - "\n\n" - "hello\n" - "42\n" - "3.14\n" - "true\n" - "null\n" - '["a", "b", "c"]\n' - '{"nested": "value"}\n' - "\n" - ), - empty_arguments_output="\n\n\n", - surrounding_text_output=( - "Let me check the weather for you.\n\n" - "\n\n" - "Tokyo\n" - "\n\n\n" - "I will get that information." - ), - escaped_strings_output=( - "\n\n" - 'He said "hello"\n' - "C:\\Users\\file.txt\n" - "line1\nline2\n" - "\n" - ), - malformed_input_outputs=[ - "", - "", - ], - # Expected results - single_tool_call_expected_name="get_weather", - single_tool_call_expected_args={"city": "Tokyo"}, - parallel_tool_calls_count=2, - parallel_tool_calls_names=["get_weather", "get_time"], - # xfail markers - Qwen3XML has systematic streaming issues - xfail_streaming={ - "test_single_tool_call_simple_args": ( - "Qwen3XML streaming has systematic issues" - ), - "test_parallel_tool_calls": "Qwen3XML streaming has systematic issues", - "test_various_data_types": "Qwen3XML streaming has systematic issues", - "test_empty_arguments": "Qwen3XML streaming has systematic issues", - "test_surrounding_text": "Qwen3XML streaming has systematic issues", - "test_escaped_strings": "Qwen3XML streaming has systematic issues", - "test_streaming_reconstruction": ( - "Qwen3XML streaming reconstruction has known issues" - ), - }, - supports_typed_arguments=False, - ) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py new file mode 100644 index 00000000000..645603d2303 --- /dev/null +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from xgrammar import StructuralTag + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.parser.abstract_parser import DelegatingParser +from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.tool_parsers.deepseekv3_tool_parser import DeepSeekV3ToolParser +from vllm.tool_parsers.deepseekv4_tool_parser import DeepSeekV4ToolParser +from vllm.tool_parsers.deepseekv31_tool_parser import DeepSeekV31ToolParser +from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser +from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser +from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser +from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser +from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser +from vllm.tool_parsers.qwen3coder_tool_parser import Qwen3CoderToolParser +from vllm.tool_parsers.structural_tag_registry import ( + SUPPORTED_STRUCTURAL_TAG_MODELS, + VLLM_BUILTIN_STRUCTURAL_TAG_MODELS, + XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS, + _get_function_parameters, + get_model_structural_tag, +) + + +@pytest.fixture +def sample_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + ) + ] + + +def test_supported_structural_tag_models_include_vllm_builtins(): + assert SUPPORTED_STRUCTURAL_TAG_MODELS == ( + XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS + ) + assert "hermes" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS + + +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_get_model_structural_tag_supports_all_xgrammar_builtins( + model: str, + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice="auto", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + + +def test_get_model_structural_tag_supports_vllm_hermes( + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model="hermes", + tools=sample_tools, + tool_choice="required", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + assert tag.model_dump() == { + "type": "structural_tag", + "format": { + "type": "tags_with_separator", + "tags": [ + { + "type": "tag", + "begin": '\n{"name": "get_weather", "arguments": ', + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + "style": "json", + }, + "end": "}\n", + }, + { + "type": "tag", + "begin": '{"name": "get_weather", "arguments": ', + "content": { + "type": "json_schema", + "json_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + "style": "json", + }, + "end": "}", + }, + ], + "separator": "", + "at_least_one": True, + "stop_after_first": False, + }, + } + + +def test_hermes_required_tool_calls_use_empty_separator(): + tools = [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_time", + "parameters": {"type": "object", "properties": {}}, + }, + ), + ] + + tag = get_model_structural_tag( + model="hermes", + tools=tools, + tool_choice="required", + reasoning=False, + ) + + assert tag is not None + assert tag.format.separator == "" + + +@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) +def test_get_model_structural_tag_supports_named_tool_choice( + model: str, + sample_tools: list[ChatCompletionToolsParam], +): + tag = get_model_structural_tag( + model=model, + tools=sample_tools, + tool_choice=ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name="get_weather") + ), + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + + +@pytest.mark.parametrize( + ("parser_cls", "model"), + [ + (DeepSeekV3ToolParser, "deepseek_r1"), + (DeepSeekV31ToolParser, "deepseek_v3_1"), + (DeepSeekV32ToolParser, "deepseek_v3_2"), + (DeepSeekV4ToolParser, "deepseek_v4"), + (Glm47MoeModelToolParser, "glm_4_7"), + (Hermes2ProToolParser, "hermes"), + (KimiK2ToolParser, "kimi"), + (Llama3JsonToolParser, "llama"), + (MinimaxM2ToolParser, "minimax"), + (Qwen3CoderToolParser, "qwen_3_coder"), + ], +) +def test_tool_parsers_declare_matching_xgrammar_builtin_model(parser_cls, model): + assert parser_cls.structural_tag_model == model + assert not parser_cls.supports_required_and_named + + +def test_tool_parsers_without_structural_tag_support_required_and_named(): + class NonStructuralTagToolParser(ToolParser): + pass + + assert NonStructuralTagToolParser.structural_tag_model is None + assert NonStructuralTagToolParser.supports_required_and_named + + +def test_non_structural_tag_parser_uses_schema_constraints( + sample_tools: list[ChatCompletionToolsParam], +): + parser = ToolParser(MagicMock()) + request = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="required", + ) + + out = parser.adjust_request(request) + + assert out.structured_outputs is not None + assert out.structured_outputs.json is not None + assert out.structured_outputs.structural_tag is None + + +def test_get_structural_tag_disables_reasoning( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + captured: list[bool] = [] + + def fake_get_model_structural_tag(*, reasoning: bool, **kwargs): + captured.append(reasoning) + return None + + monkeypatch.setattr( + "vllm.tool_parsers.structural_tag_registry.get_model_structural_tag", + fake_get_model_structural_tag, + ) + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="auto", + ) + parser = Qwen3CoderToolParser(MagicMock(), tools=sample_tools) + + parser.get_structural_tag(request) + + assert captured == [False] + + +def test_unified_parser_get_structural_tag_disables_reasoning( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + captured: list[bool] = [] + + def fake_get_model_structural_tag(*, reasoning: bool, **kwargs): + captured.append(reasoning) + return None + + monkeypatch.setattr( + "vllm.tool_parsers.structural_tag_registry.get_model_structural_tag", + fake_get_model_structural_tag, + ) + + class TestParser(DelegatingParser): + tool_parser_cls = Qwen3CoderToolParser + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools, + tool_choice="auto", + ) + parser = TestParser(MagicMock(), tools=sample_tools) + parser.reasoning_parser = MagicMock(adjust_request=lambda request: request) + + parser.adjust_request(request) + + assert captured == [False] + + +def test_xgrammar_function_parameters_are_preserved( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + captured: list[list[dict]] = [] + + def fake_get_xgrammar_model_structural_tag(*, tools: list[dict], **kwargs): + captured.append(tools) + return None + + monkeypatch.setattr( + "vllm.tool_parsers.structural_tag_registry.get_xgrammar_model_structural_tag", + fake_get_xgrammar_model_structural_tag, + ) + + get_model_structural_tag( + model="llama", + tools=sample_tools, + tool_choice="auto", + reasoning=False, + ) + + assert ( + captured[0][0]["function"]["parameters"] == sample_tools[0].function.parameters + ) + assert sample_tools[0].function.parameters is not None + + +def test_get_function_parameters_relaxes_function_strict_false(): + function = SimpleNamespace( + parameters={"type": "object", "properties": {}}, + strict=False, + ) + + assert _get_function_parameters(function) is True diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c2eb576d895..3be24d7fb34 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -28,6 +28,7 @@ from vllm.v1.core.kv_cache_utils import ( estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, + get_kv_cache_capacity, get_kv_cache_configs, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, @@ -1459,6 +1460,11 @@ def test_get_max_concurrency_for_kv_cache_config(): vllm_config, kv_cache_config_hybrid_model ) assert max_concurrency_hybrid_model == 3 + num_tokens, max_concurrency = get_kv_cache_capacity( + vllm_config, kv_cache_config_hybrid_model + ) + assert num_tokens == max_concurrency_hybrid_model * max_model_len + assert max_concurrency == max_concurrency_hybrid_model def test_allocate_with_lookahead(): diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 4d652beec81..dc8d7152b70 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -15,6 +15,7 @@ from vllm.config import ( SpeculativeConfig, VllmConfig, ) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalKwargsItem, @@ -1849,6 +1850,8 @@ def create_scheduler_with_priority( enable_chunked_prefill=True, is_encoder_decoder=model_config.is_encoder_decoder, policy="priority", # Enable priority scheduling + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( @@ -3988,6 +3991,87 @@ def test_delayed_kv_connector_free_keeps_scheduler_active(): assert not scheduler.has_finished_requests() +def test_scheduler_kv_connector_stats(): + """Test worker-side, scheduler-side, and combined KV connector stats.""" + + class GenericKVConnectorStats(KVConnectorStats): + def reset(self): + self.data = {} + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + self.data.update(other.data) + return self + + def reduce(self) -> dict[str, int | float]: + return {} + + def is_empty(self) -> bool: + return not self.data + + test_cases = ( + ({"worker": 1}, None, {"worker": 1}), + (None, {"scheduler": 2}, {"scheduler": 2}), + ({"worker": 1}, {"scheduler": 2}, {"worker": 1, "scheduler": 2}), + ) + + for worker_data, scheduler_data, expected_data in test_cases: + scheduler = create_scheduler() + worker_stats = ( + GenericKVConnectorStats(data=worker_data) if worker_data else None + ) + scheduler_stats = ( + GenericKVConnectorStats(data=scheduler_data) if scheduler_data else None + ) + scheduler.connector = Mock() + scheduler.connector.get_kv_connector_stats.return_value = ( + scheduler_stats if worker_stats is None else None + ) + scheduler.connector.take_events.return_value = [] + + def update_connector_output( + kv_connector_output: KVConnectorOutput, + scheduler=scheduler, + scheduler_stats=scheduler_stats, + ): + scheduler.connector.get_kv_connector_stats.return_value = scheduler_stats + + scheduler.connector.update_connector_output.side_effect = ( + update_connector_output + ) + + model_output = ModelRunnerOutput( + req_ids=["req_0"], + req_id_to_index={"req_0": 0}, + sampled_token_ids=[[123]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[None], + kv_connector_output=KVConnectorOutput(kv_connector_stats=worker_stats) + if worker_stats + else None, + ) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=None, + num_scheduled_tokens={"req_0": 1}, + total_num_scheduled_tokens=1, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + engine_core_outputs = scheduler.update_from_output( + scheduler_output, model_output + ) + + final_stats = next( + iter(engine_core_outputs.values()) + ).scheduler_stats.kv_connector_stats + assert final_stats == expected_data + + # ============================================================================== # Variable-length encoder cross-attention block allocation tests # ============================================================================== diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 7213a669c53..7f34250cb21 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -90,6 +90,8 @@ def create_scheduler( enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, is_encoder_decoder=model_config.is_encoder_decoder, + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 97b5fd46a2e..c10835821f5 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -49,6 +49,7 @@ def _create_vllm_config( ) mock_config.parallel_config = ParallelConfig() mock_config.speculative_config = None # No speculative decoding + mock_config.num_speculative_tokens = 0 if not lora_config: mock_config.lora_config = None else: diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index ceae041c6f9..e857b127285 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -181,6 +181,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, ): ret = original_allocate_slots_fn( self, @@ -194,6 +195,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): num_encoder_tokens, full_sequence_must_fit, reserved_blocks, + has_scheduled_reqs, ) if cur_step_action is not None: cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[ diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 20c230a4c2a..11da73b3152 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1381,3 +1381,54 @@ def test_stale_sliding_window_block_after_prepare_store_failure( expected_stored=(2, 3), expected_flushed=(2, 3) if not async_scheduling else (), ) + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): + """When skip_reading_prefix_cache=True, the offloading connector must not + load any blocks from CPU even if a matching prefix is cached there.""" + block_size = 4 + block_size_factor = 3 + offloaded_block_size = block_size * block_size_factor + num_gpu_blocks = 100 + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + ) + + # Populate the CPU offload cache with one block. + runner.new_request(token_ids=[0] * offloaded_block_size) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + expected_flushed=(0, 1, 2) if not async_scheduling else (), + ) + + # Reset GPU prefix cache so the next request cannot hit locally. + runner.scheduler.reset_prefix_cache() + + # New request with identical tokens but skip_reading_prefix_cache=True. + # The offloading connector must not load anything from CPU, but must + # still offload the freshly computed blocks (state management intact). + runner.new_request( + token_ids=[0] * offloaded_block_size, + skip_reading_prefix_cache=True, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_loaded=(), # no CPU loads must happen + expected_stored=(0, 1, 2), # tokens still offloaded to CPU + expected_flushed=(0, 1, 2) if not async_scheduling else (), + ) + + # The external lookup must have been completely skipped. + runner.manager.lookup.assert_not_called() diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 22d00b0c834..f6a354ebd43 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -324,10 +324,14 @@ class RequestRunner: self, token_ids: list[int], kv_transfer_params: dict | None = None, + skip_reading_prefix_cache: bool = False, ): self.req_id += 1 - sampling_params = SamplingParams(max_tokens=1000) + sampling_params = SamplingParams( + max_tokens=1000, + skip_reading_prefix_cache=skip_reading_prefix_cache or None, + ) sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) req = Request( diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index ef092dfb49f..12831601cba 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -32,7 +32,7 @@ from unittest.mock import patch import pytest from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import ( +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( NixlConnector, NixlConnectorMetadata, ) @@ -436,7 +436,7 @@ def test_build_connector_meta_multiple_requests(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_p_node_pull_kv_from_d(dist_init): @@ -450,7 +450,7 @@ def test_p_node_pull_kv_from_d(dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_p_node_pull_then_send_kv(dist_init): @@ -472,7 +472,7 @@ def test_p_node_pull_then_send_kv(dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_p_node_deferred_pull_on_no_handshake(dist_init): diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index 4a2ca6d2721..0c0f9f1f899 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -104,7 +104,7 @@ def _run_engine_core_handshake( speculative_config=None, ec_transfer_config=None, max_concurrent_batches=1, - model_config=SimpleNamespace(runner_type="generate"), + model_config=SimpleNamespace(runner_type="generate", is_diffusion=False), cache_config=SimpleNamespace( enable_prefix_caching=False, prefix_caching_hash_algo="builtin", diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 69593011db9..d3992b02b68 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -614,3 +614,66 @@ def test_lookup_key_server_reset_skips_drain_when_no_send_thread(): assert call_order == ["remove_all"] assert sent == [protocol.RESP_OK] + + +def test_shutdown_closes_worker_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + connector.shutdown() + + worker.close.assert_called_once_with() + + +def test_del_invokes_shutdown_and_closes_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + # __del__ is the GC backstop; it must route through shutdown() -> close(). + connector.__del__() + + worker.close.assert_called_once_with() + + +def test_shutdown_scheduler_role_is_noop(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreScheduler" + ), + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + + # Scheduler role holds no store handle, so shutdown must be a safe no-op. + assert connector.connector_worker is None + connector.shutdown() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 8cd5e6e5358..1130a7d6a78 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1558,3 +1558,34 @@ def test_lookup_records_mooncake_metrics(): assert isinstance(stats, MooncakeStoreConnectorStats) assert len(stats.data["lookup_exists"]) == 1 assert stats.data["lookup_exists"][0]["num_keys"] == 2 + + +def test_store_worker_close_releases_store(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + + store.close.assert_called_once_with() + assert worker.store is None + + +def test_store_worker_close_is_idempotent(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + worker.close() + + # Second call short-circuits because store was already released. + store.close.assert_called_once_with() + + +def test_store_worker_close_swallows_store_errors(): + worker = _make_bare_worker() + worker.store.close.side_effect = RuntimeError("boom") + + # A failure tearing down the store must not propagate out of close(). + worker.close() + + assert worker.store is None diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index f78037a1431..2d6fa834d22 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -58,6 +58,7 @@ class MockConnector(KVConnectorBase_V1): mock = MagicMock(spec_set=KVConnectorBase_V1) # Override just build_kv_connector_stats mock.build_kv_connector_stats = cls.build_kv_connector_stats + mock.get_kv_connector_stats.return_value = None return mock @classmethod @@ -93,6 +94,7 @@ class MockHMAConnector(KVConnectorBase_V1, SupportsHMA): def __new__(cls, *args, **kwargs): mock = MagicMock(spec_set=cls) + mock.get_kv_connector_stats.return_value = None return mock def start_load_kv(self, forward_context, **kwargs): @@ -366,7 +368,10 @@ def test_multi_example_connector_consistency(): def _ignore_event_collection(events: list[str]) -> list[str]: - return [event for event in events if event != "take_events"] + # Filter out per-step polling hooks that the scheduler calls repeatedly + # and which are not meaningful state transitions for these assertions. + ignored = {"get_kv_connector_stats", "has_pending_push_work", "take_events"} + return [event for event in events if event not in ignored] def get_connector_events() -> dict[str, list[str]]: @@ -1072,7 +1077,7 @@ def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch): ) with patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ): llm = LLM( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index a2a46684bb7..32652118d52 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -344,7 +344,7 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_kv_transfer_handshake(dist_init): @@ -560,7 +560,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): class TestNixlHandshake: @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_multi_xfer_one_engine( @@ -643,7 +643,7 @@ class TestNixlHandshake: connector.clear_connector_metadata() @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) @pytest.mark.parametrize( @@ -713,7 +713,7 @@ class TestNixlHandshake: raise TimeoutError("Took too long to complete async handshake.") @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) @pytest.mark.parametrize("local_tp_size", [1, 2]) @@ -725,7 +725,7 @@ class TestNixlHandshake: remote configurations. """ monkeypatch.setattr( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", lambda: local_tp_size, ) @@ -784,7 +784,7 @@ class TestNixlHandshake: check_handshake(6) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_prefill_tp_size_greater_than_decode_tp_size_mla( @@ -887,7 +887,7 @@ class TestNixlHandshake: assert req_id not in conn_p1.connector_worker._reqs_to_process @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_concurrent_load_kv( @@ -952,7 +952,7 @@ class TestNixlHandshake: raise TimeoutError("Took too long to complete async handshake.") @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_fails_on_kv_cache_layout_mismatch( @@ -967,7 +967,7 @@ class TestNixlHandshake: # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 with patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): # Initialize connector and worker (with fake NIXL wrapper) @@ -1007,7 +1007,7 @@ class TestNixlHandshake: worker.add_remote_agent(meta, remote_tp_size=1) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( @@ -1022,7 +1022,7 @@ class TestNixlHandshake: # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 with patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): # Initialize connector and worker (with fake NIXL wrapper) @@ -1063,12 +1063,93 @@ class TestNixlHandshake: # whole block is moved. worker.add_remote_agent(meta, remote_tp_size=1) + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): + """Mixed full-attn (SPLIT) + MLA (REPLICATE) single KV group under + heterogeneous TP must NOT raise (previously a NotImplementedError), + and the per-region gate must still reject a wrong block_len. + """ + vllm_config = create_vllm_config() + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501 + return_value=2, + ): + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + + # Region 0: full-attn (SPLIT). Region 1: MLA (REPLICATE). + fa_len = 4096 * worker.block_size + idx_len = 512 * worker.block_size + worker.slot_size_per_layer = [4096, 512] + worker.block_len_per_layer = [fa_len, idx_len] + worker._region_is_mla = [False, True] + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + worker.src_blocks_data = [ + (0, fa_len, worker.tp_rank), + (0, idx_len, worker.tp_rank), + ] + worker.num_descs = len(worker.src_blocks_data) + + # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio; + # REPLICATE region is unchanged. + tp_ratio = 2 + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0, 0], + device_id=0, + num_blocks=1, + block_lens=[fa_len * tp_ratio, idx_len], + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + worker.add_remote_agent(meta, remote_tp_size=1) + assert ( + FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles + ) + # Gate rejects an MLA region wrongly scaled by tp_ratio. + worker2 = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker2.block_len_per_layer = [fa_len, idx_len] + worker2._region_is_mla = [False, True] + worker2.num_blocks = 1 + worker2.dst_num_blocks[worker2.engine_id] = worker2.num_blocks + bad_meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0, 0], + device_id=0, + num_blocks=1, + # WRONG: MLA region scaled by tp_ratio (it should be replicated). + block_lens=[fa_len * tp_ratio, idx_len * tp_ratio], + kv_cache_layout=worker2.kv_cache_layout, + block_size=worker2.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker2.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + with pytest.raises(AssertionError): + worker2.add_remote_agent(bad_meta, remote_tp_size=1) + # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which # we put here is important. First run ray, it will clean up the resources, then # the rest of the tests. @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_kv_connector_stats(default_vllm_config, dist_init): @@ -1282,7 +1363,7 @@ def test_multi_kv_connector_stats_aggregation(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_scheduler_kv_connector_stats_aggregation(): @@ -1347,7 +1428,7 @@ def test_scheduler_kv_connector_stats_aggregation(): @pytest.mark.parametrize("distributed_executor_backend", ["ray", None]) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): @@ -1534,7 +1615,7 @@ def test_register_kv_caches( backend_cls = TritonAttentionBackend - nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker" + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector" with ( patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper, @@ -1784,15 +1865,17 @@ def test_kv_buffer_to_nixl_memory_types( _NIXL_SUPPORTED_DEVICE.update(FakePlatform.get_nixl_supported_devices()) with ( - patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Event" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper" ), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Thread" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Event" ), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.current_platform", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Thread" + ), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.current_platform", FakePlatform, ), patch( @@ -1811,7 +1894,7 @@ def test_kv_buffer_to_nixl_memory_types( @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): @@ -1910,7 +1993,7 @@ def _setup_worker_with_remote_engine( @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_engine_ttl_eviction(default_vllm_config, dist_init): @@ -1945,7 +2028,7 @@ def test_engine_ttl_eviction(default_vllm_config, dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_engine_ttl_disabled(default_vllm_config, dist_init): @@ -1993,7 +2076,7 @@ def test_transfer_topology_unregister(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init): @@ -2113,7 +2196,7 @@ class FailingNixlWrapper(FakeNixlWrapper): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) @pytest.mark.parametrize( @@ -2203,10 +2286,13 @@ def test_transfer_failure_logging( slot_mapping={}, ) - # Capture logs from the nixl.worker logger specifically + # Capture logs from the nixl connector loggers # vLLM loggers have propagate=False, so we need to capture directly nixl_logger = logging.getLogger( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" + ) + pull_logger = logging.getLogger( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker" ) captured_logs: list[logging.LogRecord] = [] @@ -2217,6 +2303,7 @@ def test_transfer_failure_logging( handler = LogCapture() handler.setLevel(logging.ERROR) nixl_logger.addHandler(handler) + pull_logger.addHandler(handler) try: connector.start_load_kv(dummy_ctx) @@ -2232,6 +2319,7 @@ def test_transfer_failure_logging( connector.get_finished(finished_req_ids=set()) finally: nixl_logger.removeHandler(handler) + pull_logger.removeHandler(handler) # Print logs for manual comparison between commits error_logs = [r for r in captured_logs if r.levelno >= logging.ERROR] @@ -2268,7 +2356,7 @@ def test_transfer_failure_logging( @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) def test_handshake_failure_returns_finished(default_vllm_config, dist_init): @@ -2319,7 +2407,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init): @@ -2373,7 +2461,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, ) @pytest.mark.parametrize( @@ -2516,7 +2604,7 @@ def test_failed_request_skips_kv_postprocessing( ], ) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_compatibility_hash_validation( @@ -2625,7 +2713,7 @@ def test_compatibility_hash_validation( # Patch zmq_ctx to return our mock socket with ( patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"), - patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx, + patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx, ): mock_zmq_ctx.return_value.__enter__.return_value = mock_socket @@ -2659,7 +2747,7 @@ def test_compatibility_hash_validation( ], ) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario): @@ -2725,7 +2813,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) mock_socket.recv.return_value = msg_bytes with ( patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"), - patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx, + patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx, ): mock_zmq_ctx.return_value.__enter__.return_value = mock_socket @@ -2738,7 +2826,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) ) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, ) def test_mla_broadcast_notif_uses_remote_request_id( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index af043113ed1..eed20e03668 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -34,7 +34,9 @@ from .utils import ( (False, [0]), ], ) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform" +) def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes): """Test sw_sizes is correctly computed based on SWA enabled/disabled.""" from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( @@ -782,7 +784,9 @@ def test_mamba_n1_p_side_truncation(): ], ids=["fa_swa_mamba", "fa_swa_only", "fa_only"], ) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform" +) def test_has_mamba_init( mock_platform, swa_enabled, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py new file mode 100644 index 00000000000..fe67c1ac73a --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -0,0 +1,815 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for NixlPushConnector (scheduler + worker). + +These tests cover the end-to-end mechanics of the push design without +requiring a real NIXL agent or network: + +* Scheduler stages D registrations on ``update_state_after_alloc`` and + P finished blocks on ``request_finished``. +* ``build_connector_meta`` drains them onto + ``meta.push_registrations`` / ``meta.push_finished_blocks``. +* ``has_pending_push_work`` reports True/False over the lifecycle. +* ``update_connector_output`` clears state on ``finished_sending`` and + ``finished_recving``. +* The worker matches D registrations against P finished blocks (both + scenario directions) and forwards non-PUSH_REG NIXL notifs to the main + thread's ``_get_new_notifs``. +* ``get_finished`` enqueues evictions for the writer. +""" + +from __future__ import annotations + +import logging +import queue +import threading +import time +from collections import defaultdict +from typing import Any +from unittest.mock import MagicMock, patch + +import msgspec + +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + PUSH_REG_NOTIF_PREFIX, + NixlConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( + get_base_request_id, +) +from vllm.v1.outputs import KVConnectorOutput + +from .utils import make_nixl_push_scheduler + +# ----------------------------------------------------------------- # +# Helpers / fakes # +# ----------------------------------------------------------------- # + + +def _make_request( + *, + request_id: str, + is_d_side: bool = True, + remote_engine_id: str = "prefill-engine", + remote_request_id: str | None = None, + remote_host: str = "10.0.0.1", + remote_port: int = 5601, + tp_size: int = 1, + finished: bool = True, +) -> MagicMock: + """Build a minimal Request mock used by request_finished.""" + from vllm.v1.request import RequestStatus + + req = MagicMock() + req.request_id = request_id + req.num_computed_tokens = 64 + + if is_d_side: + # D-side request: do_remote_prefill=True -> prefill on a remote P. + params: dict[str, Any] = { + "do_remote_prefill": True, + "do_remote_decode": False, + "remote_engine_id": remote_engine_id, + "remote_request_id": remote_request_id or f"prefill-{request_id}", + "remote_host": remote_host, + "remote_port": remote_port, + "tp_size": tp_size, + } + else: + # P-side request: do_remote_decode=True (we are the prefiller). + params = { + "do_remote_prefill": False, + "do_remote_decode": True, + } + req.kv_transfer_params = params + req.status = ( + RequestStatus.FINISHED_LENGTH_CAPPED if finished else RequestStatus.RUNNING + ) + return req + + +class _BlocksMock: + """Minimal stand-in for ``KVCacheBlocks`` used in update_state_after_alloc.""" + + def __init__(self, block_ids: tuple[list[int], ...]): + self._block_ids = block_ids + + def get_unhashed_block_ids_all_groups(self) -> tuple[list[int], ...]: + return self._block_ids + + +def _stub_sw_clipping(scheduler) -> None: + """Make ``get_sw_clipped_blocks`` a passthrough so tests don't need + the full sliding-window machinery.""" + scheduler.get_sw_clipped_blocks = lambda block_ids: block_ids + + +# ----------------------------------------------------------------- # +# Scheduler-side tests # +# ----------------------------------------------------------------- # + + +class TestPushScheduler: + def test_d_side_update_state_after_alloc_stages_registration(self): + """D scheduler stashes registration data + arms watchdog deadline.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = _make_request(request_id="req-d-1") + blocks = _BlocksMock(block_ids=([10, 11, 12],)) + + sched.update_state_after_alloc(request, blocks, num_external_tokens=48) + + assert request.request_id in sched._push_pending_registrations + reg = sched._push_pending_registrations[request.request_id] + # ``request_id`` is D's own vLLM request id; plus our own (D) coords. + assert reg["request_id"] == request.request_id + assert reg["decode_engine_id"] == sched.engine_id + assert reg["decode_host"] == sched.side_channel_host + assert reg["decode_port"] == sched.side_channel_port + assert reg["local_block_ids"] == ([10, 11, 12],) + assert reg["remote_engine_id"] == "prefill-engine" + + # Watchdog deadline set in the future. + deadline = sched._push_registration_deadlines[request.request_id] + assert deadline > time.perf_counter() + # do_remote_prefill flipped off so the request isn't reprocessed. + assert request.kv_transfer_params["do_remote_prefill"] is False + # Tracked as awaiting a recv. + assert request.request_id in sched._reqs_need_recv + + def test_p_side_request_finished_stages_blocks(self): + """P scheduler pushes blocks into both _finished_request_blocks (lease) + and _newly_finished_push_blocks (metadata for next step).""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = _make_request(request_id="req-p-1", is_d_side=False) + block_ids = ([20, 21, 22, 23],) + + delay, ret_params = sched.request_finished(request, block_ids) + + assert delay is True + assert ret_params is not None + assert ret_params["do_remote_prefill"] is True + assert ret_params["do_remote_decode"] is False + assert request.request_id in sched._finished_request_blocks + assert request.request_id in sched._newly_finished_push_blocks + assert request.request_id in sched._reqs_need_send # lease armed + + def test_build_connector_meta_drains_both_sides(self): + """meta.push_registrations and meta.push_finished_blocks are filled + from the staging dicts and the staging dicts are cleared.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + # Stage one D registration and one P finished entry. + d_req = _make_request(request_id="req-d-9") + sched.update_state_after_alloc( + d_req, _BlocksMock(([1, 2, 3],)), num_external_tokens=48 + ) + p_req = _make_request(request_id="req-p-9", is_d_side=False) + sched.request_finished(p_req, ([4, 5, 6],)) + + scheduler_output = MagicMock() + scheduler_output.scheduled_new_reqs = [] + scheduler_output.scheduled_cached_reqs = MagicMock( + req_ids=[], resumed_req_ids=set() + ) + + # Patch parent build_connector_meta so we don't have to set up + # all the base scheduler plumbing. + with patch.object( + sched.__class__.__mro__[1], + "build_connector_meta", + return_value=NixlConnectorMetadata(), + ): + meta = sched.build_connector_meta(scheduler_output) + + assert isinstance(meta, NixlConnectorMetadata) + assert "req-d-9" in meta.push_registrations + assert "req-p-9" in meta.push_finished_blocks + # Staging dicts cleared. + assert sched._push_pending_registrations == {} + assert sched._newly_finished_push_blocks == {} + # Lease bookkeeping kept until the WRITE completes. + assert "req-p-9" in sched._finished_request_blocks + + def test_has_pending_push_work_lifecycle(self): + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + assert sched.has_pending_push_work() is False + + # P finished blocks waiting for WRITE completion. + p_req = _make_request(request_id="req-p-7", is_d_side=False) + sched.request_finished(p_req, ([0, 1],)) + assert sched.has_pending_push_work() is True + + # Drain via build_connector_meta - lease still pending until WRITE. + scheduler_output = MagicMock() + scheduler_output.scheduled_new_reqs = [] + scheduler_output.scheduled_cached_reqs = MagicMock( + req_ids=[], resumed_req_ids=set() + ) + with patch.object( + sched.__class__.__mro__[1], + "build_connector_meta", + return_value=NixlConnectorMetadata(), + ): + sched.build_connector_meta(scheduler_output) + # Lease is pending until WRITE completes -> still True. + assert sched.has_pending_push_work() is True + + # Simulate WRITE completion via update_connector_output. + sched.update_connector_output( + KVConnectorOutput( + finished_sending={"req-p-7"}, + finished_recving=set(), + invalid_block_ids=set(), + ) + ) + assert sched.has_pending_push_work() is False + + def test_update_connector_output_clears_lease_and_watchdog(self): + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + d_req = _make_request(request_id="req-d-x") + sched.update_state_after_alloc( + d_req, _BlocksMock(([1, 2],)), num_external_tokens=32 + ) + p_req = _make_request(request_id="req-p-x", is_d_side=False) + sched.request_finished(p_req, ([3, 4],)) + + sched.update_connector_output( + KVConnectorOutput( + finished_sending={"req-p-x"}, + finished_recving={"req-d-x"}, + invalid_block_ids=set(), + ) + ) + assert "req-p-x" not in sched._finished_request_blocks + assert "req-d-x" not in sched._push_registration_deadlines + + def test_registration_watchdog_expires(self, caplog): + """Stale D registrations whose deadline has passed are dropped at + ``build_connector_meta`` time.""" + # Watchdog logs a WARNING when it drops the stale entry; that's + # what this test is verifying, so silence it in the test report. + caplog.set_level( + logging.CRITICAL, + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler"), + ) + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + d_req = _make_request(request_id="req-d-stale") + sched.update_state_after_alloc( + d_req, _BlocksMock(([7, 8],)), num_external_tokens=32 + ) + # Force the deadline into the past. + sched._push_registration_deadlines[d_req.request_id] = time.perf_counter() - 1.0 + + scheduler_output = MagicMock() + scheduler_output.scheduled_new_reqs = [] + scheduler_output.scheduled_cached_reqs = MagicMock( + req_ids=[], resumed_req_ids=set() + ) + with patch.object( + sched.__class__.__mro__[1], + "build_connector_meta", + return_value=NixlConnectorMetadata(), + ): + meta = sched.build_connector_meta(scheduler_output) + + assert d_req.request_id not in sched._push_registration_deadlines + assert d_req.request_id not in sched._push_pending_registrations + assert d_req.request_id not in meta.push_registrations + + +# ----------------------------------------------------------------- # +# Worker-side tests # +# ----------------------------------------------------------------- # + + +class _StubWriterWorker(NixlPushConnectorWorker): + """Construct a worker without invoking ``__init__`` so we can drive + the matching/notif logic without bringing up NIXL or torch.""" + + @classmethod + def fresh(cls) -> _StubWriterWorker: + w = object.__new__(cls) + + # Push-specific state managed by NixlPushConnectorWorker. + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + ReqId, + TransferHandle, + ) + + w._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list) + w._sending_transfers_lock = threading.Lock() + w._push_finished_blocks = {} + w._pending_d_registrations = {} + w._reg_send_inbox = queue.Queue() + w._finished_blocks_inbox = queue.Queue() + w._pending_completion_notifs = queue.Queue() + w._evict_finished_inbox = queue.Queue() + w._push_writer_wake = threading.Event() + w._push_writer_stop = threading.Event() + w._push_writer_thread = None + + # Base worker fields touched by start_load_kv / _get_new_notifs. + w._recving_metadata = {} + w._recving_transfers = defaultdict(list) + w._reqs_to_process = set() + w._reqs_to_send = {} + w.consumer_notification_counts_by_req = defaultdict(int) + w.tp_rank = 0 + w.world_size = 1 + w.engine_id = "test-decode-engine" + w._remote_agents = {} + + # Track _do_start_push_kv invocations. + calls: list[tuple[str, Any, dict[str, Any]]] = [] + w.start_push_calls = calls + return w + + def _do_start_push_kv( + self, + request_id: str, + local_block_ids, + registration_data: dict[str, Any], + ) -> None: # pragma: no cover - exercised through tests + # Track the call instead of issuing real WRITEs. + self.start_push_calls.append((request_id, local_block_ids, registration_data)) + + +def _registration_data( + request_id: str, + *, + decode_engine_id: str = "decode-engine", + decode_host: str = "10.0.0.2", + decode_port: int = 5602, + decode_tp_size: int = 1, + local_block_ids=((100, 101, 102),), + remote_engine_id: str = "prefill-engine", + remote_host: str = "10.0.0.1", + remote_port: int = 5601, + remote_tp_size: int = 1, +) -> dict[str, Any]: + return { + "request_id": request_id, + "decode_engine_id": decode_engine_id, + "decode_host": decode_host, + "decode_port": decode_port, + "decode_tp_size": decode_tp_size, + "local_block_ids": local_block_ids, + "remote_engine_id": remote_engine_id, + "remote_host": remote_host, + "remote_port": remote_port, + "remote_tp_size": remote_tp_size, + } + + +class TestPushWriterMatching: + def test_handle_push_reg_matches_existing_finished_blocks(self): + """PUSH_REG arrives second (P finished first): match + fire.""" + w = _StubWriterWorker.fresh() + # P had already finished; its blocks were stashed via metadata. + w._push_finished_blocks["req-A"] = ([200, 201, 202],) + + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode( + _registration_data("req-A") + ) + w._handle_push_reg_notif(notif) + + assert len(w.start_push_calls) == 1 + rid, blocks, reg = w.start_push_calls[0] + assert rid == "req-A" + assert blocks == ([200, 201, 202],) + assert reg["decode_engine_id"] == "decode-engine" + # Finished blocks consumed. + assert "req-A" not in w._push_finished_blocks + assert w._pending_d_registrations == {} + + def test_handle_push_reg_stashes_when_no_finished_blocks_yet(self): + """PUSH_REG arrives first (D registered first): stash, no fire.""" + w = _StubWriterWorker.fresh() + + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode( + _registration_data("req-B") + ) + w._handle_push_reg_notif(notif) + + assert len(w.start_push_calls) == 0 + assert "req-B" in w._pending_d_registrations + + def test_handle_push_reg_matches_after_stripping_random_suffix(self): + """P and D assign the same logical request the same + ``cmpl--`` but different per-engine random suffixes; + the writer should still match P's finished blocks via the + suffix-stripping fallback in ``_pop_matching_finished_blocks``. + """ + w = _StubWriterWorker.fresh() + # Same base id + completion index; differ only in the trailing + # ``-<8 hex>`` randomization suffix. + p_id = "cmpl-12345678-aaaa-bbbb-cccc-1234567890ab-0-aaaaaaaa" + d_id = "cmpl-12345678-aaaa-bbbb-cccc-1234567890ab-0-bbbbbbbb" + # Sanity: same base id under the helper used by the connector. + assert get_base_request_id(p_id) == get_base_request_id(d_id) + + w._push_finished_blocks[p_id] = ([1, 2, 3],) + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(_registration_data(d_id)) + w._handle_push_reg_notif(notif) + + # Suffix-stripped fallback matched and fired. + assert len(w.start_push_calls) == 1 + assert w.start_push_calls[0][0] == p_id + assert p_id not in w._push_finished_blocks + + def test_handle_push_reg_drops_malformed(self, caplog): + # The writer logs WARNING/ERROR when it sees these bad payloads; + # that's the desired behavior, so suppress the noise from test + # output rather than letting it look like a failure. + caplog.set_level( + logging.CRITICAL, + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), + ) + w = _StubWriterWorker.fresh() + # Missing request_id -> should drop without raising. + bad = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode({"decode_engine_id": "x"}) + w._handle_push_reg_notif(bad) + assert w._pending_d_registrations == {} + assert w.start_push_calls == [] + + # Undecodable payload also dropped. + w._handle_push_reg_notif(PUSH_REG_NOTIF_PREFIX + b"\xff\xff\xff") + assert w.start_push_calls == [] + + +class TestPushWriterStartLoadKv: + def test_finished_blocks_inbox_matches_stashed_registration(self): + """Run the writer-loop's finished-blocks drain against a + pre-populated _pending_d_registrations entry.""" + w = _StubWriterWorker.fresh() + w._pending_d_registrations["req-C"] = _registration_data("req-C") + + # Simulate start_load_kv enqueuing finished blocks. + w._finished_blocks_inbox.put(("req-C", ([10, 11, 12],))) + + # Drain like the writer loop does. + while True: + try: + rid, blocks = w._finished_blocks_inbox.get_nowait() + except queue.Empty: + break + matched = w._pop_matching_registration(rid) + if matched is not None: + w._do_start_push_kv(rid, blocks, matched) + else: + w._push_finished_blocks[rid] = blocks + + assert len(w.start_push_calls) == 1 + assert w.start_push_calls[0][0] == "req-C" + assert "req-C" not in w._pending_d_registrations + + def test_start_load_kv_enqueues_to_writer(self): + """``start_load_kv`` should hand registrations + finished blocks + to the writer queues without doing matching itself.""" + w = _StubWriterWorker.fresh() + # Stub heartbeats to a no-op; tests don't exercise the heartbeat + # path here. + w._send_heartbeats = lambda metadata: None + # Stub logical-to-kernel mapping used by reqs_to_recv. + w._logical_to_kernel_block_ids = lambda x: x + + meta = NixlConnectorMetadata() + meta.push_registrations = { + "req-D": _registration_data("req-D"), + } + meta.push_finished_blocks = { + "req-E": ([5, 6, 7],), + } + + w.start_load_kv(meta) + + # Things are queued for the writer; nothing fires yet. + assert w._reg_send_inbox.qsize() == 1 + assert w._finished_blocks_inbox.qsize() == 1 + assert w._push_writer_wake.is_set() + assert w.start_push_calls == [] + + +class TestPushWriterNotifs: + def test_get_new_notifs_processes_forwarded_completion_notif(self): + """Non-PUSH_REG notifs forwarded by the writer thread are drained + on the engine main thread inside ``_get_new_notifs``.""" + w = _StubWriterWorker.fresh() + # Pretend the writer thread already forwarded a completion notif + # for a request whose KV is being received. + request_id = "req-recv-1" + w._recving_metadata[request_id] = MagicMock() + # Compose the standard completion notif: req_id:tp_size. + notif_msg = f"{request_id}:1".encode() + w._pending_completion_notifs.put(notif_msg) + + # transfer_topo is consulted only for the producer-side path; we + # make it a MagicMock because the D-side branch returns early. + w.transfer_topo = MagicMock() + + notified = w._get_new_notifs() + + # Notif consumed; D-side just touches _recving_transfers. + assert notified == set() + assert request_id in w._recving_transfers + + def test_get_finished_evicts_completed_state(self): + """``get_finished`` should enqueue evictions and wake the writer.""" + w = _StubWriterWorker.fresh() + + # Stub the base ``get_finished`` to return one done_sending entry. + # Patch via the MRO's parent class. + with patch.object( + NixlPushConnectorWorker.__mro__[1], + "get_finished", + return_value=({"req-done"}, set()), + ): + done_sending, done_recving = w.get_finished() + + assert "req-done" in done_sending + assert done_recving == set() + # Eviction enqueued for the writer. + evicted = [] + while True: + try: + evicted.append(w._evict_finished_inbox.get_nowait()) + except queue.Empty: + break + assert evicted == ["req-done"] + assert w._push_writer_wake.is_set() + + +# ----------------------------------------------------------------- # +# Negative / error-path tests # +# ----------------------------------------------------------------- # + + +class TestPushSchedulerNegative: + """Failure / no-op paths on the scheduler side.""" + + def test_update_state_after_alloc_no_kv_transfer_params_is_noop(self): + """Requests without kv_transfer_params must not register anything.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = MagicMock() + request.request_id = "req-no-params" + request.kv_transfer_params = None + + sched.update_state_after_alloc( + request, _BlocksMock(([1, 2, 3],)), num_external_tokens=64 + ) + + assert sched._push_pending_registrations == {} + assert sched._push_registration_deadlines == {} + assert sched._reqs_need_recv == {} + + def test_update_state_after_alloc_zero_external_tokens_does_not_register(self): + """num_external_tokens=0 should not stage a D registration.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = _make_request(request_id="req-zero-ext") + sched.update_state_after_alloc( + request, _BlocksMock(([1, 2, 3],)), num_external_tokens=0 + ) + + assert sched._push_pending_registrations == {} + assert sched._push_registration_deadlines == {} + + def test_request_finished_unfinished_status_does_not_stage(self): + """If a request is still RUNNING, request_finished must not stash + blocks for the worker (no push needed).""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = _make_request( + request_id="req-running", is_d_side=False, finished=False + ) + + delay, ret = sched.request_finished(request, ([1, 2, 3],)) + + assert delay is False + assert ret is None + assert sched._finished_request_blocks == {} + assert sched._newly_finished_push_blocks == {} + + def test_request_finished_empty_blocks_does_not_arm_lease(self): + """Empty block-id groups should still complete cleanly without + arming the lease/finished maps.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + request = _make_request(request_id="req-empty", is_d_side=False) + delay, ret = sched.request_finished(request, ((),)) + + assert delay is False + assert ret is not None + assert "req-empty" not in sched._finished_request_blocks + assert "req-empty" not in sched._newly_finished_push_blocks + assert "req-empty" not in sched._reqs_need_send + + def test_update_connector_output_unknown_request_is_noop(self): + """Idempotent cleanup: clearing a request that was never staged + must not raise or mutate other state.""" + sched = make_nixl_push_scheduler() + _stub_sw_clipping(sched) + + # Stage one real request to ensure it's NOT touched. + live = _make_request(request_id="req-live", is_d_side=False) + sched.request_finished(live, ([1],)) + + sched.update_connector_output( + KVConnectorOutput( + finished_sending={"unknown-1"}, + finished_recving={"unknown-2"}, + invalid_block_ids=set(), + ) + ) + + # Live entry untouched. + assert "req-live" in sched._finished_request_blocks + + +class TestPushWriterNegative: + """Failure / drop / idempotence paths in the writer thread.""" + + def test_pop_matching_registration_returns_none_when_empty(self): + w = _StubWriterWorker.fresh() + assert w._pop_matching_registration("nope") is None + + def test_pop_matching_finished_blocks_returns_none_when_empty(self): + w = _StubWriterWorker.fresh() + assert w._pop_matching_finished_blocks("nope") is None + + def test_pop_matching_registration_no_match_when_base_ids_differ(self): + """A registration whose base id (after stripping the random suffix) + does NOT match the lookup request_id must not be popped.""" + w = _StubWriterWorker.fresh() + # Two unrelated requests: different base UUIDs, so stripping the + # trailing ``-<8 hex>`` suffix still yields different base ids. + unrelated_d = "cmpl-aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa-0-11111111" + lookup = "cmpl-bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb-0-22222222" + assert get_base_request_id(unrelated_d) != get_base_request_id(lookup) + + w._pending_d_registrations[unrelated_d] = _registration_data(unrelated_d) + result = w._pop_matching_registration(lookup) + assert result is None + # Original entry untouched. + assert unrelated_d in w._pending_d_registrations + + def test_handle_push_reg_with_non_dict_payload_is_dropped(self, caplog): + """msgpack-encoded non-dict payload (e.g. a list) should be + dropped without raising.""" + caplog.set_level( + logging.CRITICAL, + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), + ) + w = _StubWriterWorker.fresh() + bad = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode([1, 2, 3]) + w._handle_push_reg_notif(bad) + assert w._pending_d_registrations == {} + assert w.start_push_calls == [] + + def test_handle_push_reg_with_non_string_request_id_is_dropped(self, caplog): + """request_id must be a str; integers, None, etc. must drop.""" + caplog.set_level( + logging.CRITICAL, + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), + ) + w = _StubWriterWorker.fresh() + for bogus_rid in (123, None, 4.5, b"bytes-not-str"): + payload = _registration_data("placeholder") + payload["request_id"] = bogus_rid # type: ignore[assignment] + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(payload) + w._handle_push_reg_notif(notif) + assert w._pending_d_registrations == {} + assert w.start_push_calls == [] + + def test_handle_push_reg_idempotent_for_same_request_id(self): + """Receiving the same PUSH_REG twice (e.g. P retries after a + flake) keeps the entry staged exactly once and never fires.""" + w = _StubWriterWorker.fresh() + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode( + _registration_data("req-dup") + ) + w._handle_push_reg_notif(notif) + w._handle_push_reg_notif(notif) + assert "req-dup" in w._pending_d_registrations + assert len(w._pending_d_registrations) == 1 + assert w.start_push_calls == [] + + def test_get_finished_enqueues_eviction_for_each_done_request(self): + """``get_finished`` must enqueue an eviction for every request + in ``done_sending`` so the writer can drop stale matching state. + Unlike the happy-path test, this verifies the *cardinality*: N + completed requests -> N evictions, in order.""" + w = _StubWriterWorker.fresh() + with patch.object( + NixlPushConnectorWorker.__mro__[1], + "get_finished", + return_value=({"req-1", "req-2", "req-3"}, set()), + ): + done_sending, _ = w.get_finished() + assert done_sending == {"req-1", "req-2", "req-3"} + + evicted: list[str] = [] + while True: + try: + evicted.append(w._evict_finished_inbox.get_nowait()) + except queue.Empty: + break + assert sorted(evicted) == ["req-1", "req-2", "req-3"] + + def test_get_finished_with_no_completions_does_not_enqueue_eviction(self): + """If there's nothing newly done, no eviction should be enqueued. + The wake event IS still set because ``get_finished`` always wakes + the writer to drain notifs.""" + w = _StubWriterWorker.fresh() + with patch.object( + NixlPushConnectorWorker.__mro__[1], + "get_finished", + return_value=(set(), set()), + ): + done_sending, done_recving = w.get_finished() + assert done_sending == set() + assert done_recving == set() + assert w._evict_finished_inbox.qsize() == 0 + # Wake set so the writer drains NIXL notifs even when idle. + assert w._push_writer_wake.is_set() + + def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): + """A completion notif for a request the worker doesn't know + about should be logged but not crash.""" + caplog.set_level( + logging.CRITICAL, + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"), + ) + w = _StubWriterWorker.fresh() + w.transfer_topo = MagicMock() + # Forward a completion notif for an unknown request_id. + w._pending_completion_notifs.put(b"never-heard-of-you:1") + + notified = w._get_new_notifs() + assert notified == set() + # Did not register anywhere. + assert "never-heard-of-you" not in w._recving_transfers + + def test_start_load_kv_with_empty_metadata_is_noop(self): + """Empty metadata must not wake the writer or enqueue anything.""" + w = _StubWriterWorker.fresh() + w._send_heartbeats = lambda metadata: None + w._logical_to_kernel_block_ids = lambda x: x + + meta = NixlConnectorMetadata() + w.start_load_kv(meta) + + assert w._reg_send_inbox.qsize() == 0 + assert w._finished_blocks_inbox.qsize() == 0 + # Wake should NOT be set if there was nothing to push. + assert not w._push_writer_wake.is_set() + + def test_get_new_notifs_extends_lease_on_heartbeat(self): + """``HB:`` notifs forwarded by the writer thread must extend the + leases of tracked P-side requests on the engine main thread, and + ignore request IDs that aren't being tracked.""" + w = _StubWriterWorker.fresh() + w.transfer_topo = MagicMock() + # _handle_heartbeat reads ``self._lease_extension`` (set in the + # real ``__init__``). + w._lease_extension = 10 + + # Tracked P-side requests with a lease about to expire. + old_expiry = time.perf_counter() - 5.0 + w._reqs_to_send["req-a"] = old_expiry + w._reqs_to_send["req-b"] = old_expiry + + # Forwarded heartbeat covers a tracked request, an unknown one, + # and another tracked one. + w._pending_completion_notifs.put(b"HB:req-a,req-unknown,req-b") + + notified = w._get_new_notifs() + assert notified == set() + + # Tracked leases were renewed strictly forward in time. + now = time.perf_counter() + for rid in ("req-a", "req-b"): + assert w._reqs_to_send[rid] > old_expiry + # New expiry must be roughly now + _lease_extension. + assert w._reqs_to_send[rid] >= now + # Unknown request must not be inserted by the heartbeat path. + assert "req-unknown" not in w._reqs_to_send diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py index 0760d7141ec..78e9e1196fd 100644 --- a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -44,7 +44,7 @@ from vllm.v1.simple_kv_offload.metadata import ( ) NIXL_WRAPPER_PATCH = ( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper" ) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index c432b1b20ed..34a8ec57281 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -554,3 +554,91 @@ def test_fs_tiering_offloading(tmp_path) -> None: finally: subscriber.close() del llm + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="HMA mamba-align CPU offload test is CUDA-only", +) +@pytest.mark.parametrize( + "model,block_size,tp_size", + [ + # ("Qwen/Qwen3.6-35B-A3B", 1056, 2), + # ("tiiuae/falcon-mamba-7b", 16, 1), + ("state-spaces/mamba-1.4b-hf", 16, 1) + ], +) +def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int): + kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "cpu_bytes_to_use": 4 << 30, + "block_size": block_size, + }, + ) + llm = LLM( + model=model, + max_model_len=block_size * 10, + gpu_memory_utilization=0.85, + tensor_parallel_size=tp_size, + kv_transfer_config=kv_transfer_config, + language_model_only=True, + enable_prefix_caching=True, + mamba_cache_mode="align", + disable_hybrid_kv_cache_manager=False, + ) + + _PROMPT_SIZE: int = block_size * 2 + _PROMPT_TEXT = "Hi. Give me a set of trivia questions and their answers " + + # build prompt ids to match prompt_size + tokenizer = llm.get_tokenizer() + raw_ids: list[int] = tokenizer.encode(_PROMPT_TEXT) + while len(raw_ids) < _PROMPT_SIZE: + raw_ids = tokenizer.encode("....") + raw_ids + initial_ids: list[int] = raw_ids[:_PROMPT_SIZE] + + sampling_params = SamplingParams(max_tokens=128, temperature=0, ignore_eos=True) + + failures: list[str] = [] + + def _get_output_str(outputs): + return outputs[0].outputs[0].text + + def _verify(llm, prompt, label: str): + cold_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) + _wait_for_prefix_cache_reset(llm) + cpu_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) + + cold_text = _get_output_str(cold_outputs) + cpu_text = _get_output_str(cpu_outputs) + print(f"{label} : cold outputs\n{cold_text}") + print(f"{label} : cpu outputs\n{cpu_text}") + + if cold_text != cpu_text: + failures.append( + f"{label}: mismatch\n cold: {cold_text!r}\n cpu: {cpu_text!r}" + ) + + try: + # Mamba has only a single state. The CPU cache stores are triggered + # at offload block boundaries. When the prompt is exactly at the boundary, + # The CPU offload should not load the cached block. + # This is because we'd use that state to recompute the last token. This + # does not work for mamba as there is only one KV value and that is for + # for the token at the boundary. + # This is fine for other attention types as we have all the necessary + # token KV values in the hit blocks. + prompt = TokensPrompt(prompt_token_ids=initial_ids) + _verify(llm, prompt, "block-boundary-prompt") + + # Test for prompt token ids at non-block boundaries. + # Reuse is okay for this case. + prompt = TokensPrompt(prompt_token_ids=[0] + initial_ids) + _verify(llm, prompt, "block-mid-prompt") + + assert not failures, "\n\n".join(failures) + + finally: + del llm diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index d92b6326763..95e8254fe40 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -587,7 +587,9 @@ def test_cannot_recv(): assert_scheduler_empty(scheduler) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform" +) def test_p_side_chunked_prefill_mamba(mock_platform): """P-side integration: Mamba N-1 truncation + chunked prefill completes. diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 95d49faf042..5ab6b68400c 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -73,9 +73,19 @@ class TestTPMappingStructure: def _make_mock_worker_for_splits(group_spec_types): - """Build a mock NixlConnectorWorker with _group_spec_types for split tests.""" + """Build a mock NixlConnectorWorker with _group_spec_types for split tests. + + No per-region replicate flags are configured (``block_len_per_layer`` empty + and ``num_regions == 0``), so ``_fa_desc_replicated`` takes its early-return + path and treats every FA descriptor as SPLIT, matching the legacy behavior + these tests assert. + """ worker = object.__new__(NixlConnectorWorker) worker._group_spec_types = group_spec_types + worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=False) + worker.block_len_per_layer = [] + worker.num_regions = 0 + worker._region_is_mla = [] return worker diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index c5411be6207..7df9e20e6a5 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -524,3 +524,66 @@ def make_nixl_scheduler( sched.blocks_per_sw = [] sched.is_bidirectional_kv_xfer_enabled = False return sched + + +def make_nixl_push_scheduler( + *, + decoder_kv_blocks_ttl: float = 30.0, + push_registration_timeout: float | None = None, + is_bidirectional_kv_xfer_enabled: bool = False, + has_mamba: bool = False, +): + """Create a NixlPushConnectorScheduler via __new__ (skipping __init__). + + The push scheduler can't reuse :func:`make_nixl_scheduler` because it + is a different class (``NixlPushConnectorScheduler`` vs + ``NixlConnectorScheduler``) and carries push-specific state. Only the + fields touched by the unit tests are populated. + """ + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( + NixlPushConnectorScheduler, + ) + + sched = object.__new__(NixlPushConnectorScheduler) + + # Base scheduler fields (shared with pull / heartbeat path). + sched._reqs_need_recv = {} + sched._reqs_need_send = {} + sched._reqs_in_batch = set() + sched._reqs_not_processed = set() + sched._reqs_need_save = {} + sched._kv_lease_duration = 30 + sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl + sched.use_host_buffer = False + sched.engine_id = "decode-engine" + sched.side_channel_host = "127.0.0.1" + sched.side_channel_port = 5600 + sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled + sched._has_mamba = has_mamba + + # vllm_config is consulted for parallel_config.tensor_parallel_size. + vllm_config = MagicMock() + vllm_config.parallel_config.tensor_parallel_size = 1 + sched.vllm_config = vllm_config + + # Push-specific state. + sched._push_pending_registrations = {} + sched._push_registration_deadlines = {} + sched._finished_request_blocks = {} + sched._newly_finished_push_blocks = {} + sched._push_registration_timeout = ( + push_registration_timeout + if push_registration_timeout is not None + else decoder_kv_blocks_ttl + ) + + # Heartbeat fields touched by base request_finished / + # update_connector_output. + sched._heartbeat_by_engine = {} + sched._heartbeat_req_engine = {} + sched._last_heartbeat_time = 0.0 + sched.blocks_per_sw = [] + + return sched diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index bd77fbe91fa..ab30f1bb9e2 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -28,6 +28,7 @@ from vllm.v1.metrics.perf import ( ExecutionContext, FfnMetrics, InvalidComponent, + MLAAttentionMetrics, ModelMetrics, ParsedArgs, UnembedMetrics, @@ -1021,3 +1022,317 @@ def test_quantized_model_metrics_aggregation(): assert total_flops > 0 assert total_flops == sum(breakdown.values()) + + +#### MLA Attention Tests #### + + +def test_mla_config_parser(): + """Test MLAConfigParser extracts MLA-specific fields from DeepseekV3Config.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=61, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + + parser_chain = MLAAttentionMetrics.get_parser() + result = parser_chain.parse(vllm_config) + + assert result.kv_lora_rank == 512 + assert result.qk_nope_head_dim == 128 + assert result.qk_rope_head_dim == 64 + assert result.v_head_dim == 128 + assert result.q_lora_rank == 1536 + assert result.num_attention_heads == 128 + assert result.hidden_size == 7168 + + +def test_mla_attention_metrics_decode(): + """Test MLA decode metrics use compressed KV cache, not standard head_dim.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + # Single decode token with 1024 context + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=1024, is_prefill=False + ) + + write_breakdown = metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + + # KV cache write should be 1 * (512 + 64) * cache_byte_size * 1 layer + # = 576 * 2 = 1152 bytes (for bfloat16 cache) + kv_compressed_dim = 512 + 64 # kv_lora_rank + qk_rope_head_dim + expected_kv_cache_write = 1 * kv_compressed_dim * 2 * 1 # T * dim * bytes * L + assert write_breakdown["kv_cache"] == expected_kv_cache_write + + # Verify read bytes include compressed KV cache reads for context + read_breakdown = metrics.get_read_bytes_breakdown(ctx, per_gpu=False) + assert "attn_input" in read_breakdown + assert read_breakdown["attn_input"] > 0 + + +def test_mla_attention_metrics_prefill(): + """Test MLA prefill metrics account for low-rank Q and KV projections.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=2048, context_len=2048, is_prefill=True + ) + + flops_breakdown = metrics.get_num_flops_breakdown(ctx, per_gpu=False) + + # Should have two-stage Q projection (q_a and q_b) + assert "q_a_proj" in flops_breakdown + assert "q_b_proj" in flops_breakdown + assert "q_proj" not in flops_breakdown # Since q_lora_rank is not None + + # Should have KV projections + assert "kv_a_proj" in flops_breakdown + assert "kv_b_proj" in flops_breakdown + + # Should have attention and output + assert "attn_qk" in flops_breakdown + assert "attn_av" in flops_breakdown + assert "out_proj" in flops_breakdown + + # Verify q_a_proj: 2 * T * D * q_lora_rank * L + expected_q_a = 2 * 2048 * 7168 * 1536 * 1 + assert flops_breakdown["q_a_proj"] == expected_q_a + + # Verify kv_a_proj: 2 * T * D * (kv_lora_rank + qk_rope_head_dim) * L + expected_kv_a = 2 * 2048 * 7168 * (512 + 64) * 1 + assert flops_breakdown["kv_a_proj"] == expected_kv_a + + +def test_mla_kv_cache_vs_standard_attention(): + """Test MLA KV cache writes are dramatically smaller than standard MHA.""" + # MLA config (DeepSeek-V3 style) + mla_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=1, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ) + mla_vllm_config = create_mock_vllm_config(mla_config) + mla_metrics = MLAAttentionMetrics.from_vllm_config(mla_vllm_config) + + # Standard MHA config with same num_heads and head_dim + standard_config = Qwen3Config( + hidden_size=7168, + num_attention_heads=128, + num_key_value_heads=128, # MHA: same as num_heads + num_hidden_layers=1, + head_dim=128, + ) + standard_vllm_config = create_mock_vllm_config(standard_config) + standard_metrics = AttentionMetrics.from_vllm_config(standard_vllm_config) + + # Compare KV cache write for 100 tokens + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=100, is_prefill=True + ) + + mla_write = mla_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + standard_write = standard_metrics.get_write_bytes_breakdown(ctx, per_gpu=False) + + # MLA: T * (kv_lora_rank + qk_rope_head_dim) * cache_bytes * L + # = 100 * 576 * 2 * 1 = 115,200 + mla_kv_cache = mla_write["kv_cache"] + + # Standard: 2 * T * num_kv_heads * head_dim * cache_bytes * L + # = 2 * 100 * 128 * 128 * 2 * 1 = 6,553,600 + standard_kv_cache = standard_write["kv_cache"] + + # MLA KV cache should be dramatically smaller (about 57x) + assert mla_kv_cache < standard_kv_cache + ratio = standard_kv_cache / mla_kv_cache + assert ratio > 50 # Should be ~56.9x + + +def test_mla_per_gpu_with_tensor_parallelism(): + """Test MLA metrics with tensor parallelism.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=8, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + # Test with TP=8 + vllm_config = create_mock_vllm_config(hf_config, tensor_parallel_size=8) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=64, context_len=1024, is_prefill=True + ) + + global_flops = metrics.get_num_flops(ctx, per_gpu=False) + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) + + # Both should be positive + assert global_flops > 0 + assert per_gpu_flops > 0 + # Global should exceed per-GPU + assert global_flops > per_gpu_flops + + +def test_mla_per_gpu_with_pipeline_parallelism(): + """Test MLA metrics with pipeline parallelism.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=16, # Divisible by PP + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=4) + metrics = MLAAttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=512, is_prefill=False + ) + + global_flops = metrics.get_num_flops(ctx, per_gpu=False) + per_gpu_flops = metrics.get_num_flops(ctx, per_gpu=True) + + # With PP=4, layers are divided by 4 + assert global_flops == 4 * per_gpu_flops + + +def test_mla_model_metrics_excludes_standard_attention(): + """Test that ModelMetrics uses MLAAttentionMetrics, not AttentionMetrics, + for DeepSeek MLA models.""" + hf_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=4, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + vllm_config = create_mock_vllm_config(hf_config) + model_metrics = ModelMetrics(vllm_config) + + # Should have MLAAttentionMetrics but NOT standard AttentionMetrics + component_types = [m.component_type() for m in model_metrics.metrics] + assert "mla_attn" in component_types + assert "attn" not in component_types + + # Should still have FFN and unembed + assert "ffn" in component_types + assert "unembed" in component_types + + # Breakdowns should work end-to-end + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + total_flops = model_metrics.get_num_flops(ctx) + breakdown = model_metrics.get_num_flops_breakdown(ctx) + assert total_flops == sum(breakdown.values()) + assert total_flops > 0 + + # Verify MLA-specific keys in breakdown + assert any(k.startswith("mla_attn.") for k in breakdown) + assert not any(k.startswith("attn.") for k in breakdown) + + +def test_standard_attention_still_works_for_non_mla(): + """Regression test: non-MLA models still use standard AttentionMetrics.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=12, + vocab_size=32000, + intermediate_size=8192, + ) + vllm_config = create_mock_vllm_config(hf_config) + model_metrics = ModelMetrics(vllm_config) + + component_types = [m.component_type() for m in model_metrics.metrics] + assert "attn" in component_types + assert "mla_attn" not in component_types + + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + total_flops = model_metrics.get_num_flops(ctx) + assert total_flops > 0 + + +def test_mla_attention_scaling_with_layers(): + """Test that MLA attention metrics scale proportionally with layers.""" + base_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=8, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + double_config = DeepseekV3Config( + hidden_size=7168, + num_attention_heads=128, + num_hidden_layers=16, # Double layers + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + q_lora_rank=1536, + ) + + base_vllm = create_mock_vllm_config(base_config) + double_vllm = create_mock_vllm_config(double_config) + + base_metrics = MLAAttentionMetrics.from_vllm_config(base_vllm) + double_metrics = MLAAttentionMetrics.from_vllm_config(double_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=100, context_len=512, is_prefill=True + ) + + # All metrics should double with double layers + assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx) + assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx) + assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 1db07baf93d..9d39621f4fa 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -70,6 +70,7 @@ def _make_runner(**overrides: Any) -> Any: runner.use_aux_hidden_state_outputs = False runner.speculative_config = None runner.speculator = None + runner.num_speculative_steps = 0 runner.encoder_cache = None runner.is_pooling_model = False runner.is_last_pp_rank = True @@ -102,18 +103,22 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch): lambda load_config: SimpleNamespace(load_model=lambda **_: model), ) monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append) - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state") + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) monkeypatch.setattr( eplb, "is_mixture_of_experts", lambda loaded_model: getattr(loaded_model, "is_moe", False), ) - runner = _make_runner() + runner = _make_runner(is_last_pp_rank=False) mrv2.GPUModelRunner.load_model(runner) assert runner.model is model - assert runner.model_state == "model-state" + assert runner.model_state is not None assert prepared == [model] assert runner.eplb_state is not None assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] @@ -133,10 +138,14 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch): lambda load_config: SimpleNamespace(load_model=lambda **_: model), ) monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append) - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state") + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True) - runner = _make_runner() + runner = _make_runner(is_last_pp_rank=False) mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True) assert runner.load_config.load_format == "dummy" diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 1d75b7c7628..d744da0b89b 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -167,6 +167,7 @@ def _rocm_aiter_fused_moe_impl( output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -177,6 +178,10 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + extra_kwargs: dict = {} + if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): + extra_kwargs["gate_mode"] = gate_mode + return fused_moe( hidden_states, w1, @@ -198,6 +203,7 @@ def _rocm_aiter_fused_moe_impl( bias1=bias1, bias2=bias2, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + **extra_kwargs, ) @@ -219,6 +225,7 @@ def _rocm_aiter_fused_moe_fake( output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -1804,6 +1811,21 @@ class rocm_aiter_ops: except (ImportError, ModuleNotFoundError): return False + @classmethod + @if_aiter_supported + @functools.cache + def fused_moe_supports_gate_mode(cls) -> bool: + """Probe whether the installed aiter.fused_moe accepts `gate_mode`. + + Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14). + Builds with older AITER must omit this argument. + """ + import inspect + + from aiter.fused_moe import fused_moe + + return "gate_mode" in inspect.signature(fused_moe).parameters + @staticmethod @if_aiter_supported def register_ops_once() -> None: @@ -2172,6 +2194,7 @@ class rocm_aiter_ops: output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -2194,6 +2217,7 @@ class rocm_aiter_ops: output_dtype, hidden_pad, intermediate_pad, + gate_mode, bias1, bias2, moe_sorting_dispatch_policy, diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 962efd7724a..8875ed49f6e 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -784,8 +784,10 @@ class xpu_ops: return_softmax_lse: bool | None = False, s_aux: torch.Tensor | None = None, return_attn_probs: bool | None = False, + dynamic_causal: torch.Tensor | None = None, mask_mod: Callable | None = None, aux_tensors: list | None = None, + **kwargs, ): assert cu_seqlens_k is not None or seqused_k is not None, ( "cu_seqlens_k or seqused_k must be provided" diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index abdcedd12be..25ceadc41a1 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -4001,20 +4001,27 @@ class ASRDataset(HuggingFaceDataset): Dataset class for processing a ASR dataset for transcription. Tested on the following set: - +----------------+----------------------------------------+--------------------------+-----------------------------+ - | Dataset | Domain | Speaking Style | hf-subset | - +----------------+----------------------------------------+--------------------------+-----------------------------+ - | TED-LIUM | TED talks | Oratory | release1, release2, release3| - | | | | release3-speaker-adaptation | - | VoxPopuli | European Parliament | Oratory | en, de, it, fr, ... | - | LibriSpeech | Audiobook | Narrated | "LIUM/tedlium" | - | GigaSpeech | Audiobook, podcast, YouTube | Narrated, spontaneous | xs, s, m, l, xl, dev, test | - | SPGISpeech | Financial meetings | Oratory, spontaneous | S, M, L, dev, test | - | AMI | Meetings | Spontaneous | ihm, sdm | - +----------------+----------------------------------------+--------------------------+-----------------------------+ + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ + | Dataset | Domain | Speaking Style | hf-subset | + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ + | TED-LIUM | TED talks | Oratory | release1, release2, release3| + | | | | release3-speaker-adaptation | + | VoxPopuli | European Parliament | Oratory | en, de, it, fr, ... | + | LibriSpeech | Audiobook | Narrated | "LIUM/tedlium" | + | GigaSpeech | Audiobook, podcast, YouTube | Narrated, spontaneous | xs, s, m, l, xl, dev, test | + | SPGISpeech | Financial meetings | Oratory, spontaneous | S, M, L, dev, test | + | Earnings22-Cleaned-AA | Long form earnings calls | Prepared remarks, Q&A | test | + | Earnings22-Tiny-Filtered | Earnings calls | Prepared remarks, Q&A | validation | + | AMI | Meetings | Spontaneous | ihm, sdm | + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ """ # noqa: E501 + EARNINGS22_CLEANED_DATASET = "ArtificialAnalysis/Earnings22-Cleaned-AA" + EARNINGS22_TINY_FILTERED_DATASET = ( + "D4nt3/esb-datasets-earnings22-validation-tiny-filtered" + ) + SUPPORTED_DATASET_PATHS = { "openslr/librispeech_asr", "facebook/voxpopuli", @@ -4022,11 +4029,52 @@ class ASRDataset(HuggingFaceDataset): "edinburghcstr/ami", "speechcolab/gigaspeech", "kensho/spgispeech", + EARNINGS22_CLEANED_DATASET, + EARNINGS22_TINY_FILTERED_DATASET, } DEFAULT_OUTPUT_LEN = 1024 IS_MULTIMODAL = True + def load_data(self) -> None: + if self.hf_name == self.EARNINGS22_CLEANED_DATASET: + # This subset stores repo-local MP3 paths instead of a HF `Audio` + # column, so eagerly materialize it back into the common schema. + self.data = load_dataset( + self.dataset_path, + name=self.dataset_subset, + split=self.dataset_split, + streaming=False, + trust_remote_code=self.trust_remote_code, + ) + if not getattr(self, "disable_shuffle", False): + self.data = self.data.shuffle(seed=self.random_seed) + self._materialize_local_audio_column() + return + if self.hf_name == self.EARNINGS22_TINY_FILTERED_DATASET: + super().load_data() + self._disable_audio_decode() + return + + super().load_data() + + def _disable_audio_decode(self) -> None: + from datasets import Audio + + self.data = self.data.cast_column("audio", Audio(decode=False)) + + def _materialize_local_audio_column(self) -> None: + local_path_root = Path( + hf_api().snapshot_download(self.hf_name, repo_type="dataset") + ) + self.data = self.data.map( + lambda item: { + "audio": str(local_path_root / item["url"]), + "text": item["transcript"], + } + ) + self._disable_audio_decode() + def sample( self, tokenizer: TokenizerLike, @@ -4052,14 +4100,35 @@ class ASRDataset(HuggingFaceDataset): if len(sampled_requests) >= num_requests: break audio = item["audio"] - y, sr = audio["array"], audio["sampling_rate"] - duration_s = get_audio_duration(y=y, sr=sr) + if ( + isinstance(audio, dict) + and "array" in audio + and "sampling_rate" in audio + ): + y, sr = audio["array"], audio["sampling_rate"] + duration_s = get_audio_duration(y=y, sr=sr) + mm_content = {"audio": (y, sr)} + elif isinstance(audio, str): + duration_s = sf.info(audio).duration + mm_content = {"audio_path": audio} + elif isinstance(audio, dict) and audio.get("path"): + duration_s = sf.info(audio["path"]).duration + mm_content = {"audio_path": audio["path"]} + elif isinstance(audio, dict) and audio.get("bytes") is not None: + with BytesIO(audio["bytes"]) as audio_buffer: + y, sr = sf.read(audio_buffer, dtype="float32") + duration_s = get_audio_duration(y=y, sr=sr) + mm_content = {"audio": (y, sr)} + else: + raise ValueError( + "ASR samples must provide decoded audio arrays, " + "embedded audio bytes, or a local audio path." + ) if duration_s < asr_min_audio_len_sec or duration_s > asr_max_audio_len_sec: skipped += 1 continue durations.append(duration_s) - mm_content = {"audio": (y, sr)} sampled_requests.append( SampleRequest( prompt=prompt, diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index d282033ba1f..db58f422b80 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -445,7 +445,6 @@ async def async_request_openai_audio( api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Audio API", {"transcriptions", "translations"}) - content = [{"type": "text", "text": request_func_input.prompt}] payload = { "model": request_func_input.model_name if request_func_input.model_name @@ -469,19 +468,26 @@ async def async_request_openai_audio( buffer.seek(0) return buffer - mm_audio = request_func_input.multi_modal_content - if not isinstance(mm_audio, dict) or "audio" not in mm_audio: - raise TypeError("multi_modal_content must be a dict containing 'audio'") - with to_bytes(*mm_audio["audio"]) as f: + async def send_audio_file( + audio_file: io.BytesIO | Any, + *, + input_audio_duration: float, + filename: str | None = None, + content_type: str | None = None, + ) -> RequestFuncOutput: form = aiohttp.FormData() - form.add_field("file", f, content_type="audio/wav") + add_field_kwargs: dict[str, str] = {} + if filename is not None: + add_field_kwargs["filename"] = filename + if content_type is not None: + add_field_kwargs["content_type"] = content_type + form.add_field("file", audio_file, **add_field_kwargs) for key, value in payload.items(): form.add_field(key, str(value)) output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len - output.input_audio_duration = soundfile.info(f).duration - f.seek(0) + output.input_audio_duration = input_audio_duration generated_text = "" ttft = 0.0 @@ -541,9 +547,36 @@ async def async_request_openai_audio( exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) - if pbar: - pbar.update(1) - return output + if pbar: + pbar.update(1) + return output + + mm_audio = request_func_input.multi_modal_content + if not isinstance(mm_audio, dict): + raise TypeError( + "multi_modal_content must be a dict containing 'audio' or 'audio_path'" + ) + if "audio" in mm_audio: + with to_bytes(*mm_audio["audio"]) as f: + input_audio_duration = soundfile.info(f).duration + f.seek(0) + return await send_audio_file( + f, + input_audio_duration=input_audio_duration, + filename="audio.wav", + content_type="audio/wav", + ) + if "audio_path" in mm_audio: + audio_path = mm_audio["audio_path"] + with open(audio_path, "rb") as f: + return await send_audio_file( + f, + input_audio_duration=soundfile.info(audio_path).duration, + filename=os.path.basename(audio_path), + ) + raise TypeError( + "multi_modal_content must be a dict containing 'audio' or 'audio_path'" + ) async def _run_pooling_request( diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index cbf7be44ae9..4d6fdbe22af 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -248,6 +248,68 @@ async def fetch_spec_decode_metrics( return None +@dataclass +class DiffusionMetrics: + """Diffusion (dLLM) decoding metrics from the server's Prometheus endpoint.""" + + num_denoising_steps: int + num_canvas_positions: int + num_committed_tokens: int + + +async def fetch_diffusion_metrics( + base_url: str, session: aiohttp.ClientSession +) -> DiffusionMetrics | None: + """Fetch diffusion decoding metrics from the server's Prometheus endpoint. + + Returns None if the model is not a diffusion model or metrics are not + available. + """ + metrics_url = f"{base_url}/metrics" + try: + async with session.get(metrics_url) as response: + if response.status != 200: + return None + text = await response.text() + + num_denoising_steps = 0 + num_canvas_positions = 0 + num_committed_tokens = 0 + found_diffusion = False + + for line in text.split("\n"): + line = line.strip() + if not line or line.startswith("#"): + continue + + if line.startswith("vllm:diffusion"): + # Extract metric name (before labels) to avoid matching + # substrings inside label values. + parts = line.split(None, 1) + metric_name = parts[0].split("{")[0] + if not metric_name.endswith("_total"): + continue + found_diffusion = True + with contextlib.suppress(ValueError): + if "num_denoising_steps" in metric_name: + num_denoising_steps += int(float(parts[-1])) + elif "num_canvas_positions" in metric_name: + num_canvas_positions += int(float(parts[-1])) + elif "num_committed_tokens" in metric_name: + num_committed_tokens += int(float(parts[-1])) + + if not found_diffusion: + return None + + return DiffusionMetrics( + num_denoising_steps=num_denoising_steps, + num_canvas_positions=num_canvas_positions, + num_committed_tokens=num_committed_tokens, + ) + except (aiohttp.ClientError, asyncio.TimeoutError): + return None + + class TaskType(Enum): GENERATION = "generation" POOLING = "pooling" @@ -887,6 +949,7 @@ async def benchmark( print("Self timing is set, using the timestamps from the trace file.") spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session) + diffusion_metrics_before = await fetch_diffusion_metrics(base_url, session) pbar = None if disable_tqdm else tqdm(total=len(input_requests)) @@ -1016,6 +1079,34 @@ async def benchmark( "per_position_acceptance_rates": per_pos_rates, } + diffusion_metrics_after = await fetch_diffusion_metrics(base_url, session) + diffusion_stats: dict[str, Any] | None = None + if diffusion_metrics_before is not None and diffusion_metrics_after is not None: + delta_steps = ( + diffusion_metrics_after.num_denoising_steps + - diffusion_metrics_before.num_denoising_steps + ) + delta_positions = ( + diffusion_metrics_after.num_canvas_positions + - diffusion_metrics_before.num_canvas_positions + ) + delta_committed = ( + diffusion_metrics_after.num_committed_tokens + - diffusion_metrics_before.num_committed_tokens + ) + if delta_steps > 0 and delta_committed > 0: + block_size = delta_positions / delta_steps # canvas length (CL) + num_canvases = delta_committed / block_size # = number of commit steps + denoising_steps = delta_steps - num_canvases # exclude commit steps + diffusion_stats = { + "denoising_steps": denoising_steps, + "canvas_positions": delta_positions, + "committed_tokens": delta_committed, + "committed_throughput": delta_committed / benchmark_duration, + "steps_per_canvas": denoising_steps / num_canvases, + "committed_per_step": delta_committed / denoising_steps, + } + if task_type == TaskType.GENERATION: metrics, actual_output_lens = calculate_metrics( input_requests=input_requests, @@ -1134,6 +1225,16 @@ async def benchmark( "per_position_acceptance_rates", [] ) + if diffusion_stats is not None: + result["diffusion_committed_throughput"] = diffusion_stats[ + "committed_throughput" + ] + result["diffusion_steps_per_canvas"] = diffusion_stats["steps_per_canvas"] + result["diffusion_committed_per_step"] = diffusion_stats["committed_per_step"] + result["diffusion_committed_tokens"] = int(diffusion_stats["committed_tokens"]) + result["diffusion_denoising_steps"] = int(diffusion_stats["denoising_steps"]) + result["diffusion_canvas_positions"] = int(diffusion_stats["canvas_positions"]) + def process_one_metric( # E.g., "ttft" metric_attribute_name: str, @@ -1179,7 +1280,22 @@ async def benchmark( process_one_metric("itl", "ITL", "Inter-token Latency") process_one_metric("e2el", "E2EL", "End-to-end Latency") - if spec_decode_stats is not None: + if diffusion_stats is not None: + print("{s:{c}^{n}}".format(s="Diffusion Decoding", n=50, c="-")) + for label, key, value_fmt in ( + ("Committed throughput (tok/s):", "committed_throughput", "{:<10.2f}"), + ("Denoising steps per canvas:", "steps_per_canvas", "{:<10.2f}"), + ("Committed per denoising step:", "committed_per_step", "{:<10.2f}"), + ("Committed tokens:", "committed_tokens", "{:<10d}"), + ("Denoising steps:", "denoising_steps", "{:<10d}"), + ("Canvas positions evaluated:", "canvas_positions", "{:<10d}"), + ): + value = diffusion_stats[key] + if value_fmt.endswith("d}"): + value = int(value) + print("{:<40} ".format(label) + value_fmt.format(value)) + + if spec_decode_stats is not None and diffusion_stats is None: print("{s:{c}^{n}}".format(s="Speculative Decoding", n=50, c="-")) print( "{:<40} {:<10.2f}".format( diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index b189c45c8d7..82ab1842fe9 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -10,6 +10,7 @@ from vllm.config.compilation import ( PassConfig, ) from vllm.config.device import DeviceConfig +from vllm.config.diffusion import DiffusionConfig from vllm.config.ec_transfer import ECTransferConfig from vllm.config.kernel import KernelConfig from vllm.config.kv_events import KVEventsConfig @@ -72,6 +73,8 @@ __all__ = [ "PassConfig", # From vllm.config.device "DeviceConfig", + # From vllm.config.diffusion + "DiffusionConfig", # From vllm.config.ec_transfer "ECTransferConfig", # From vllm.config.kernel diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 352ccec3202..9b96c64513b 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -146,6 +146,14 @@ class CacheConfig: num_cpu_blocks: int | None = field(default=None, init=False) """The number of blocks to allocate for CPU memory.""" + # Set after KV cache initialization. + kv_cache_size_tokens: int | None = field(default=None, init=False) + """Per-DP-engine KV cache capacity in tokens (group-aware). Uses + group-aware capacity since num_gpu_blocks * block_size can be wrong + for hybrid models where requests occupy multiple KV cache groups.""" + kv_cache_max_concurrency: float | None = field(default=None, init=False) + """Per-DP-engine maximum concurrency at max_model_len tokens.""" + kv_sharing_fast_prefill: bool = False """This feature is work in progress and no prefill optimization takes place with this flag enabled currently. @@ -204,6 +212,8 @@ class CacheConfig: # Post-init/derived counters "num_gpu_blocks", "num_cpu_blocks", + "kv_cache_size_tokens", + "kv_cache_max_concurrency", # WIP feature toggle not impacting compiled graph shape "kv_sharing_fast_prefill", } diff --git a/vllm/config/diffusion.py b/vllm/config/diffusion.py new file mode 100644 index 00000000000..6f59c40a836 --- /dev/null +++ b/vllm/config/diffusion.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for discrete diffusion (dLLM) models.""" + +from pydantic import Field + +from vllm.config.utils import config + + +@config +class DiffusionConfig: + """Configuration for discrete diffusion language models (dLLMs). + + dLLMs generate tokens via iterative denoising over a fixed-length canvas + rather than left-to-right autoregressive decoding. They reuse the + speculative-decoding data path (draft token ids, scheduled spec decode + tokens) with overloaded semantics for block-based generation. + """ + + canvas_length: int = Field(default=None, gt=0) # type: ignore[assignment] + """Length of the denoising canvas (block). Also determines the number of + speculative tokens scheduled per step.""" + + max_denoising_steps: int | None = None + """Maximum number of denoising iterations per canvas block. + If not set, read from the model's generation_config.json.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 015e75afac2..42c11eacd46 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1546,6 +1546,11 @@ class ModelConfig: """Extract the HF encoder/decoder model flag.""" return is_encoder_decoder(self.hf_config) + @cached_property + def is_diffusion(self) -> bool: + """Detect discrete diffusion (dLLM) models from HF config.""" + return getattr(self.hf_config, "canvas_length", None) is not None + @property def uses_alibi(self) -> bool: cfg = self.hf_text_config diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 9669bd1cc41..95f3ed48d47 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -143,6 +143,13 @@ class SchedulerConfig: checking the first chunk. Prevents over-admission and KV cache thrashing with chunked prefill.""" + watermark: float = Field(default=0.0, ge=0.0, lt=1.0) + """Fraction of total KV cache blocks to keep free (the watermark) when + admitting waiting or preempted requests into the running queue. This headroom + helps avoid frequent KV cache eviction and the resulting repeated preemption + of requests when GPU memory is scarce. Must be in the range [0.0, 1.0); 0.0 + (the default) disables the watermark.""" + async_scheduling: bool | None = None """If set to False, disable async scheduling. Async scheduling helps to avoid gaps in GPU utilization, leading to better latency and throughput. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 86a2f4d09e0..890d2b72e31 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -31,6 +31,7 @@ from .attention import AttentionConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig +from .diffusion import DiffusionConfig from .ec_transfer import ECTransferConfig from .kernel import KernelConfig from .kv_events import KVEventsConfig @@ -323,6 +324,9 @@ class VllmConfig: """LoRA configuration.""" speculative_config: SpeculativeConfig | None = None """Speculative decoding configuration.""" + diffusion_config: DiffusionConfig | None = None + """Diffusion LLM (dLLM) configuration.""" + structured_outputs_config: StructuredOutputsConfig = Field( default_factory=StructuredOutputsConfig ) @@ -511,6 +515,11 @@ class VllmConfig: and self.speculative_config.num_speculative_tokens is not None ): return self.speculative_config.num_speculative_tokens + if ( + self.diffusion_config is not None + and self.diffusion_config.canvas_length is not None + ): + return self.diffusion_config.canvas_length return 0 @property @@ -519,6 +528,9 @@ class VllmConfig: if use_v2_model_runner is not None: return use_v2_model_runner + if self.model_config is not None and self.model_config.is_diffusion: + return True + if not self._is_default_v2_model_runner_model(): return False @@ -1654,12 +1666,7 @@ class VllmConfig: self.compilation_config.max_cudagraph_capture_size ) if max_cudagraph_capture_size is None: - decode_query_len = 1 - if ( - self.speculative_config - and self.speculative_config.num_speculative_tokens - ): - decode_query_len += self.speculative_config.num_speculative_tokens + decode_query_len = 1 + self.num_speculative_tokens max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index fd1c826322c..967ce5d75c3 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -9,6 +9,7 @@ import torch.distributed as dist import vllm.envs as envs from vllm.distributed import get_dp_group, get_ep_group +from vllm.distributed.utils import StatelessProcessGroup from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.utils.flashinfer import ( @@ -342,7 +343,12 @@ class NixlEPAll2AllManager(All2AllManagerBase): _lock = threading.RLock() def __init__(self, cpu_group, tcp_store_group=None): - assert tcp_store_group is not None + if tcp_store_group is None: + tcp_store_group = StatelessProcessGroup( + rank=cpu_group.rank(), + world_size=cpu_group.size(), + store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()), + ) super().__init__(cpu_group, tcp_store_group) self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index 75290f6a012..aad7999d08a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -179,6 +179,18 @@ KVConnectorFactory.register_connector( "NixlConnector", ) +KVConnectorFactory.register_connector( + "NixlPullConnector", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl", + "NixlPullConnector", +) + +KVConnectorFactory.register_connector( + "NixlPushConnector", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl", + "NixlPushConnector", +) + KVConnectorFactory.register_connector( "MultiConnector", "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 71d89f43a79..954fedafe89 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -569,6 +569,18 @@ class KVConnectorBase_V1(ABC): """ return () + def has_pending_push_work(self) -> bool: + """Return True if the connector has push-mode work that requires + the engine main loop to keep stepping (e.g. a P-side request whose + KV blocks are waiting to be WRITTEN to a D node). + + Connectors that don't implement push-based KV transfer should + leave this as False. + """ + # TODO: replace with a more general connector hook for keeping the + # scheduler alive (e.g. extend has_unfinished_requests). + return False + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 14d4b381a3c..d53cd13c2e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -153,6 +153,21 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): else: self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config) + def shutdown(self): + """Release connector resources on teardown. + + Closes the worker's MooncakeDistributedStore handle so its + TransferEngine and RDMA registrations are released. Invoked from the + engine's explicit shutdown path and as a backstop from ``__del__``; + a no-op on the scheduler role, which holds no store handle. + """ + worker = getattr(self, "connector_worker", None) + if worker is not None: + worker.close() + + def __del__(self): + self.shutdown() + # ============================================================ # Scheduler-side methods # ============================================================ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 9c3ac83e06a..105762ccfcf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1426,6 +1426,22 @@ class MooncakeStoreWorker: return self.kv_send_thread.get_kv_events() return [] + def close(self) -> None: + """Release the MooncakeDistributedStore handle on teardown. + + Closing the store frees its TransferEngine, the registered RDMA + buffers, and the connection to the master server. Idempotent so it is + safe to call from both the explicit shutdown path and ``__del__``. + """ + store = getattr(self, "store", None) + if store is None: + return + self.store = None + try: + store.close() + except Exception as e: + logger.warning("Error closing MooncakeDistributedStore: %s", e) + # ============================================================ # Lookup Key Server diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 46354337e65..bfb6ee466ad 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -538,6 +538,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): for c in self._connectors: yield from c.take_events() + def has_pending_push_work(self) -> bool: + return any(c.has_pending_push_work() for c in self._connectors) + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py index ed5c892fb9d..fd5996f64bc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py @@ -2,14 +2,35 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """NIXL KV-cache transfer connector (disaggregated prefill / decode).""" +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import ( + NixlBaseConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import ( + NixlBaseConnector, NixlConnector, + NixlPullConnector, + NixlPushConnector, ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlAgentMetadata, NixlConnectorMetadata, NixlHandshakePayload, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import ( + NixlPullConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( + NixlPushConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) @@ -22,10 +43,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( __all__ = [ "NixlAgentMetadata", + "NixlBaseConnector", + "NixlBaseConnectorScheduler", + "NixlBaseConnectorWorker", "NixlConnector", "NixlConnectorMetadata", "NixlConnectorScheduler", "NixlConnectorWorker", "NixlHandshakePayload", "NixlKVConnectorStats", + "NixlPullConnector", + "NixlPullConnectorScheduler", + "NixlPullConnectorWorker", + "NixlPushConnector", + "NixlPushConnectorScheduler", + "NixlPushConnectorWorker", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py new file mode 100644 index 00000000000..cba81cadd84 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base scheduler-side logic for the NIXL connector.""" + +import threading +import time +from typing import TYPE_CHECKING, Any + +import msgspec +import zmq + +from vllm import envs +from vllm.distributed.kv_transfer.kv_connector.utils import ( + BlockIds, + EngineId, + yield_req_data, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, + KVConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + GET_META_MSG, + HeartbeatInfo, + NixlConnectorMetadata, + NixlHandshakePayload, + ReqId, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.utils.network_utils import make_zmq_path +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + SlidingWindowSpec, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.outputs import KVConnectorOutput + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class NixlBaseConnectorScheduler: + """Base implementation of Scheduler side methods shared by pull and push.""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + self.vllm_config = vllm_config + self.block_size = vllm_config.cache_config.block_size + self.engine_id: EngineId = engine_id + self.kv_cache_config = kv_cache_config + self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST + self.side_channel_port = ( + envs.VLLM_NIXL_SIDE_CHANNEL_PORT + + vllm_config.parallel_config.data_parallel_index + ) + assert vllm_config.kv_transfer_config is not None + self._kv_lease_duration: int = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "kv_lease_duration", 30 + ) + ) + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. + self._heartbeat_interval = self._kv_lease_duration // 6 + if current_platform.device_type == "cpu": + self.use_host_buffer = False + else: + self.use_host_buffer = ( + vllm_config.kv_transfer_config.kv_buffer_device == "cpu" + ) + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + # Also handle unlikely SW-only model case instead of checking num_groups>1. + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.kv_cache_groups + ) + ) + self._has_mamba = any( + isinstance(g.kv_cache_spec, MambaSpec) + for g in kv_cache_config.kv_cache_groups + ) + + logger.info("Initializing NIXL Scheduler %s", engine_id) + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + logger.info("Hybrid Memory Allocator is enabled with NIXL") + + # Background thread for handling new handshake requests. + self._nixl_handshake_listener_t: threading.Thread | None = None + self._stop_event = threading.Event() + + # Requests that need to start recv/send. + # New requests are added by update_state_after_alloc in + # the scheduler. Used to make metadata passed to Worker. + self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} + self._reqs_need_save: dict[ReqId, Request] = {} + # Reqs to send and their expiration time + self._reqs_need_send: dict[ReqId, float] = {} + self._reqs_in_batch: set[ReqId] = set() + # Reqs to remove from processed set because they're not to send after + # remote prefill or aborted. + self._reqs_not_processed: set[ReqId] = set() + + # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to + # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine + self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} + # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal + self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {} + self._last_heartbeat_time: float = 0.0 + + # Gather Sliding Window sizes for each kv cache group (if any) in number of + # blocks per KV cache group. This is used to clip the local attention window. + sw_sizes_tokens: list[tuple[int, int]] = [ + (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) + if isinstance(g.kv_cache_spec, SlidingWindowSpec) + else (0, self.block_size) + for g in kv_cache_config.kv_cache_groups + ] + # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively + # account for boundary overlap eg window isn't fully aligned with blocks. + self.blocks_per_sw = [ + cdiv(n_tokens, block_size) + 1 if n_tokens else 0 + for n_tokens, block_size in sw_sizes_tokens + ] + + # Threshold to decide whether to compute kv cache locally + # or pull from a remote node: minimum number of remote + # tokens to amortize the xfer latencies + self.kv_recompute_threshold: int = int( + vllm_config.kv_transfer_config.get_from_extra_config( + "kv_recompute_threshold", 64 + ) + ) + + # Bi-directional KV transfer feature supports KV block + # transfers from D node to P node + self.is_bidirectional_kv_xfer_enabled = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "bidirectional_kv_xfer", False + ) + ) + self.decoder_kv_blocks_ttl = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "decoder_kv_blocks_ttl", 480 + ) + ) + + if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0: + logger.info( + "Bidirectional KV transfer is enabled and the kv " + "recompute threshold is set to %d tokens." + "KV blocks on D are released after a TTL of %d seconds.", + self.kv_recompute_threshold, + self.decoder_kv_blocks_ttl, + ) + + def shutdown(self): + self._stop_event.set() + if self._nixl_handshake_listener_t is not None: + self._nixl_handshake_listener_t.join() + self._nixl_handshake_listener_t = None + + def on_new_request(self, request: "Request") -> None: + """Track a request that may need heartbeats.""" + params = request.kv_transfer_params + # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are + # effectively disabled for Bidirectional KV transfer. + if params is None or not params.get("do_remote_prefill"): + return + # Only track if all required remote fields are present. + remote_engine_id = params.get("remote_engine_id") + remote_request_id = params.get("remote_request_id") + host = params.get("remote_host") + port = params.get("remote_port") + tp_size = params.get("tp_size") + if ( + remote_engine_id is None + or remote_request_id is None + or host is None + or port is None + or tp_size is None + ): + return + if remote_engine_id not in self._heartbeat_by_engine: + self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo( + req_ids=set(), + host=host, + port=port, + tp_size=tp_size, + ) + self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id) + self._heartbeat_req_engine[request.request_id] = ( + remote_engine_id, + remote_request_id, + ) + + def _stop_heartbeat(self, req_id: ReqId) -> None: + """Remove *req_id* from heartbeat tracking (if tracked).""" + if key := self._heartbeat_req_engine.pop(req_id, None): + engine_id, remote_id = key + if info := self._heartbeat_by_engine.get(engine_id): + info.req_ids.discard(remote_id) + if not info.req_ids: + # Clean up empty engines so we don't leak a key when remote dies. + del self._heartbeat_by_engine[engine_id] + + def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: + """ + Clip the number of blocks to the sliding window size for each kv cache group + that employs SWA. + This is necessary because the KV Cache manager initially allocates blocks for + the entire sequence length, and successively cleans up blocks that are outside + the window prior to the `request_finished_all_groups` hook. + """ + if len(block_ids) == 0 or not self._is_hma_required: + # No blocks to clip eg Full prefix cache hit or not a hybrid model. + return block_ids + # NOTE (NickLucche) This logic is currently handled at the connector level + # because offloading connectors might want to receive the whole sequence even + # for SWA groups. We will abstract this logic once the interface is more stable + assert len(block_ids) == len(self.blocks_per_sw), ( + "Number of KV cache groups must match" + ) + # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged + return tuple( + [ + blocks[-self.blocks_per_sw[i] :] + if self.blocks_per_sw[i] > 0 + else blocks + for i, blocks in enumerate(block_ids) + ] + ) + + def set_xfer_handshake_metadata( + self, metadata: dict[int, KVConnectorHandshakeMetadata] + ) -> None: + """ + Set the KV connector handshake metadata for this connector. + + Args: + metadata (dict): the handshake metadata to set. + """ + encoded_data: dict[int, bytes] = {} + encoder = msgspec.msgpack.Encoder() + for tp_rank, rank_metadata in metadata.items(): + if not isinstance(rank_metadata, NixlHandshakePayload): + raise ValueError( + "NixlConnectorScheduler expects NixlHandshakePayload for " + "handshake metadata." + ) + encoded_data[tp_rank] = encoder.encode(rank_metadata) + logger.debug( + "Tp rank %d: encoded NixlHandshakePayload size: %s bytes", + tp_rank, + str(len(encoded_data[tp_rank])), + ) + + # Only start the listener when we have metadata to serve. + if self._nixl_handshake_listener_t is None: + ready_event = threading.Event() + self._nixl_handshake_listener_t = threading.Thread( + target=self._nixl_handshake_listener, + args=( + encoded_data, + ready_event, + self._stop_event, + self.side_channel_host, + self.side_channel_port, + ), + daemon=True, + name="nixl_handshake_listener", + ) + self._nixl_handshake_listener_t.start() + ready_event.wait() # Wait for listener ZMQ socket to be ready. + + @staticmethod + def _nixl_handshake_listener( + encoded_data: dict[int, Any], + ready_event: threading.Event, + stop_event: threading.Event, + host: str, + port: int, + ): + """Background thread for getting new NIXL handshakes.""" + # NOTE(rob): this is a simple implementation. We will move + # to a better approach via HTTP endpoint soon. + + # Listen for new requests for metadata. + path = make_zmq_path("tcp", host, port) + logger.debug("Starting listening on path: %s", path) + with zmq_ctx(zmq.ROUTER, path) as sock: + sock.setsockopt(zmq.RCVTIMEO, 1000) + ready_event.set() + while True: + try: + identity, _, msg = sock.recv_multipart() + except zmq.Again: + if stop_event.is_set(): + break + continue + # Decode the message which contains (GET_META_MSG, rank) + msg, target_tp_rank = msgspec.msgpack.decode(msg) + logger.debug( + "Received message for tp rank %s", + target_tp_rank, + ) + if msg != GET_META_MSG: + logger.warning("Connection listener got unexpected message %s", msg) + sock.send_multipart((identity, b"", encoded_data[target_tp_rank])) + + def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int: + """D-side only. Returns N-1 for Mamba models since the decoder + always recomputes the last token and must start from h(N-1).""" + if self._has_mamba and num_prompt_tokens > 1: + return num_prompt_tokens - 1 + return num_prompt_tokens + + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the last prompt token so the prefiller computes + h(N-1) instead of h(N). The decoder recomputes the last token to + derive h(N) correctly. + + Guarded by ``_p_side_truncated`` to avoid repeated truncation if the + request is preempted and rescheduled.""" + params = request.kv_transfer_params + if ( + params is not None + # Guard against repeated truncation after preemption/reschedule. + and not params.get("_p_side_truncated") + and request.num_prompt_tokens > 1 + ): + if request.prompt_token_ids is not None: + request.prompt_token_ids.pop() + elif request.prompt_embeds is not None: + request.prompt_embeds = request.prompt_embeds[:-1] + else: + return + + request._all_token_ids.pop() + request.num_prompt_tokens -= 1 + request.max_tokens = 1 + params["_p_side_truncated"] = True + + def _build_save_meta( + self, + meta: NixlConnectorMetadata, + scheduler_output: SchedulerOutput, + ) -> None: + # only called when use_host_buffer is True to build the save metadata + + # NOTE: For the prefill side, there might be a chance that an early added + # request is a chunked prefill, so we need to check if new blocks are added + for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output): + req_to_save = self._reqs_need_save.get(req_id) + if req_to_save is None or new_block_id_groups is None: + continue + req = req_to_save + + assert req.kv_transfer_params is not None + clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups) + meta.add_new_req_to_save( + request_id=req_id, + local_block_ids=clipped_block_id_groups, + kv_transfer_params=req.kv_transfer_params, + ) + assert scheduler_output.num_scheduled_tokens is not None + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + is_partial = ( + req.num_computed_tokens + num_scheduled_tokens + ) < req.num_prompt_tokens + if not is_partial: + # For non-partial prefills, once new req_meta is scheduled, it + # can be removed from _reqs_need_save. + # For partial prefill case, we will retain the request in + # _reqs_need_save until all blocks are scheduled with req_meta. + # Therefore, only pop if `not is_partial`. + self._reqs_need_save.pop(req_id) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> KVConnectorMetadata: + meta = NixlConnectorMetadata() + + # Loop through scheduled reqs and convert to ReqMeta. + for req_id, (req, block_ids) in self._reqs_need_recv.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_recv( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + if self.use_host_buffer: + self._build_save_meta(meta, scheduler_output) + + meta.reqs_to_send = self._reqs_need_send + meta.reqs_in_batch = self._reqs_in_batch + meta.reqs_not_processed = self._reqs_not_processed + + # Package heartbeats, throttled by heartbeat_interval. + if self._heartbeat_by_engine: + now = time.perf_counter() + if now - self._last_heartbeat_time >= self._heartbeat_interval: + self._last_heartbeat_time = now + meta.heartbeat_by_engine = self._heartbeat_by_engine + + # Clear the list once workers start the transfers + self._reqs_need_recv.clear() + self._reqs_in_batch = set() + self._reqs_not_processed = set() + self._reqs_need_send = {} + + return meta + + def update_connector_output(self, connector_output: "KVConnectorOutput") -> None: + """Stop heartbeating for requests whose KV transfer completed.""" + for req_id in connector_output.finished_recving or (): + self._stop_heartbeat(req_id) + + def has_pending_push_work(self) -> bool: + return False + + ############################################################ + # Abstract methods that subclasses must implement + ############################################################ + + def get_num_new_matched_tokens( + self, request: "Request", num_computed_tokens: int + ) -> tuple[int, bool]: + raise NotImplementedError + + def update_state_after_alloc( + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int + ): + raise NotImplementedError + + def request_finished( + self, + request: "Request", + block_ids: BlockIds, + ) -> tuple[bool, dict[str, Any] | None]: + raise NotImplementedError diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py new file mode 100644 index 00000000000..e587b0cd1fa --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -0,0 +1,2286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base worker-side logic for the NIXL connector.""" + +import logging +import os +import queue +import threading +import time +import uuid +from collections import defaultdict +from collections.abc import Iterator +from concurrent.futures import Future, ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, cast + +import msgspec +import numpy as np +import torch +import zmq + +from vllm.distributed.kv_transfer.kv_connector.utils import ( + BlockIds, + EngineId, + EngineTransferInfo, + TransferTopology, + get_current_attn_backends, + kv_postprocess_blksize_and_layout_on_receive, + kv_postprocess_blksize_on_receive, + kv_postprocess_layout_on_receive, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + GET_META_MSG, + NixlAgentMetadata, + NixlConnectorMetadata, + NixlHandshakePayload, + ReqId, + ReqMeta, + TransferHandle, + compute_nixl_compatibility_hash, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlKVConnectorStats, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + _is_attention_spec, + _is_ssm_spec, + compute_tp_mapping, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( + _NIXL_SUPPORTED_DEVICE, + get_representative_spec_type, + zmq_ctx, +) +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( + MambaConvSplitInfo, + derive_mamba_conv_split, +) +from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.network_utils import make_zmq_path +from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.worker.block_table import BlockTable +from vllm.v1.worker.utils import select_common_block_size + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import KVCacheConfig + +logger = init_logger(__name__) + + +class NixlBaseConnectorWorker: + """Base implementation of Worker side methods shared by pull and push.""" + + def _compute_desc_ids( + self, + block_ids: BlockIds, + dst_num_blocks: int, + block_size_ratio: float | None, + physical_blocks_per_logical: int, + ) -> np.ndarray: + """Compute NIXL descriptor IDs for given block IDs.""" + num_fa_regions = self.num_regions + num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0 + + num_blocks = dst_num_blocks + if block_size_ratio is not None: + num_blocks = int(num_blocks * block_size_ratio) + num_fa_descs = num_fa_regions * num_blocks + + # All-attention fast path: single vectorized broadcast. + if num_ssm_regions == 0: + # NOTE (NickLucche) With HMA, every kv group has the same number of layers + # and layers from different groups share the same kv tensor. + # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be + # read across all regions, same for [3], but group0-group1 blocks will + # always differ (different areas). Therefore we can just flatten the + # block_ids and compute the descs ids for all groups at once. + block_arr = np.concatenate(block_ids)[None, :] + region_ids = np.arange(num_fa_regions)[:, None] + return (region_ids * num_blocks + block_arr).flatten() + + # Compute desc ids per group using the right stride: FA descs have + # num_blocks entries per region (kernel granularity), SSM descs have + # logical_blocks entries per region (no kernel splitting). + logical_blocks = num_blocks // physical_blocks_per_logical + all_descs: list[np.ndarray] = [] + for i, group in enumerate(block_ids): + group_arr = np.asarray(group) + if _is_attention_spec(self._group_spec_types[i]): + fa_region_ids = np.arange(num_fa_regions)[:, None] + all_descs.append( + (fa_region_ids * num_blocks + group_arr[None, :]).flatten() + ) + elif _is_ssm_spec(self._group_spec_types[i]): + # NOTE (NickLucche) SSM and Attention block regions can + # be exchanged arbitrarily by manager. Therefore, descs + # are laid out as: + # [descs_fa (all regions) | descs_ssm (all regions)]. + # num_fa_descs offset must be computed per-engine since + # P and D can have different num_blocks (and thus + # different FA desc counts). + ssm_region_ids = np.arange(num_ssm_regions)[:, None] + all_descs.append( + ( + ssm_region_ids * logical_blocks + + group_arr[None, :] + + num_fa_descs + ).flatten() + ) + else: + raise ValueError( + f"Unknown spec type {self._group_spec_types[i]} at index {i}" + ) + + return np.concatenate(all_descs) + + def _build_local_splits_from_plan( + self, + plan: TPMapping, + src_blocks_data: list[tuple[int, int, int]], + num_fa_descs: int, + ) -> Iterator[list[tuple[int, int, int]]]: + """Build split handle data for P_TP > D_TP scenario. + + num_fa_descs is the boundary between FA and SSM descriptors. + Split counts are derived from source_ranks_per_group lengths. + FA uses rank_to_attention_slot for the slot offset; + SSM uses the rank's positional index. + """ + fa_idx = next( + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) + ) + fa_num_splits = len(plan.source_ranks_per_group[fa_idx]) + + has_ssm_descs = num_fa_descs < len(src_blocks_data) + ssm_idx = next( + (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)), + None, + ) + ssm_num_splits = ( + len(plan.source_ranks_per_group[ssm_idx]) + if has_ssm_descs and ssm_idx is not None + else 0 + ) + + # Per-FA-descriptor replicate flag, in _build_fa_local emission order. + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + + for p_idx, p_rank in enumerate(plan.all_source_ranks): + fa_slot = plan.rank_to_attention_slot.get(p_rank, 0) + + handle: list[tuple[int, int, int]] = [] + for j, (addr, local_len, dev) in enumerate(src_blocks_data): + if j < num_fa_descs: + if fa_desc_replicated[j]: + # REPLICATE (MLA): whole block written on every rank. + handle.append((addr, local_len, dev)) + else: + # SPLIT (full-attn): this rank's head slice. + chunk = local_len // fa_num_splits + handle.append((addr + fa_slot * chunk, chunk, dev)) + else: + chunk = local_len // ssm_num_splits + handle.append((addr + p_idx * chunk, chunk, dev)) + yield handle + + def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: + """Per-FA-descriptor replicate flag, in _build_fa_local emission order + (region-major; K then optional V per region). Length ``num_fa_descs``. + """ + assert self.transfer_topo is not None + n_regions = len(self.block_len_per_layer) + if n_regions == 0 or self.num_regions == 0: + return [False] * num_fa_descs + nblk = num_fa_descs // self.num_regions + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks + flags: list[bool] = [] + for i in range(n_regions): + replicated = self._is_region_replicated(i) + num_streams = 1 if replicated or not virtually_split else 2 + flags.extend([replicated] * (num_streams * nblk)) + assert len(flags) == num_fa_descs, ( + f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}" + ) + return flags + + def _is_region_replicated(self, region_idx: int) -> bool: + """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT. + + REPLICATE (MLA): identical on every rank, whole block read from one + rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP. + Defaults to SPLIT when the per-region map is unset (e.g. tests that set + block_len_per_layer without register_kv_caches). + """ + return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + nixl_wrapper_cls = NixlWrapper + if nixl_wrapper_cls is None: + logger.error("NIXL is not available") + raise RuntimeError("NIXL is not available") + logger.info("Initializing NIXL wrapper") + logger.info("Initializing NIXL worker %s", engine_id) + + # Config. + self.vllm_config = vllm_config + # mypy will complain on re-assignment otherwise. + self.block_size: int = cast(int, vllm_config.cache_config.block_size) + + if vllm_config.kv_transfer_config is None: + raise ValueError("kv_transfer_config must be set for NixlConnector") + self.kv_transfer_config = vllm_config.kv_transfer_config + + self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config( + "backends", ["UCX"] + ) + kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config( + "kv_lease_duration", 30 + ) + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. + self._lease_extension = kv_lease_duration * 2 // 3 + + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.kv_cache_groups + ) + ) + self.kv_cache_config = kv_cache_config + self._layer_specs = { + layer: group.kv_cache_spec + for group in kv_cache_config.kv_cache_groups + for layer in group.layer_names + } + self.hma_group_size = len(kv_cache_config.kv_cache_tensors) + + # ---- Model state (derived from model config) ---- + mamba_ssm_size = (0, 0) + # Conv state sub-projection decomposition (None when no Mamba). + # The 3-read transfer requires DS (dim, state_len) conv layout so + # that x/B/C sub-projections are contiguous in memory. + self._conv_decomp: MambaConvSplitInfo | None = None + self._has_mamba = any( + isinstance(g.kv_cache_spec, MambaSpec) + for g in kv_cache_config.kv_cache_groups + ) + if self._has_mamba: + assert self._is_hma_required + from vllm.model_executor.layers.mamba.mamba_utils import ( + is_conv_state_dim_first, + ) + + assert is_conv_state_dim_first(), ( + "3-read Mamba conv transfer requires DS conv state layout. " + "Set VLLM_SSM_CONV_STATE_LAYOUT=DS" + ) + mamba_spec = next( + spec + for spec in self._layer_specs.values() + if isinstance(spec, MambaSpec) + ) + self._conv_decomp = derive_mamba_conv_split( + mamba_spec, + vllm_config.parallel_config.tensor_parallel_size, + ) + mamba_ssm_size = self._conv_decomp.ssm_sizes + self._mamba_ssm_size = mamba_ssm_size + + # Agent. + non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] + # Configure NIXL num_threads to avoid UAR exhaustion on Mellanox NICs. + # Each UCX thread allocates UARs (doorbell pages) via DevX, and + # excessive NIXL UAR usage can exhaust NIC UAR space. This can cause + # components like NVSHMEM (used by DeepEP kernels) to fail during RDMA + # initialization with "mlx5dv_devx_alloc_uar" errors. + # Ref: https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf#page=63 + num_threads = vllm_config.kv_transfer_config.get_from_extra_config( + "num_threads", 4 + ) + if nixl_agent_config is None: + config = None + else: + # Enable telemetry by default for NIXL 0.7.1 and above. + config = ( + nixl_agent_config(backends=self.nixl_backends, capture_telemetry=True) + if len(non_ucx_backends) > 0 + else nixl_agent_config(num_threads=num_threads, capture_telemetry=True) + ) + + self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config) + # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. + self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict) + + # Metadata. + self.engine_id: EngineId = engine_id + self.tp_rank = get_tensor_model_parallel_rank() + self.world_size = get_tensor_model_parallel_world_size() + + self.num_blocks = kv_cache_config.num_blocks + self.enable_permute_local_kv = False + self.enable_heterogeneous_attn_post_process = False + + # KV Caches and nixl tracking data. + self.device_type = current_platform.device_type + self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device + if self.device_type not in _NIXL_SUPPORTED_DEVICE: + raise RuntimeError(f"{self.device_type} is not supported.") + elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]: + raise RuntimeError( + f"{self.device_type} with {self.kv_buffer_device} kv_buffer " + "is not supported." + ) + self.device_kv_caches: dict[str, torch.Tensor] = {} + + # cpu kv buffer for xfer + # used when device memory can not be registered under nixl + self.host_xfer_buffers: dict[str, torch.Tensor] = {} + if self.device_type == "cpu": + self.use_host_buffer = False + else: + self.use_host_buffer = self.kv_buffer_device == "cpu" + + # reserve different cores for start_load_kv() from model_forward() + if self.device_type == "cpu": + numa_core_list = current_platform.discover_numa_topology() + # setup one last core in each numa for kv transfer. + rsv_cores_for_kv = [ + max(each_numa_core_list) for each_numa_core_list in numa_core_list + ] + + if rsv_cores_for_kv: + if not hasattr(os, "sched_setaffinity"): + raise NotImplementedError( + "os.sched_setaffinity is not available on this platform" + ) + os.sched_setaffinity(0, rsv_cores_for_kv) + + # support for oot platform which can't register nixl memory + # type based on kv_buffer_device + nixl_memory_type = current_platform.get_nixl_memory_type() + if nixl_memory_type is None: + if self.kv_buffer_device in ["cuda", "xpu"]: + nixl_memory_type = "VRAM" + elif self.kv_buffer_device == "cpu": + nixl_memory_type = "DRAM" + if nixl_memory_type is None: + raise RuntimeError( + f"{self.device_type} with {self.kv_buffer_device} kv_buffer " + "is not supported." + ) + self.nixl_memory_type = nixl_memory_type + + # Note: host xfer buffer ops when use_host_buffer is True + self.copy_blocks: CopyBlocksOp | None = None + + # Map of engine_id -> kv_caches_base_addr. For TP case, each local + self.device_id: int = 0 + # Current rank may pull from multiple remote TP workers. + # EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer + self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict) + + # Number of NIXL regions. Currently one region per cache + # (so 1 per layer for MLA, otherwise 2 per layer) + self.num_regions = 0 + + # nixl_prepped_dlist_handle. + self.src_xfer_handles_by_block_size: dict[int, int] = {} + # Populated dynamically during handshake based on remote configuration. + # Keep track of regions at different tp_ratio values. tp_ratio->handles + self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} + # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. + self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) + + # Map of engine_id -> num_blocks. All ranks in the same deployment will + # have the same number of blocks. + self.dst_num_blocks: dict[EngineId, int] = {} + self._registered_descs: list[Any] = [] + + # In progress transfers. + # [req_id -> list[handle]] + self._recving_metadata: dict[ReqId, ReqMeta] = {} + self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list) + # Track the expiration time of requests that are waiting to be sent. + self._reqs_to_send: dict[ReqId, float] = {} + # Set of requests that have been part of a batch, regardless of status. + self._reqs_to_process: set[ReqId] = set() + + # Invalid blocks from failed NIXL operations (thread-safe queue of block ids) + self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue() + # requests that skipped transfer (handshake or transfer failures) + # Uses Queue for thread-safe cross-thread coordination with the + # background handshake thread, matching the _ready_requests pattern. + self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() + + # Handshake metadata of this worker for NIXL transfers. + self.xfer_handshake_metadata: NixlHandshakePayload | None = None + # Background thread for initializing new NIXL handshakes. + self._handshake_initiation_executor = ThreadPoolExecutor( + # NIXL is not guaranteed to be thread-safe, limit 1 worker. + max_workers=1, + thread_name_prefix="vllm-nixl-handshake-initiator", + ) + self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]() + self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {} + # Protects _handshake_futures and _remote_agents. + self._handshake_lock = threading.RLock() + + # TTL-based eviction of stale remote engine state. + self._engine_last_active: dict[EngineId, float] = {} + self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config( + "engine_ttl", 3600.0 + ) + + self.block_size = vllm_config.cache_config.block_size + self.model_config = vllm_config.model_config + + self.use_mla = self.model_config.use_mla + + # Get the attention backend from the first layer + # NOTE (NickLucche) models with multiple backends are not supported yet + self.attn_backends = get_current_attn_backends(vllm_config) + self.backend_name = self.attn_backends[0].get_name() + + self.kv_cache_layout = get_kv_cache_layout() + self.host_buffer_kv_cache_layout = self.kv_cache_layout + logger.info( + "Detected attention backend(s) %s", + [backend.get_name() for backend in self.attn_backends], + ) + logger.info("Detected kv cache layout %s", self.kv_cache_layout) + + # lazy initialized in register_kv_caches + self.compat_hash: str | None = None + self.transfer_topo: TransferTopology | None = None + + # With heterogeneous TP, P must wait for all assigned D TP workers to + # finish reading before safely freeing the blocks. + self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int) + self.xfer_stats = NixlKVConnectorStats() + + self._physical_blocks_per_logical_kv_block = 1 + self._sync_block_size_with_kernel() + + # Unwrap UniformTypeKVCacheSpecs to get the representative spec type + self._group_spec_types = tuple( + get_representative_spec_type(g.kv_cache_spec) + for g in self.kv_cache_config.kv_cache_groups + ) + + # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE + # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models + # combining both (e.g. GQA main + MLA Eagle-3 draft). + self._region_is_mla = list[bool]() + + # Enable different block lengths for different layers *only* when MLA is used. + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. + self.block_len_per_layer = list[int]() + + # Per-engine TP mappings. Generated during handshake. + self.tp_mappings: dict[EngineId, TPMapping] = {} + + self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( + "enforce_handshake_compat", True + ) + + def _sync_block_size_with_kernel(self) -> None: + backends = get_current_attn_backends(self.vllm_config) + kernel_block_size = select_common_block_size(self.block_size, backends) + # Number of blocks not accounting for kernel block mismatches + self._logical_num_blocks = self.num_blocks + if self.block_size != kernel_block_size: + logger.info_once( + "User-specified logical block size (%s) does not match" + " physical kernel block size (%s). Using the latter.", + self.block_size, + kernel_block_size, + ) + assert self.block_size > kernel_block_size + self._physical_blocks_per_logical_kv_block = ( + self.block_size // kernel_block_size + ) + self.block_size = kernel_block_size + self.num_blocks *= self._physical_blocks_per_logical_kv_block + + def _nixl_handshake( + self, + host: str, + port: int, + remote_tp_size: int, + expected_engine_id: str, + ) -> dict[int, str]: + """Do a NIXL handshake with a remote instance.""" + + # the first time we connect to a remote agent. + # be careful, the handshake happens in a background thread. + # it does not have an active cuda context until any cuda runtime + # call is made. when UCX fails to find a valid cuda context, it will + # disable any cuda ipc communication, essentially disabling any NVLink + # communication. + # when we are using device buffers, we need to set the device + # explicitly to make sure the handshake background thread has a valid + # cuda context. + if not self.use_host_buffer: + current_platform.set_device(self.device_id) + + # When target instance TP > local TP, we need to perform multiple + # handshakes. Do it in a single background job for simplicity. + # Regardless, only handshake with the remote TP rank(s) that current + # local rank will read from. Note that With homogeneous TP, + # this happens to be the same single rank_i. + assert self.transfer_topo is not None + p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size) + remote_rank_to_agent_name = {} + path = make_zmq_path("tcp", host, port) + + with zmq_ctx(zmq.REQ, path) as sock: + for remote_rank in p_remote_ranks: + logger.debug( + "Querying metadata on path: %s at remote tp rank %s", + path, + remote_rank, + ) + + start_time = time.perf_counter() + # Send query for the request. + msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank)) + # Set receive timeout to 5 seconds to avoid hanging on dead server + sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds + sock.send(msg) + handshake_bytes = sock.recv() + + # Decode handshake payload to get compatibility hash + handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload) + try: + handshake_payload = handshake_decoder.decode(handshake_bytes) + except (msgspec.DecodeError, msgspec.ValidationError) as e: + raise RuntimeError( + f"Failed to decode NixlHandshakePayload. This likely indicates " + f"an incompatibility between connector version. Error: {e}" + ) from e + + got_metadata_time = time.perf_counter() + logger.debug( + "NIXL handshake: get metadata took: %s", + got_metadata_time - start_time, + ) + + # Check compatibility hash BEFORE decoding agent metadata + assert self.compat_hash is not None + if ( + self.enforce_compat_hash + and handshake_payload.compatibility_hash != self.compat_hash + ): + raise RuntimeError( + f"NIXL compatibility hash mismatch. " + f"Local: {self.compat_hash}, " + f"Remote: {handshake_payload.compatibility_hash}. " + f"Prefill and decode instances have incompatible " + f"configurations. This may be due to: different vLLM versions," + f" models, dtypes, KV cache layouts, attention backends, etc. " + f"Both instances must use identical configurations." + f"Disable this check using " + f'--kv-transfer-config \'{{"kv_connector_extra_config": ' + f'{{"enforce_handshake_compat": false}}}}\'' + ) + + logger.info( + "NIXL compatibility check passed (hash: %s)", + handshake_payload.compatibility_hash, + ) + + # Decode agent metadata + metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata) + try: + metadata = metadata_decoder.decode( + handshake_payload.agent_metadata_bytes + ) + except (msgspec.DecodeError, msgspec.ValidationError) as e: + # This should not happen if hash matched + raise RuntimeError( + f"Failed to decode NixlAgentMetadata. Error: {e}" + ) from e + + # Ensure engine id matches. + if metadata.engine_id != expected_engine_id: + raise RuntimeError( + f"Remote NIXL agent engine ID mismatch. " + f"Expected {expected_engine_id}," + f"received {metadata.engine_id}." + ) + + # Register Remote agent. + remote_agent_name = self.add_remote_agent( + metadata, remote_rank, remote_tp_size + ) + setup_agent_time = time.perf_counter() + logger.debug( + "NIXL handshake: add agent took: %s", + setup_agent_time - got_metadata_time, + ) + remote_rank_to_agent_name[remote_rank] = remote_agent_name + return remote_rank_to_agent_name + + def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None: + """ + Initialize transfer buffer in CPU mem for accelerators + NOT directly supported by NIXL (e.g., tpu) + """ + xfer_buffers: dict[str, torch.Tensor] = {} + inv_order = [0, 1, 3, 2, 4] + try: + for layer_name, kv_cache in kv_caches.items(): + kv_shape = kv_cache.shape + kv_dtype = kv_cache.dtype + permute_shape = False + if ( + self.kv_cache_layout == "NHD" + and self.vllm_config.kv_transfer_config is not None + and self.vllm_config.kv_transfer_config.enable_permute_local_kv + ): + logger.info_once( + "'enable_permute_local_kv' flag is enabled while " + "device KV Layout is NHD. Init host buffer with" + " HND to better support Decode/Prefill TP_ratio > 1." + ) + # Since NHD will not support Decode/Prefill TP_ratio > 1, + # we can leverage host_buffer for permute + self.host_buffer_kv_cache_layout = "HND" + kv_shape = ( + tuple(kv_shape[i] for i in inv_order) + if not self.use_mla + else kv_shape + ) + permute_shape = not self.use_mla + + xfer_buffers[layer_name] = torch.empty( + kv_shape, dtype=kv_dtype, device="cpu" + ) + if permute_shape: + xfer_buffers[layer_name] = xfer_buffers[layer_name].permute( + inv_order + ) + except MemoryError as e: + logger.error("NIXLConnectorWorker gets %s.", e) + raise + + self.host_xfer_buffers = xfer_buffers + + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): + """Assign copy (d2h, h2d) operations when host buffer is used.""" + # Set a no-op if the host buffer is not cpu. + if self.kv_buffer_device != "cpu": + return + # Set a no-op if self.device_type is 'cpu'. + if self.device_type == "cpu": + return + assert self.use_host_buffer + self.copy_blocks = copy_operation + + def _log_failure( + self, + failure_type: str, + req_id: str | None, + msg: str = "", + error: Exception | None = None, + meta: ReqMeta | None = None, + **extra_context, + ): + """Log transfer failure with structured context for easier debugging.""" + context: dict[str, Any] = { + "failure_type": failure_type, + "request_id": req_id, + "engine_id": self.engine_id, + } + if meta is None and req_id is not None: + # Try to get metadata from in progress transfers when not provided + meta = self._recving_metadata.get(req_id) + + if meta and meta.remote: + context.update( + { + "remote_engine_id": meta.remote.engine_id, + "remote_request_id": meta.remote.request_id, + "remote_host": meta.remote.host, + "remote_port": meta.remote.port, + "num_local_blocks": sum( + len(group) for group in meta.local_block_ids + ), + "num_remote_blocks": sum( + len(group) for group in meta.remote.block_ids + ), + "local_block_ids_sample": meta.local_block_ids[0][:10] + if meta.local_block_ids + else [], + } + ) + + context.update(extra_context) + if msg: + failure_type = f"{failure_type}. {msg}" + + logger.error( + "NIXL transfer failure: %s | Context: %s", + failure_type, + context, + exc_info=error is not None, + stacklevel=2, + ) + + def _ensure_handshake( + self, + engine_id: EngineId, + host: str, + port: int, + tp_size: int, + ) -> Future[dict[int, str]] | None: + """ + Ensure a handshake is in-flight (or already done) for *engine_id*. + + Returns the ``Future`` if a handshake is pending (or was just + started), or ``None`` if the handshake already completed + successfully. Callers can attach per-request callbacks to the + returned future. + Failures to handshake are logged and the request is marked as failed. + """ + self._evict_stale_engines() + with self._handshake_lock: + if engine_id in self._remote_agents: + return None + fut = self._handshake_futures.get(engine_id) + if fut is not None: + return fut + fut = self._handshake_initiation_executor.submit( + self._nixl_handshake, + host, + port, + tp_size, + engine_id, + ) + self._handshake_futures[engine_id] = fut + + def done_callback(f: Future[dict[int, str]], eid=engine_id): + with self._handshake_lock: + del self._handshake_futures[eid] + try: + self._remote_agents[eid] = f.result() + self._engine_last_active[eid] = time.perf_counter() + except Exception as e: + self._log_failure( + failure_type="handshake_setup_failed", + req_id=None, + error=e, + remote_engine_id=eid, + ) + + fut.add_done_callback(done_callback) + return fut + + def _background_nixl_handshake( + self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta + ): + # Do NIXL handshake in background and add to _ready_requests when done. + assert meta.remote is not None + fut = self._ensure_handshake( + remote_engine_id, + meta.remote.host, + meta.remote.port, + meta.tp_size, + ) + if fut is None: + # Already handshaked — only happens if caller does not pre-check. + self._ready_requests.put((req_id, meta)) + return + + # Check handshake success before proceeding with request. + def request_ready(f: Future[Any], entry=(req_id, meta)): + try: + f.result() + self._ready_requests.put(entry) + except Exception as e: + self._log_failure( + failure_type="handshake_failed", + req_id=req_id, + error=e, + meta=meta, + ) + self._handle_failed_transfer(req_id, None) + + fut.add_done_callback(request_ready) + + def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: + """Register a cross-layers KV cache tensor with NIXL. + + `use_uniform_kv_cache()` guarantees a single KV cache group whose + layers all share the same `AttentionSpec`, so any layer name from + `_layer_specs` yields the correct per-layer spec for `page_size_bytes`. + """ + first_layer = next(iter(self._layer_specs)) + # Forwarding a real layer name rather than a synthetic key + self.register_kv_caches({first_layer: kv_cache}) + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): + """Register the KV Cache data in nixl.""" + self.transfer_topo = TransferTopology( + tp_rank=self.tp_rank, + tp_size=self.world_size, + block_size=self.block_size, + engine_id=self.engine_id, + is_mla=self.use_mla, + total_num_kv_heads=self.model_config.get_total_num_kv_heads(), + attn_backends=self.attn_backends, + # SSM States come in tuples (ssm, conv) + tensor_shape=next(iter(kv_caches.values())).shape + if not self._has_mamba + else None, + is_mamba=self._has_mamba, + ) + self.compat_hash = compute_nixl_compatibility_hash( + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks + ) + + if self.use_host_buffer: + self.initialize_host_xfer_buffer(kv_caches=kv_caches) + assert len(self.host_xfer_buffers) == len(kv_caches), ( + f"host_buffer: {len(self.host_xfer_buffers)}, " + f"kv_caches: {len(kv_caches)}" + ) + xfer_buffers = self.host_xfer_buffers + else: + xfer_buffers = kv_caches + assert not self.host_xfer_buffers, ( + "host_xfer_buffer should not be initialized when " + f"kv_buffer_device is {self.kv_buffer_device}" + ) + + logger.info( + "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, " + "use_host_buffer: %s", + self.use_mla, + self.kv_buffer_device, + self.use_host_buffer, + ) + + caches_data = [] + # With hybrid allocator, layers can share a kv cache tensor + seen_base_addresses = [] + + # Note(tms): I modified this from the original region setup code. + # K and V are now in different regions. Advantage is that we can + # elegantly support MLA and any cases where the K and V tensors + # are non-contiguous (it's not locally guaranteed that they will be) + # Disadvantage is that the encoded NixlAgentMetadata is now larger + # (roughly 8KB vs 5KB). + # Conversely for FlashInfer, K and V are registered in the same region + # to better exploit the memory layout (ie num_blocks is the first dim). + tensor_size_bytes = None + + for layer_name, cache_or_caches in xfer_buffers.items(): + # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to + # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. + # However, physical page_size may differ when kernel requires a specific + # block size. This leads to SSM and FA layers having different num_blocks. + # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. + layer_spec = self._layer_specs.get(layer_name) + if layer_spec is None: + logger.debug( + "Skipping layer %s as no KVCache spec is present. " + "This is likely because the layer is sharing its KV cache", + layer_name, + ) + continue + if isinstance(layer_spec, UniformTypeKVCacheSpecs): + # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs + layer_spec = layer_spec.kv_cache_specs[layer_name] + cache_list = self.transfer_topo.get_transfer_cache_regions( + cache_or_caches, layer_spec + ) + # `layer_spec.page_size_bytes` only accounts for logical page_size, that is + # the page_size assuming constant `self._logical_num_blocks`. + physical_page_size = ( + layer_spec.page_size_bytes + if isinstance(layer_spec, MambaSpec) + else layer_spec.page_size_bytes + // self._physical_blocks_per_logical_kv_block + ) + # For when registering multiple tensors eg K/V in separate regions. + physical_page_size = physical_page_size // len(cache_list) + if self.transfer_topo._cross_layers_blocks: + # When cross-layers blocks are used, multiply by number of layers + physical_page_size = physical_page_size * len( + self.kv_cache_config.kv_cache_tensors + ) + num_blocks = ( + self._logical_num_blocks + if isinstance(layer_spec, MambaSpec) + else self.num_blocks + ) + # `page_size` accounts for physical blocks, st KVCache is always + # [`num_blocks` * `page_size`] + curr_tensor_size_bytes = num_blocks * physical_page_size + + # TODO (NickLucche) we could eventually unify how we handle FA/FI regions, + # registering a single tensor for both K/V and splitting logically like FI. + for cache in cache_list: + base_addr = cache.data_ptr() + if base_addr in seen_base_addresses: + # NOTE (NickLucche) HMA employs memory pooling to share tensors + # across groups. This results in skipping all tensors but the ones + # pointed to by group0. Also, generally we will have more blocks + # per tensor but fewer regions. + logger.debug("Skipping %s because it's already seen", layer_name) + continue + logger.debug( + "Registering layer %s with cache shape: %s", layer_name, cache.shape + ) + seen_base_addresses.append(base_addr) + # Only record non-Mamba page sizes. + if isinstance(layer_spec, MambaSpec): + self.block_len_per_layer.append( + physical_page_size // self._physical_blocks_per_logical_kv_block + ) + else: + self.block_len_per_layer.append(physical_page_size) + is_mla_region = isinstance(layer_spec, MLAAttentionSpec) + self._region_is_mla.append(is_mla_region) + + if not is_mla_region: + if tensor_size_bytes is None: + tensor_size_bytes = curr_tensor_size_bytes + assert tensor_size_bytes == curr_tensor_size_bytes, ( + "All non-MLA kv cache tensors must have the same size" + ) + + if cache.shape[0] != num_blocks: + raise AssertionError( + "All kv cache tensors must have the same number of " + f"blocks; layer={layer_name}, " + f"expected_num_blocks={num_blocks}, " + f"cache_shape={tuple(cache.shape)}, " + f"cache_stride={tuple(cache.stride())}, " + f"layer_spec={type(layer_spec).__name__}, " + f"backend={self.backend_name}, " + "all_backends=" + f"{[backend.get_name() for backend in self.attn_backends]}, " + f"kv_cache_layout={self.kv_cache_layout}, " + "blocks_first=" + f"{self.transfer_topo.is_kv_layout_blocks_first}" + ) + + # Need to make sure the device ID is non-negative for NIXL, + # Torch uses -1 to indicate CPU tensors. + self.device_id = max(cache.get_device(), 0) + caches_data.append( + (base_addr, curr_tensor_size_bytes, self.device_id, "") + ) + + logger.debug( + "Different block lengths collected: %s", set(self.block_len_per_layer) + ) + assert ( + len(self.block_len_per_layer) + == len(seen_base_addresses) + == len(self._region_is_mla) + ) + + self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses + self.num_regions = len(caches_data) + + if self.transfer_topo.virtually_split_kv_in_blocks: + # NOTE (NickLucche) When FlashInfer is used, memory is registered + # with joint KV for each block. This minimizes the overhead in + # registerMem allowing faster descs queries. In order to be able to + # split on kv_heads dim as required by heterogeneous TP, one must + # be able to index K/V separately. Hence we double the number + # of 'virtual' regions here and halve `block_len` below. + # Similarly for Mamba layers, we register SSM+Conv as a single region and + # then duplicate it logically to be able to index SSM/Conv separately. + # Exception: key-only REPLICATE regions (MLA) have no V half, so + # they contribute a single desc stream and are not doubled. + self.num_regions = sum( + 1 if self._is_region_replicated(i) else 2 + for i in range(len(self._region_is_mla)) + ) + + # Total local FA descriptors (boundary between FA and mamba descs). + self.num_descs = self.num_regions * self.num_blocks + + descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) + logger.debug("Registering descs: %s", caches_data) + self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) + logger.debug("Done registering descs") + self._registered_descs.append(descs) + + self.device_kv_caches = kv_caches + self.dst_num_blocks[self.engine_id] = self.num_blocks + + if self._has_mamba: + logger.info( + "Hybrid SSM registration: num_blocks=%s, " + "logical_num_blocks=%s, ratio=%s, num_regions=%s, " + "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s", + self.num_blocks, + self._logical_num_blocks, + self._physical_blocks_per_logical_kv_block, + self.num_regions, + self.num_descs, + self._mamba_ssm_size, + set(self.block_len_per_layer), + ) + + # Register local/src descr for NIXL xfer. + self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = ( + self.register_local_xfer_handler(self.block_size) + ) + + # After KV Caches registered, listen for new connections. + agent_metadata = NixlAgentMetadata( + engine_id=self.engine_id, + agent_metadata=self.nixl_wrapper.get_agent_metadata(), + device_id=self.device_id, + kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank], + num_blocks=self.num_blocks, + block_lens=self.block_len_per_layer, + kv_cache_layout=self.kv_cache_layout + if not self.use_host_buffer + else self.host_buffer_kv_cache_layout, + block_size=self.block_size, + ssm_sizes=self._mamba_ssm_size, + attn_backend_name=self.backend_name, + physical_blocks_per_logical_kv_block=( + self._physical_blocks_per_logical_kv_block + ), + ) + # Wrap metadata in payload with hash for defensive decoding + assert self.compat_hash is not None + encoder = msgspec.msgpack.Encoder() + self.xfer_handshake_metadata = NixlHandshakePayload( + compatibility_hash=self.compat_hash, + agent_metadata_bytes=encoder.encode(agent_metadata), + ) + + def _build_mamba_local( + self, + base_addresses: list[int], + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build 4 desc regions (x, B, C, ssm) per layer for local mamba + blocks, enabling the 3-read transfer with DS conv layout.""" + assert block_size_ratio == 1, ( + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " + f"Got block_size_ratio={block_size_ratio}." + ) + assert self._conv_decomp is not None + conv_offsets = self._conv_decomp.local_conv_offsets + conv_size, ssm_size = self._mamba_ssm_size + num_blocks = self._logical_num_blocks * block_size_ratio + physical_per_logical = self._physical_blocks_per_logical_kv_block + + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(base_addresses): + # Jump one page_size, but ssm page_size may be bigger when kernel + # locks block size to a specific value (physical_per_logical scale). + page_stride = ( + self.block_len_per_layer[i] // block_size_ratio * physical_per_logical + ) + for off, sz in conv_offsets: + for blk in range(num_blocks): + result.append( + (base_addr + blk * page_stride + off, sz, self.device_id) + ) + # SSM temporal state follows the conv state. + for blk in range(num_blocks): + result.append( + ( + base_addr + blk * page_stride + conv_size, + ssm_size, + self.device_id, + ) + ) + return result + + def _build_mamba_remote( + self, + nixl_agent_meta: NixlAgentMetadata, + tp_ratio: int, + transfer_info: EngineTransferInfo, + ) -> list[tuple[int, int, int]]: + """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer + for the 3-read transfer. For hetero-TP, each D rank reads only its + sub-projection slice from the P rank.""" + assert self._conv_decomp is not None + effective_ratio = max(tp_ratio, 1) + # Mamba conv state is always TP-sharded, even when attention KV + # is replicated (num_kv_heads < tp_size). + local_offset = self.tp_rank % effective_ratio + conv_size_remote = nixl_agent_meta.ssm_sizes[0] + + conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio) + if tp_ratio >= 1: + ssm_read_size = self._mamba_ssm_size[1] + else: + ssm_read_size = nixl_agent_meta.ssm_sizes[1] + + remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical + num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical + device_id = nixl_agent_meta.device_id + + result: list[tuple[int, int, int]] = [] + # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case + # block lengths vary across layers (e.g. MLA). + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical + for off, sz in conv_offsets: + for blk in range(num_blocks): + result.append((base_addr + blk * page_stride + off, sz, device_id)) + # SSM temporal state is also TP-sharded on the heads dimension. + for blk in range(num_blocks): + ssm_addr = ( + base_addr + + blk * page_stride + + conv_size_remote + + local_offset * ssm_read_size + ) + result.append((ssm_addr, ssm_read_size, device_id)) + return result + + def _build_fa_local( + self, + base_addresses: list[int], + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build local FA descriptors for all layers.""" + assert self.transfer_topo is not None + num_blocks = self.num_blocks * block_size_ratio + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(base_addresses): + kv_block_len = ( + self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=False + ) + // block_size_ratio + ) + page_stride = self.block_len_per_layer[i] // block_size_ratio + for block_id in range(num_blocks): + block_offset = block_id * page_stride + addr = base_addr + block_offset + result.append((addr, kv_block_len, self.device_id)) + + if ( + self.transfer_topo.virtually_split_kv_in_blocks + and not self._is_region_replicated(i) + ): + # Separate and interleave K/V regions to maintain the same + # descs ordering. This is needed for selecting contiguous heads + # when split across TP ranks. (Skipped for key-only REPLICATE.) + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=False + ) + for block_id in range(num_blocks): + block_offset = block_id * page_stride + addr = base_addr + block_offset + v_addr = addr + kv_block_len + result.append((v_addr, second_split, self.device_id)) + return result + + def _build_fa_remote( + self, + plan: TPMapping, + nixl_agent_meta: NixlAgentMetadata, + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build remote FA descriptors for all layers.""" + assert self.transfer_topo is not None + fa_group_idx = next( + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) + ) + # SPLIT regions read their head slice from this many remote ranks at a + # per-rank offset; REPLICATE regions read the whole block once. + split_reads = len(plan.source_ranks_per_group[fa_group_idx]) + num_blocks = nixl_agent_meta.num_blocks + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + replicated = self._is_region_replicated(i) + # Read our whole local region size from remote.. + local_block_len = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=False + ) + remote_kv_block_len = local_block_len // block_size_ratio + if block_size_ratio > 1: + # ..using remote kv_block_len as transfer unit + local_block_len = remote_kv_block_len + + # REPLICATE reads the whole block once at offset 0; SPLIT gathers + # its head slice from `split_reads` remote ranks at a per-rank offset. + num_reads = 1 if replicated else split_reads + rank_offset = ( + 0 if replicated else plan.rank_offset_factor * remote_kv_block_len + ) + local_block_len = local_block_len // num_reads + + page_size = nixl_agent_meta.block_lens[i] + for block_id in range(num_blocks): + block_offset = block_id * page_size + # For each block, grab the kv heads chunk belonging to current local + # tp rank of size local_block_len. + addr = base_addr + block_offset + rank_offset + result.append((addr, local_block_len, nixl_agent_meta.device_id)) + + emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated + if emits_v: + # With FlashInfer index V separately to allow head splitting. + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=False + ) + second_split = second_split // num_reads + for block_id in range(num_blocks): + block_offset = block_id * page_size + addr = base_addr + block_offset + rank_offset + # Hop over the first split of remote page, K, to read V. + v_addr = addr + nixl_agent_meta.block_lens[i] // 2 + result.append((v_addr, second_split, nixl_agent_meta.device_id)) + return result + + def register_local_xfer_handler( + self, + block_size: int, + ) -> tuple[int, list[tuple[int, int, int]]]: + """ + Function used for register local xfer handler with local block_size or + Remote block_size. + + When local block_size is same as remote block_size, we use local block_size + to register local_xfer_handler during init. + + When remote block size is less than local block size, we need to use + register another local_xfer_handler using remote block len to ensure + data copy correctness. + """ + assert self.transfer_topo is not None + block_size_ratio = self.block_size // block_size + local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank] + + blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio) + logger.debug( + "Created %s blocks for src engine %s and rank %s on device id %s", + len(blocks_data), + self.engine_id, + self.tp_rank, + self.device_id, + ) + if self._has_mamba: + assert self.num_descs == len(blocks_data) + # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split + # is unnecessary — a single conv desc per block suffices. Consider + # adding a fast path that falls back to the standard 2-region + # registration (_build_fa_local mamba=True) when no hetero-TP + # remote has been seen. Currently we always register 4 regions + # because local descs are created before knowing the remote TP. + logger.debug("Registering local Mamba descriptors (4 regions/layer)") + blocks_data.extend( + self._build_mamba_local(local_base_addresses, block_size_ratio) + ) + + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) + # NIXL_INIT_AGENT to be used for preparations of local descs. + return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data + + def add_remote_agent( + self, + nixl_agent_meta: NixlAgentMetadata, + remote_tp_rank: int = 0, + remote_tp_size: int = 1, + ) -> str: + """ + Add the remote NIXL agent and prepare the descriptors for reading cache + blocks from remote. + + In particular, handle both homogeneous and heterogeneous TP. The former + requires local rank_i to read from remote rank_i. + The latter, in the case of D.world_size < P.world_size, requires that a + local (D) TP worker reads from multiple remote (P) TP workers. + Conversely, assuming D.world_size > P.world_size, two or more local TP + workers will read from a single remote TP worker. + + Here's an example for the last case described above (non-MLA): + + rank_offset p_remote_tp_rank + (kv split no) + -------------------------------- + 0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ] + / + 1 0 Worker1 ---- 2nd half of KV -----/ + + 0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ] + / + 1 1 Worker3 ---- 2nd half of KV -----/ + + + Decoder TP workers Prefix TP workers + (world_size=4) (world_size=2) + tp_ratio = 4 // 2 = 2 + + Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim] + then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format. + Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio + first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split + along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0. + + Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1. + + Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0 + so that the whole cache is shared by "tp_ratio" D TP workers. + + For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and + tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer. + """ # noqa: E501 + engine_id = nixl_agent_meta.engine_id + # TODO re-evaluate refreshing for scaling/recovery + if remote_tp_rank in self._remote_agents.get(engine_id, {}): + logger.debug( + "Remote agent with engine_id %s and rank" + "%s already exchanged metadata, skip handshake.", + engine_id, + remote_tp_rank, + ) + return self._remote_agents[engine_id][remote_tp_rank] + + ### Register remote engine in TransferTopology (idempotent). + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo + physical_blocks_per_logical = ( + nixl_agent_meta.physical_blocks_per_logical_kv_block + ) + transfer_info = EngineTransferInfo( + remote_tp_size=remote_tp_size, + remote_block_size=nixl_agent_meta.block_size, + remote_block_len=nixl_agent_meta.block_lens[0], + remote_physical_blocks_per_logical=physical_blocks_per_logical, + ) + transfer_topo.register_remote_engine(engine_id, transfer_info) + logger.info("Transfer plan: %s", transfer_topo.describe(engine_id)) + + self.tp_mappings[engine_id] = compute_tp_mapping( + transfer_topology=transfer_topo, + remote_tp_size=remote_tp_size, + group_spec_types=self._group_spec_types, + ) + + remote_agent_name = self.nixl_wrapper.add_remote_agent( + nixl_agent_meta.agent_metadata + ) + + # Create dst descs and xfer side handles. TP workers have same #blocks + # so we only register once per engine_id. + # Example: + # block_size_ratio > 1: + # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12| + # local origin:| 0| 1| 8| 12| + # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15| + block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size) + + if engine_id not in self.dst_num_blocks: + self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks + + # Keep track of remote agent kv caches base addresses. + self.kv_caches_base_addr[engine_id][remote_tp_rank] = ( + nixl_agent_meta.kv_caches_base_addr + ) + self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size) + + # This is 1 when P and D `--tensor-parallel-size` match. Otherwise, + # this is the ratio between the two sizes. + tp_ratio = transfer_topo.tp_ratio(remote_tp_size) + + logger.debug( + "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s", + engine_id, + remote_tp_rank, + tp_ratio, + ) + + plan = self.tp_mappings[engine_id] + + ### (Optional) Register local agent memory regions. MLA is not split. + if ( + tp_ratio < 0 + and not self.use_mla + and tp_ratio not in self.src_xfer_handles_by_tp_ratio + ): + # Remote tp_size > local tp_size: read from multiple remote ranks. + # Logically "split" own regions into |tp_ratio| chunks. Mind that + # we only do this once per remote tp_size (replica-friendly). + self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] + + for handle_data in self._build_local_splits_from_plan( + plan, + self.src_blocks_data, + self.num_descs, + ): + descs = self.nixl_wrapper.get_xfer_descs( + handle_data, self.nixl_memory_type + ) + handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + + ### Register remote agent memory regions + # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With + # heterogeneous TP, prepare the descriptors by splitting the P KV cache along + # kv_head dim, of D worker's kv_head size (D>P). + # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..]. + + # Register all remote blocks, but only the corresponding kv heads. + blocks_data = self._build_fa_remote( + plan, + nixl_agent_meta, + block_size_ratio, + ) + logger.debug( + "Created %s blocks for dst engine %s with remote rank %s and local rank %s", + len(blocks_data), + engine_id, + remote_tp_rank, + self.tp_rank, + ) + if self._has_mamba: + logger.debug( + "Registering remote Mamba blocks for engine %s rank %s", + engine_id, + remote_tp_rank, + ) + blocks_data.extend( + self._build_mamba_remote( + nixl_agent_meta, + tp_ratio, + transfer_info, + ) + ) + + # Register with NIXL. + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) + self.dst_xfer_side_handles[engine_id][remote_tp_rank] = ( + self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) + ) + + if block_size_ratio > 1: + # when prefill with smaller block_size, we need to init a + # new handler with same block_len to match + self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( + self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] + ) + + return remote_agent_name + + def _validate_remote_agent_handshake( + self, nixl_agent_meta: NixlAgentMetadata, remote_tp_size: int + ): + """ + Validate the remote agent handshake metadata ensuring the + invariants hold true. + """ + remote_engine_id = nixl_agent_meta.engine_id + + assert self.transfer_topo is not None + remote_info = self.transfer_topo.get_engine_info(remote_engine_id) + assert remote_info.remote_tp_size == remote_tp_size + + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size) + block_size_ratio = self.transfer_topo.block_size_ratio( + nixl_agent_meta.block_size + ) + # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba. + # Mamba models can have replicated FA KV with tp_ratio < 0. + # MLA models do not need to handle kv replication. + if not self.use_mla and not self._has_mamba: + assert not ( + tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) + ) + + remote_physical_per_logical = ( + nixl_agent_meta.physical_blocks_per_logical_kv_block + ) + if ( + self._has_mamba + and remote_physical_per_logical + != self._physical_blocks_per_logical_kv_block + and self.vllm_config.cache_config.enable_prefix_caching + ): + raise RuntimeError( + "Prefix caching with heterogeneous physical_blocks_per_logical " + "is not supported for Mamba hybrid models. " + f"Local: {self._physical_blocks_per_logical_kv_block}, " + f"Remote: {remote_physical_per_logical}. " + "Disable prefix caching with --no-enable-prefix-caching." + ) + + if self._is_hma_required: + assert block_size_ratio == 1, ( + "HMA does not support different remote block size yet" + ) + kv_cache_layout = ( + self.kv_cache_layout + if not self.use_host_buffer + else self.host_buffer_kv_cache_layout + ) + if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout: + if ( + self.kv_transfer_config.enable_permute_local_kv + and nixl_agent_meta.kv_cache_layout == "HND" + ): + logger.info( + "Remote is HND and local is NHD, enabled additional permute " + "on local device KV." + ) + assert not self._is_hma_required, ( + "HMA does not support block size post processing" + ) + self.enable_permute_local_kv = True + else: + raise RuntimeError( + "Heterogeneous TP expects same kv_cache_layout. " + "Or enable experimental feature to use HND to NHD support by " + "setting 'enable_permute_local_kv'=True in --kv-transfer-config." + ) + # if remote_agent used attn is not same as local, + # hint heterogenuous attn post process + if ( + nixl_agent_meta.attn_backend_name != self.backend_name + and self.backend_name in ["CPU_ATTN"] + ): + if self._is_hma_required: + raise RuntimeError( + "heterogeneous attn post process is not supported with HMA" + ) + logger.info( + "[Experimental] CPU_ATTN backend is used, " + "hint heterogeneous attn post process" + ) + self.enable_heterogeneous_attn_post_process = True + + # Heterogeneous TP requires head-splitting, which only works with + # HND layout. MLA and replicated-KV cases don't split on heads. + # Mamba doesn't support heterogeneous TP. + if ( + abs(tp_ratio) != 1 + and not self.use_mla + and not self.transfer_topo.is_kv_replicated(remote_engine_id) + and kv_cache_layout != "HND" + and not self.enable_permute_local_kv + ): + raise RuntimeError( + "Heterogeneous TP head-dimension splitting requires contiguous heads. " + "Use HND layout on the prefill side." + ) + + # Per-region block_len validation enforcing the P/D invariant. + # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer) + # only allow the number of blocks to differ; SPLIT regions scale with + # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here. + if not self._has_mamba: + assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( + "Number of KV layers must match between prefill and decode" + ) + model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated( + remote_engine_id + ) + for i, local_len in enumerate(self.block_len_per_layer): + replicated = model_replicated or self._is_region_replicated(i) + remote_len = nixl_agent_meta.block_lens[i] + if replicated: + assert local_len // block_size_ratio == remote_len, ( + "KV cache sizes must match between P and D when " + f"replicated (region {i}: local={local_len}, " + f"remote={remote_len}, bsr={block_size_ratio})." + ) + elif tp_ratio > 0: + assert remote_len == (local_len * tp_ratio) // block_size_ratio, ( + f"SPLIT region {i}: remote P KV block_len {remote_len} " + f"must equal local {local_len} * tp_ratio {tp_ratio} " + f"// block_size_ratio {block_size_ratio}." + ) + else: + assert block_size_ratio == 1, ( + "Different local/remote block sizes are not supported " + "when P TP > D TP." + ) + assert remote_len == local_len // (-tp_ratio), ( + f"SPLIT region {i}: remote P KV block_len " + f"{remote_len} must equal local {local_len} " + f"// |tp_ratio| {-tp_ratio}." + ) + + # TP workers that handhshake with same remote have same #blocks. + assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks + # Same number of regions/~layers. + assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) + + def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): + """copy recved kv from host buffer to device.""" + assert self.use_host_buffer + assert self.copy_blocks is not None + + local_block_ids = meta.local_physical_block_ids + # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups + for group_block_ids in local_block_ids: + self.copy_blocks( + self.host_xfer_buffers, + self.device_kv_caches, + group_block_ids, + group_block_ids, + "h2d", + ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "synced recved kv of request[%s] to device kv buffer," + "local_block_ids: %s. ", + req_id, + ",".join(map(str, local_block_ids)), + ) + + def save_kv_to_host(self, metadata: NixlConnectorMetadata): + """copy kv from device to host buffer.""" + assert self.use_host_buffer + assert self.copy_blocks is not None + + for req_id, meta in metadata.reqs_to_save.items(): + meta.local_physical_block_ids = self._logical_to_kernel_block_ids( + meta.local_block_ids + ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "save_load_kv for request[%s] to host xfer buffer." + "local_block_ids: %s. ", + req_id, + ",".join(map(str, meta.local_physical_block_ids)), + ) + # blocking + for group_block_ids in meta.local_physical_block_ids: + self.copy_blocks( + self.device_kv_caches, + self.host_xfer_buffers, + group_block_ids, + group_block_ids, + "d2h", + ) + + def post_process_device_kv_on_receive( + self, + block_size_ratio: int, + block_ids_list: list[list[int]], + ): + """ + Post process device kv cache after receiving from remote. + + 3 types of post processing supported: + * kv_cache_postprocess_layout => convert from HND to NHD + * kv_cache_postprocess_blksize => convert from small block size + to large block size + * kv_cache_postprocess_blksize_and_layout => convert from small + block size to large block size and convert from HND to NHD + + """ + if len(self.device_kv_caches) == 0: + return + assert block_size_ratio >= 1, "Only nP < nD supported currently." + assert self.transfer_topo is not None + if self.enable_permute_local_kv and block_size_ratio > 1: + logger.debug( + "Post-processing device kv cache on receive by converting " + "block_size with %sx bigger and permuting layout from HND" + " to NHD.", + block_size_ratio, + ) + elif self.enable_permute_local_kv: + logger.debug( + "Post-processing device kv cache on receive by permuting layout" + "from HND to NHD." + ) + else: + logger.debug( + "Post-processing device kv cache on receive by converting " + "block_size with %sx bigger.", + block_size_ratio, + ) + + split_k_and_v = self.transfer_topo.split_k_and_v + + for block_ids in block_ids_list: + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + + for _, cache_or_caches in self.device_kv_caches.items(): + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches] + for cache in cache_list: + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive( + cache, indices, block_size_ratio + ) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive( + cache, indices, block_size_ratio + ) + + def post_process_device_kv_on_receive_heterogeneous_attn( + self, block_ids: list[int] + ): + """ + Post process device kv cache after receiving from remote + for heterogeneous attention. + """ + assert self.enable_heterogeneous_attn_post_process + + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + + for _, cache_or_caches in self.device_kv_caches.items(): + blocks_to_update = cache_or_caches.index_select(1, indices) + current_platform.pack_kv_cache( + key=blocks_to_update[0], + value=blocks_to_update[1], + key_cache=cache_or_caches[0], + value_cache=cache_or_caches[1], + block_ids=block_ids, + indices=indices, + ) + + def get_finished(self) -> tuple[set[str], set[str]]: + """ + Get requests that are done sending or recving on this specific worker. + The scheduler process (via the MultiprocExecutor) will use this output + to track which workers are done. + """ + assert self.transfer_topo is not None + done_sending = self._get_new_notifs() + done_recving = self._pop_done_transfers(self._recving_transfers) + + # Drain queue of requests where handshake or transfer setup failed. + failed_recv_reqs = set[ReqId]() + while not self._failed_recv_reqs.empty(): + try: + failed_recv_reqs.add(self._failed_recv_reqs.get_nowait()) + except queue.Empty: + break + + # Add failed requests to done_recving for scheduler tracking + # (blocks are already marked invalid, scheduler will handle recompute) + done_recving.update(failed_recv_reqs) + + if len(done_sending) > 0 or len(done_recving) > 0: + logger.debug( + "Rank %s, get_finished: %s requests done sending " + "and %s requests done recving (%s failed)", + self.tp_rank, + len(done_sending), + len(done_recving), + len(failed_recv_reqs), + ) + + block_ids_for_blocksize_post_process = defaultdict(list) + block_ids_for_heterogeneous_attn_post_process = list[list[int]]() + for req_id in done_recving: + # clean up metadata for completed requests + meta = self._recving_metadata.pop(req_id, None) + assert meta is not None, f"{req_id} not found in recving_metadata list" + + # Skip KV sync and post-processing for failed requests + if req_id in failed_recv_reqs: + logger.warning( + "Skipping KV post-processing for failed request %s", + req_id, + ) + continue + + assert meta.remote is not None + if self.use_host_buffer: + self.sync_recved_kv_to_device(req_id, meta) + + # post processing for heteroblocksize + remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) + block_size_ratio = self.transfer_topo.block_size_ratio( + remote_info.remote_block_size + ) + if not self.use_mla and ( + block_size_ratio > 1 or self.enable_permute_local_kv + ): + assert not self._is_hma_required + block_ids_for_blocksize_post_process[block_size_ratio].append( + meta.local_physical_block_ids[0] + ) + # post processing for heterogeneous attention + if self.enable_heterogeneous_attn_post_process: + block_ids_for_heterogeneous_attn_post_process.append( + meta.local_physical_block_ids[0] + ) + for ( + block_size_ratio, + block_ids_list, + ) in block_ids_for_blocksize_post_process.items(): + self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + + for block_ids in block_ids_for_heterogeneous_attn_post_process: + self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) + + # Handle timeout to avoid stranding blocks on remote. + now = time.perf_counter() + while self._reqs_to_send: + req_id, expires = next(iter(self._reqs_to_send.items())) + # Sorted dict, oldest requests are put first so we can exit early. + if now < expires: + break + count = self.consumer_notification_counts_by_req.pop(req_id, 0) + self.xfer_stats.record_kv_expired_req() + logger.warning( + "Releasing expired KV blocks for request %s which were " + "retrieved by %d remote worker(s) before lease expired.", + req_id, + count, + ) + self._reqs_to_process.remove(req_id) + del self._reqs_to_send[req_id] + done_sending.add(req_id) + + return done_sending, done_recving + + def _get_new_notifs(self) -> set[str]: + """Get req_ids which got a remote xfer notification. + + Subclasses must implement this to handle mode-specific notifications. + """ + raise NotImplementedError + + def _handle_heartbeat(self, payload: str) -> None: + """Extend leases for requests referenced in a heartbeat. + + Args: + payload: comma-separated P-side request IDs, e.g. + "req_abc,req_def". + """ + new_expiry = time.perf_counter() + self._lease_extension + for req_id in payload.split(","): + if req_id in self._reqs_to_send: + old = self._reqs_to_send[req_id] + self._reqs_to_send[req_id] = max(old, new_expiry) + logger.debug( + "Heartbeat extended lease for request %s " + "by %ds (old_expiry=%.1f, new_expiry=%.1f)", + req_id, + self._lease_extension, + old, + new_expiry, + ) + + def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: + """ + Pop completed xfers by checking for DONE state. + Args: + transfers: dict of req_id -> list[running_xfer] + Returns: + set of req_ids that have all done xfers + """ + done_req_ids: set[str] = set() + for req_id, handles in list(transfers.items()): + in_progress = [] + for handle in handles: + try: + xfer_state = self.nixl_wrapper.check_xfer_state(handle) + if xfer_state == "DONE": + # Get telemetry from NIXL + res = self.nixl_wrapper.get_xfer_telemetry(handle) + self.xfer_stats.record_transfer(res) + self.nixl_wrapper.release_xfer_handle(handle) + elif xfer_state == "PROC": + in_progress.append(handle) + continue + else: + self._log_failure( + failure_type="transfer_failed", + msg="Marking blocks as invalid", + req_id=req_id, + xfer_state=xfer_state, + ) + self._handle_failed_transfer(req_id, handle) + except Exception as e: + self._log_failure( + failure_type="transfer_exception", + msg="Marking blocks as invalid", + req_id=req_id, + error=e, + ) + self._handle_failed_transfer(req_id, handle) + + if not in_progress: + # Only report request as completed when all transfers are done. + done_req_ids.add(req_id) + del transfers[req_id] + else: + transfers[req_id] = in_progress + return done_req_ids + + def _handle_failed_transfer(self, req_id: str, handle: int | None): + """ + Handle a failed transfer by marking all (logical) blocks as invalid and + recording the failure. + + Args: + req_id: The request ID. + handle: The transfer handle. + """ + # Use .get() here as the metadata cleanup is handled by get_finished() + # TODO (NickLucche) handle failed transfer for HMA. + if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required: + self._invalid_block_ids.put(set(meta.local_block_ids[0])) + self._failed_recv_reqs.put(req_id) + if handle is not None: + self.nixl_wrapper.release_xfer_handle(handle) + self.xfer_stats.record_failed_transfer() + + def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None: + """ + Send heartbeat notifications to remote engines, extending lease on KV blocks. + """ + for engine_id, hb_info in metadata.heartbeat_by_engine.items(): + # Proactive handshake (this request may still be in waiting queue) so + # the **next** heartbeat for this remote can go through. + if ( + self._ensure_handshake( + engine_id, hb_info.host, hb_info.port, hb_info.tp_size + ) + is not None + ): + continue # handshake is still pending + + # Build the heartbeat message: "HB:req1,req2,..." + hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode() + for agent_name in self._remote_agents[engine_id].values(): + try: + self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg) + except Exception: + logger.debug( + "Failed to send heartbeat to engine %s", + engine_id, + exc_info=True, + ) + + def get_mapped_blocks( + self, block_ids: np.ndarray, block_size_ratio: int + ) -> np.ndarray: + """ + Calculates the new set of block IDs by mapping every element + in the (potentially sparse) input array. + Example: block_ids=[0, 2], block_size_ratio=2 + get_mapped_blocks 0 1 [2 3] 4 5 + # remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1|| + # local is |h0-b0......||h1-b0......||h2-b0........ + local_block_ids 0 [1] 2 + """ + if block_ids.size == 0: + return np.array([], dtype=np.int64) + + start_ids = block_ids * block_size_ratio + offsets = np.arange(block_size_ratio) + mapped_2d = start_ids[:, None] + offsets[None, :] + + return mapped_2d.flatten().astype(np.int64) + + def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: + """ + Convert logical block ids to kernel physical block ids. + This is required when the logical block size (the one set by the user) + does not match the one required by the attn backend. + """ + if self._physical_blocks_per_logical_kv_block == 1: + # Noop when physical and logical block sizes are the same + return block_ids + block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape( + 1, -1 + ) + # Mamba blocks have no logical<>physical discrepancy + group_specs = self.kv_cache_config.kv_cache_groups + return [ + BlockTable.map_to_kernel_blocks( + np.array(group), + self._physical_blocks_per_logical_kv_block, + block_arange, + ).tolist() + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) + else group + for i, group in enumerate(block_ids) + ] + + def _apply_prefix_caching( + self, + local_block_ids: BlockIds, + remote_block_ids: BlockIds, + remote_physical_per_logical: int, + ) -> tuple[BlockIds, list]: + """Apply prefix caching by trimming local/remote block ID lists. + + For non-Mamba models: end-trim remote to match local count, so that + already-cached prefix blocks are skipped in the transfer. + + For Mamba hybrid (prefix caching not yet supported): front-trim both + to the minimum count to handle kernel block count discrepancies from + logical block rounding in heterogeneous TP. + """ + # Partial prefix cache hit: just read uncomputed blocks. + # Skip mamba groups — their blocks represent full state (conv+ssm), + # not per-token data, so trimming would corrupt the transfer. + remote_block_ids = list(remote_block_ids) + if not self._has_mamba: + for i, remote_group in enumerate(remote_block_ids): + num_local_blocks = len(local_block_ids[i]) + assert num_local_blocks <= len(remote_group) + if num_local_blocks < len(remote_group): + remote_block_ids[i] = remote_group[-num_local_blocks:] + else: + # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP + # can cause different kernel block counts due to logical block rounding. + # Example: 640 prompt tokens, kernel_block_size=64 + # remote physical_per_logical=10, local physical_per_logical=6 + # remote logical ids from kv_transfer_params = [0] + # local logical ids allocated = [0, 1] + # remote kernel blocks: [0..9] (1*10=10) + # local kernel blocks: [0..11] (2*6=12) + # actual data blocks = ceil(640/64) = 10, trim both to 10 + # Vice versa (remote physical_per_logical=6, local=10): + # remote logical ids = [0, 1], local logical ids = [0] + # remote kernel blocks: [0..11] (2*6=12) + # local kernel blocks: [0..9] (1*10=10) + # actual data blocks = ceil(640/64) = 10, trim both to 10 + local_block_ids = list(local_block_ids) + for i, remote_group in enumerate(remote_block_ids): + num_local_blocks = len(local_block_ids[i]) + num_remote_blocks = len(remote_group) + if ( + _is_ssm_spec(self._group_spec_types[i]) + and num_local_blocks < num_remote_blocks + ): + # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks + # prior to the last one are placeholders (null blocks). Mind that + # this doesn't really impact transfer, as we only still care about + # the last "block", the full in-place state. + assert num_local_blocks == 1, "SSM can only have one local block" + remote_block_ids[i] = remote_group[-num_local_blocks:] + elif ( + self._physical_blocks_per_logical_kv_block + == remote_physical_per_logical + and num_local_blocks < num_remote_blocks + ): + # Partial prefix cache hit for FA group. + remote_block_ids[i] = remote_group[-num_local_blocks:] + else: + # TODO Handle prefix caching with different block_sizes + max_padding = max( + self._physical_blocks_per_logical_kv_block, + remote_physical_per_logical, + ) + assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( + f"Group {i}: |{num_local_blocks} - " + f"{num_remote_blocks}| >= {max_padding}" + ) + num_blocks = min(num_local_blocks, num_remote_blocks) + local_block_ids[i] = local_block_ids[i][:num_blocks] + remote_block_ids[i] = remote_group[:num_blocks] + return local_block_ids, remote_block_ids + + def _logical_to_remote_kernel_block_ids( + self, block_ids: BlockIds, remote_physical_per_logical: int + ) -> BlockIds: + """Map logical block IDs to physical kernel block IDs on the remote. + + Args: + block_ids: per-group lists of logical block IDs. + remote_physical_per_logical: remote engine's physical blocks + per logical block. + + Returns: + Same structure with FA groups expanded (each logical block L + becomes kernel blocks [L*remote_physical_per_logical, .. + L*remote_physical_per_logical + + remote_physical_per_logical - 1]). + Mamba groups are passed through unchanged. + """ + if remote_physical_per_logical == 1: + return block_ids + remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1) + group_specs = self.kv_cache_config.kv_cache_groups + result = [ + BlockTable.map_to_kernel_blocks( + np.array(group), + remote_physical_per_logical, + remote_arange, + ).tolist() + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) + else group + for i, group in enumerate(block_ids) + ] + return result + + def get_backend_aware_kv_block_len( + self, layer_idx: int, first_split: bool = True, mamba_view: bool = False + ) -> int: + """ + Get the block length for one K/V element (K and V have the same size). + + For FA and other backends, this is equal to the length of the whole + block, as K and V are in separate regions. + For FlashInfer, this is half the length of the whole block, as K and V + share the same region. + Similarly, for SSM-based models, state and conv are interleaved, but crucially + the their size differs. + Reference diagram: + KVCacheTensor (Shared) + / \\ + / \\ + / \\ + Attention (FlashInfer) View Mamba View + | | + | | + +-------------------+ +-------------------+ + | KVCacheTensor | | KVCacheTensor | + | | | | + |<----- page ------>| |<----- page ------->| + | size | | size | + | Key 0 | Val 0 | |Conv 0 | SSM 0 | + | Key 1 | Val 1 | |Conv 1 | SSM 1 | + | ... | ... | | ... | ... | + | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | + | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | + +-------------------+ +--------------------+ + |1st_split-2nd_split| |1st_split-2nd_split | + """ + assert self.transfer_topo is not None + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks + if virtually_split and mamba_view: + block_len = self._mamba_ssm_size[not first_split] + else: + half_block = virtually_split and not self._is_region_replicated(layer_idx) + block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1) + return block_len + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """ + Get the KV transfer stats for the connector. + """ + # Clear stats for next iteration + if not self.xfer_stats.is_empty(): + return self.xfer_stats.clone_and_reset() + return None + + def get_block_ids_with_load_errors(self) -> set[int]: + """ + Return and clear the set of block IDs that failed to load. + + This is called by the scheduler to identify blocks that need + to be retried after a NIXL transfer failure. + """ + # Drain the queue (thread-safe, no lock needed). + result: set[int] = set() + while not self._invalid_block_ids.empty(): + try: + result.update(self._invalid_block_ids.get_nowait()) + except queue.Empty: + break + return result + + def _evict_stale_engines(self) -> None: + """Scan for and evict remote engines that have exceeded their TTL. + + Called from the main thread in when a new remote engine appears. + We can only go OOM as we discover and register a new remote, therefore we make + sure we clean up stale engine data structures before then. This invariant + prevents us from using background threads, though memory usage is not guaranteed + to be "optimal" until a new handshake is performed. + + Engines with active transfers or pending handshakes cannot be stale: + - Active transfers touch _engine_last_active in start_load_kv. + - Pending handshakes don't have an _engine_last_active entry yet + """ + # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number + # of remote engines is registered all at once (adding a background cleanup + # thread wouldnt help either). + # If that scenario is plausible, we can follow up with an LRU eviction policy. + if self._engine_ttl <= 0: + return + + now = time.perf_counter() + for eid, last_active in list(self._engine_last_active.items()): + if now - last_active > self._engine_ttl: + self._cleanup_remote_engine(eid) + + def _cleanup_remote_engine( + self, engine_id: EngineId, *, log_eviction: bool = True + ) -> None: + """Remove all state for a single remote engine. + + Releases NIXL resources (dlist handles, remote agents) and clears + all per-engine data structures. Used by both TTL eviction and + shutdown. + """ + assert engine_id in self._remote_agents + + for handle in self.dst_xfer_side_handles.pop(engine_id).values(): + self.nixl_wrapper.release_dlist_handle(handle) + for agent_name in self._remote_agents.pop(engine_id).values(): + self.nixl_wrapper.remove_remote_agent(agent_name) + + del self.kv_caches_base_addr[engine_id] + del self.dst_num_blocks[engine_id] + del self.tp_mappings[engine_id] + if self.transfer_topo is not None: + self.transfer_topo.unregister_remote_engine(engine_id) + + last_active = self._engine_last_active.pop(engine_id) + if log_eviction: + logger.info( + "Evicted stale remote engine %s (inactive for %.1fs).", + engine_id, + time.perf_counter() - last_active, + ) + + def __del__(self): + self.shutdown() + + def shutdown(self): + """Shutdown the connector worker.""" + if not hasattr(self, "_handshake_initiation_executor"): + # error happens during init, no need to shutdown + return + self._handshake_initiation_executor.shutdown(wait=False) + for handles in self._recving_transfers.values(): + for handle in handles: + self.nixl_wrapper.release_xfer_handle(handle) + self._recving_transfers.clear() + for handle in self.src_xfer_handles_by_block_size.values(): + self.nixl_wrapper.release_dlist_handle(handle) + self.src_xfer_handles_by_block_size.clear() + for handles in self.src_xfer_handles_by_tp_ratio.values(): + for handle in handles: + self.nixl_wrapper.release_dlist_handle(handle) + self.src_xfer_handles_by_tp_ratio.clear() + for engine_id in list(self._remote_agents): + self._cleanup_remote_engine(engine_id, log_eviction=False) + for desc in self._registered_descs: + self.nixl_wrapper.deregister_memory(desc) + self._registered_descs.clear() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index dad81e84c45..b3214505309 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -1,6 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NixlConnector – thin facade that delegates to scheduler / worker.""" +"""NIXL connector facades. + +This module hosts the thin facade classes that vLLM's KV-connector layer +instantiates. Almost all the real work lives in the per-mode scheduler +and worker classes; the connector classes here only forward calls. + +* :class:`NixlBaseConnector` – common logic shared by pull and push. +* :class:`NixlPullConnector` – pull-based (READ) KV transfer. +* :class:`NixlPushConnector` – push-based (WRITE) KV transfer. +* ``NixlConnector`` – backward-compatible alias for :class:`NixlPullConnector`. +""" from typing import TYPE_CHECKING, Any @@ -28,16 +38,22 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( - NixlConnectorScheduler, +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import ( + NixlPullConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( + NixlPushConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( NixlKVConnectorStats, NixlPromMetrics, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( - NixlConnectorWorker, -) from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata @@ -47,6 +63,12 @@ from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import ( + NixlBaseConnectorScheduler, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, + ) from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request @@ -54,7 +76,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class NixlConnector(KVConnectorBase_V1, SupportsHMA): +class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): + """Base connector with common logic shared by pull and push modes.""" + @property def prefer_cross_layer_blocks(self) -> bool: if any( @@ -106,16 +130,9 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA): self.kv_cache_config = kv_cache_config self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id self.kv_transfer_config = vllm_config.kv_transfer_config - if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler: NixlConnectorScheduler | None = ( - NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) - ) - self.connector_worker: NixlConnectorWorker | None = None - elif role == KVConnectorRole.WORKER: - self.connector_scheduler = None - self.connector_worker = NixlConnectorWorker( - vllm_config, self.engine_id, kv_cache_config - ) + # Subclasses must set self.connector_scheduler and self.connector_worker + self.connector_scheduler: NixlBaseConnectorScheduler | None = None + self.connector_worker: NixlBaseConnectorWorker | None = None ############################################################ # Class Methods @@ -256,11 +273,6 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA): vllm_config, metric_types, labelnames, per_engine_labelvalues ) - def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: - assert self.connector_worker is not None - assert isinstance(self._connector_metadata, NixlConnectorMetadata) - self.connector_worker.start_load_kv(self._connector_metadata) - def wait_for_layer_load(self, layer_name: str) -> None: """NixlConnector does not do layerwise saving.""" pass @@ -281,6 +293,11 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA): if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks: self.connector_worker.save_kv_to_host(self._connector_metadata) + def has_pending_push_work(self) -> bool: + if self.connector_scheduler is not None: + return self.connector_scheduler.has_pending_push_work() + return False + def shutdown(self): if self.connector_worker is not None: self.connector_worker.shutdown() @@ -299,3 +316,79 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA): """ assert self.connector_worker is not None return self.connector_worker.xfer_handshake_metadata + + +class NixlPullConnector(NixlBaseConnector): + """Pull-based (READ) NIXL KV transfer connector.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, role, kv_cache_config) + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = NixlPullConnectorScheduler( + vllm_config, self.engine_id, kv_cache_config + ) + self.connector_worker = None + elif role == KVConnectorRole.WORKER: + self.connector_scheduler = None + self.connector_worker = NixlPullConnectorWorker( + vllm_config, self.engine_id, kv_cache_config + ) + + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: + assert self.connector_worker is not None + assert isinstance(self.connector_worker, NixlPullConnectorWorker) + assert isinstance(self._connector_metadata, NixlConnectorMetadata) + self.connector_worker.start_load_kv(self._connector_metadata) + + +class NixlPushConnector(NixlBaseConnector): + """Push-based (WRITE) NIXL KV transfer connector.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, role, kv_cache_config) + self.connector_scheduler: NixlPushConnectorScheduler | None = None + self.connector_worker: NixlPushConnectorWorker | None = None + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = NixlPushConnectorScheduler( + vllm_config, self.engine_id, kv_cache_config + ) + elif role == KVConnectorRole.WORKER: + self.connector_worker = NixlPushConnectorWorker( + vllm_config, self.engine_id, kv_cache_config + ) + else: + raise ValueError(f"Unsupported KVConnectorRole: {role}") + + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: + """Drive push processing on the worker. + + The worker enqueues registrations / finished blocks for the + background ``nixl-push-writer`` thread; the writer issues the + WRITE transfers and polls NIXL notifs without further + engine-thread involvement. + """ + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, NixlConnectorMetadata) + self.connector_worker.start_load_kv(self._connector_metadata) + + +# Backward compatibility: NixlConnector is the pull-based connector. +NixlConnector = NixlPullConnector + + +__all__ = [ + "NixlBaseConnector", + "NixlConnector", + "NixlPullConnector", + "NixlPushConnector", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index b9e3436f501..c120f939aff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -19,6 +19,11 @@ TransferHandle = int ReqId = str GET_META_MSG = b"get_meta_msg" + +# Push-mode (WRITE-based) registration notification. +# Sent worker-to-worker over NIXL: D worker -> P worker, encoded as +# PUSH_REG_NOTIF_PREFIX + msgpack(registration_data). +PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # # NIXL Connector Version # @@ -160,6 +165,8 @@ class ReqMeta: local_physical_block_ids: BlockIds tp_size: int remote: RemoteMeta | None = None + # Remote block size, discovered during NIXL handshake (push mode). + remote_block_size: int | None = None class NixlConnectorMetadata(KVConnectorMetadata): @@ -171,6 +178,12 @@ class NixlConnectorMetadata(KVConnectorMetadata): self.reqs_not_processed: set[ReqId] = set() # Heartbeat data grouped by remote engine, sent by D worker to P. self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} + # Push mode (D side): registration data the D worker should send to + # P workers via NIXL notification on this step. + self.push_registrations: dict[ReqId, dict[str, Any]] = {} + # Push mode (P side): newly finished request blocks to be matched + # against pending D registrations on the P worker. + self.push_finished_blocks: dict[ReqId, BlockIds] = {} def _add_new_req( self, @@ -182,6 +195,7 @@ class NixlConnectorMetadata(KVConnectorMetadata): local_physical_block_ids=local_block_ids, # P workers don't need to receive tp_size from proxy here. tp_size=kv_transfer_params.get("tp_size", 1), + remote_block_size=kv_transfer_params.get("remote_block_size"), ) def add_new_req_to_save( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py new file mode 100644 index 00000000000..f13e2160566 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pull-specific scheduler-side logic for the NIXL connector.""" + +import time +from typing import TYPE_CHECKING, Any + +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import ( + NixlBaseConnectorScheduler, +) +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): + """Pull-specific scheduler logic (READ-based KV transfer).""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, engine_id, kv_cache_config) + + def get_num_new_matched_tokens( + self, request: "Request", num_computed_tokens: int + ) -> tuple[int, bool]: + """ + For remote prefill, pull all prompt blocks from remote + asynchronously relative to engine execution. + + Args: + request (Request): the request object. + num_computed_tokens (int): the number of locally + computed tokens for this request + Returns: + * the number of tokens that can be loaded from the + external KV cache beyond what is already computed. + * true if the external KV cache tokens will be loaded + asynchronously (between scheduler steps). + """ + + params = request.kv_transfer_params + logger.debug( + "NIXLConnector get_num_new_matched_tokens: " + "num_computed_tokens=%s, kv_transfer_params=%s", + num_computed_tokens, + params, + ) + + if params is not None and params.get("do_remote_prefill"): + # Remote prefill: get all prompt blocks from remote. + token_ids = request.prompt_token_ids or [] + actual = self._mamba_prefill_token_count(len(token_ids)) + count = actual - num_computed_tokens + if count > 0: + return count, True + + if params is not None and params.get("do_remote_decode") and self._has_mamba: + self._truncate_mamba_request_for_prefill(request) + + if ( + params is not None + and params.get("do_remote_decode") + and params.get("remote_block_ids") + and all( + p in params + for p in ( + "remote_engine_id", + "remote_request_id", + "remote_host", + "remote_port", + ) + ) + ): + # Decode node has kv blocks for part of prefill request, so, provide them + # as an external token count to scheduler. + # The tokens will be loaded if not already present + # in the prefill node local cache + remote_num_tokens = params.get("remote_num_tokens") or 0 + count = ( + min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens + ) + if count > 0: + # Check kv_recompute_threshold: skip pull if + # remote tokens are below the threshold. + if ( + self.kv_recompute_threshold > 0 + and count < self.kv_recompute_threshold + ): + logger.debug( + "Skipping remote pull for %s: %d remote tokens < threshold %d", + request.request_id, + count, + self.kv_recompute_threshold, + ) + return 0, False + return count, True + + # No remote prefill for this request. + return 0, False + + def update_state_after_alloc( + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int + ): + params = request.kv_transfer_params + logger.debug( + "NIXLConnector update_state_after_alloc: " + "num_external_tokens=%s, kv_transfer_params=%s", + num_external_tokens, + params, + ) + + if not params: + return + + if params.get("do_remote_decode") or ( + params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled + ): + self._reqs_in_batch.add(request.request_id) + if self.use_host_buffer and params.get("do_remote_decode"): + # NOTE: when accelerator is not directly supported by Nixl, + # prefilled blocks need to be saved to host memory before transfer. + self._reqs_need_save[request.request_id] = request + elif params.get("do_remote_prefill") or ( + params.get("do_remote_decode") + and self.is_bidirectional_kv_xfer_enabled + and not params.get("_remote_blocks_processed") + ): + if params.get("remote_block_ids"): + if all( + p in params + for p in ( + "remote_engine_id", + "remote_request_id", + "remote_host", + "remote_port", + ) + ): + # If remote_blocks and num_external_tokens = 0, we have + # a full prefix cache hit on the local node. We need to call + # send_notif in _read_blocks to free the memory on the remote node. + + unhashed_local_block_ids: BlockIds = ( + blocks.get_unhashed_block_ids_all_groups() + if num_external_tokens > 0 + else () + ) + local_block_ids = self.get_sw_clipped_blocks( + unhashed_local_block_ids + ) + + # Get unhashed blocks to pull from remote. Mind that a full prefix + # cache hit is indicated with an empty list. + self._reqs_need_recv[request.request_id] = ( + request, + local_block_ids, + ) + + else: + logger.warning( + "Got invalid KVTransferParams: %s. This " + "request will not utilize KVTransfer", + params, + ) + else: + assert num_external_tokens == 0 + # Only trigger 1 KV transfer per request. + params["do_remote_prefill"] = False + params["_remote_blocks_processed"] = True + + def request_finished( + self, + request: "Request", + block_ids: "BlockIds", + ) -> tuple[bool, dict[str, Any] | None]: + """ + Once a request is finished, determine whether request blocks + should be freed now or will be sent asynchronously and freed later. + """ + from vllm.v1.request import RequestStatus + + params = request.kv_transfer_params + logger.debug( + "NIXLConnector request_finished(%s), request_status=%s, " + "kv_transfer_params=%s", + request.request_id, + request.status, + params, + ) + if not params: + return False, None + + is_p_node = bool(params.get("do_remote_decode")) + is_d_node = not is_p_node + + # Stop heartbeating for aborted requests that never reached finished_recving: + # normal path cleans up in update_connector_output. + self._stop_heartbeat(request.request_id) + + if params.get("do_remote_prefill"): + # If do_remote_prefill is still True when the request is finished, + # update_state_after_alloc must not have been called (the request + # must have been aborted before it was scheduled, e.g. via the + # abort_immediately path used to clean up KV-transfer requests + # rejected at the D-side serving layer). + # To avoid stranding the prefill blocks in the prefill instance, + # we must add empty block_ids to _reqs_need_recv so that our + # worker side will notify and free blocks in the prefill instance. + self._reqs_need_recv[request.request_id] = (request, []) + params["do_remote_prefill"] = False + return False, None + + if is_d_node and not self.is_bidirectional_kv_xfer_enabled: + return False, None + + if request.status not in ( + RequestStatus.FINISHED_LENGTH_CAPPED, + RequestStatus.FINISHED_STOPPED, + ): + # Also include the case of a P/D Prefill request with immediate + # block free (eg abort). Stop tracking this request. + self._reqs_not_processed.add(request.request_id) + # Clear _reqs_need_save if a request is aborted as partial prefill. + self._reqs_need_save.pop(request.request_id, None) + return False, None + + # TODO: check whether block_ids actually ever be 0. If not we could + # remove the conditional below + delay_free_blocks = any(len(group) > 0 for group in block_ids) + remote_num_tokens = 0 + if delay_free_blocks: + # Prefill request on remote. It will be read from D upon completion + request_kv_blocks_ttl = self._kv_lease_duration + if is_d_node: + # For blocks pinned on D, use a simpler timeout for now instead of a + # lease mechanism as turn2 request is client-driven. + request_kv_blocks_ttl = self.decoder_kv_blocks_ttl + logger.debug( + "NIXLConnector request_finished(%s) waiting for %d seconds " + "before releasing blocks", + request.request_id, + request_kv_blocks_ttl, + ) + self._reqs_need_send[request.request_id] = ( + time.perf_counter() + request_kv_blocks_ttl + ) + # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), + # trimming down after allocating for the whole sequence length. Empty + # blocks are always at the start of the list. + # Here we "unpad" blocks to send the actual remote blocks to be read. + block_ids = self.get_sw_clipped_blocks(block_ids) + + remote_num_tokens = request.num_computed_tokens + + return delay_free_blocks, dict( + do_remote_prefill=is_p_node, + do_remote_decode=is_d_node, + remote_block_ids=block_ids, + remote_engine_id=self.engine_id, + remote_request_id=request.request_id, + remote_host=self.side_channel_host, + remote_port=self.side_channel_port, + tp_size=self.vllm_config.parallel_config.tensor_parallel_size, + remote_num_tokens=remote_num_tokens, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py new file mode 100644 index 00000000000..26f5fde24d8 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pull-specific (READ) worker-side logic for the NIXL connector.""" + +import time +from typing import TYPE_CHECKING + +import numpy as np + +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ReqMeta, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + ReadSpec, +) +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import KVCacheConfig + +logger = init_logger(__name__) + + +class NixlPullConnectorWorker(NixlBaseConnectorWorker): + """Pull-specific (READ) worker logic.""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, engine_id, kv_cache_config) + + def start_load_kv(self, metadata: NixlConnectorMetadata): + """ + Start loading by triggering non-blocking nixl_xfer. + We check for these trnxs to complete in each step(). + """ + for req_id, meta in metadata.reqs_to_recv.items(): + meta.local_physical_block_ids = self._logical_to_kernel_block_ids( + meta.local_block_ids + ) + assert meta.remote is not None + # Remote block IDs are kept logical here; expanded in + # _read_blocks_for_req using the remote engine's phys ratio. + remote_engine_id = meta.remote.engine_id + logger.debug( + "start_load_kv for request %s from remote engine %s. " + "Num local_block_ids: %s. Num remote_block_ids: %s. ", + req_id, + remote_engine_id, + len(meta.local_physical_block_ids), + len(meta.remote.block_ids), + ) + # always store metadata for failure recovery + self._recving_metadata[req_id] = meta + if remote_engine_id not in self._remote_agents: + # Initiate handshake with remote engine to exchange metadata. + with self._handshake_lock: + if remote_engine_id not in self._remote_agents: + self._background_nixl_handshake(req_id, remote_engine_id, meta) + continue + + # Handshake already completed, start async read xfer. + self._read_blocks_for_req(req_id, meta) + + # Start transfers for requests whose handshakes have now finished. + while not self._ready_requests.empty(): + self._read_blocks_for_req(*self._ready_requests.get_nowait()) + + # Keep around the requests that have been part of a batch. This is + # needed because async scheduling pushes the misalignment between the + # moment in which requests expiration is set (P side) and the moment in + # which blocks are read from D. As P can now more easily lag behind D + # while processing the next batch, we make sure to only set an + # expiration for requests that have not been read from D yet. + for req_id in metadata.reqs_in_batch: + self._reqs_to_process.add(req_id) + + # Remove all requests that are not to be processed (eg aborted). + for req_id in metadata.reqs_not_processed: + self._reqs_to_process.discard(req_id) + # We should never get an abort after setting an expiry timer + assert req_id not in self._reqs_to_send + + # Add to requests that are waiting to be read and track expiration. + for req_id, expiration_time in metadata.reqs_to_send.items(): + if req_id in self._reqs_to_process: + self._reqs_to_send[req_id] = expiration_time + + # Send heartbeats to P-side engines to keep KV blocks alive while + # requests sit in the D scheduler WAITING queue. + self._send_heartbeats(metadata) + + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): + assert meta.remote is not None and self.transfer_topo is not None + engine_id = meta.remote.engine_id + # Update last activity from this remote. Mind that cleanup is done on main + # thread (this one), so we don't race on this structure. + self._engine_last_active[engine_id] = time.perf_counter() + plan = self.tp_mappings[engine_id] + remote_info = self.transfer_topo.get_engine_info(engine_id) + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) + + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids, + remote_info.remote_physical_blocks_per_logical, + ) + remote_block_ids = meta.remote.block_ids + local_block_ids = meta.local_physical_block_ids + num_groups = len(local_block_ids) + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=[ + list(local_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], + remote_block_ids=[ + list(remote_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], + ) + for rank in plan.all_source_ranks + ] + + # D may have to perform multiple reads from different remote ranks. + # MLA opt: when P TP > D TP, only a single read is executed for + # the first remote rank (cache is duplicated).. + if self.use_mla and tp_ratio < 0: + assert len(read_specs) == 1 + + for i, spec in enumerate(read_specs): + remote_block_size = remote_info.remote_block_size + logger.debug( + "Remote agent %s available, calling _read_blocks" + " on remote rank %s with remote block size %s for req %s", + meta.remote.engine_id, + spec.remote_rank, + remote_block_size, + req_id, + ) + # Get side handles. + if tp_ratio < 0 and not self.use_mla: + assert remote_block_size == self.block_size + # Remote tp_size > local tp_size: we must perform multiple + # reads. Get the memory chunk onto which we will write to. + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + else: + # Single read from remote, we write to the whole memory region. + # Also handle remote block size different from local block size. + local_xfer_side_handle = self.src_xfer_handles_by_block_size[ + remote_block_size + ] + + # Destination handle: remote_engine_id -> remote_rank -> handle. + remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][ + spec.remote_rank + ] + + self._read_blocks( + read_spec=spec, + request_id=req_id, + dst_engine_id=meta.remote.engine_id, + remote_request_id=meta.remote.request_id, + local_xfer_side_handle=local_xfer_side_handle, + remote_xfer_side_handle=remote_xfer_side_handle, + ) + + if self.use_mla and tp_ratio < 0 and read_specs: + # ..but we still need to notify the other remote ranks that we + # have the blocks we need so they can update the request state. + notif_id = f"{meta.remote.request_id}:{self.world_size}".encode() + remote_agents = self._remote_agents[meta.remote.engine_id] + for rank_to_notify, agent in remote_agents.items(): + if rank_to_notify != read_specs[0].remote_rank: + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) + + def _read_blocks( + self, + read_spec: ReadSpec, + dst_engine_id: str, + request_id: str, + remote_request_id: str, + local_xfer_side_handle: int, + remote_xfer_side_handle: int, + ): + """ + Post a READ point-to-point xfer request from a single local worker to + a single remote worker. + """ + assert self.transfer_topo is not None + remote_rank = read_spec.remote_rank + local_block_ids = read_spec.local_block_ids + remote_block_ids = read_spec.remote_block_ids + + remote_info = self.transfer_topo.get_engine_info(dst_engine_id) + block_size_ratio = self.transfer_topo.block_size_ratio( + remote_info.remote_block_size + ) + if block_size_ratio > 1: + # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. + assert not self._is_hma_required + local_block_ids0 = local_block_ids[0] if local_block_ids else [] + remote_block_ids0 = remote_block_ids[0] + local_block_ids_mapped = self.get_mapped_blocks( + np.asarray(local_block_ids0), block_size_ratio + ).tolist() + if len(local_block_ids_mapped) > len(remote_block_ids0): + # NOTE: + # get_mapped_blocks will always expand block_ids for n times. + # ex: + # prefill block_ids with block_size as 4: + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + # Local decode block_ids with block_size as 16: [1, 2, 3] + # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + # Then we clip local to align with prefill + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + local_block_ids_mapped = local_block_ids_mapped[ + : len(remote_block_ids0) + ] + local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] + remote_block_ids = [remote_block_ids0] + # NOTE(rob): having the staging blocks be on the READER side is + # not going to work well (since we will have to call rearrange tensors). + # after we detect the txn is complete (which means we cannot make the + # read trxn async easily). If we want to make "READ" happen cleanly, + # then we will need to have the staging blocks on the remote side. + + # NOTE(rob): according to nvidia the staging blocks are used to + # saturate IB with heterogeneous TP sizes. + + # Number of D TP workers that will read from dst P. Propagate info + # on notification so that dst worker can wait before freeing blocks. + notif_id = f"{remote_request_id}:{self.world_size}".encode() + + # Full prefix cache hit: do not need to read remote blocks, + # just notify P worker that we have the blocks we need. + if len(local_block_ids) == 0: + # A full prefix cache hit is indicated with an empty list. + agent_name = self._remote_agents[dst_engine_id][remote_rank] + try: + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id) + except Exception as e: + self._log_failure( + failure_type="notification_failed", + msg="P worker blocks will be freed after timeout. " + "This may indicate network issues.", + req_id=request_id, + error=e, + dst_engine_id=dst_engine_id, + remote_rank=remote_rank, + remote_agent_name=agent_name, + ) + self.xfer_stats.record_failed_notification() + return + + assert ( + len(remote_block_ids) + == len(local_block_ids) + == len(self.kv_cache_config.kv_cache_groups) + ) + remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical + local_block_ids, remote_block_ids = self._apply_prefix_caching( + local_block_ids, remote_block_ids, remote_physical_per_logical + ) + + # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from + # corresponding rank. With heterogeneous TP, fixing D>P, the D tp + # workers will issue xfers to parts of the P worker remote kv caches. + + # Get descs ids. + remote_block_descs_ids = self._compute_desc_ids( + block_ids=remote_block_ids, + dst_num_blocks=self.dst_num_blocks[dst_engine_id], + block_size_ratio=None, + physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical, + ) + local_block_descs_ids = self._compute_desc_ids( + block_ids=local_block_ids, + dst_num_blocks=self.dst_num_blocks[self.engine_id], + block_size_ratio=block_size_ratio, + physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block, + ) + + assert len(local_block_descs_ids) == len(remote_block_descs_ids) + + # Prepare transfer with Nixl. + handle = None + try: + handle = self.nixl_wrapper.make_prepped_xfer( + "READ", + local_xfer_side_handle, + local_block_descs_ids, + remote_xfer_side_handle, + remote_block_descs_ids, + notif_msg=notif_id, + ) + + # Begin async xfer. + self.nixl_wrapper.transfer(handle) + + # Use handle to check completion in future step(). + self._recving_transfers[request_id].append(handle) + except Exception as e: + # mark all (logical) blocks for this request as invalid + self._log_failure( + failure_type="transfer_setup_failed", + req_id=request_id, + msg="Marking blocks as invalid", + error=e, + dst_engine_id=dst_engine_id, + remote_rank=remote_rank, + ) + self._handle_failed_transfer(request_id, handle) + + def _get_new_notifs(self) -> set[str]: + """ + Get req_ids which got a remote xfer message. When multiple consumers + are reading from the same producer (heterogeneous TP scenario), wait + for all consumers to be done pulling. + + Also handles heartbeat notifications ("HB:req1,req2,...") by + extending the lease on the referenced requests. + """ + assert self.transfer_topo is not None + notified_req_ids: set[str] = set() + for notifs in self.nixl_wrapper.get_new_notifs().values(): + for notif in notifs: + msg = notif.decode("utf-8") + + # Handle heartbeat messages from D-side. + if msg.startswith("HB:"): + self._handle_heartbeat(msg[3:]) + continue + + req_id, tp_size = msg.rsplit(":", 1) + if ( + req_id not in self._reqs_to_send + and req_id not in self._reqs_to_process + ): + logger.error( + "Potentially invalid KV blocks for " + "unrecognized request %s were retrieved by " + "a decode worker. They may have expired.", + req_id, + ) + continue + + # NOTE: `tp_ratio` is the opposite when swapping local<>remote + n_consumers = int(tp_size) + tp_ratio = self.transfer_topo.tp_ratio(n_consumers) + + # Number of reads *per producer* to wait for. + # When remote D TP > local P TP we expect `tp_ratio` reads. + consumers_per_producer = ( + -tp_ratio if n_consumers > self.world_size else 1 + ) + + self.consumer_notification_counts_by_req[req_id] += 1 + # Wait all consumers (D) to be done reading before freeing. + if ( + self.consumer_notification_counts_by_req[req_id] + == consumers_per_producer + ): + notified_req_ids.add(req_id) + del self.consumer_notification_counts_by_req[req_id] + self._reqs_to_process.remove(req_id) + self._reqs_to_send.pop(req_id, None) + return notified_req_ids diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py new file mode 100644 index 00000000000..dc976ae3a39 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Push-specific scheduler-side logic for the NIXL connector. + +In push mode, scheduler-side responsibilities are: + +* D side (decode): on ``update_state_after_alloc``, stash registration data + (D's identity + locally allocated block IDs) into + ``_push_pending_registrations``. The D worker drains it from + ``meta.push_registrations`` next step and sends a NIXL notification to the + P worker (no scheduler-level networking). +* P side (prefill): on ``request_finished``, stash the finished block IDs + into ``_finished_request_blocks`` for the lease, and into + ``_newly_finished_push_blocks`` so the P worker picks them up via + ``meta.push_finished_blocks`` and matches against any D registrations + it already received via NIXL notifications. +* Both sides: ``has_pending_push_work`` keeps the engine main loop stepping + while pushes are in flight. ``update_connector_output`` cleans up + ``_finished_request_blocks`` once the WRITE completes. + +A soft per-registration watchdog on the D scheduler fails requests that have +been registered but not fulfilled within a configurable timeout. +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import ( + NixlBaseConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ReqId, +) +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.outputs import KVConnectorOutput + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): + """Push-specific scheduler logic (WRITE-based KV transfer). + + All P2P communication is deferred to the worker level via NIXL + notifications. The scheduler communicates with workers only through + the standard ``build_connector_meta`` / ``update_connector_output`` + hooks. + """ + + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: KVCacheConfig, + ): + super().__init__(vllm_config, engine_id, kv_cache_config) + + # D-side: registration data to pass to D workers via metadata on + # the next ``build_connector_meta`` call. + self._push_pending_registrations: dict[ReqId, dict[str, Any]] = {} + + # D-side: track the wall-clock deadline for each registered request + # to detect "registered but never fulfilled" failures (e.g. the P + # node disappeared after registration). Keyed by D request_id. + self._push_registration_deadlines: dict[ReqId, float] = {} + + # P-side: block IDs for finished requests, kept for the lease and + # used to drive ``has_pending_push_work``. + self._finished_request_blocks: dict[ReqId, BlockIds] = {} + # P-side: newly finished blocks to ship to P workers on next step. + self._newly_finished_push_blocks: dict[ReqId, BlockIds] = {} + + # Soft watchdog timeout (seconds) for D-side registrations that + # never receive a push completion. Defaults to the existing + # decoder KV blocks TTL so behaviour matches the lease. + assert vllm_config.kv_transfer_config is not None + self._push_registration_timeout: float = float( + vllm_config.kv_transfer_config.get_from_extra_config( + "push_registration_timeout", + self.decoder_kv_blocks_ttl, + ) + ) + + def get_num_new_matched_tokens( + self, request: Request, num_computed_tokens: int + ) -> tuple[int, bool]: + """In push mode, D doesn't pull — it registers blocks and waits. + + However, we still need to handle the do_remote_prefill case where D + needs to know how many tokens will be pushed. + """ + params = request.kv_transfer_params + logger.debug( + "NixlPushConnector get_num_new_matched_tokens: " + "num_computed_tokens=%s, kv_transfer_params=%s", + num_computed_tokens, + params, + ) + + if params is not None and params.get("do_remote_prefill"): + token_ids = request.prompt_token_ids or [] + actual = self._mamba_prefill_token_count(len(token_ids)) + count = actual - num_computed_tokens + if count > 0: + return count, True + + if params is not None and params.get("do_remote_decode") and self._has_mamba: + self._truncate_mamba_request_for_prefill(request) + + return 0, False + + def update_state_after_alloc( + self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int + ): + """In push mode, D stores registration data for the worker to send + to P via NIXL notification (deferred to ``build_connector_meta``). + """ + params = request.kv_transfer_params + logger.debug( + "NixlPushConnector update_state_after_alloc: " + "num_external_tokens=%s, kv_transfer_params=%s", + num_external_tokens, + params, + ) + + if not params: + return + + # P side: track the request as in-batch so the lease accounting + # matches what the worker expects on the next step. + if params.get("do_remote_decode"): + self._reqs_in_batch.add(request.request_id) + + # P side with host-buffer offload: defer save to the worker. + if self.use_host_buffer and params.get("do_remote_decode"): + self._reqs_need_save[request.request_id] = request + return + + # D side: only act on the first call (``do_remote_prefill`` is + # unset on re-entry by the marker below). + if not params.get("do_remote_prefill"): + return + + if num_external_tokens <= 0: + # Nothing to receive: full prefix-cache hit on D, no + # registration to stage. + return + + # First-pass D path: stash registration data the worker will + # ship to P on the next ``build_connector_meta`` cycle. + logger.debug( + "KV PUSH mode: D node storing registration for request %s", + request.request_id, + ) + local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups() + local_block_ids = self.get_sw_clipped_blocks(local_block_ids) + + # ``remote_*`` fields are P's coordinates (from D's perspective). + # ``decode_*`` fields are D's own info that P needs for the + # reverse handshake before WRITE-ing. + self._push_pending_registrations[request.request_id] = { + "request_id": request.request_id, + "decode_engine_id": self.engine_id, + "decode_host": self.side_channel_host, + "decode_port": self.side_channel_port, + "decode_tp_size": (self.vllm_config.parallel_config.tensor_parallel_size), + "local_block_ids": local_block_ids, + "remote_engine_id": params["remote_engine_id"], + "remote_host": params["remote_host"], + "remote_port": params["remote_port"], + "remote_tp_size": params["tp_size"], + } + self._push_registration_deadlines[request.request_id] = ( + time.perf_counter() + self._push_registration_timeout + ) + # In push mode D doesn't know P's blocks; P determines them + # from the registration. We still track the request as + # needing recv so the engine waits for P's WRITE completion. + # ``remote_block_ids`` is also seeded to an empty tuple so the + # base scheduler's ``add_new_req_to_recv`` can build the + # ReqMeta without a KeyError — the actual remote block IDs are + # learned by P over the NIXL handshake at WRITE time. + params["remote_block_ids"] = () + self._reqs_need_recv[request.request_id] = (request, local_block_ids) + + # Mark as processed so a re-entry (e.g. preemption + reschedule) + # doesn't re-stage the registration. + params["do_remote_prefill"] = False + + def request_finished( + self, + request: Request, + block_ids: BlockIds, + ) -> tuple[bool, dict[str, Any] | None]: + """Push-mode request_finished: stores blocks for workers.""" + from vllm.v1.request import RequestStatus + + params = request.kv_transfer_params + logger.debug( + "NixlPushConnector request_finished(%s), request_status=%s, " + "kv_transfer_params=%s", + request.request_id, + request.status, + params, + ) + if not params: + return False, None + + is_p_node = bool(params.get("do_remote_decode")) + + self._stop_heartbeat(request.request_id) + # Drop any pending registration deadline; the request either + # completed or was cancelled. + self._push_registration_deadlines.pop(request.request_id, None) + + if params.get("do_remote_prefill"): + # ``do_remote_prefill`` is still set, which means + # ``update_state_after_alloc`` never ran (it would have + # flipped this flag to False). The request was aborted + # before it could be scheduled — e.g. rejected at the D + # serving layer via abort_immediately. To keep P from + # stranding the prefill blocks, we still register an empty + # recv so the worker emits a notif that lets P free them. + self._reqs_need_recv[request.request_id] = (request, []) + params["do_remote_prefill"] = False + return False, None + + # Push connector only acts on the P-side terminal path; D-side + # finishing without a remote prefill is a no-op. + if not is_p_node: + return False, None + + if request.status not in ( + RequestStatus.FINISHED_LENGTH_CAPPED, + RequestStatus.FINISHED_STOPPED, + ): + self._reqs_not_processed.add(request.request_id) + self._reqs_need_save.pop(request.request_id, None) + return False, None + + delay_free_blocks = any(len(group) > 0 for group in block_ids) + remote_num_tokens = 0 + if delay_free_blocks: + logger.debug( + "NixlPushConnector request_finished(%s) waiting for %d seconds " + "before releasing blocks", + request.request_id, + self._kv_lease_duration, + ) + self._reqs_need_send[request.request_id] = ( + time.perf_counter() + self._kv_lease_duration + ) + + block_ids = self.get_sw_clipped_blocks(block_ids) + remote_num_tokens = request.num_computed_tokens + + # Store finished blocks for worker-level matching with D + # registrations (via NIXL notifications). + self._finished_request_blocks[request.request_id] = block_ids + self._newly_finished_push_blocks[request.request_id] = block_ids + + return delay_free_blocks, dict( + do_remote_prefill=True, + do_remote_decode=False, + remote_block_ids=block_ids, + remote_engine_id=self.engine_id, + remote_request_id=request.request_id, + remote_host=self.side_channel_host, + remote_port=self.side_channel_port, + tp_size=self.vllm_config.parallel_config.tensor_parallel_size, + remote_num_tokens=remote_num_tokens, + ) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> KVConnectorMetadata: + meta = super().build_connector_meta(scheduler_output) + assert isinstance(meta, NixlConnectorMetadata) + + # Watchdog: any D-side registration whose deadline has passed without + # a corresponding push completion is treated as failed and cleaned up. + # The corresponding request is already tracked via _reqs_need_recv; + # the engine layer will eventually time it out via the lease, but we + # at least drop the stale registration so we don't keep retrying. + now = time.perf_counter() + # Deadlines are inserted in non-decreasing order (monotonic clock + + # constant timeout, armed once per request), and dict insertion order + # is preserved across key deletions, so we can stop at the first + # not-yet-expired entry instead of scanning the whole dict. + expired = [] + for rid, deadline in self._push_registration_deadlines.items(): + if deadline > now: + break + expired.append(rid) + for rid in expired: + self._push_registration_deadlines.pop(rid, None) + # Avoid resending a registration that already timed out. + self._push_pending_registrations.pop(rid, None) + logger.warning( + "NixlPushConnector: registration for request %s timed out " + "after %.1fs without a push completion", + rid, + self._push_registration_timeout, + ) + + # D side: package pending registrations for D workers to send out. + if self._push_pending_registrations: + meta.push_registrations = dict(self._push_pending_registrations) + self._push_pending_registrations.clear() + + # P side: package newly finished blocks for P workers to match against + # any D registrations they have received via NIXL notifications. + if self._newly_finished_push_blocks: + meta.push_finished_blocks = dict(self._newly_finished_push_blocks) + self._newly_finished_push_blocks.clear() + + return meta + + def has_pending_push_work(self) -> bool: + # Keep the engine main loop alive while we have: + # - finished P blocks awaiting WRITE completion, or + # - pending D registrations the worker has not yet shipped, or + # - newly finished blocks not yet shipped to P workers. + return bool(self._finished_request_blocks or self._push_pending_registrations) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + """Clean up finished request blocks after push completes.""" + super().update_connector_output(connector_output) + for req_id in connector_output.finished_sending or (): + self._finished_request_blocks.pop(req_id, None) + # On D side, finished_recving means the push completed; clear the + # watchdog so we don't trip an expiration on a fulfilled request. + for req_id in connector_output.finished_recving or (): + self._push_registration_deadlines.pop(req_id, None) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py new file mode 100644 index 00000000000..a15fc204d26 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -0,0 +1,742 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Push-specific (WRITE) worker-side logic for the NIXL connector. + +A dedicated ``nixl-push-writer`` thread owns all push-related NIXL ops: +calls ``get_new_notifs`` (routing PUSH_REG internally; HB / completion +notifs are forwarded to the engine main thread), sends PUSH_REG via +``send_notif``, matches D registrations with P finished blocks, and +issues WRITE transfers via ``make_prepped_xfer`` / ``transfer``. + +The engine main thread feeds the writer through three queues: +``_reg_send_inbox`` (D-side regs to send), ``_finished_blocks_inbox`` +(P-side blocks from metadata) and ``_pending_completion_notifs`` +(non-PUSH_REG notifs forwarded back for HB / completion accounting). + +Wake model: the writer self-polls every +``_PUSH_WRITER_POLL_INTERVAL_MS`` only while it has unmatched +``_push_finished_blocks`` (i.e. P-side blocks waiting for a D PUSH_REG +notif that has no other wake source). All other progress is +event-driven: the engine main thread sets ``_push_writer_wake`` from +``start_load_kv`` (when handing it new work) and from ``get_finished`` +(so each engine step gives the writer a chance to drain NIXL notifs); +the handshake-completion callback sets the same event after a deferred +PUSH_REG send has been queued. When a request's lease expires (the base +worker reports it via ``done_sending``) or the WRITE completes, +``get_finished`` enqueues an eviction onto ``_evict_finished_inbox`` so +the writer drops any leftover ``_push_finished_blocks`` / +``_pending_d_registrations`` and stops self-polling. +""" + +import queue +import threading +import time +from collections import defaultdict +from concurrent.futures import Future +from typing import TYPE_CHECKING, Any + +import msgspec +import numpy as np + +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + PUSH_REG_NOTIF_PREFIX, + NixlConnectorMetadata, + RemoteMeta, + ReqId, + ReqMeta, + TransferHandle, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id +from vllm.logger import init_logger + +if TYPE_CHECKING: + import torch + + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import KVCacheConfig + +logger = init_logger(__name__) + +# Writer-thread poll cadence while there is in-flight push state. When +# fully idle, the writer blocks on a wake event signalled by the engine +# main thread (start_load_kv / get_finished). Smaller -> lower latency +# while active, slightly more CPU. +_PUSH_WRITER_POLL_INTERVAL_MS = 1.0 + + +class NixlPushConnectorWorker(NixlBaseConnectorWorker): + """Push-specific (WRITE) worker logic. See module docstring.""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, engine_id, kv_cache_config) + + # Push-specific state. + # P-side: outgoing WRITE handles awaiting completion, keyed by + # request_id. Mutated by writer (submit) and main thread + # (``_pop_done_transfers``); guarded by + # ``_sending_transfers_lock``. + self._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list) + self._sending_transfers_lock = threading.Lock() + + # Writer-thread owned matching state. + # P-side: finished request blocks received from scheduler metadata + # that have not yet been matched with an incoming D registration. + self._push_finished_blocks: dict[ReqId, BlockIds] = {} + # P-side: D registrations received via NIXL notification that have + # not yet been matched with a finished P request. + self._pending_d_registrations: dict[ReqId, dict[str, Any]] = {} + + # Cross-thread channels. + self._reg_send_inbox: queue.Queue[tuple[str, dict[str, Any]]] = queue.Queue() + self._finished_blocks_inbox: queue.Queue[tuple[str, BlockIds]] = queue.Queue() + self._pending_completion_notifs: queue.Queue[bytes] = queue.Queue() + # Main thread → writer: req_ids whose lease has expired or whose + # WRITE has completed. Writer drops them from + # ``_push_finished_blocks`` so an unmatched entry doesn't keep the + # writer busy-polling forever. + self._evict_finished_inbox: queue.Queue[str] = queue.Queue() + + # Wake signal from engine main thread (start_load_kv / get_finished). + # Writer self-polls at _PUSH_WRITER_POLL_INTERVAL_MS while it has + # active in-flight state; otherwise it blocks until signalled. + self._push_writer_wake = threading.Event() + + self._push_writer_stop = threading.Event() + self._push_writer_thread: threading.Thread | None = None + + # --- Lifecycle ----------------------------------------------------- # + + def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]): + super().register_kv_caches(kv_caches) + if self._push_writer_thread is None: + self._push_writer_thread = threading.Thread( + target=self._push_writer_loop, + daemon=True, + name="nixl-push-writer", + ) + self._push_writer_thread.start() + logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank) + + def shutdown(self): + self._push_writer_stop.set() + # Unblock the writer if it's waiting in the no-active-state branch. + self._push_writer_wake.set() + if self._push_writer_thread is not None: + self._push_writer_thread.join(timeout=2) + self._push_writer_thread = None + with self._sending_transfers_lock: + for handles in self._sending_transfers.values(): + for handle in handles: + self.nixl_wrapper.release_xfer_handle(handle) + self._sending_transfers.clear() + super().shutdown() + + # --- Engine-main-thread entry point -------------------------------- # + + def start_load_kv(self, metadata: NixlConnectorMetadata): + """Pre-process metadata; defer NIXL ops to the writer thread.""" + # D-side: track reqs waiting for P to push. + for req_id, meta in metadata.reqs_to_recv.items(): + meta.local_physical_block_ids = self._logical_to_kernel_block_ids( + meta.local_block_ids + ) + assert meta.remote is not None + remote_engine_id = meta.remote.engine_id + logger.debug( + "start_load_kv (push) for request %s from remote engine %s. " + "Num local_block_ids: %s. Num remote_block_ids: %s. ", + req_id, + remote_engine_id, + len(meta.local_physical_block_ids), + len(meta.remote.block_ids), + ) + self._recving_metadata[req_id] = meta + + # --- D-side: registrations to send to P via NIXL --- + if metadata.push_registrations: + for req_id, reg_data in metadata.push_registrations.items(): + self._reg_send_inbox.put((req_id, reg_data)) + self._push_writer_wake.set() + + # --- P-side: newly finished blocks awaiting a D registration match --- + if metadata.push_finished_blocks: + for req_id, block_ids in metadata.push_finished_blocks.items(): + self._finished_blocks_inbox.put((req_id, block_ids)) + self._push_writer_wake.set() + + # Batch + lease tracking (same as pull). + for req_id in metadata.reqs_in_batch: + self._reqs_to_process.add(req_id) + for req_id in metadata.reqs_not_processed: + self._reqs_to_process.discard(req_id) + assert req_id not in self._reqs_to_send + for req_id, expiration_time in metadata.reqs_to_send.items(): + if req_id in self._reqs_to_process: + self._reqs_to_send[req_id] = expiration_time + + # Heartbeats still leave from the main thread (base worker behaviour). + self._send_heartbeats(metadata) + + # --- Writer thread ------------------------------------------------- # + + def _push_writer_loop(self) -> None: + sleep_s = _PUSH_WRITER_POLL_INTERVAL_MS / 1000.0 + + while not self._push_writer_stop.is_set(): + try: + # 1. D registrations to send. + while True: + try: + rid, rd = self._reg_send_inbox.get_nowait() + except queue.Empty: + break + self._send_registration_to_p(rid, rd) + + # 2. P-side finished blocks; match against pending regs. + while True: + try: + rid, blocks = self._finished_blocks_inbox.get_nowait() + except queue.Empty: + break + matched = self._pop_matching_registration(rid) + if matched is not None: + self._do_start_push_kv(rid, blocks, matched) + else: + self._push_finished_blocks[rid] = blocks + + # 2b. Evict finished blocks for requests that have either + # completed (WRITE acknowledged) or whose lease expired + # without a D registration. Drop pending registrations + # for the same reason so we don't leak state. + while True: + try: + rid = self._evict_finished_inbox.get_nowait() + except queue.Empty: + break + self._push_finished_blocks.pop(rid, None) + self._pending_d_registrations.pop(rid, None) + + # 3. NIXL notifs: route PUSH_REG; forward the rest. + for notifs in self.nixl_wrapper.get_new_notifs().values(): + for notif in notifs: + if notif.startswith(PUSH_REG_NOTIF_PREFIX): + self._handle_push_reg_notif(notif) + else: + self._pending_completion_notifs.put(notif) + except Exception: + logger.exception("nixl-push-writer error; continuing") + + # Self-poll only while there is no other wake source: P-side + # finished blocks waiting for a D PUSH_REG match. All other + # progress is event-driven (see module docstring). + if self._push_finished_blocks: + self._push_writer_stop.wait(timeout=sleep_s) + else: + self._push_writer_wake.wait() + self._push_writer_wake.clear() + + def _handle_push_reg_notif(self, notif: bytes) -> None: + try: + reg_data = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :]) + except Exception: + logger.exception("Failed to decode PUSH_REG notification payload") + return + rid = reg_data.get("request_id") if isinstance(reg_data, dict) else None + if not isinstance(rid, str): + logger.warning("PUSH_REG notif missing request_id; dropping") + return + + match = self._pop_matching_finished_blocks(rid) + if match is not None: + fin_id, blocks = match + self._do_start_push_kv(fin_id, blocks, reg_data) + else: + self._pending_d_registrations[rid] = reg_data + + # --- D-side registration send (writer thread) ---------------------- # + + def _send_registration_to_p( + self, + req_id: str, + reg_data: dict[str, Any], + ) -> None: + """Handshake (if needed) then send PUSH_REG. ``send_notif`` always + executes on the writer; the handshake runs on the background executor + and the request is re-queued onto ``_reg_send_inbox`` once it + completes (at which point ``_ensure_handshake`` returns ``None`` and we + send directly).""" + fut = self._ensure_handshake( + reg_data["remote_engine_id"], + reg_data["remote_host"], + reg_data["remote_port"], + reg_data["remote_tp_size"], + ) + if fut is None: + self._do_send_reg_notif(req_id, reg_data) + return + + def _on_handshake( + f: Future[dict[int, str]], + rid: str = req_id, + rd: dict[str, Any] = reg_data, + ) -> None: + try: + f.result() + except Exception as e: + self._log_failure( + failure_type="push_reg_handshake_failed", req_id=rid, error=e + ) + self._handle_failed_transfer(rid, None) + return + # Re-queue for the writer to send now that the handshake is done. + self._reg_send_inbox.put((rid, rd)) + # Wake the writer so it sends the PUSH_REG promptly even if + # otherwise parked. + self._push_writer_wake.set() + + fut.add_done_callback(_on_handshake) + + def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: + engine_id = reg_data["remote_engine_id"] + notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data) + agents = self._remote_agents.get(engine_id) + if not agents: + logger.error( + "No remote agents for engine %s; cannot send registration for %s", + engine_id, + req_id, + ) + self._handle_failed_transfer(req_id, None) + return + for rank, agent_name in agents.items(): + try: + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_msg) + except Exception as e: + self._log_failure( + failure_type="push_reg_notif_failed", + req_id=req_id, + error=e, + remote_rank=rank, + ) + logger.debug( + "Sent PUSH_REG for %s to engine %s (%dB)", req_id, engine_id, len(notif_msg) + ) + + # --- Matching helpers --------------------------------------------- # + + def _pop_matching_registration(self, request_id: str) -> dict[str, Any] | None: + """Pop the D-side registration matching *request_id*. + + Exact key first, then a match after stripping the random suffix from + both sides. No match leaves the request unmatched (push not started). + """ + data = self._pending_d_registrations.pop(request_id, None) + if data is not None: + return data + base_id = get_base_request_id(request_id) + for reg_id in list(self._pending_d_registrations): + if get_base_request_id(reg_id) == base_id: + return self._pending_d_registrations.pop(reg_id) + return None + + def _pop_matching_finished_blocks( + self, request_id: str + ) -> tuple[str, BlockIds] | None: + """Pop the P-side finished blocks matching *request_id*. + + Same lookup as ``_pop_matching_registration``: exact key, then a + match after stripping the random suffix from both sides. + """ + blocks = self._push_finished_blocks.pop(request_id, None) + if blocks is not None: + return request_id, blocks + base_id = get_base_request_id(request_id) + for fin_id in list(self._push_finished_blocks): + if get_base_request_id(fin_id) == base_id: + return fin_id, self._push_finished_blocks.pop(fin_id) + return None + + # --- WRITE transfer logic (writer thread) ------------------------- # + + def _do_start_push_kv( + self, + request_id: str, + local_block_ids: BlockIds, + registration_data: dict[str, Any], + ) -> None: + """Start push-based KV transfer from P worker to D node. + + ``local_block_ids`` are P's *logical* block IDs (from the P + scheduler's metadata). ``registration_data["local_block_ids"]`` + are D's *logical* block IDs (from D's scheduler, sent over the + PUSH_REG notif). All conversion to physical block IDs is + deferred to ``_xfer_blocks_for_req`` so each side uses its own + physical-blocks-per-logical ratio (P uses + ``self._physical_blocks_per_logical_kv_block``; D's ratio is + learned during the NIXL handshake).""" + decode_engine_id = registration_data["decode_engine_id"] + remote_block_ids = registration_data["local_block_ids"] + decode_host = registration_data["decode_host"] + decode_port = registration_data["decode_port"] + decode_request_id = registration_data["request_id"] + if not local_block_ids: + logger.warning("No local blocks to push for request %s", request_id) + return + + if not self._ensure_d_handshake( + decode_engine_id, + decode_host, + decode_port, + registration_data["decode_tp_size"], + request_id, + ): + return + + # Both sides are kept in logical form here; ``_xfer_blocks_for_req`` + # expands each side using the appropriate ratio. + logical_local = self._as_grouped_block_ids(local_block_ids) + logical_remote = self._as_grouped_block_ids(remote_block_ids) + physical_local = self._logical_to_kernel_block_ids(logical_local) + + push_meta = ReqMeta( + local_block_ids=logical_local, + local_physical_block_ids=physical_local, + tp_size=self.world_size, + remote=RemoteMeta( + block_ids=logical_remote, + host="", + port=0, + engine_id=decode_engine_id, + request_id=decode_request_id, + ), + ) + + t0 = time.perf_counter() + self._xfer_blocks_for_req(req_id=request_id, meta=push_meta) + elapsed_ms = (time.perf_counter() - t0) * 1000.0 + if elapsed_ms > 200.0: + logger.warning( + "_do_start_push_kv for %s took %.1fms (slow NIXL submission)", + request_id, + elapsed_ms, + ) + + def _ensure_d_handshake( + self, + decode_engine_id: str, + decode_host: str, + decode_port: int, + decode_tp_size: int, + request_id: str, + ) -> bool: + """First-time P→D handshake. Blocking call on the writer thread. + + Returns True iff the handshake succeeded (or had already been + completed). Returns False if the handshake raised; the request is + skipped in that case (the engine layer will reschedule or fail it + via the standard lease/timeout path).""" + if decode_engine_id in self._remote_agents: + return True + try: + remote_agents = self._nixl_handshake( + decode_host, + decode_port, + decode_tp_size, + decode_engine_id, + ) + except Exception: + logger.exception( + "Failed handshake to D %s for push %s", + decode_engine_id, + request_id, + ) + return False + with self._handshake_lock: + self._remote_agents[decode_engine_id] = remote_agents + logger.info( + "Push handshake to D %s done (%d agents)", + decode_engine_id, + len(remote_agents), + ) + return True + + @staticmethod + def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds: + """Normalise a sequence of block IDs to a tuple-of-groups shape. + + ``BlockIds`` is canonically a tuple of per-group lists, but some + registration payloads collapse a single-group case to a flat + list. Re-wrap that case so downstream group-aware helpers see a + consistent shape.""" + if block_ids and not isinstance(block_ids[0], (list, tuple)): + return (list(block_ids),) + return block_ids + + def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): + """Issue WRITE transfers to one or more remote TP ranks.""" + assert meta.remote is not None and self.transfer_topo is not None + engine_id = meta.remote.engine_id + plan = self.tp_mappings[engine_id] + remote_info = self.transfer_topo.get_engine_info(engine_id) + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) + + # Expand D's logical IDs using the ratio learned during the + # NIXL handshake. ``meta`` is freshly built by + # ``_do_start_push_kv`` so mutating it here is safe. + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids, + remote_info.remote_physical_blocks_per_logical, + ) + remote_block_ids = meta.remote.block_ids + local_block_ids = meta.local_physical_block_ids + num_groups = len(local_block_ids) + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=[ + list(local_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], + remote_block_ids=[ + list(remote_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], + ) + for rank in plan.all_source_ranks + ] + + if self.use_mla and tp_ratio < 0: + assert len(read_specs) == 1 + + for i, spec in enumerate(read_specs): + remote_block_size = remote_info.remote_block_size + logger.debug( + "Remote agent %s available, calling _xfer_blocks" + " on remote rank %s with remote block size %s for req %s", + meta.remote.engine_id, + spec.remote_rank, + remote_block_size, + req_id, + ) + if tp_ratio < 0 and not self.use_mla: + assert remote_block_size == self.block_size + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + else: + local_xfer_side_handle = self.src_xfer_handles_by_block_size[ + remote_block_size + ] + + remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][ + spec.remote_rank + ] + + self._xfer_blocks( + read_spec=spec, + request_id=req_id, + dst_engine_id=meta.remote.engine_id, + remote_request_id=meta.remote.request_id, + local_xfer_side_handle=local_xfer_side_handle, + remote_xfer_side_handle=remote_xfer_side_handle, + ) + + if self.use_mla and tp_ratio < 0 and read_specs: + notif_id = f"{meta.remote.request_id}:{self.world_size}".encode() + remote_agents = self._remote_agents[meta.remote.engine_id] + for rank_to_notify, agent in remote_agents.items(): + if rank_to_notify != read_specs[0].remote_rank: + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) + + def _xfer_blocks( + self, + read_spec: ReadSpec, + dst_engine_id: str, + request_id: str, + remote_request_id: str, + local_xfer_side_handle: int, + remote_xfer_side_handle: int, + ): + """Post a WRITE point-to-point xfer request.""" + assert self.transfer_topo is not None + remote_rank = read_spec.remote_rank + local_block_ids = read_spec.local_block_ids + remote_block_ids = read_spec.remote_block_ids + + remote_info = self.transfer_topo.get_engine_info(dst_engine_id) + block_size_ratio = self.transfer_topo.block_size_ratio( + remote_info.remote_block_size + ) + if block_size_ratio > 1: + assert not self._is_hma_required + local_block_ids0 = local_block_ids[0] if local_block_ids else [] + remote_block_ids0 = remote_block_ids[0] + local_block_ids_mapped = self.get_mapped_blocks( + np.asarray(local_block_ids0), block_size_ratio + ).tolist() + if len(local_block_ids_mapped) > len(remote_block_ids0): + local_block_ids_mapped = local_block_ids_mapped[ + : len(remote_block_ids0) + ] + local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] + remote_block_ids = [remote_block_ids0] + + notif_id = f"{remote_request_id}:{self.world_size}".encode() + + if len(local_block_ids) == 0: + logger.warning("No blocks to push for request %s", request_id) + return + + # Align per-group block counts for push. + local_block_ids = list(local_block_ids) + remote_block_ids = list(remote_block_ids) + for i in range(min(len(local_block_ids), len(remote_block_ids))): + num_local = len(local_block_ids[i]) + num_remote = len(remote_block_ids[i]) + if num_local > num_remote: + local_block_ids[i] = local_block_ids[i][:num_remote] + elif num_local < num_remote: + remote_block_ids[i] = remote_block_ids[i][:num_local] + + # Get descs ids. + remote_block_descs_ids = self._compute_desc_ids( + block_ids=remote_block_ids, + dst_num_blocks=self.dst_num_blocks[dst_engine_id], + block_size_ratio=None, + physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical, + ) + local_block_descs_ids = self._compute_desc_ids( + block_ids=local_block_ids, + dst_num_blocks=self.dst_num_blocks[self.engine_id], + block_size_ratio=block_size_ratio, + physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block, + ) + + assert len(local_block_descs_ids) == len(remote_block_descs_ids) + + handle = None + try: + handle = self.nixl_wrapper.make_prepped_xfer( + "WRITE", + local_xfer_side_handle, + local_block_descs_ids, + remote_xfer_side_handle, + remote_block_descs_ids, + notif_msg=notif_id, + ) + self.nixl_wrapper.transfer(handle) + # Track push WRITE handles so P can free blocks once done. + with self._sending_transfers_lock: + self._sending_transfers[request_id].append(handle) + except Exception as e: + self._log_failure( + failure_type="transfer_setup_failed", + req_id=request_id, + msg="Push WRITE submission failed; releasing handle", + error=e, + dst_engine_id=dst_engine_id, + remote_rank=remote_rank, + ) + # On the P side this WRITE failure is purely outbound; we + # don't have a ``_recving_metadata`` entry to invalidate, so + # we just release the handle and let the engine reschedule + # via the lease / watchdog. + if handle is not None: + self.nixl_wrapper.release_xfer_handle(handle) + self.xfer_stats.record_failed_transfer() + + # --- Notification handling on engine main thread ------------------ # + + def _get_new_notifs(self) -> set[str]: + """Drain HB / completion notifs forwarded by the writer thread. + + The writer owns ``nixl_wrapper.get_new_notifs`` for push; PUSH_REG + notifs are handled there. Everything else is forwarded here for + existing accounting. + """ + assert self.transfer_topo is not None + notified_req_ids: set[str] = set() + while True: + try: + notif = self._pending_completion_notifs.get_nowait() + except queue.Empty: + break + + msg = notif.decode("utf-8") + if msg.startswith("HB:"): + self._handle_heartbeat(msg[3:]) + continue + + req_id, tp_size = msg.rsplit(":", 1) + + # Not tracked as a P-side send/process for this notif. + if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process: + if req_id in self._recving_metadata: + # D-side: P signalled push completion. The transfer was + # driven entirely by P (we don't own a NIXL handle here), + # so materialise an empty entry in ``_recving_transfers`` + # and let ``_pop_done_transfers`` report it done on the + # next ``get_finished``. + self._recving_transfers.setdefault(req_id, []) + else: + # Not tracked on either side (lease may have expired + # before the notif arrived). Log and skip. + logger.error( + "Unrecognized request %s notif (may have expired).", + req_id, + ) + continue + + n_consumers = int(tp_size) + tp_ratio = self.transfer_topo.tp_ratio(n_consumers) + consumers_per_producer = -tp_ratio if n_consumers > self.world_size else 1 + self.consumer_notification_counts_by_req[req_id] += 1 + if ( + self.consumer_notification_counts_by_req[req_id] + == consumers_per_producer + ): + notified_req_ids.add(req_id) + del self.consumer_notification_counts_by_req[req_id] + self._reqs_to_process.remove(req_id) + self._reqs_to_send.pop(req_id, None) + return notified_req_ids + + def get_finished(self) -> tuple[set[str], set[str]]: + # Engine main thread asking for completions: also wake the writer + # so it gets a chance to drain NIXL notifs (heartbeats, completion + # notifs, late PUSH_REGs) even if it had been parked. + self._push_writer_wake.set() + + done_sending, done_recving = super().get_finished() + + # ``_pop_done_transfers`` mutates ``_sending_transfers``; the + # writer thread also appends to it, so guard the pop. + with self._sending_transfers_lock: + done_pushing = self._pop_done_transfers(self._sending_transfers) + for req_id in done_pushing: + self._reqs_to_send.pop(req_id, None) + self._reqs_to_process.discard(req_id) + self.consumer_notification_counts_by_req.pop(req_id, None) + done_sending.add(req_id) + + # Tell the writer to drop any state it still holds for any + # request that just finished (push completed) or expired + # (lease ran out without a D registration ever arriving). + for req_id in done_sending: + self._evict_finished_inbox.put(req_id) + if done_sending: + self._push_writer_wake.set() + + return done_sending, done_recving diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py index b2122ed0d30..3da8e28a749 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py @@ -1,674 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Scheduler-side logic for the NIXL connector.""" +"""Backward-compatible re-export of NixlPullConnectorScheduler.""" -import threading -import time -from typing import TYPE_CHECKING, Any - -import msgspec -import zmq - -from vllm import envs -from vllm.distributed.kv_transfer.kv_connector.utils import ( - BlockIds, - EngineId, - yield_req_data, -) -from vllm.distributed.kv_transfer.kv_connector.v1.base import ( - KVConnectorHandshakeMetadata, - KVConnectorMetadata, -) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( - GET_META_MSG, - HeartbeatInfo, - NixlConnectorMetadata, - NixlHandshakePayload, - ReqId, -) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.network_utils import make_zmq_path -from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_cache_interface import ( - FullAttentionSpec, - MambaSpec, - SlidingWindowSpec, +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, ) -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.core.kv_cache_manager import KVCacheBlocks - from vllm.v1.kv_cache_interface import KVCacheConfig - from vllm.v1.outputs import KVConnectorOutput - from vllm.v1.request import Request +# Backward compatibility: NixlConnectorScheduler is the pull-based scheduler. +NixlConnectorScheduler = NixlPullConnectorScheduler -logger = init_logger(__name__) - - -class NixlConnectorScheduler: - """Implementation of Scheduler side methods""" - - def __init__( - self, - vllm_config: "VllmConfig", - engine_id: str, - kv_cache_config: "KVCacheConfig", - ): - self.vllm_config = vllm_config - self.block_size = vllm_config.cache_config.block_size - self.engine_id: EngineId = engine_id - self.kv_cache_config = kv_cache_config - self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST - self.side_channel_port = ( - envs.VLLM_NIXL_SIDE_CHANNEL_PORT - + vllm_config.parallel_config.data_parallel_index - ) - assert vllm_config.kv_transfer_config is not None - self._kv_lease_duration: int = ( - vllm_config.kv_transfer_config.get_from_extra_config( - "kv_lease_duration", 30 - ) - ) - # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. - self._heartbeat_interval = self._kv_lease_duration // 6 - if current_platform.device_type == "cpu": - self.use_host_buffer = False - else: - self.use_host_buffer = ( - vllm_config.kv_transfer_config.kv_buffer_device == "cpu" - ) - self._is_hma_required = ( - not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager - # Also handle unlikely SW-only model case instead of checking num_groups>1. - and any( - not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups - ) - ) - self._has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) - for g in kv_cache_config.kv_cache_groups - ) - - logger.info("Initializing NIXL Scheduler %s", engine_id) - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - logger.info("Hybrid Memory Allocator is enabled with NIXL") - - # Background thread for handling new handshake requests. - self._nixl_handshake_listener_t: threading.Thread | None = None - self._stop_event = threading.Event() - - # Requests that need to start recv/send. - # New requests are added by update_state_after_alloc in - # the scheduler. Used to make metadata passed to Worker. - self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} - self._reqs_need_save: dict[ReqId, Request] = {} - # Reqs to send and their expiration time - self._reqs_need_send: dict[ReqId, float] = {} - self._reqs_in_batch: set[ReqId] = set() - # Reqs to remove from processed set because they're not to send after - # remote prefill or aborted. - self._reqs_not_processed: set[ReqId] = set() - - # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to - # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine - self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} - # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal - self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {} - self._last_heartbeat_time: float = 0.0 - - # Gather Sliding Window sizes for each kv cache group (if any) in number of - # blocks per KV cache group. This is used to clip the local attention window. - sw_sizes_tokens: list[tuple[int, int]] = [ - (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) - if isinstance(g.kv_cache_spec, SlidingWindowSpec) - else (0, self.block_size) - for g in kv_cache_config.kv_cache_groups - ] - # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively - # account for boundary overlap eg window isn't fully aligned with blocks. - self.blocks_per_sw = [ - cdiv(n_tokens, block_size) + 1 if n_tokens else 0 - for n_tokens, block_size in sw_sizes_tokens - ] - - # Threshold to decide whether to compute kv cache locally - # or pull from a remote node: minimum number of remote - # tokens to amortize the xfer latencies - self.kv_recompute_threshold: int = int( - vllm_config.kv_transfer_config.get_from_extra_config( - "kv_recompute_threshold", 64 - ) - ) - - # Bi-directional KV transfer feature supports KV block - # transfers from D node to P node - self.is_bidirectional_kv_xfer_enabled = ( - vllm_config.kv_transfer_config.get_from_extra_config( - "bidirectional_kv_xfer", False - ) - ) - self.decoder_kv_blocks_ttl = ( - vllm_config.kv_transfer_config.get_from_extra_config( - "decoder_kv_blocks_ttl", 480 - ) - ) - - if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0: - logger.info( - "Bidirectional KV transfer is enabled and the kv " - "recompute threshold is set to %d tokens." - "KV blocks on D are released after a TTL of %d seconds.", - self.kv_recompute_threshold, - self.decoder_kv_blocks_ttl, - ) - - def shutdown(self): - self._stop_event.set() - if self._nixl_handshake_listener_t is not None: - self._nixl_handshake_listener_t.join() - self._nixl_handshake_listener_t = None - - def on_new_request(self, request: "Request") -> None: - """Track a request that may need heartbeats.""" - params = request.kv_transfer_params - # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are - # effectively disabled for Bidirectional KV transfer. - if params is None or not params.get("do_remote_prefill"): - return - # Only track if all required remote fields are present. - remote_engine_id = params.get("remote_engine_id") - remote_request_id = params.get("remote_request_id") - host = params.get("remote_host") - port = params.get("remote_port") - tp_size = params.get("tp_size") - if ( - remote_engine_id is None - or remote_request_id is None - or host is None - or port is None - or tp_size is None - ): - return - if remote_engine_id not in self._heartbeat_by_engine: - self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo( - req_ids=set(), - host=host, - port=port, - tp_size=tp_size, - ) - self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id) - self._heartbeat_req_engine[request.request_id] = ( - remote_engine_id, - remote_request_id, - ) - - def _stop_heartbeat(self, req_id: ReqId) -> None: - """Remove *req_id* from heartbeat tracking (if tracked).""" - if key := self._heartbeat_req_engine.pop(req_id, None): - engine_id, remote_id = key - if info := self._heartbeat_by_engine.get(engine_id): - info.req_ids.discard(remote_id) - if not info.req_ids: - # Clean up empty engines so we don't leak a key when remote dies. - del self._heartbeat_by_engine[engine_id] - - def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: - """ - Clip the number of blocks to the sliding window size for each kv cache group - that employs SWA. - This is necessary because the KV Cache manager initially allocates blocks for - the entire sequence length, and successively cleans up blocks that are outside - the window prior to the `request_finished_all_groups` hook. - """ - if len(block_ids) == 0 or not self._is_hma_required: - # No blocks to clip eg Full prefix cache hit or not a hybrid model. - return block_ids - # NOTE (NickLucche) This logic is currently handled at the connector level - # because offloading connectors might want to receive the whole sequence even - # for SWA groups. We will abstract this logic once the interface is more stable - assert len(block_ids) == len(self.blocks_per_sw), ( - "Number of KV cache groups must match" - ) - # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged - return tuple( - [ - blocks[-self.blocks_per_sw[i] :] - if self.blocks_per_sw[i] > 0 - else blocks - for i, blocks in enumerate(block_ids) - ] - ) - - def set_xfer_handshake_metadata( - self, metadata: dict[int, KVConnectorHandshakeMetadata] - ) -> None: - """ - Set the KV connector handshake metadata for this connector. - - Args: - metadata (dict): the handshake metadata to set. - """ - encoded_data: dict[int, bytes] = {} - encoder = msgspec.msgpack.Encoder() - for tp_rank, rank_metadata in metadata.items(): - if not isinstance(rank_metadata, NixlHandshakePayload): - raise ValueError( - "NixlConnectorScheduler expects NixlHandshakePayload for " - "handshake metadata." - ) - encoded_data[tp_rank] = encoder.encode(rank_metadata) - logger.debug( - "Tp rank %d: encoded NixlHandshakePayload size: %s bytes", - tp_rank, - str(len(encoded_data[tp_rank])), - ) - - # Only start the listener when we have metadata to serve. - if self._nixl_handshake_listener_t is None: - ready_event = threading.Event() - self._nixl_handshake_listener_t = threading.Thread( - target=self._nixl_handshake_listener, - args=( - encoded_data, - ready_event, - self._stop_event, - self.side_channel_host, - self.side_channel_port, - ), - daemon=True, - name="nixl_handshake_listener", - ) - self._nixl_handshake_listener_t.start() - ready_event.wait() # Wait for listener ZMQ socket to be ready. - - @staticmethod - def _nixl_handshake_listener( - encoded_data: dict[int, Any], - ready_event: threading.Event, - stop_event: threading.Event, - host: str, - port: int, - ): - """Background thread for getting new NIXL handshakes.""" - # NOTE(rob): this is a simple implementation. We will move - # to a better approach via HTTP endpoint soon. - - # Listen for new requests for metadata. - path = make_zmq_path("tcp", host, port) - logger.debug("Starting listening on path: %s", path) - with zmq_ctx(zmq.ROUTER, path) as sock: - sock.setsockopt(zmq.RCVTIMEO, 1000) - ready_event.set() - while True: - try: - identity, _, msg = sock.recv_multipart() - except zmq.Again: - if stop_event.is_set(): - break - continue - # Decode the message which contains (GET_META_MSG, rank) - msg, target_tp_rank = msgspec.msgpack.decode(msg) - logger.debug( - "Received message for tp rank %s", - target_tp_rank, - ) - if msg != GET_META_MSG: - logger.warning("Connection listener got unexpected message %s", msg) - sock.send_multipart((identity, b"", encoded_data[target_tp_rank])) - - def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: - return num_prompt_tokens - 1 - return num_prompt_tokens - - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. - - Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled.""" - params = request.kv_transfer_params - if ( - params is not None - # Guard against repeated truncation after preemption/reschedule. - and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 - ): - if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() - elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] - else: - return - - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 - request.max_tokens = 1 - params["_p_side_truncated"] = True - - def get_num_new_matched_tokens( - self, request: "Request", num_computed_tokens: int - ) -> tuple[int, bool]: - """ - For remote prefill, pull all prompt blocks from remote - asynchronously relative to engine execution. - - Args: - request (Request): the request object. - num_computed_tokens (int): the number of locally - computed tokens for this request - Returns: - * the number of tokens that can be loaded from the - external KV cache beyond what is already computed. - * true if the external KV cache tokens will be loaded - asynchronously (between scheduler steps). - """ - - params = request.kv_transfer_params - logger.debug( - "NIXLConnector get_num_new_matched_tokens: " - "num_computed_tokens=%s, kv_transfer_params=%s", - num_computed_tokens, - params, - ) - - if params is not None and params.get("do_remote_prefill"): - # Remote prefill: get all prompt blocks from remote. - token_ids = request.prompt_token_ids or [] - actual = self._mamba_prefill_token_count(len(token_ids)) - count = actual - num_computed_tokens - if count > 0: - return count, True - - if params is not None and params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) - - if ( - params is not None - and params.get("do_remote_decode") - and params.get("remote_block_ids") - and all( - p in params - for p in ( - "remote_engine_id", - "remote_request_id", - "remote_host", - "remote_port", - ) - ) - ): - # Decode node has kv blocks for part of prefill request, so, provide them - # as an external token count to scheduler. - # The tokens will be loaded if not already present - # in the prefill node local cache - remote_num_tokens = params.get("remote_num_tokens") or 0 - count = ( - min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens - ) - if count > 0: - # Check kv_recompute_threshold: skip pull if - # remote tokens are below the threshold. - if ( - self.kv_recompute_threshold > 0 - and count < self.kv_recompute_threshold - ): - logger.debug( - "Skipping remote pull for %s: %d remote tokens < threshold %d", - request.request_id, - count, - self.kv_recompute_threshold, - ) - return 0, False - return count, True - - # No remote prefill for this request. - return 0, False - - def update_state_after_alloc( - self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int - ): - params = request.kv_transfer_params - logger.debug( - "NIXLConnector update_state_after_alloc: " - "num_external_tokens=%s, kv_transfer_params=%s", - num_external_tokens, - params, - ) - - if not params: - return - - if params.get("do_remote_decode") or ( - params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled - ): - self._reqs_in_batch.add(request.request_id) - if self.use_host_buffer and params.get("do_remote_decode"): - # NOTE: when accelerator is not directly supported by Nixl, - # prefilled blocks need to be saved to host memory before transfer. - self._reqs_need_save[request.request_id] = request - elif params.get("do_remote_prefill") or ( - params.get("do_remote_decode") - and self.is_bidirectional_kv_xfer_enabled - and not params.get("_remote_blocks_processed") - ): - if params.get("remote_block_ids"): - if all( - p in params - for p in ( - "remote_engine_id", - "remote_request_id", - "remote_host", - "remote_port", - ) - ): - # If remote_blocks and num_external_tokens = 0, we have - # a full prefix cache hit on the local node. We need to call - # send_notif in _read_blocks to free the memory on the remote node. - - unhashed_local_block_ids: BlockIds = ( - blocks.get_unhashed_block_ids_all_groups() - if num_external_tokens > 0 - else () - ) - local_block_ids = self.get_sw_clipped_blocks( - unhashed_local_block_ids - ) - - # Get unhashed blocks to pull from remote. Mind that a full prefix - # cache hit is indicated with an empty list. - self._reqs_need_recv[request.request_id] = ( - request, - local_block_ids, - ) - - else: - logger.warning( - "Got invalid KVTransferParams: %s. This " - "request will not utilize KVTransfer", - params, - ) - else: - assert num_external_tokens == 0 - # Only trigger 1 KV transfer per request. - params["do_remote_prefill"] = False - params["_remote_blocks_processed"] = True - - def _build_save_meta( - self, - meta: NixlConnectorMetadata, - scheduler_output: SchedulerOutput, - ) -> None: - # only called when use_host_buffer is True to build the save metadata - - # NOTE: For the prefill side, there might be a chance that an early added - # request is a chunked prefill, so we need to check if new blocks are added - for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output): - req_to_save = self._reqs_need_save.get(req_id) - if req_to_save is None or new_block_id_groups is None: - continue - req = req_to_save - - assert req.kv_transfer_params is not None - clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups) - meta.add_new_req_to_save( - request_id=req_id, - local_block_ids=clipped_block_id_groups, - kv_transfer_params=req.kv_transfer_params, - ) - assert scheduler_output.num_scheduled_tokens is not None - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] - is_partial = ( - req.num_computed_tokens + num_scheduled_tokens - ) < req.num_prompt_tokens - if not is_partial: - # For non-partial prefills, once new req_meta is scheduled, it - # can be removed from _reqs_need_save. - # For partial prefill case, we will retain the request in - # _reqs_need_save until all blocks are scheduled with req_meta. - # Therefore, only pop if `not is_partial`. - self._reqs_need_save.pop(req_id) - - def build_connector_meta( - self, - scheduler_output: SchedulerOutput, - ) -> KVConnectorMetadata: - meta = NixlConnectorMetadata() - - # Loop through scheduled reqs and convert to ReqMeta. - for req_id, (req, block_ids) in self._reqs_need_recv.items(): - assert req.kv_transfer_params is not None - meta.add_new_req_to_recv( - request_id=req_id, - local_block_ids=block_ids, - kv_transfer_params=req.kv_transfer_params, - ) - - if self.use_host_buffer: - self._build_save_meta(meta, scheduler_output) - - meta.reqs_to_send = self._reqs_need_send - meta.reqs_in_batch = self._reqs_in_batch - meta.reqs_not_processed = self._reqs_not_processed - - # Package heartbeats, throttled by heartbeat_interval. - if self._heartbeat_by_engine: - now = time.perf_counter() - if now - self._last_heartbeat_time >= self._heartbeat_interval: - self._last_heartbeat_time = now - meta.heartbeat_by_engine = self._heartbeat_by_engine - - # Clear the list once workers start the transfers - self._reqs_need_recv.clear() - self._reqs_in_batch = set() - self._reqs_not_processed = set() - self._reqs_need_send = {} - - return meta - - def update_connector_output(self, connector_output: "KVConnectorOutput") -> None: - """Stop heartbeating for requests whose KV transfer completed.""" - for req_id in connector_output.finished_recving or (): - self._stop_heartbeat(req_id) - - def request_finished( - self, - request: "Request", - block_ids: BlockIds, - ) -> tuple[bool, dict[str, Any] | None]: - """ - Once a request is finished, determine whether request blocks - should be freed now or will be sent asynchronously and freed later. - """ - from vllm.v1.request import RequestStatus - - params = request.kv_transfer_params - logger.debug( - "NIXLConnector request_finished(%s), request_status=%s, " - "kv_transfer_params=%s", - request.request_id, - request.status, - params, - ) - if not params: - return False, None - - is_p_node = bool(params.get("do_remote_decode")) - is_d_node = not is_p_node - - # Stop heartbeating for aborted requests that never reached finished_recving: - # normal path cleans up in update_connector_output. - self._stop_heartbeat(request.request_id) - - if params.get("do_remote_prefill"): - # If do_remote_prefill is still True when the request is finished, - # update_state_after_alloc must not have been called (the request - # must have been aborted before it was scheduled, e.g. via the - # abort_immediately path used to clean up KV-transfer requests - # rejected at the D-side serving layer). - # To avoid stranding the prefill blocks in the prefill instance, - # we must add empty block_ids to _reqs_need_recv so that our - # worker side will notify and free blocks in the prefill instance. - self._reqs_need_recv[request.request_id] = (request, []) - params["do_remote_prefill"] = False - return False, None - - if is_d_node and not self.is_bidirectional_kv_xfer_enabled: - return False, None - - if request.status not in ( - RequestStatus.FINISHED_LENGTH_CAPPED, - RequestStatus.FINISHED_STOPPED, - ): - # Also include the case of a P/D Prefill request with immediate - # block free (eg abort). Stop tracking this request. - self._reqs_not_processed.add(request.request_id) - # Clear _reqs_need_save if a request is aborted as partial prefill. - self._reqs_need_save.pop(request.request_id, None) - return False, None - - # TODO: check whether block_ids actually ever be 0. If not we could - # remove the conditional below - delay_free_blocks = any(len(group) > 0 for group in block_ids) - remote_num_tokens = 0 - if delay_free_blocks: - # Prefill request on remote. It will be read from D upon completion - request_kv_blocks_ttl = self._kv_lease_duration - if is_d_node: - # For blocks pinned on D, use a simpler timeout for now instead of a - # lease mechanism as turn2 request is client-driven. - request_kv_blocks_ttl = self.decoder_kv_blocks_ttl - logger.debug( - "NIXLConnector request_finished(%s) waiting for %d seconds " - "before releasing blocks", - request.request_id, - request_kv_blocks_ttl, - ) - self._reqs_need_send[request.request_id] = ( - time.perf_counter() + request_kv_blocks_ttl - ) - # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), - # trimming down after allocating for the whole sequence length. Empty - # blocks are always at the start of the list. - # Here we "unpad" blocks to send the actual remote blocks to be read. - block_ids = self.get_sw_clipped_blocks(block_ids) - - remote_num_tokens = request.num_computed_tokens - - return delay_free_blocks, dict( - do_remote_prefill=is_p_node, - do_remote_decode=is_d_node, - remote_block_ids=block_ids, - remote_engine_id=self.engine_id, - remote_request_id=request.request_id, - remote_host=self.side_channel_host, - remote_port=self.side_channel_port, - tp_size=self.vllm_config.parallel_config.tensor_parallel_size, - remote_num_tokens=remote_num_tokens, - ) +__all__ = ["NixlConnectorScheduler", "NixlPullConnectorScheduler"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py index 2fa3829eaec..b8606167348 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py @@ -6,6 +6,7 @@ import contextlib from collections.abc import Iterator from typing import Any +import regex as re import zmq from vllm.platforms import current_platform @@ -55,3 +56,13 @@ def get_representative_spec_type(spec: KVCacheSpec) -> type[KVCacheSpec]: inner = next(iter(spec.kv_cache_specs.values())) return type(inner) return type(spec) + + +# Trailing 8-hex randomization suffix appended by +# ``input_processor.assign_request_id`` as ``-{random_uuid():.8}``. +_RANDOM_SUFFIX_RE = re.compile(r"-[0-9a-f]{8}$", re.IGNORECASE) + + +def get_base_request_id(request_id: str) -> str: + """Strip the per-request ``-<8 hex>`` randomization suffix, if present.""" + return _RANDOM_SUFFIX_RE.sub("", request_id) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index e4b20c01f4d..66ad155bdae 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -1,2566 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Worker-side logic for the NIXL connector.""" +"""Backward-compatible re-export of NixlPullConnectorWorker.""" -import logging -import os -import queue -import threading -import time -import uuid -from collections import defaultdict -from collections.abc import Iterator -from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING, Any, cast - -import msgspec -import numpy as np -import torch -import zmq - -from vllm.distributed.kv_transfer.kv_connector.utils import ( - BlockIds, - EngineId, - EngineTransferInfo, - TransferTopology, - get_current_attn_backends, - kv_postprocess_blksize_and_layout_on_receive, - kv_postprocess_blksize_on_receive, - kv_postprocess_layout_on_receive, +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import ( + NixlPullConnectorWorker, ) -from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( - GET_META_MSG, - NixlAgentMetadata, - NixlConnectorMetadata, - NixlHandshakePayload, - ReqId, - ReqMeta, - TransferHandle, - compute_nixl_compatibility_hash, -) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( - NixlKVConnectorStats, -) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( - ReadSpec, - TPMapping, - _is_attention_spec, - _is_ssm_spec, - compute_tp_mapping, -) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( - _NIXL_SUPPORTED_DEVICE, - get_representative_spec_type, - zmq_ctx, -) -from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( - MambaConvSplitInfo, - derive_mamba_conv_split, -) -from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config -from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.logger import init_logger -from vllm.platforms import current_platform -from vllm.utils.network_utils import make_zmq_path -from vllm.v1.attention.backends.utils import get_kv_cache_layout -from vllm.v1.kv_cache_interface import ( - FullAttentionSpec, - MambaSpec, - UniformTypeKVCacheSpecs, -) -from vllm.v1.worker.block_table import BlockTable -from vllm.v1.worker.utils import select_common_block_size -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig +# Backward compatibility: NixlConnectorWorker is the pull-based worker. +NixlConnectorWorker = NixlPullConnectorWorker -logger = init_logger(__name__) - -class NixlConnectorWorker: - """Implementation of Worker side methods""" - - def _compute_desc_ids( - self, - block_ids: BlockIds, - dst_num_blocks: int, - block_size_ratio: float | None, - physical_blocks_per_logical: int, - ) -> np.ndarray: - """Compute NIXL descriptor IDs for given block IDs.""" - num_fa_regions = self.num_regions - num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0 - - num_blocks = dst_num_blocks - if block_size_ratio is not None: - num_blocks = int(num_blocks * block_size_ratio) - num_fa_descs = num_fa_regions * num_blocks - - # All-attention fast path: single vectorized broadcast. - if num_ssm_regions == 0: - # NOTE (NickLucche) With HMA, every kv group has the same number of layers - # and layers from different groups share the same kv tensor. - # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be - # read across all regions, same for [3], but group0-group1 blocks will - # always differ (different areas). Therefore we can just flatten the - # block_ids and compute the descs ids for all groups at once. - block_arr = np.concatenate(block_ids)[None, :] - region_ids = np.arange(num_fa_regions)[:, None] - return (region_ids * num_blocks + block_arr).flatten() - - # Compute desc ids per group using the right stride: FA descs have - # num_blocks entries per region (kernel granularity), SSM descs have - # logical_blocks entries per region (no kernel splitting). - logical_blocks = num_blocks // physical_blocks_per_logical - all_descs: list[np.ndarray] = [] - for i, group in enumerate(block_ids): - group_arr = np.asarray(group) - if _is_attention_spec(self._group_spec_types[i]): - fa_region_ids = np.arange(num_fa_regions)[:, None] - all_descs.append( - (fa_region_ids * num_blocks + group_arr[None, :]).flatten() - ) - elif _is_ssm_spec(self._group_spec_types[i]): - # NOTE (NickLucche) SSM and Attention block regions can - # be exchanged arbitrarily by manager. Therefore, descs - # are laid out as: - # [descs_fa (all regions) | descs_ssm (all regions)]. - # num_fa_descs offset must be computed per-engine since - # P and D can have different num_blocks (and thus - # different FA desc counts). - ssm_region_ids = np.arange(num_ssm_regions)[:, None] - all_descs.append( - ( - ssm_region_ids * logical_blocks - + group_arr[None, :] - + num_fa_descs - ).flatten() - ) - else: - raise ValueError( - f"Unknown spec type {self._group_spec_types[i]} at index {i}" - ) - - return np.concatenate(all_descs) - - def _build_local_splits_from_plan( - self, - plan: TPMapping, - src_blocks_data: list[tuple[int, int, int]], - num_fa_descs: int, - ) -> Iterator[list[tuple[int, int, int]]]: - """Build split handle data for P_TP > D_TP scenario. - - num_fa_descs is the boundary between FA and SSM descriptors. - Split counts are derived from source_ranks_per_group lengths. - FA uses rank_to_attention_slot for the slot offset; - SSM uses the rank's positional index. - """ - fa_idx = next( - i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) - ) - fa_num_splits = len(plan.source_ranks_per_group[fa_idx]) - - has_ssm_descs = num_fa_descs < len(src_blocks_data) - ssm_idx = next( - (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)), - None, - ) - ssm_num_splits = ( - len(plan.source_ranks_per_group[ssm_idx]) - if has_ssm_descs and ssm_idx is not None - else 0 - ) - - for p_idx, p_rank in enumerate(plan.all_source_ranks): - fa_slot = plan.rank_to_attention_slot.get(p_rank, 0) - - handle: list[tuple[int, int, int]] = [] - for j, (addr, local_len, dev) in enumerate(src_blocks_data): - if j < num_fa_descs: - chunk = local_len // fa_num_splits - handle.append((addr + fa_slot * chunk, chunk, dev)) - else: - chunk = local_len // ssm_num_splits - handle.append((addr + p_idx * chunk, chunk, dev)) - yield handle - - def __init__( - self, - vllm_config: "VllmConfig", - engine_id: str, - kv_cache_config: "KVCacheConfig", - ): - nixl_wrapper_cls = NixlWrapper - if nixl_wrapper_cls is None: - logger.error("NIXL is not available") - raise RuntimeError("NIXL is not available") - logger.info("Initializing NIXL wrapper") - logger.info("Initializing NIXL worker %s", engine_id) - - # Config. - self.vllm_config = vllm_config - # mypy will complain on re-assignment otherwise. - self.block_size: int = cast(int, vllm_config.cache_config.block_size) - - if vllm_config.kv_transfer_config is None: - raise ValueError("kv_transfer_config must be set for NixlConnector") - self.kv_transfer_config = vllm_config.kv_transfer_config - - self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config( - "backends", ["UCX"] - ) - kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config( - "kv_lease_duration", 30 - ) - # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. - self._lease_extension = kv_lease_duration * 2 // 3 - - self._is_hma_required = ( - not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager - and any( - not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups - ) - ) - self.kv_cache_config = kv_cache_config - self._layer_specs = { - layer: group.kv_cache_spec - for group in kv_cache_config.kv_cache_groups - for layer in group.layer_names - } - self.hma_group_size = len(kv_cache_config.kv_cache_tensors) - - # ---- Model state (derived from model config) ---- - mamba_ssm_size = (0, 0) - # Conv state sub-projection decomposition (None when no Mamba). - # The 3-read transfer requires DS (dim, state_len) conv layout so - # that x/B/C sub-projections are contiguous in memory. - self._conv_decomp: MambaConvSplitInfo | None = None - self._has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) - for g in kv_cache_config.kv_cache_groups - ) - if self._has_mamba: - assert self._is_hma_required - from vllm.model_executor.layers.mamba.mamba_utils import ( - is_conv_state_dim_first, - ) - - assert is_conv_state_dim_first(), ( - "3-read Mamba conv transfer requires DS conv state layout. " - "Set VLLM_SSM_CONV_STATE_LAYOUT=DS" - ) - mamba_spec = next( - spec - for spec in self._layer_specs.values() - if isinstance(spec, MambaSpec) - ) - self._conv_decomp = derive_mamba_conv_split( - mamba_spec, - vllm_config.parallel_config.tensor_parallel_size, - ) - mamba_ssm_size = self._conv_decomp.ssm_sizes - self._mamba_ssm_size = mamba_ssm_size - - # Agent. - non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] - # Configure NIXL num_threads to avoid UAR exhaustion on Mellanox NICs. - # Each UCX thread allocates UARs (doorbell pages) via DevX, and - # excessive NIXL UAR usage can exhaust NIC UAR space. This can cause - # components like NVSHMEM (used by DeepEP kernels) to fail during RDMA - # initialization with "mlx5dv_devx_alloc_uar" errors. - # Ref: https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf#page=63 - num_threads = vllm_config.kv_transfer_config.get_from_extra_config( - "num_threads", 4 - ) - if nixl_agent_config is None: - config = None - else: - # Enable telemetry by default for NIXL 0.7.1 and above. - config = ( - nixl_agent_config(backends=self.nixl_backends, capture_telemetry=True) - if len(non_ucx_backends) > 0 - else nixl_agent_config(num_threads=num_threads, capture_telemetry=True) - ) - - self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config) - # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. - self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict) - - # Metadata. - self.engine_id: EngineId = engine_id - self.tp_rank = get_tensor_model_parallel_rank() - self.world_size = get_tensor_model_parallel_world_size() - - self.num_blocks = kv_cache_config.num_blocks - self.enable_permute_local_kv = False - self.enable_heterogeneous_attn_post_process = False - - # KV Caches and nixl tracking data. - self.device_type = current_platform.device_type - self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device - if self.device_type not in _NIXL_SUPPORTED_DEVICE: - raise RuntimeError(f"{self.device_type} is not supported.") - elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]: - raise RuntimeError( - f"{self.device_type} with {self.kv_buffer_device} kv_buffer " - "is not supported." - ) - self.device_kv_caches: dict[str, torch.Tensor] = {} - - # cpu kv buffer for xfer - # used when device memory can not be registered under nixl - self.host_xfer_buffers: dict[str, torch.Tensor] = {} - if self.device_type == "cpu": - self.use_host_buffer = False - else: - self.use_host_buffer = self.kv_buffer_device == "cpu" - - # reserve different cores for start_load_kv() from model_forward() - if self.device_type == "cpu": - numa_core_list = current_platform.discover_numa_topology() - # setup one last core in each numa for kv transfer. - rsv_cores_for_kv = [ - max(each_numa_core_list) for each_numa_core_list in numa_core_list - ] - - if rsv_cores_for_kv: - if not hasattr(os, "sched_setaffinity"): - raise NotImplementedError( - "os.sched_setaffinity is not available on this platform" - ) - os.sched_setaffinity(0, rsv_cores_for_kv) - - # support for oot platform which can't register nixl memory - # type based on kv_buffer_device - nixl_memory_type = current_platform.get_nixl_memory_type() - if nixl_memory_type is None: - if self.kv_buffer_device in ["cuda", "xpu"]: - nixl_memory_type = "VRAM" - elif self.kv_buffer_device == "cpu": - nixl_memory_type = "DRAM" - if nixl_memory_type is None: - raise RuntimeError( - f"{self.device_type} with {self.kv_buffer_device} kv_buffer " - "is not supported." - ) - self.nixl_memory_type = nixl_memory_type - - # Note: host xfer buffer ops when use_host_buffer is True - self.copy_blocks: CopyBlocksOp | None = None - - # Map of engine_id -> kv_caches_base_addr. For TP case, each local - self.device_id: int = 0 - # Current rank may pull from multiple remote TP workers. - # EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer - self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict) - - # Number of NIXL regions. Currently one region per cache - # (so 1 per layer for MLA, otherwise 2 per layer) - self.num_regions = 0 - - # nixl_prepped_dlist_handle. - self.src_xfer_handles_by_block_size: dict[int, int] = {} - # Populated dynamically during handshake based on remote configuration. - # Keep track of regions at different tp_ratio values. tp_ratio->handles - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} - # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}. - self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict) - - # Map of engine_id -> num_blocks. All ranks in the same deployment will - # have the same number of blocks. - self.dst_num_blocks: dict[EngineId, int] = {} - self._registered_descs: list[Any] = [] - - # In progress transfers. - # [req_id -> list[handle]] - self._recving_metadata: dict[ReqId, ReqMeta] = {} - self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list) - # Track the expiration time of requests that are waiting to be sent. - self._reqs_to_send: dict[ReqId, float] = {} - # Set of requests that have been part of a batch, regardless of status. - self._reqs_to_process: set[ReqId] = set() - - # Invalid blocks from failed NIXL operations (thread-safe queue of block ids) - self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue() - # requests that skipped transfer (handshake or transfer failures) - # Uses Queue for thread-safe cross-thread coordination with the - # background handshake thread, matching the _ready_requests pattern. - self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() - - # Handshake metadata of this worker for NIXL transfers. - self.xfer_handshake_metadata: NixlHandshakePayload | None = None - # Background thread for initializing new NIXL handshakes. - self._handshake_initiation_executor = ThreadPoolExecutor( - # NIXL is not guaranteed to be thread-safe, limit 1 worker. - max_workers=1, - thread_name_prefix="vllm-nixl-handshake-initiator", - ) - self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]() - self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {} - # Protects _handshake_futures and _remote_agents. - self._handshake_lock = threading.RLock() - - # TTL-based eviction of stale remote engine state. - self._engine_last_active: dict[EngineId, float] = {} - self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config( - "engine_ttl", 3600.0 - ) - - self.block_size = vllm_config.cache_config.block_size - self.model_config = vllm_config.model_config - - self.use_mla = self.model_config.use_mla - - # Get the attention backend from the first layer - # NOTE (NickLucche) models with multiple backends are not supported yet - self.attn_backends = get_current_attn_backends(vllm_config) - self.backend_name = self.attn_backends[0].get_name() - - self.kv_cache_layout = get_kv_cache_layout() - self.host_buffer_kv_cache_layout = self.kv_cache_layout - logger.info( - "Detected attention backend(s) %s", - [backend.get_name() for backend in self.attn_backends], - ) - logger.info("Detected kv cache layout %s", self.kv_cache_layout) - - # lazy initialized in register_kv_caches - self.compat_hash: str | None = None - self.transfer_topo: TransferTopology | None = None - - # With heterogeneous TP, P must wait for all assigned D TP workers to - # finish reading before safely freeing the blocks. - self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int) - self.xfer_stats = NixlKVConnectorStats() - - self._physical_blocks_per_logical_kv_block = 1 - self._sync_block_size_with_kernel() - - # Unwrap UniformTypeKVCacheSpecs to get the representative spec type - self._group_spec_types = tuple( - get_representative_spec_type(g.kv_cache_spec) - for g in self.kv_cache_config.kv_cache_groups - ) - - # Per-engine TP mappings. Generated during handshake. - self.tp_mappings: dict[EngineId, TPMapping] = {} - - self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( - "enforce_handshake_compat", True - ) - - def _sync_block_size_with_kernel(self) -> None: - backends = get_current_attn_backends(self.vllm_config) - kernel_block_size = select_common_block_size(self.block_size, backends) - # Number of blocks not accounting for kernel block mismatches - self._logical_num_blocks = self.num_blocks - if self.block_size != kernel_block_size: - logger.info_once( - "User-specified logical block size (%s) does not match" - " physical kernel block size (%s). Using the latter.", - self.block_size, - kernel_block_size, - ) - assert self.block_size > kernel_block_size - self._physical_blocks_per_logical_kv_block = ( - self.block_size // kernel_block_size - ) - self.block_size = kernel_block_size - self.num_blocks *= self._physical_blocks_per_logical_kv_block - - def _nixl_handshake( - self, - host: str, - port: int, - remote_tp_size: int, - expected_engine_id: str, - ) -> dict[int, str]: - """Do a NIXL handshake with a remote instance.""" - - # the first time we connect to a remote agent. - # be careful, the handshake happens in a background thread. - # it does not have an active cuda context until any cuda runtime - # call is made. when UCX fails to find a valid cuda context, it will - # disable any cuda ipc communication, essentially disabling any NVLink - # communication. - # when we are using device buffers, we need to set the device - # explicitly to make sure the handshake background thread has a valid - # cuda context. - if not self.use_host_buffer: - current_platform.set_device(self.device_id) - - # When target instance TP > local TP, we need to perform multiple - # handshakes. Do it in a single background job for simplicity. - # Regardless, only handshake with the remote TP rank(s) that current - # local rank will read from. Note that With homogeneous TP, - # this happens to be the same single rank_i. - assert self.transfer_topo is not None - p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size) - remote_rank_to_agent_name = {} - path = make_zmq_path("tcp", host, port) - - with zmq_ctx(zmq.REQ, path) as sock: - for remote_rank in p_remote_ranks: - logger.debug( - "Querying metadata on path: %s at remote tp rank %s", - path, - remote_rank, - ) - - start_time = time.perf_counter() - # Send query for the request. - msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank)) - # Set receive timeout to 5 seconds to avoid hanging on dead server - sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds - sock.send(msg) - handshake_bytes = sock.recv() - - # Decode handshake payload to get compatibility hash - handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload) - try: - handshake_payload = handshake_decoder.decode(handshake_bytes) - except (msgspec.DecodeError, msgspec.ValidationError) as e: - raise RuntimeError( - f"Failed to decode NixlHandshakePayload. This likely indicates " - f"an incompatibility between connector version. Error: {e}" - ) from e - - got_metadata_time = time.perf_counter() - logger.debug( - "NIXL handshake: get metadata took: %s", - got_metadata_time - start_time, - ) - - # Check compatibility hash BEFORE decoding agent metadata - assert self.compat_hash is not None - if ( - self.enforce_compat_hash - and handshake_payload.compatibility_hash != self.compat_hash - ): - raise RuntimeError( - f"NIXL compatibility hash mismatch. " - f"Local: {self.compat_hash}, " - f"Remote: {handshake_payload.compatibility_hash}. " - f"Prefill and decode instances have incompatible " - f"configurations. This may be due to: different vLLM versions," - f" models, dtypes, KV cache layouts, attention backends, etc. " - f"Both instances must use identical configurations." - f"Disable this check using " - f'--kv-transfer-config \'{{"kv_connector_extra_config": ' - f'{{"enforce_handshake_compat": false}}}}\'' - ) - - logger.info( - "NIXL compatibility check passed (hash: %s)", - handshake_payload.compatibility_hash, - ) - - # Decode agent metadata - metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata) - try: - metadata = metadata_decoder.decode( - handshake_payload.agent_metadata_bytes - ) - except (msgspec.DecodeError, msgspec.ValidationError) as e: - # This should not happen if hash matched - raise RuntimeError( - f"Failed to decode NixlAgentMetadata. Error: {e}" - ) from e - - # Ensure engine id matches. - if metadata.engine_id != expected_engine_id: - raise RuntimeError( - f"Remote NIXL agent engine ID mismatch. " - f"Expected {expected_engine_id}," - f"received {metadata.engine_id}." - ) - - # Register Remote agent. - remote_agent_name = self.add_remote_agent( - metadata, remote_rank, remote_tp_size - ) - setup_agent_time = time.perf_counter() - logger.debug( - "NIXL handshake: add agent took: %s", - setup_agent_time - got_metadata_time, - ) - remote_rank_to_agent_name[remote_rank] = remote_agent_name - return remote_rank_to_agent_name - - def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None: - """ - Initialize transfer buffer in CPU mem for accelerators - NOT directly supported by NIXL (e.g., tpu) - """ - xfer_buffers: dict[str, torch.Tensor] = {} - inv_order = [0, 1, 3, 2, 4] - try: - for layer_name, kv_cache in kv_caches.items(): - kv_shape = kv_cache.shape - kv_dtype = kv_cache.dtype - permute_shape = False - if ( - self.kv_cache_layout == "NHD" - and self.vllm_config.kv_transfer_config is not None - and self.vllm_config.kv_transfer_config.enable_permute_local_kv - ): - logger.info_once( - "'enable_permute_local_kv' flag is enabled while " - "device KV Layout is NHD. Init host buffer with" - " HND to better support Decode/Prefill TP_ratio > 1." - ) - # Since NHD will not support Decode/Prefill TP_ratio > 1, - # we can leverage host_buffer for permute - self.host_buffer_kv_cache_layout = "HND" - kv_shape = ( - tuple(kv_shape[i] for i in inv_order) - if not self.use_mla - else kv_shape - ) - permute_shape = not self.use_mla - - xfer_buffers[layer_name] = torch.empty( - kv_shape, dtype=kv_dtype, device="cpu" - ) - if permute_shape: - xfer_buffers[layer_name] = xfer_buffers[layer_name].permute( - inv_order - ) - except MemoryError as e: - logger.error("NIXLConnectorWorker gets %s.", e) - raise - - self.host_xfer_buffers = xfer_buffers - - def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): - """Assign copy (d2h, h2d) operations when host buffer is used.""" - # Set a no-op if the host buffer is not cpu. - if self.kv_buffer_device != "cpu": - return - # Set a no-op if self.device_type is 'cpu'. - if self.device_type == "cpu": - return - assert self.use_host_buffer - self.copy_blocks = copy_operation - - def _log_failure( - self, - failure_type: str, - req_id: str | None, - msg: str = "", - error: Exception | None = None, - meta: ReqMeta | None = None, - **extra_context, - ): - """Log transfer failure with structured context for easier debugging.""" - context: dict[str, Any] = { - "failure_type": failure_type, - "request_id": req_id, - "engine_id": self.engine_id, - } - if meta is None and req_id is not None: - # Try to get metadata from in progress transfers when not provided - meta = self._recving_metadata.get(req_id) - - if meta and meta.remote: - context.update( - { - "remote_engine_id": meta.remote.engine_id, - "remote_request_id": meta.remote.request_id, - "remote_host": meta.remote.host, - "remote_port": meta.remote.port, - "num_local_blocks": sum( - len(group) for group in meta.local_block_ids - ), - "num_remote_blocks": sum( - len(group) for group in meta.remote.block_ids - ), - "local_block_ids_sample": meta.local_block_ids[0][:10] - if meta.local_block_ids - else [], - } - ) - - context.update(extra_context) - if msg: - failure_type = f"{failure_type}. {msg}" - - logger.error( - "NIXL transfer failure: %s | Context: %s", - failure_type, - context, - exc_info=error is not None, - stacklevel=2, - ) - - def _ensure_handshake( - self, - engine_id: EngineId, - host: str, - port: int, - tp_size: int, - ) -> Future[dict[int, str]] | None: - """ - Ensure a handshake is in-flight (or already done) for *engine_id*. - - Returns the ``Future`` if a handshake is pending (or was just - started), or ``None`` if the handshake already completed - successfully. Callers can attach per-request callbacks to the - returned future. - Failures to handshake are logged and the request is marked as failed. - """ - self._evict_stale_engines() - with self._handshake_lock: - if engine_id in self._remote_agents: - return None - fut = self._handshake_futures.get(engine_id) - if fut is not None: - return fut - fut = self._handshake_initiation_executor.submit( - self._nixl_handshake, - host, - port, - tp_size, - engine_id, - ) - self._handshake_futures[engine_id] = fut - - def done_callback(f: Future[dict[int, str]], eid=engine_id): - with self._handshake_lock: - del self._handshake_futures[eid] - try: - self._remote_agents[eid] = f.result() - self._engine_last_active[eid] = time.perf_counter() - except Exception as e: - self._log_failure( - failure_type="handshake_setup_failed", - req_id=None, - error=e, - remote_engine_id=eid, - ) - - fut.add_done_callback(done_callback) - return fut - - def _background_nixl_handshake( - self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta - ): - # Do NIXL handshake in background and add to _ready_requests when done. - assert meta.remote is not None - fut = self._ensure_handshake( - remote_engine_id, - meta.remote.host, - meta.remote.port, - meta.tp_size, - ) - if fut is None: - # Already handshaked — only happens if caller does not pre-check. - self._ready_requests.put((req_id, meta)) - return - - # Check handshake success before proceeding with request. - def request_ready(f: Future[Any], entry=(req_id, meta)): - try: - f.result() - self._ready_requests.put(entry) - except Exception as e: - self._log_failure( - failure_type="handshake_failed", - req_id=req_id, - error=e, - meta=meta, - ) - self._handle_failed_transfer(req_id, None) - - fut.add_done_callback(request_ready) - - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: - """Register a cross-layers KV cache tensor with NIXL. - - `use_uniform_kv_cache()` guarantees a single KV cache group whose - layers all share the same `AttentionSpec`, so any layer name from - `_layer_specs` yields the correct per-layer spec for `page_size_bytes`. - """ - first_layer = next(iter(self._layer_specs)) - # Forwarding a real layer name rather than a synthetic key - self.register_kv_caches({first_layer: kv_cache}) - - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - """Register the KV Cache data in nixl.""" - self.transfer_topo = TransferTopology( - tp_rank=self.tp_rank, - tp_size=self.world_size, - block_size=self.block_size, - engine_id=self.engine_id, - is_mla=self.use_mla, - total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backends=self.attn_backends, - # SSM States come in tuples (ssm, conv) - tensor_shape=next(iter(kv_caches.values())).shape - if not self._has_mamba - else None, - is_mamba=self._has_mamba, - ) - self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks - ) - - if self.use_host_buffer: - self.initialize_host_xfer_buffer(kv_caches=kv_caches) - assert len(self.host_xfer_buffers) == len(kv_caches), ( - f"host_buffer: {len(self.host_xfer_buffers)}, " - f"kv_caches: {len(kv_caches)}" - ) - xfer_buffers = self.host_xfer_buffers - else: - xfer_buffers = kv_caches - assert not self.host_xfer_buffers, ( - "host_xfer_buffer should not be initialized when " - f"kv_buffer_device is {self.kv_buffer_device}" - ) - - logger.info( - "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, " - "use_host_buffer: %s", - self.use_mla, - self.kv_buffer_device, - self.use_host_buffer, - ) - - caches_data = [] - # With hybrid allocator, layers can share a kv cache tensor - seen_base_addresses = [] - - # Note(tms): I modified this from the original region setup code. - # K and V are now in different regions. Advantage is that we can - # elegantly support MLA and any cases where the K and V tensors - # are non-contiguous (it's not locally guaranteed that they will be) - # Disadvantage is that the encoded NixlAgentMetadata is now larger - # (roughly 8KB vs 5KB). - # Conversely for FlashInfer, K and V are registered in the same region - # to better exploit the memory layout (ie num_blocks is the first dim). - tensor_size_bytes = None - - # Enable different block lengths for different layers *only* when MLA is used. - # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. - self.block_len_per_layer = list[int]() - for layer_name, cache_or_caches in xfer_buffers.items(): - # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to - # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. - # However, physical page_size may differ when kernel requires a specific - # block size. This leads to SSM and FA layers having different num_blocks. - # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. - layer_spec = self._layer_specs.get(layer_name) - if layer_spec is None: - logger.debug( - "Skipping layer %s as no KVCache spec is present. " - "This is likely because the layer is sharing its KV cache", - layer_name, - ) - continue - if isinstance(layer_spec, UniformTypeKVCacheSpecs): - # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs - layer_spec = layer_spec.kv_cache_specs[layer_name] - cache_list = self.transfer_topo.get_transfer_cache_regions( - cache_or_caches, layer_spec - ) - # `layer_spec.page_size_bytes` only accounts for logical page_size, that is - # the page_size assuming constant `self._logical_num_blocks`. - physical_page_size = ( - layer_spec.page_size_bytes - if isinstance(layer_spec, MambaSpec) - else layer_spec.page_size_bytes - // self._physical_blocks_per_logical_kv_block - ) - # For when registering multiple tensors eg K/V in separate regions. - physical_page_size = physical_page_size // len(cache_list) - if self.transfer_topo._cross_layers_blocks: - # When cross-layers blocks are used, multiply by number of layers - physical_page_size = physical_page_size * len( - self.kv_cache_config.kv_cache_tensors - ) - num_blocks = ( - self._logical_num_blocks - if isinstance(layer_spec, MambaSpec) - else self.num_blocks - ) - # `page_size` accounts for physical blocks, st KVCache is always - # [`num_blocks` * `page_size`] - curr_tensor_size_bytes = num_blocks * physical_page_size - if tensor_size_bytes is None: - tensor_size_bytes = curr_tensor_size_bytes - - # TODO (NickLucche) we could eventually unify how we handle FA/FI regions, - # registering a single tensor for both K/V and splitting logically like FI. - for cache in cache_list: - base_addr = cache.data_ptr() - if base_addr in seen_base_addresses: - # NOTE (NickLucche) HMA employs memory pooling to share tensors - # across groups. This results in skipping all tensors but the ones - # pointed to by group0. Also, generally we will have more blocks - # per tensor but fewer regions. - logger.debug("Skipping %s because it's already seen", layer_name) - continue - logger.debug( - "Registering layer %s with cache shape: %s", layer_name, cache.shape - ) - seen_base_addresses.append(base_addr) - # Only record non-Mamba page sizes. - if isinstance(layer_spec, MambaSpec): - self.block_len_per_layer.append( - physical_page_size // self._physical_blocks_per_logical_kv_block - ) - else: - self.block_len_per_layer.append(physical_page_size) - - if cache.shape[0] != num_blocks: - raise AssertionError( - "All kv cache tensors must have the same number of " - f"blocks; layer={layer_name}, " - f"expected_num_blocks={num_blocks}, " - f"cache_shape={tuple(cache.shape)}, " - f"cache_stride={tuple(cache.stride())}, " - f"layer_spec={type(layer_spec).__name__}, " - f"backend={self.backend_name}, " - "all_backends=" - f"{[backend.get_name() for backend in self.attn_backends]}, " - f"kv_cache_layout={self.kv_cache_layout}, " - "blocks_first=" - f"{self.transfer_topo.is_kv_layout_blocks_first}" - ) - - if not self.use_mla: - # Different kv cache shape is not supported by HeteroTP. - # This must also hold true for Mamba-like models. - assert tensor_size_bytes == curr_tensor_size_bytes, ( - "All kv cache tensors must have the same size" - ) - # Need to make sure the device ID is non-negative for NIXL, - # Torch uses -1 to indicate CPU tensors. - self.device_id = max(cache.get_device(), 0) - caches_data.append( - (base_addr, curr_tensor_size_bytes, self.device_id, "") - ) - - logger.debug( - "Different block lengths collected: %s", set(self.block_len_per_layer) - ) - assert len(self.block_len_per_layer) == len(seen_base_addresses) - - self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses - self.num_regions = len(caches_data) - - if self.transfer_topo.virtually_split_kv_in_blocks: - # NOTE (NickLucche) When FlashInfer is used, memory is registered - # with joint KV for each block. This minimizes the overhead in - # registerMem allowing faster descs queries. In order to be able to - # split on kv_heads dim as required by heterogeneous TP, one must - # be able to index K/V separately. Hence we double the number - # of 'virtual' regions here and halve `block_len` below. - # Similarly for Mamba layers, we register SSM+Conv as a single region and - # then duplicate it logically to be able to index SSM/Conv separately. - self.num_regions *= 2 - - # Total local FA descriptors (boundary between FA and mamba descs). - self.num_descs = self.num_regions * self.num_blocks - - descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) - logger.debug("Registering descs: %s", caches_data) - self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) - logger.debug("Done registering descs") - self._registered_descs.append(descs) - - self.device_kv_caches = kv_caches - self.dst_num_blocks[self.engine_id] = self.num_blocks - - if self._has_mamba: - logger.info( - "Hybrid SSM registration: num_blocks=%s, " - "logical_num_blocks=%s, ratio=%s, num_regions=%s, " - "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s", - self.num_blocks, - self._logical_num_blocks, - self._physical_blocks_per_logical_kv_block, - self.num_regions, - self.num_descs, - self._mamba_ssm_size, - set(self.block_len_per_layer), - ) - - # Register local/src descr for NIXL xfer. - self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = ( - self.register_local_xfer_handler(self.block_size) - ) - - # After KV Caches registered, listen for new connections. - agent_metadata = NixlAgentMetadata( - engine_id=self.engine_id, - agent_metadata=self.nixl_wrapper.get_agent_metadata(), - device_id=self.device_id, - kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank], - num_blocks=self.num_blocks, - block_lens=self.block_len_per_layer, - kv_cache_layout=self.kv_cache_layout - if not self.use_host_buffer - else self.host_buffer_kv_cache_layout, - block_size=self.block_size, - ssm_sizes=self._mamba_ssm_size, - attn_backend_name=self.backend_name, - physical_blocks_per_logical_kv_block=( - self._physical_blocks_per_logical_kv_block - ), - ) - # Wrap metadata in payload with hash for defensive decoding - assert self.compat_hash is not None - encoder = msgspec.msgpack.Encoder() - self.xfer_handshake_metadata = NixlHandshakePayload( - compatibility_hash=self.compat_hash, - agent_metadata_bytes=encoder.encode(agent_metadata), - ) - - def _build_mamba_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> list[tuple[int, int, int]]: - """Build 4 desc regions (x, B, C, ssm) per layer for local mamba - blocks, enabling the 3-read transfer with DS conv layout.""" - assert block_size_ratio == 1, ( - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " - f"Got block_size_ratio={block_size_ratio}." - ) - assert self._conv_decomp is not None - conv_offsets = self._conv_decomp.local_conv_offsets - conv_size, ssm_size = self._mamba_ssm_size - num_blocks = self._logical_num_blocks * block_size_ratio - physical_per_logical = self._physical_blocks_per_logical_kv_block - - result: list[tuple[int, int, int]] = [] - for i, base_addr in enumerate(base_addresses): - # Jump one page_size, but ssm page_size may be bigger when kernel - # locks block size to a specific value (physical_per_logical scale). - page_stride = ( - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical - ) - for off, sz in conv_offsets: - for blk in range(num_blocks): - result.append( - (base_addr + blk * page_stride + off, sz, self.device_id) - ) - # SSM temporal state follows the conv state. - for blk in range(num_blocks): - result.append( - ( - base_addr + blk * page_stride + conv_size, - ssm_size, - self.device_id, - ) - ) - return result - - def _build_mamba_remote( - self, - nixl_agent_meta: NixlAgentMetadata, - tp_ratio: int, - transfer_info: EngineTransferInfo, - ) -> list[tuple[int, int, int]]: - """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer - for the 3-read transfer. For hetero-TP, each D rank reads only its - sub-projection slice from the P rank.""" - assert self._conv_decomp is not None - effective_ratio = max(tp_ratio, 1) - # Mamba conv state is always TP-sharded, even when attention KV - # is replicated (num_kv_heads < tp_size). - local_offset = self.tp_rank % effective_ratio - conv_size_remote = nixl_agent_meta.ssm_sizes[0] - - conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio) - if tp_ratio >= 1: - ssm_read_size = self._mamba_ssm_size[1] - else: - ssm_read_size = nixl_agent_meta.ssm_sizes[1] - - remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical - num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical - device_id = nixl_agent_meta.device_id - - result: list[tuple[int, int, int]] = [] - # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case - # block lengths vary across layers (e.g. MLA). - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical - for off, sz in conv_offsets: - for blk in range(num_blocks): - result.append((base_addr + blk * page_stride + off, sz, device_id)) - # SSM temporal state is also TP-sharded on the heads dimension. - for blk in range(num_blocks): - ssm_addr = ( - base_addr - + blk * page_stride - + conv_size_remote - + local_offset * ssm_read_size - ) - result.append((ssm_addr, ssm_read_size, device_id)) - return result - - def _build_fa_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> list[tuple[int, int, int]]: - """Build local FA descriptors for all layers.""" - assert self.transfer_topo is not None - num_blocks = self.num_blocks * block_size_ratio - result: list[tuple[int, int, int]] = [] - for i, base_addr in enumerate(base_addresses): - kv_block_len = ( - self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) - // block_size_ratio - ) - page_stride = self.block_len_per_layer[i] // block_size_ratio - for block_id in range(num_blocks): - block_offset = block_id * page_stride - addr = base_addr + block_offset - result.append((addr, kv_block_len, self.device_id)) - - if self.transfer_topo.virtually_split_kv_in_blocks: - # Separate and interleave K/V regions to maintain the same - # descs ordering. This is needed for selecting contiguous heads - # when split across TP ranks. - second_split = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=False, mamba_view=False - ) - for block_id in range(num_blocks): - block_offset = block_id * page_stride - addr = base_addr + block_offset - v_addr = addr + kv_block_len - result.append((v_addr, second_split, self.device_id)) - return result - - def _build_fa_remote( - self, - plan: TPMapping, - nixl_agent_meta: NixlAgentMetadata, - block_size_ratio: int, - ) -> list[tuple[int, int, int]]: - """Build remote FA descriptors for all layers.""" - assert self.transfer_topo is not None - fa_group_idx = next( - i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) - ) - num_attn_reads = len(plan.source_ranks_per_group[fa_group_idx]) - num_blocks = nixl_agent_meta.num_blocks - result: list[tuple[int, int, int]] = [] - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - # Read our whole local region size from remote.. - local_block_len = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) - remote_kv_block_len = local_block_len // block_size_ratio - if block_size_ratio > 1: - # ..using remote kv_block_len as transfer unit - local_block_len = remote_kv_block_len - - local_block_len = local_block_len // num_attn_reads - rank_offset = plan.rank_offset_factor * remote_kv_block_len - - page_size = nixl_agent_meta.block_lens[i] - for block_id in range(num_blocks): - block_offset = block_id * page_size - # For each block, grab the kv heads chunk belonging to current local - # tp rank of size local_block_len. - addr = base_addr + block_offset + rank_offset - result.append((addr, local_block_len, nixl_agent_meta.device_id)) - - if self.transfer_topo.virtually_split_kv_in_blocks: - # With FlashInfer index V separately to allow head splitting. - second_split = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=False, mamba_view=False - ) - second_split = second_split // num_attn_reads - for block_id in range(num_blocks): - block_offset = block_id * page_size - addr = base_addr + block_offset + rank_offset - # Hop over the first split of remote page, K, to read V. - v_addr = addr + nixl_agent_meta.block_lens[i] // 2 - result.append((v_addr, second_split, nixl_agent_meta.device_id)) - return result - - def register_local_xfer_handler( - self, - block_size: int, - ) -> tuple[int, list[tuple[int, int, int]]]: - """ - Function used for register local xfer handler with local block_size or - Remote block_size. - - When local block_size is same as remote block_size, we use local block_size - to register local_xfer_handler during init. - - When remote block size is less than local block size, we need to use - register another local_xfer_handler using remote block len to ensure - data copy correctness. - """ - assert self.transfer_topo is not None - block_size_ratio = self.block_size // block_size - local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank] - - blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio) - logger.debug( - "Created %s blocks for src engine %s and rank %s on device id %s", - len(blocks_data), - self.engine_id, - self.tp_rank, - self.device_id, - ) - if self._has_mamba: - assert self.num_descs == len(blocks_data) - # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split - # is unnecessary — a single conv desc per block suffices. Consider - # adding a fast path that falls back to the standard 2-region - # registration (_build_fa_local mamba=True) when no hetero-TP - # remote has been seen. Currently we always register 4 regions - # because local descs are created before knowing the remote TP. - logger.debug("Registering local Mamba descriptors (4 regions/layer)") - blocks_data.extend( - self._build_mamba_local(local_base_addresses, block_size_ratio) - ) - - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) - # NIXL_INIT_AGENT to be used for preparations of local descs. - return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data - - def add_remote_agent( - self, - nixl_agent_meta: NixlAgentMetadata, - remote_tp_rank: int = 0, - remote_tp_size: int = 1, - ) -> str: - """ - Add the remote NIXL agent and prepare the descriptors for reading cache - blocks from remote. - - In particular, handle both homogeneous and heterogeneous TP. The former - requires local rank_i to read from remote rank_i. - The latter, in the case of D.world_size < P.world_size, requires that a - local (D) TP worker reads from multiple remote (P) TP workers. - Conversely, assuming D.world_size > P.world_size, two or more local TP - workers will read from a single remote TP worker. - - Here's an example for the last case described above (non-MLA): - - rank_offset p_remote_tp_rank - (kv split no) - -------------------------------- - 0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ] - / - 1 0 Worker1 ---- 2nd half of KV -----/ - - 0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ] - / - 1 1 Worker3 ---- 2nd half of KV -----/ - - - Decoder TP workers Prefix TP workers - (world_size=4) (world_size=2) - tp_ratio = 4 // 2 = 2 - - Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim] - then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format. - Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio - first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split - along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0. - - Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1. - - Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0 - so that the whole cache is shared by "tp_ratio" D TP workers. - - For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and - tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer. - """ # noqa: E501 - engine_id = nixl_agent_meta.engine_id - # TODO re-evaluate refreshing for scaling/recovery - if remote_tp_rank in self._remote_agents.get(engine_id, {}): - logger.debug( - "Remote agent with engine_id %s and rank" - "%s already exchanged metadata, skip handshake.", - engine_id, - remote_tp_rank, - ) - return self._remote_agents[engine_id][remote_tp_rank] - - ### Register remote engine in TransferTopology (idempotent). - assert self.transfer_topo is not None - transfer_topo = self.transfer_topo - physical_blocks_per_logical = ( - nixl_agent_meta.physical_blocks_per_logical_kv_block - ) - transfer_info = EngineTransferInfo( - remote_tp_size=remote_tp_size, - remote_block_size=nixl_agent_meta.block_size, - remote_block_len=nixl_agent_meta.block_lens[0], - remote_physical_blocks_per_logical=physical_blocks_per_logical, - ) - transfer_topo.register_remote_engine(engine_id, transfer_info) - logger.info("Transfer plan: %s", transfer_topo.describe(engine_id)) - - self.tp_mappings[engine_id] = compute_tp_mapping( - transfer_topology=transfer_topo, - remote_tp_size=remote_tp_size, - group_spec_types=self._group_spec_types, - ) - - remote_agent_name = self.nixl_wrapper.add_remote_agent( - nixl_agent_meta.agent_metadata - ) - - # Create dst descs and xfer side handles. TP workers have same #blocks - # so we only register once per engine_id. - # Example: - # block_size_ratio > 1: - # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12| - # local origin:| 0| 1| 8| 12| - # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15| - block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size) - - if engine_id not in self.dst_num_blocks: - self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks - - # Keep track of remote agent kv caches base addresses. - self.kv_caches_base_addr[engine_id][remote_tp_rank] = ( - nixl_agent_meta.kv_caches_base_addr - ) - self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size) - - # This is 1 when P and D `--tensor-parallel-size` match. Otherwise, - # this is the ratio between the two sizes. - tp_ratio = transfer_topo.tp_ratio(remote_tp_size) - - logger.debug( - "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s", - engine_id, - remote_tp_rank, - tp_ratio, - ) - - plan = self.tp_mappings[engine_id] - - ### (Optional) Register local agent memory regions. MLA is not split. - if ( - tp_ratio < 0 - and not self.use_mla - and tp_ratio not in self.src_xfer_handles_by_tp_ratio - ): - # Remote tp_size > local tp_size: read from multiple remote ranks. - # Logically "split" own regions into |tp_ratio| chunks. Mind that - # we only do this once per remote tp_size (replica-friendly). - self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] - - for handle_data in self._build_local_splits_from_plan( - plan, - self.src_blocks_data, - self.num_descs, - ): - descs = self.nixl_wrapper.get_xfer_descs( - handle_data, self.nixl_memory_type - ) - handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) - - ### Register remote agent memory regions - # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With - # heterogeneous TP, prepare the descriptors by splitting the P KV cache along - # kv_head dim, of D worker's kv_head size (D>P). - # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..]. - - # Register all remote blocks, but only the corresponding kv heads. - blocks_data = self._build_fa_remote( - plan, - nixl_agent_meta, - block_size_ratio, - ) - logger.debug( - "Created %s blocks for dst engine %s with remote rank %s and local rank %s", - len(blocks_data), - engine_id, - remote_tp_rank, - self.tp_rank, - ) - if self._has_mamba: - logger.debug( - "Registering remote Mamba blocks for engine %s rank %s", - engine_id, - remote_tp_rank, - ) - blocks_data.extend( - self._build_mamba_remote( - nixl_agent_meta, - tp_ratio, - transfer_info, - ) - ) - - # Register with NIXL. - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) - self.dst_xfer_side_handles[engine_id][remote_tp_rank] = ( - self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) - ) - - if block_size_ratio > 1: - # when prefill with smaller block_size, we need to init a - # new handler with same block_len to match - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] - ) - - return remote_agent_name - - def _validate_remote_agent_handshake( - self, nixl_agent_meta: NixlAgentMetadata, remote_tp_size: int - ): - """ - Validate the remote agent handshake metadata ensuring the - invariants hold true. - """ - remote_engine_id = nixl_agent_meta.engine_id - - assert self.transfer_topo is not None - remote_info = self.transfer_topo.get_engine_info(remote_engine_id) - assert remote_info.remote_tp_size == remote_tp_size - - tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size) - block_size_ratio = self.transfer_topo.block_size_ratio( - nixl_agent_meta.block_size - ) - # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba. - # Mamba models can have replicated FA KV with tp_ratio < 0. - # MLA models do not need to handle kv replication. - if not self.use_mla and not self._has_mamba: - assert not ( - tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) - ) - - remote_physical_per_logical = ( - nixl_agent_meta.physical_blocks_per_logical_kv_block - ) - if ( - self._has_mamba - and remote_physical_per_logical - != self._physical_blocks_per_logical_kv_block - and self.vllm_config.cache_config.enable_prefix_caching - ): - raise RuntimeError( - "Prefix caching with heterogeneous physical_blocks_per_logical " - "is not supported for Mamba hybrid models. " - f"Local: {self._physical_blocks_per_logical_kv_block}, " - f"Remote: {remote_physical_per_logical}. " - "Disable prefix caching with --no-enable-prefix-caching." - ) - - if self._is_hma_required: - assert block_size_ratio == 1, ( - "HMA does not support different remote block size yet" - ) - kv_cache_layout = ( - self.kv_cache_layout - if not self.use_host_buffer - else self.host_buffer_kv_cache_layout - ) - if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout: - if ( - self.kv_transfer_config.enable_permute_local_kv - and nixl_agent_meta.kv_cache_layout == "HND" - ): - logger.info( - "Remote is HND and local is NHD, enabled additional permute " - "on local device KV." - ) - assert not self._is_hma_required, ( - "HMA does not support block size post processing" - ) - self.enable_permute_local_kv = True - else: - raise RuntimeError( - "Heterogeneous TP expects same kv_cache_layout. " - "Or enable experimental feature to use HND to NHD support by " - "setting 'enable_permute_local_kv'=True in --kv-transfer-config." - ) - # if remote_agent used attn is not same as local, - # hint heterogenuous attn post process - if ( - nixl_agent_meta.attn_backend_name != self.backend_name - and self.backend_name in ["CPU_ATTN"] - ): - if self._is_hma_required: - raise RuntimeError( - "heterogeneous attn post process is not supported with HMA" - ) - logger.info( - "[Experimental] CPU_ATTN backend is used, " - "hint heterogeneous attn post process" - ) - self.enable_heterogeneous_attn_post_process = True - - # Heterogeneous TP requires head-splitting, which only works with - # HND layout. MLA and replicated-KV cases don't split on heads. - # Mamba doesn't support heterogeneous TP. - if ( - abs(tp_ratio) != 1 - and not self.use_mla - and not self.transfer_topo.is_kv_replicated(remote_engine_id) - and kv_cache_layout != "HND" - and not self.enable_permute_local_kv - ): - raise RuntimeError( - "Heterogeneous TP head-dimension splitting requires contiguous heads. " - "Use HND layout on the prefill side." - ) - - # Block len can only vary across layers when using MLA. - remote_block_len = nixl_agent_meta.block_lens[0] - if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id): - # With replicated KV cache, only the number of blocks can differ. - # TODO (ZhanqiuHu): For mamba models, validate FA and mamba - # block_lens separately. - if not self._has_mamba: - for i in range(len(self.block_len_per_layer)): - assert ( - self.block_len_per_layer[i] // block_size_ratio - == nixl_agent_meta.block_lens[i] - ), "KV cache sizes must match between P and D when replicated" - else: - # When MLA is not used, this is a list of the same block length - for block_len in nixl_agent_meta.block_lens: - assert block_len == remote_block_len, ( - "All remote layers must have the same block size" - ) - - # HMA hybrid models (mamba+attention) pad block_len to - # max(attn_page, mamba_page), so the linear tp_ratio scaling - # assumption only holds for pure-attention models. - if not self._has_mamba: - if tp_ratio > 0: - assert ( - remote_block_len - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio - ), ( - "Remote P worker KV layer cache must be of shape [2, N," - " local_kv_heads*tp_ratio, page_size, head_dim] and " - "same dtype." - ) - else: - assert block_size_ratio == 1, ( - "Different local/remote block sizes are not supported" - " when P TP > D TP." - ) - assert remote_block_len == self.block_len_per_layer[0] // ( - -tp_ratio - ), ( - "Remote P worker KV layer cache must be of shape [2, N," - " local_kv_heads/tp_ratio, page_size, head_dim] and " - "same dtype." - ) - - # TP workers that handhshake with same remote have same #blocks. - assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks - # Same number of regions/~layers. - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) - - def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): - """copy recved kv from host buffer to device.""" - assert self.use_host_buffer - assert self.copy_blocks is not None - - local_block_ids = meta.local_physical_block_ids - # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups - for group_block_ids in local_block_ids: - self.copy_blocks( - self.host_xfer_buffers, - self.device_kv_caches, - group_block_ids, - group_block_ids, - "h2d", - ) - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "synced recved kv of request[%s] to device kv buffer," - "local_block_ids: %s. ", - req_id, - ",".join(map(str, local_block_ids)), - ) - - def save_kv_to_host(self, metadata: NixlConnectorMetadata): - """copy kv from device to host buffer.""" - assert self.use_host_buffer - assert self.copy_blocks is not None - - for req_id, meta in metadata.reqs_to_save.items(): - meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids - ) - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "save_load_kv for request[%s] to host xfer buffer." - "local_block_ids: %s. ", - req_id, - ",".join(map(str, meta.local_physical_block_ids)), - ) - # blocking - for group_block_ids in meta.local_physical_block_ids: - self.copy_blocks( - self.device_kv_caches, - self.host_xfer_buffers, - group_block_ids, - group_block_ids, - "d2h", - ) - - def post_process_device_kv_on_receive( - self, - block_size_ratio: int, - block_ids_list: list[list[int]], - ): - """ - Post process device kv cache after receiving from remote. - - 3 types of post processing supported: - * kv_cache_postprocess_layout => convert from HND to NHD - * kv_cache_postprocess_blksize => convert from small block size - to large block size - * kv_cache_postprocess_blksize_and_layout => convert from small - block size to large block size and convert from HND to NHD - - """ - if len(self.device_kv_caches) == 0: - return - assert block_size_ratio >= 1, "Only nP < nD supported currently." - assert self.transfer_topo is not None - if self.enable_permute_local_kv and block_size_ratio > 1: - logger.debug( - "Post-processing device kv cache on receive by converting " - "block_size with %sx bigger and permuting layout from HND" - " to NHD.", - block_size_ratio, - ) - elif self.enable_permute_local_kv: - logger.debug( - "Post-processing device kv cache on receive by permuting layout" - "from HND to NHD." - ) - else: - logger.debug( - "Post-processing device kv cache on receive by converting " - "block_size with %sx bigger.", - block_size_ratio, - ) - - split_k_and_v = self.transfer_topo.split_k_and_v - - for block_ids in block_ids_list: - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) - - for _, cache_or_caches in self.device_kv_caches.items(): - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches] - for cache in cache_list: - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive( - cache, indices, block_size_ratio - ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive( - cache, indices, block_size_ratio - ) - - def post_process_device_kv_on_receive_heterogeneous_attn( - self, block_ids: list[int] - ): - """ - Post process device kv cache after receiving from remote - for heterogeneous attention. - """ - assert self.enable_heterogeneous_attn_post_process - - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) - - for _, cache_or_caches in self.device_kv_caches.items(): - blocks_to_update = cache_or_caches.index_select(1, indices) - current_platform.pack_kv_cache( - key=blocks_to_update[0], - value=blocks_to_update[1], - key_cache=cache_or_caches[0], - value_cache=cache_or_caches[1], - block_ids=block_ids, - indices=indices, - ) - - def get_finished(self) -> tuple[set[str], set[str]]: - """ - Get requests that are done sending or recving on this specific worker. - The scheduler process (via the MultiprocExecutor) will use this output - to track which workers are done. - """ - assert self.transfer_topo is not None - done_sending = self._get_new_notifs() - done_recving = self._pop_done_transfers(self._recving_transfers) - - # Drain queue of requests where handshake or transfer setup failed. - failed_recv_reqs = set[ReqId]() - while not self._failed_recv_reqs.empty(): - try: - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait()) - except queue.Empty: - break - - # Add failed requests to done_recving for scheduler tracking - # (blocks are already marked invalid, scheduler will handle recompute) - done_recving.update(failed_recv_reqs) - - if len(done_sending) > 0 or len(done_recving) > 0: - logger.debug( - "Rank %s, get_finished: %s requests done sending " - "and %s requests done recving (%s failed)", - self.tp_rank, - len(done_sending), - len(done_recving), - len(failed_recv_reqs), - ) - - block_ids_for_blocksize_post_process = defaultdict(list) - block_ids_for_heterogeneous_attn_post_process = list[list[int]]() - for req_id in done_recving: - # clean up metadata for completed requests - meta = self._recving_metadata.pop(req_id, None) - assert meta is not None, f"{req_id} not found in recving_metadata list" - - # Skip KV sync and post-processing for failed requests - if req_id in failed_recv_reqs: - logger.warning( - "Skipping KV post-processing for failed request %s", - req_id, - ) - continue - - assert meta.remote is not None - if self.use_host_buffer: - self.sync_recved_kv_to_device(req_id, meta) - - # post processing for heteroblocksize - remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) - block_size_ratio = self.transfer_topo.block_size_ratio( - remote_info.remote_block_size - ) - if not self.use_mla and ( - block_size_ratio > 1 or self.enable_permute_local_kv - ): - assert not self._is_hma_required - block_ids_for_blocksize_post_process[block_size_ratio].append( - meta.local_physical_block_ids[0] - ) - # post processing for heterogeneous attention - if self.enable_heterogeneous_attn_post_process: - block_ids_for_heterogeneous_attn_post_process.append( - meta.local_physical_block_ids[0] - ) - for ( - block_size_ratio, - block_ids_list, - ) in block_ids_for_blocksize_post_process.items(): - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) - - for block_ids in block_ids_for_heterogeneous_attn_post_process: - self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) - - # Handle timeout to avoid stranding blocks on remote. - now = time.perf_counter() - while self._reqs_to_send: - req_id, expires = next(iter(self._reqs_to_send.items())) - # Sorted dict, oldest requests are put first so we can exit early. - if now < expires: - break - count = self.consumer_notification_counts_by_req.pop(req_id, 0) - self.xfer_stats.record_kv_expired_req() - logger.warning( - "Releasing expired KV blocks for request %s which were " - "retrieved by %d remote worker(s) before lease expired.", - req_id, - count, - ) - self._reqs_to_process.remove(req_id) - del self._reqs_to_send[req_id] - done_sending.add(req_id) - - return done_sending, done_recving - - def _get_new_notifs(self) -> set[str]: - """ - Get req_ids which got a remote xfer message. When multiple consumers - are reading from the same producer (heterogeneous TP scenario), wait - for all consumers to be done pulling. - - Also handles heartbeat notifications ("HB:req1,req2,...") by - extending the lease on the referenced requests. - """ - assert self.transfer_topo is not None - notified_req_ids: set[str] = set() - for notifs in self.nixl_wrapper.get_new_notifs().values(): - for notif in notifs: - msg = notif.decode("utf-8") - - # Handle heartbeat messages from D-side. - if msg.startswith("HB:"): - self._handle_heartbeat(msg[3:]) - continue - - req_id, tp_size = msg.rsplit(":", 1) - if ( - req_id not in self._reqs_to_send - and req_id not in self._reqs_to_process - ): - logger.error( - "Potentially invalid KV blocks for " - "unrecognized request %s were retrieved by " - "a decode worker. They may have expired.", - req_id, - ) - continue - - # NOTE: `tp_ratio` is the opposite when swapping local<>remote - n_consumers = int(tp_size) - tp_ratio = self.transfer_topo.tp_ratio(n_consumers) - - # Number of reads *per producer* to wait for. - # When remote D TP > local P TP we expect `tp_ratio` reads. - consumers_per_producer = ( - -tp_ratio if n_consumers > self.world_size else 1 - ) - - self.consumer_notification_counts_by_req[req_id] += 1 - # Wait all consumers (D) to be done reading before freeing. - if ( - self.consumer_notification_counts_by_req[req_id] - == consumers_per_producer - ): - notified_req_ids.add(req_id) - del self.consumer_notification_counts_by_req[req_id] - self._reqs_to_process.remove(req_id) - self._reqs_to_send.pop(req_id, None) - return notified_req_ids - - def _handle_heartbeat(self, payload: str) -> None: - """Extend leases for requests referenced in a heartbeat. - - Args: - payload: comma-separated P-side request IDs, e.g. - "req_abc,req_def". - """ - new_expiry = time.perf_counter() + self._lease_extension - for req_id in payload.split(","): - if req_id in self._reqs_to_send: - old = self._reqs_to_send[req_id] - self._reqs_to_send[req_id] = max(old, new_expiry) - logger.debug( - "Heartbeat extended lease for request %s " - "by %ds (old_expiry=%.1f, new_expiry=%.1f)", - req_id, - self._lease_extension, - old, - new_expiry, - ) - - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: - """ - Pop completed xfers by checking for DONE state. - Args: - transfers: dict of req_id -> list[running_xfer] - Returns: - set of req_ids that have all done xfers - """ - done_req_ids: set[str] = set() - for req_id, handles in list(transfers.items()): - in_progress = [] - for handle in handles: - try: - xfer_state = self.nixl_wrapper.check_xfer_state(handle) - if xfer_state == "DONE": - # Get telemetry from NIXL - res = self.nixl_wrapper.get_xfer_telemetry(handle) - self.xfer_stats.record_transfer(res) - self.nixl_wrapper.release_xfer_handle(handle) - elif xfer_state == "PROC": - in_progress.append(handle) - continue - else: - self._log_failure( - failure_type="transfer_failed", - msg="Marking blocks as invalid", - req_id=req_id, - xfer_state=xfer_state, - ) - self._handle_failed_transfer(req_id, handle) - except Exception as e: - self._log_failure( - failure_type="transfer_exception", - msg="Marking blocks as invalid", - req_id=req_id, - error=e, - ) - self._handle_failed_transfer(req_id, handle) - - if not in_progress: - # Only report request as completed when all transfers are done. - done_req_ids.add(req_id) - del transfers[req_id] - else: - transfers[req_id] = in_progress - return done_req_ids - - def _handle_failed_transfer(self, req_id: str, handle: int | None): - """ - Handle a failed transfer by marking all (logical) blocks as invalid and - recording the failure. - - Args: - req_id: The request ID. - handle: The transfer handle. - """ - # Use .get() here as the metadata cleanup is handled by get_finished() - # TODO (NickLucche) handle failed transfer for HMA. - if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required: - self._invalid_block_ids.put(set(meta.local_block_ids[0])) - self._failed_recv_reqs.put(req_id) - if handle is not None: - self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() - - def start_load_kv(self, metadata: NixlConnectorMetadata): - """ - Start loading by triggering non-blocking nixl_xfer. - We check for these trnxs to complete in each step(). - """ - for req_id, meta in metadata.reqs_to_recv.items(): - meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids - ) - assert meta.remote is not None - # Remote block IDs are kept logical here; expanded in - # _read_blocks_for_req using the remote engine's phys ratio. - remote_engine_id = meta.remote.engine_id - logger.debug( - "start_load_kv for request %s from remote engine %s. " - "Num local_block_ids: %s. Num remote_block_ids: %s. ", - req_id, - remote_engine_id, - len(meta.local_physical_block_ids), - len(meta.remote.block_ids), - ) - # always store metadata for failure recovery - self._recving_metadata[req_id] = meta - if remote_engine_id not in self._remote_agents: - # Initiate handshake with remote engine to exchange metadata. - with self._handshake_lock: - if remote_engine_id not in self._remote_agents: - self._background_nixl_handshake(req_id, remote_engine_id, meta) - continue - - # Handshake already completed, start async read xfer. - self._read_blocks_for_req(req_id, meta) - - # Start transfers for requests whose handshakes have now finished. - while not self._ready_requests.empty(): - self._read_blocks_for_req(*self._ready_requests.get_nowait()) - - # Keep around the requests that have been part of a batch. This is - # needed because async scheduling pushes the misalignment between the - # moment in which requests expiration is set (P side) and the moment in - # which blocks are read from D. As P can now more easily lag behind D - # while processing the next batch, we make sure to only set an - # expiration for requests that have not been read from D yet. - for req_id in metadata.reqs_in_batch: - self._reqs_to_process.add(req_id) - - # Remove all requests that are not to be processed (eg aborted). - for req_id in metadata.reqs_not_processed: - self._reqs_to_process.discard(req_id) - # We should never get an abort after setting an expiry timer - assert req_id not in self._reqs_to_send - - # Add to requests that are waiting to be read and track expiration. - for req_id, expiration_time in metadata.reqs_to_send.items(): - if req_id in self._reqs_to_process: - self._reqs_to_send[req_id] = expiration_time - - # Send heartbeats to P-side engines to keep KV blocks alive while - # requests sit in the D scheduler WAITING queue. - self._send_heartbeats(metadata) - - def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None: - """ - Send heartbeat notifications to remote engines, extending lease on KV blocks. - """ - for engine_id, hb_info in metadata.heartbeat_by_engine.items(): - # Proactive handshake (this request may still be in waiting queue) so - # the **next** heartbeat for this remote can go through. - if ( - self._ensure_handshake( - engine_id, hb_info.host, hb_info.port, hb_info.tp_size - ) - is not None - ): - continue # handshake is still pending - - # Build the heartbeat message: "HB:req1,req2,..." - hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode() - for agent_name in self._remote_agents[engine_id].values(): - try: - self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg) - except Exception: - logger.debug( - "Failed to send heartbeat to engine %s", - engine_id, - exc_info=True, - ) - - def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): - assert meta.remote is not None and self.transfer_topo is not None - engine_id = meta.remote.engine_id - # Update last activity from this remote. Mind that cleanup is done on main - # thread (this one), so we don't race on this structure. - self._engine_last_active[engine_id] = time.perf_counter() - plan = self.tp_mappings[engine_id] - remote_info = self.transfer_topo.get_engine_info(engine_id) - tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) - - meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( - meta.remote.block_ids, - remote_info.remote_physical_blocks_per_logical, - ) - remote_block_ids = meta.remote.block_ids - local_block_ids = meta.local_physical_block_ids - num_groups = len(local_block_ids) - read_specs = [ - ReadSpec( - remote_rank=rank, - local_block_ids=[ - list(local_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - remote_block_ids=[ - list(remote_block_ids[g]) - if rank in plan.source_ranks_per_group[g] - else [] - for g in range(num_groups) - ], - ) - for rank in plan.all_source_ranks - ] - - # D may have to perform multiple reads from different remote ranks. - # MLA opt: when P TP > D TP, only a single read is executed for - # the first remote rank (cache is duplicated).. - if self.use_mla and tp_ratio < 0: - assert len(read_specs) == 1 - - for i, spec in enumerate(read_specs): - remote_block_size = remote_info.remote_block_size - logger.debug( - "Remote agent %s available, calling _read_blocks" - " on remote rank %s with remote block size %s for req %s", - meta.remote.engine_id, - spec.remote_rank, - remote_block_size, - req_id, - ) - # Get side handles. - if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size - # Remote tp_size > local tp_size: we must perform multiple - # reads. Get the memory chunk onto which we will write to. - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] - else: - # Single read from remote, we write to the whole memory region. - # Also handle remote block size different from local block size. - local_xfer_side_handle = self.src_xfer_handles_by_block_size[ - remote_block_size - ] - - # Destination handle: remote_engine_id -> remote_rank -> handle. - remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][ - spec.remote_rank - ] - - self._read_blocks( - read_spec=spec, - request_id=req_id, - dst_engine_id=meta.remote.engine_id, - remote_request_id=meta.remote.request_id, - local_xfer_side_handle=local_xfer_side_handle, - remote_xfer_side_handle=remote_xfer_side_handle, - ) - - if self.use_mla and tp_ratio < 0 and read_specs: - # ..but we still need to notify the other remote ranks that we - # have the blocks we need so they can update the request state. - notif_id = f"{meta.remote.request_id}:{self.world_size}".encode() - remote_agents = self._remote_agents[meta.remote.engine_id] - for rank_to_notify, agent in remote_agents.items(): - if rank_to_notify != read_specs[0].remote_rank: - self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) - - def _read_blocks( - self, - read_spec: ReadSpec, - dst_engine_id: str, - request_id: str, - remote_request_id: str, - local_xfer_side_handle: int, - remote_xfer_side_handle: int, - ): - """ - Post a READ point-to-point xfer request from a single local worker to - a single remote worker. - """ - assert self.transfer_topo is not None - remote_rank = read_spec.remote_rank - local_block_ids = read_spec.local_block_ids - remote_block_ids = read_spec.remote_block_ids - - remote_info = self.transfer_topo.get_engine_info(dst_engine_id) - block_size_ratio = self.transfer_topo.block_size_ratio( - remote_info.remote_block_size - ) - if block_size_ratio > 1: - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - # NOTE: - # get_mapped_blocks will always expand block_ids for n times. - # ex: - # prefill block_ids with block_size as 4: - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - # Local decode block_ids with block_size as 16: [1, 2, 3] - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # Then we clip local to align with prefill - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] - # NOTE(rob): having the staging blocks be on the READER side is - # not going to work well (since we will have to call rearrange tensors). - # after we detect the txn is complete (which means we cannot make the - # read trxn async easily). If we want to make "READ" happen cleanly, - # then we will need to have the staging blocks on the remote side. - - # NOTE(rob): according to nvidia the staging blocks are used to - # saturate IB with heterogeneous TP sizes. - - # Number of D TP workers that will read from dst P. Propagate info - # on notification so that dst worker can wait before freeing blocks. - notif_id = f"{remote_request_id}:{self.world_size}".encode() - - # Full prefix cache hit: do not need to read remote blocks, - # just notify P worker that we have the blocks we need. - if len(local_block_ids) == 0: - # A full prefix cache hit is indicated with an empty list. - agent_name = self._remote_agents[dst_engine_id][remote_rank] - try: - self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id) - except Exception as e: - self._log_failure( - failure_type="notification_failed", - msg="P worker blocks will be freed after timeout. " - "This may indicate network issues.", - req_id=request_id, - error=e, - dst_engine_id=dst_engine_id, - remote_rank=remote_rank, - remote_agent_name=agent_name, - ) - self.xfer_stats.record_failed_notification() - return - - assert ( - len(remote_block_ids) - == len(local_block_ids) - == len(self.kv_cache_config.kv_cache_groups) - ) - remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical - local_block_ids, remote_block_ids = self._apply_prefix_caching( - local_block_ids, remote_block_ids, remote_physical_per_logical - ) - - # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from - # corresponding rank. With heterogeneous TP, fixing D>P, the D tp - # workers will issue xfers to parts of the P worker remote kv caches. - - # Get descs ids. - remote_block_descs_ids = self._compute_desc_ids( - block_ids=remote_block_ids, - dst_num_blocks=self.dst_num_blocks[dst_engine_id], - block_size_ratio=None, - physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical, - ) - local_block_descs_ids = self._compute_desc_ids( - block_ids=local_block_ids, - dst_num_blocks=self.dst_num_blocks[self.engine_id], - block_size_ratio=block_size_ratio, - physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block, - ) - - assert len(local_block_descs_ids) == len(remote_block_descs_ids) - - # Prepare transfer with Nixl. - handle = None - try: - handle = self.nixl_wrapper.make_prepped_xfer( - "READ", - local_xfer_side_handle, - local_block_descs_ids, - remote_xfer_side_handle, - remote_block_descs_ids, - notif_msg=notif_id, - ) - - # Begin async xfer. - self.nixl_wrapper.transfer(handle) - - # Use handle to check completion in future step(). - self._recving_transfers[request_id].append(handle) - except Exception as e: - # mark all (logical) blocks for this request as invalid - self._log_failure( - failure_type="transfer_setup_failed", - req_id=request_id, - msg="Marking blocks as invalid", - error=e, - dst_engine_id=dst_engine_id, - remote_rank=remote_rank, - ) - self._handle_failed_transfer(request_id, handle) - - def get_mapped_blocks( - self, block_ids: np.ndarray, block_size_ratio: int - ) -> np.ndarray: - """ - Calculates the new set of block IDs by mapping every element - in the (potentially sparse) input array. - Example: block_ids=[0, 2], block_size_ratio=2 - get_mapped_blocks 0 1 [2 3] 4 5 - # remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1|| - # local is |h0-b0......||h1-b0......||h2-b0........ - local_block_ids 0 [1] 2 - """ - if block_ids.size == 0: - return np.array([], dtype=np.int64) - - start_ids = block_ids * block_size_ratio - offsets = np.arange(block_size_ratio) - mapped_2d = start_ids[:, None] + offsets[None, :] - - return mapped_2d.flatten().astype(np.int64) - - def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: - """ - Convert logical block ids to kernel physical block ids. - This is required when the logical block size (the one set by the user) - does not match the one required by the attn backend. - """ - if self._physical_blocks_per_logical_kv_block == 1: - # Noop when physical and logical block sizes are the same - return block_ids - block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape( - 1, -1 - ) - # Mamba blocks have no logical<>physical discrepancy - group_specs = self.kv_cache_config.kv_cache_groups - return [ - BlockTable.map_to_kernel_blocks( - np.array(group), - self._physical_blocks_per_logical_kv_block, - block_arange, - ).tolist() - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) - else group - for i, group in enumerate(block_ids) - ] - - def _apply_prefix_caching( - self, - local_block_ids: BlockIds, - remote_block_ids: BlockIds, - remote_physical_per_logical: int, - ) -> tuple[BlockIds, list]: - """Apply prefix caching by trimming local/remote block ID lists. - - For non-Mamba models: end-trim remote to match local count, so that - already-cached prefix blocks are skipped in the transfer. - - For Mamba hybrid (prefix caching not yet supported): front-trim both - to the minimum count to handle kernel block count discrepancies from - logical block rounding in heterogeneous TP. - """ - # Partial prefix cache hit: just read uncomputed blocks. - # Skip mamba groups — their blocks represent full state (conv+ssm), - # not per-token data, so trimming would corrupt the transfer. - remote_block_ids = list(remote_block_ids) - if not self._has_mamba: - for i, remote_group in enumerate(remote_block_ids): - num_local_blocks = len(local_block_ids[i]) - assert num_local_blocks <= len(remote_group) - if num_local_blocks < len(remote_group): - remote_block_ids[i] = remote_group[-num_local_blocks:] - else: - # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP - # can cause different kernel block counts due to logical block rounding. - # Example: 640 prompt tokens, kernel_block_size=64 - # remote physical_per_logical=10, local physical_per_logical=6 - # remote logical ids from kv_transfer_params = [0] - # local logical ids allocated = [0, 1] - # remote kernel blocks: [0..9] (1*10=10) - # local kernel blocks: [0..11] (2*6=12) - # actual data blocks = ceil(640/64) = 10, trim both to 10 - # Vice versa (remote physical_per_logical=6, local=10): - # remote logical ids = [0, 1], local logical ids = [0] - # remote kernel blocks: [0..11] (2*6=12) - # local kernel blocks: [0..9] (1*10=10) - # actual data blocks = ceil(640/64) = 10, trim both to 10 - local_block_ids = list(local_block_ids) - for i, remote_group in enumerate(remote_block_ids): - num_local_blocks = len(local_block_ids[i]) - num_remote_blocks = len(remote_group) - if ( - _is_ssm_spec(self._group_spec_types[i]) - and num_local_blocks < num_remote_blocks - ): - # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks - # prior to the last one are placeholders (null blocks). Mind that - # this doesn't really impact transfer, as we only still care about - # the last "block", the full in-place state. - assert num_local_blocks == 1, "SSM can only have one local block" - remote_block_ids[i] = remote_group[-num_local_blocks:] - elif ( - self._physical_blocks_per_logical_kv_block - == remote_physical_per_logical - and num_local_blocks < num_remote_blocks - ): - # Partial prefix cache hit for FA group. - remote_block_ids[i] = remote_group[-num_local_blocks:] - else: - # TODO Handle prefix caching with different block_sizes - max_padding = max( - self._physical_blocks_per_logical_kv_block, - remote_physical_per_logical, - ) - assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( - f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| >= {max_padding}" - ) - num_blocks = min(num_local_blocks, num_remote_blocks) - local_block_ids[i] = local_block_ids[i][:num_blocks] - remote_block_ids[i] = remote_group[:num_blocks] - return local_block_ids, remote_block_ids - - def _logical_to_remote_kernel_block_ids( - self, block_ids: BlockIds, remote_physical_per_logical: int - ) -> BlockIds: - """Map logical block IDs to physical kernel block IDs on the remote. - - Args: - block_ids: per-group lists of logical block IDs. - remote_physical_per_logical: remote engine's physical blocks - per logical block. - - Returns: - Same structure with FA groups expanded (each logical block L - becomes kernel blocks [L*remote_physical_per_logical, .. - L*remote_physical_per_logical + - remote_physical_per_logical - 1]). - Mamba groups are passed through unchanged. - """ - if remote_physical_per_logical == 1: - return block_ids - remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1) - group_specs = self.kv_cache_config.kv_cache_groups - result = [ - BlockTable.map_to_kernel_blocks( - np.array(group), - remote_physical_per_logical, - remote_arange, - ).tolist() - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) - else group - for i, group in enumerate(block_ids) - ] - return result - - def get_backend_aware_kv_block_len( - self, layer_idx: int, first_split: bool = True, mamba_view: bool = False - ) -> int: - """ - Get the block length for one K/V element (K and V have the same size). - - For FA and other backends, this is equal to the length of the whole - block, as K and V are in separate regions. - For FlashInfer, this is half the length of the whole block, as K and V - share the same region. - Similarly, for SSM-based models, state and conv are interleaved, but crucially - the their size differs. - Reference diagram: - KVCacheTensor (Shared) - / \\ - / \\ - / \\ - Attention (FlashInfer) View Mamba View - | | - | | - +-------------------+ +-------------------+ - | KVCacheTensor | | KVCacheTensor | - | | | | - |<----- page ------>| |<----- page ------->| - | size | | size | - | Key 0 | Val 0 | |Conv 0 | SSM 0 | - | Key 1 | Val 1 | |Conv 1 | SSM 1 | - | ... | ... | | ... | ... | - | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | - | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | - +-------------------+ +--------------------+ - |1st_split-2nd_split| |1st_split-2nd_split | - """ - assert self.transfer_topo is not None - if self.transfer_topo.virtually_split_kv_in_blocks: - if mamba_view: - block_len = self._mamba_ssm_size[not first_split] - else: - block_len = self.block_len_per_layer[layer_idx] // 2 - else: - block_len = self.block_len_per_layer[layer_idx] - return block_len - - def get_kv_connector_stats(self) -> KVConnectorStats | None: - """ - Get the KV transfer stats for the connector. - """ - # Clear stats for next iteration - if not self.xfer_stats.is_empty(): - return self.xfer_stats.clone_and_reset() - return None - - def get_block_ids_with_load_errors(self) -> set[int]: - """ - Return and clear the set of block IDs that failed to load. - - This is called by the scheduler to identify blocks that need - to be retried after a NIXL transfer failure. - """ - # Drain the queue (thread-safe, no lock needed). - result: set[int] = set() - while not self._invalid_block_ids.empty(): - try: - result.update(self._invalid_block_ids.get_nowait()) - except queue.Empty: - break - return result - - def _evict_stale_engines(self) -> None: - """Scan for and evict remote engines that have exceeded their TTL. - - Called from the main thread in when a new remote engine appears. - We can only go OOM as we discover and register a new remote, therefore we make - sure we clean up stale engine data structures before then. This invariant - prevents us from using background threads, though memory usage is not guaranteed - to be "optimal" until a new handshake is performed. - - Engines with active transfers or pending handshakes cannot be stale: - - Active transfers touch _engine_last_active in start_load_kv. - - Pending handshakes don't have an _engine_last_active entry yet - """ - # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number - # of remote engines is registered all at once (adding a background cleanup - # thread wouldnt help either). - # If that scenario is plausible, we can follow up with an LRU eviction policy. - if self._engine_ttl <= 0: - return - - now = time.perf_counter() - for eid, last_active in list(self._engine_last_active.items()): - if now - last_active > self._engine_ttl: - self._cleanup_remote_engine(eid) - - def _cleanup_remote_engine( - self, engine_id: EngineId, *, log_eviction: bool = True - ) -> None: - """Remove all state for a single remote engine. - - Releases NIXL resources (dlist handles, remote agents) and clears - all per-engine data structures. Used by both TTL eviction and - shutdown. - """ - assert engine_id in self._remote_agents - - for handle in self.dst_xfer_side_handles.pop(engine_id).values(): - self.nixl_wrapper.release_dlist_handle(handle) - for agent_name in self._remote_agents.pop(engine_id).values(): - self.nixl_wrapper.remove_remote_agent(agent_name) - - del self.kv_caches_base_addr[engine_id] - del self.dst_num_blocks[engine_id] - del self.tp_mappings[engine_id] - if self.transfer_topo is not None: - self.transfer_topo.unregister_remote_engine(engine_id) - - last_active = self._engine_last_active.pop(engine_id) - if log_eviction: - logger.info( - "Evicted stale remote engine %s (inactive for %.1fs).", - engine_id, - time.perf_counter() - last_active, - ) - - def __del__(self): - self.shutdown() - - def shutdown(self): - """Shutdown the connector worker.""" - if not hasattr(self, "_handshake_initiation_executor"): - # error happens during init, no need to shutdown - return - self._handshake_initiation_executor.shutdown(wait=False) - for handles in self._recving_transfers.values(): - for handle in handles: - self.nixl_wrapper.release_xfer_handle(handle) - self._recving_transfers.clear() - for handle in self.src_xfer_handles_by_block_size.values(): - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_block_size.clear() - for handles in self.src_xfer_handles_by_tp_ratio.values(): - for handle in handles: - self.nixl_wrapper.release_dlist_handle(handle) - self.src_xfer_handles_by_tp_ratio.clear() - for engine_id in list(self._remote_agents): - self._cleanup_remote_engine(engine_id, log_eviction=False) - for desc in self._registered_descs: - self.nixl_wrapper.deregister_memory(desc) - self._registered_descs.clear() +__all__ = ["NixlConnectorWorker", "NixlPullConnectorWorker"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 24e7143e630..1d3d83709be 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -19,7 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( _TransferMetricName, ) from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( @@ -94,6 +94,24 @@ def get_sliding_window_size_in_blocks( return None +def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None: + """Scan all KV cache groups in *spec* and return the single mamba alignment + size, or None if no group requires mamba alignment. + + For MambaSpec groups in "align" cache mode the hit window must be rounded + down to a multiple of the offloaded block size. Asserts that all such + groups agree on the same value. + """ + mamba_align_size: int | None = None + for idx, gpu_block_size in enumerate(spec.gpu_block_size): + kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec + if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": + offload_block_size = gpu_block_size * spec.block_size_factor + assert mamba_align_size is None or mamba_align_size == offload_block_size + mamba_align_size = offload_block_size + return mamba_align_size + + class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] block_size_factor: int @@ -290,6 +308,7 @@ class OffloadingConnectorScheduler: # used by _lookup self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups + self._mamba_align_size: int | None = resolve_mamba_align_size(spec) self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -408,6 +427,12 @@ class OffloadingConnectorScheduler: # for sliding window attention, we must reduce by 1 to make sure # we still have a hit after reduction max_hit_size_tokens -= 1 + if self._mamba_align_size is not None: + # Constrain hit-window to the mamba block size. + max_hit_size_tokens = round_down( + max_hit_size_tokens, self._mamba_align_size + ) + num_hit_tokens: int = 0 defer_lookup = False lookup_groups = self._lookup_groups @@ -571,7 +596,11 @@ class OffloadingConnectorScheduler: req_status.update_offload_keys() req_status.num_locally_computed_tokens = num_computed_tokens - num_hit_tokens = self._lookup(req_status) + num_hit_tokens: int | None + if request.skip_reading_prefix_cache: + num_hit_tokens = 0 + else: + num_hit_tokens = self._lookup(req_status) req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0490cbc3e4b..f863fad17de 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -38,6 +38,7 @@ from vllm.config import ( CompilationConfig, ConfigType, DeviceConfig, + DiffusionConfig, ECTransferConfig, EPLBConfig, KernelConfig, @@ -600,6 +601,8 @@ class EngineArgs: scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl + watermark: float = SchedulerConfig.watermark + disable_hybrid_kv_cache_manager: bool | None = ( SchedulerConfig.disable_hybrid_kv_cache_manager ) @@ -614,6 +617,7 @@ class EngineArgs: spec_method: str | None = None spec_model: str | None = None spec_tokens: int | None = None + diffusion_config: dict[str, Any] | None = None show_hidden_metrics_for_version: str | None = ( ObservabilityConfig.show_hidden_metrics_for_version @@ -1408,6 +1412,7 @@ class EngineArgs: "--scheduler-reserve-full-isl", **scheduler_kwargs["scheduler_reserve_full_isl"], ) + scheduler_group.add_argument("--watermark", **scheduler_kwargs["watermark"]) scheduler_group.add_argument( "--disable-hybrid-kv-cache-manager", **scheduler_kwargs["disable_hybrid_kv_cache_manager"], @@ -1470,6 +1475,10 @@ class EngineArgs: vllm_group.add_argument( "--spec-tokens", **speculative_kwargs["num_speculative_tokens"] ) + vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads) + vllm_group.add_argument( + "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"] + ) vllm_group.add_argument( "--kv-transfer-config", **vllm_kwargs["kv_transfer_config"] ) @@ -1699,6 +1708,14 @@ class EngineArgs: ) return SpeculativeConfig(**self.speculative_config) + def create_diffusion_config(self) -> DiffusionConfig | None: + if self.diffusion_config is None: + return None + cfg = self.diffusion_config + if isinstance(cfg, str): + cfg = json.loads(cfg) + return DiffusionConfig(**cfg) + def create_engine_config( self, usage_context: UsageContext | None = None, @@ -2013,6 +2030,7 @@ class EngineArgs: target_model_config=model_config, target_parallel_config=parallel_config, ) + diffusion_config = self.create_diffusion_config() self._set_default_max_num_seqs_and_batched_tokens_args( usage_context, @@ -2045,6 +2063,7 @@ class EngineArgs: max_long_partial_prefills=self.max_long_partial_prefills, long_prefill_token_threshold=self.long_prefill_token_threshold, scheduler_reserve_full_isl=self.scheduler_reserve_full_isl, + watermark=self.watermark, disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager, async_scheduling=self.async_scheduling, stream_interval=self.stream_interval, @@ -2239,6 +2258,7 @@ class EngineArgs: kernel_config=kernel_config, lora_config=lora_config, speculative_config=speculative_config, + diffusion_config=diffusion_config, structured_outputs_config=self.structured_outputs_config, observability_config=observability_config, compilation_config=compilation_config, diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 8f6cccdb0fc..266a3154212 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -564,6 +564,7 @@ class AnthropicServingMessages(OpenAIServingChat): self.block_signature: str | None = None self.signature_emitted: bool = False self.tool_use_id: str | None = None + self.pending_content: list[str] = [] def reset(self) -> None: self.block_type = None @@ -571,6 +572,7 @@ class AnthropicServingMessages(OpenAIServingChat): self.block_signature = None self.signature_emitted = False self.tool_use_id = None + self.pending_content.clear() def start(self, block: AnthropicContentBlock) -> None: self.block_type = block.type @@ -635,10 +637,30 @@ class AnthropicServingMessages(OpenAIServingChat): state.start(block) return event + def stop_and_flush() -> list[str]: + buffered = list(state.pending_content) + state.pending_content.clear() + events = stop_active_block() + if not buffered: + return events + text = "".join(buffered) + events.append(start_block(AnthropicContentBlock(type="text", text=""))) + pc_chunk = AnthropicStreamEvent( + index=state.block_index, + type="content_block_delta", + delta=AnthropicDelta(type="text_delta", text=text), + ) + pc_data = pc_chunk.model_dump_json(exclude_unset=True) + events.append(wrap_data_with_event(pc_data, "content_block_delta")) + events.extend(stop_active_block()) + return events + async for item in generator: if item.startswith("data:"): data_str = item[5:].strip().rstrip("\n") if data_str == "[DONE]": + for event in stop_and_flush(): + yield event stop_message = AnthropicStreamEvent( type="message_stop", ) @@ -675,7 +697,7 @@ class AnthropicServingMessages(OpenAIServingChat): # last chunk including usage info if len(origin_chunk.choices) == 0: - for event in stop_active_block(): + for event in stop_and_flush(): yield event stop_reason = self.stop_reason_map.get( finish_reason or "stop" @@ -707,7 +729,7 @@ class AnthropicServingMessages(OpenAIServingChat): pass else: if state.block_type != "thinking": - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock( @@ -733,9 +755,13 @@ class AnthropicServingMessages(OpenAIServingChat): if origin_chunk.choices[0].delta.content is not None: if origin_chunk.choices[0].delta.content == "": pass + elif state.block_type == "tool_use": + state.pending_content.append( + origin_chunk.choices[0].delta.content + ) else: if state.block_type != "text": - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock(type="text", text="") @@ -773,7 +799,7 @@ class AnthropicServingMessages(OpenAIServingChat): state.tool_use_id != tool_call.id and tool_name is not None ): - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock( diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index 2dc37b904ae..b38830d6e41 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -2,14 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools +from collections.abc import Callable, Sequence +import torch from tqdm import tqdm from vllm import RequestOutput, TextPrompt, TokensPrompt from vllm.entrypoints.offline_utils import OfflineInferenceMixin from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.sampling_params import BeamSearchParams, SamplingParams +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import ( + BeamSearchParams, + SamplingParams, + StructuredOutputsParams, +) +from vllm.tokenizers import TokenizerLike +from vllm.v1.structured_output.backend_types import StructuredOutputBackend +from vllm.v1.structured_output.request import get_structured_output_key from .utils import ( BeamSearchInstance, @@ -20,6 +30,27 @@ from .utils import ( logger = init_logger(__name__) +# Engine-side cap on `SamplingParams.allowed_token_ids`; keep in sync with +# MAX_NUM_ALLOWED_TOKEN_IDS in vllm/v1/worker/gpu/sample/logit_bias.py. +_MAX_NUM_ALLOWED_TOKEN_IDS = 1024 + + +_bitmask_cache: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + + +def _bitmask_to_token_ids(bitmask_row: torch.Tensor, vocab_size: int) -> list[int]: + """Convert a packed int32 bitmask row to a list of allowed token IDs.""" + if vocab_size not in _bitmask_cache: + indices = torch.arange(vocab_size) + _bitmask_cache[vocab_size] = ( + indices, + indices >> 5, # i // 32 + indices & 31, # i % 32 + ) + indices, word_indices, bit_indices = _bitmask_cache[vocab_size] + mask = ((bitmask_row[word_indices] >> bit_indices) & 1).bool() + return indices[mask].tolist() + class BeamSearchOfflineMixin(OfflineInferenceMixin): """Offline inference for beam search""" @@ -69,10 +100,22 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): if concurrency_limit is None: concurrency_limit = len(engine_inputs) + structured_output_backend: StructuredOutputBackend | None = None + structured_output_key = None + structured_output_bitmask = None + if params.structured_outputs is not None: + ( + structured_output_backend, + structured_output_key, + structured_output_bitmask, + ) = self._init_beam_search_structured_output( + params.structured_outputs, tokenizer + ) + # generate 2 * beam_width candidates at each step # following the huggingface transformers implementation # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa - sampling_params = SamplingParams( + base_sampling_params = SamplingParams( logprobs=2 * beam_width, max_tokens=1, temperature=temperature, @@ -94,77 +137,43 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): ), ) - for prompt_start in range(0, len(instances), concurrency_limit): - instances_batch = instances[prompt_start : prompt_start + concurrency_limit] + try: + for prompt_start in range(0, len(instances), concurrency_limit): + instances_batch = instances[ + prompt_start : prompt_start + concurrency_limit + ] - token_iter = range(max_tokens) - if use_tqdm: - token_iter = tqdm( - token_iter, desc="Beam search", unit="token", unit_scale=False - ) - logger.warning( - "The progress bar shows the upper bound on token steps and " - "may finish early due to stopping conditions. It does not " - "reflect instance-level progress." - ) - for _ in token_iter: - all_beams: list[BeamSearchSequence] = list( - sum((instance.beams for instance in instances_batch), []) - ) - pos = [0] + list( - itertools.accumulate( - len(instance.beams) for instance in instances_batch + token_iter = range(max_tokens) + if use_tqdm: + token_iter = tqdm( + token_iter, + desc="Beam search", + unit="token", + unit_scale=False, ) - ) - instance_start_and_end: list[tuple[int, int]] = list( - zip(pos[:-1], pos[1:]) - ) - - if len(all_beams) == 0: - break - - # only runs for one step - # we don't need to use tqdm here - output = self._render_and_run_requests( - prompts=(beam.get_prompt() for beam in all_beams), - params=self._params_to_seq(sampling_params, len(all_beams)), - output_type=RequestOutput, - lora_requests=[beam.lora_request for beam in all_beams], - use_tqdm=False, - ) - - for (start, end), instance in zip( - instance_start_and_end, instances_batch - ): - instance_new_beams = [] - for i in range(start, end): - current_beam = all_beams[i] - result = output[i] - - if result.outputs[0].logprobs is not None: - # if `result.outputs[0].logprobs` is None, it means - # the sequence is completed because of the - # max-model-len or abortion. we don't need to add - # it to the new beams. - logprobs = result.outputs[0].logprobs[0] - for token_id, logprob_obj in logprobs.items(): - new_beam = BeamSearchSequence( - current_beam.orig_prompt, - tokens=current_beam.tokens + [token_id], - logprobs=current_beam.logprobs + [logprobs], - lora_request=current_beam.lora_request, - cum_logprob=current_beam.cum_logprob - + logprob_obj.logprob, - ) - - if token_id == eos_token_id and not ignore_eos: - instance.completed.append(new_beam) - else: - instance_new_beams.append(new_beam) - sorted_beams = sorted( - instance_new_beams, key=sort_beams_key, reverse=True + logger.warning( + "The progress bar shows the upper bound on token " + "steps and may finish early due to stopping " + "conditions. It does not reflect instance-level " + "progress." ) - instance.beams = sorted_beams[:beam_width] + for _ in token_iter: + should_stop = self._beam_search_step( + instances_batch=instances_batch, + base_sampling_params=base_sampling_params, + eos_token_id=eos_token_id, + ignore_eos=ignore_eos, + beam_width=beam_width, + sort_beams_key=sort_beams_key, + structured_output_backend=structured_output_backend, + structured_output_key=structured_output_key, + structured_output_bitmask=structured_output_bitmask, + ) + if should_stop: + break + finally: + if structured_output_backend is not None: + structured_output_backend.destroy() outputs = [] for instance in instances: @@ -180,3 +189,265 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): outputs.append(BeamSearchOutput(sequences=best_beams)) return outputs + + def _beam_search_step( + self, + instances_batch: list[BeamSearchInstance], + base_sampling_params: SamplingParams, + eos_token_id: int | None, + ignore_eos: bool, + beam_width: int, + sort_beams_key: Callable, + structured_output_backend: StructuredOutputBackend | None, + structured_output_key: tuple | None, + structured_output_bitmask: torch.Tensor | None, + ) -> bool: + """Run one token step of beam search across a batch of instances. + + Returns True if all beams are exhausted and search should stop. + """ + all_beams: list[BeamSearchSequence] = list( + sum((instance.beams for instance in instances_batch), []) + ) + pos = [0] + list( + itertools.accumulate(len(instance.beams) for instance in instances_batch) + ) + instance_start_and_end: list[tuple[int, int]] = list(zip(pos[:-1], pos[1:])) + + if len(all_beams) == 0: + return True + + if structured_output_backend is not None: + assert ( + structured_output_key is not None + and structured_output_bitmask is not None + ) + beam_entries = self._build_beam_sampling_params( + all_beams, + base_sampling_params, + structured_output_backend, + structured_output_key, + structured_output_bitmask, + ) + active_indices = [ + i for i, entry in enumerate(beam_entries) if entry is not None + ] + for i, entry in enumerate(beam_entries): + if entry is None: + beam = all_beams[i] + assert beam.orig_prompt["type"] != "enc_dec" + prompt_len = len(beam.orig_prompt["prompt_token_ids"]) + if len(beam.tokens) > prompt_len: + for (s, e), inst in zip( + instance_start_and_end, + instances_batch, + ): + if s <= i < e: + inst.completed.append(beam) + break + + if not active_indices: + return True + + active_beams = [all_beams[i] for i in active_indices] + active_params: Sequence[SamplingParams | PoolingParams] = [ + beam_entries[i][0] # type: ignore[index] + for i in active_indices + ] + else: + active_indices = list(range(len(all_beams))) + active_beams = all_beams + active_params = self._params_to_seq( # type: ignore[assignment] + base_sampling_params, len(all_beams) + ) + + # only runs for one step + # we don't need to use tqdm here + active_output = self._render_and_run_requests( + prompts=(beam.get_prompt() for beam in active_beams), + params=active_params, + output_type=RequestOutput, + lora_requests=[beam.lora_request for beam in active_beams], + use_tqdm=False, + ) + + output: list[RequestOutput | None] = [None] * len(all_beams) + for idx, active_idx in enumerate(active_indices): + output[active_idx] = active_output[idx] + + # Logprobs are computed from raw logits before + # allowed_token_ids masking, so they may contain + # tokens outside the grammar's allowed set. This filtering is also + # the only grammar enforcement for beams whose allowed set exceeds + # the engine-side allowed_token_ids cap. + allowed_sets: list[set[int] | None] = [None] * len(all_beams) + if structured_output_backend is not None: + for i, entry in enumerate(beam_entries): + if entry is not None: + allowed_sets[i] = set(entry[1]) + + for (start, end), instance in zip(instance_start_and_end, instances_batch): + instance_new_beams = [] + for i in range(start, end): + current_beam = all_beams[i] + result = output[i] + + if result is None: + continue + + if result.outputs[0].logprobs is not None: + # if logprobs is None, the sequence completed + # due to max-model-len or abortion. + logprobs = result.outputs[0].logprobs[0] + allowed = allowed_sets[i] + for token_id, logprob_obj in logprobs.items(): + if allowed is not None and token_id not in allowed: + continue + new_beam = BeamSearchSequence( + current_beam.orig_prompt, + tokens=current_beam.tokens + [token_id], + logprobs=current_beam.logprobs + [logprobs], + lora_request=current_beam.lora_request, + cum_logprob=current_beam.cum_logprob + logprob_obj.logprob, + ) + + if token_id == eos_token_id and not ignore_eos: + instance.completed.append(new_beam) + else: + instance_new_beams.append(new_beam) + sorted_beams = sorted( + instance_new_beams, + key=sort_beams_key, + reverse=True, + ) + instance.beams = sorted_beams[:beam_width] + + return False + + def _init_beam_search_structured_output( + self, + structured_outputs: StructuredOutputsParams, + tokenizer: TokenizerLike, + ) -> tuple[StructuredOutputBackend, tuple, torch.Tensor]: + """Initialize the structured output backend for beam search.""" + vllm_config = self.llm_engine.vllm_config + so_config = vllm_config.structured_outputs_config + if so_config is None: + raise ValueError( + "structured_outputs_config is required for beam search " + "with structured outputs" + ) + + # Resolve the backend name from engine config if not already set. + if not structured_outputs._backend: + structured_outputs._backend = so_config.backend + + backend_name = structured_outputs._backend + vocab_size = self.model_config.get_vocab_size() + + backend: StructuredOutputBackend + if backend_name == "xgrammar": + from vllm.v1.structured_output.backend_xgrammar import ( + XgrammarBackend, + ) + + backend = XgrammarBackend( + vllm_config=vllm_config, + tokenizer=tokenizer, + vocab_size=vocab_size, + ) + elif backend_name == "guidance": + from vllm.v1.structured_output.backend_guidance import ( + GuidanceBackend, + ) + + backend = GuidanceBackend( + vllm_config=vllm_config, + tokenizer=tokenizer, + vocab_size=vocab_size, + ) + elif backend_name == "outlines": + from vllm.v1.structured_output.backend_outlines import ( + OutlinesBackend, + ) + + backend = OutlinesBackend( + vllm_config=vllm_config, + tokenizer=tokenizer, + vocab_size=vocab_size, + ) + elif backend_name == "lm-format-enforcer": + from vllm.v1.structured_output.backend_lm_format_enforcer import ( + LMFormatEnforcerBackend, + ) + + backend = LMFormatEnforcerBackend( + vllm_config=vllm_config, + tokenizer=tokenizer, + vocab_size=vocab_size, + ) + else: + raise ValueError(f"Unsupported structured output backend: {backend_name}") + + structured_output_key = get_structured_output_key(structured_outputs) + bitmask = backend.allocate_token_bitmask(1) + + return backend, structured_output_key, bitmask + + def _build_beam_sampling_params( + self, + beams: list[BeamSearchSequence], + base_params: SamplingParams, + backend: StructuredOutputBackend, + structured_output_key: tuple, + bitmask: torch.Tensor, + ) -> list[tuple[SamplingParams, list[int]] | None]: + """Build per-beam SamplingParams and allowed token IDs from grammar. + + Returns None for beams where the grammar has terminated. + """ + vocab_size = self.model_config.get_vocab_size() + request_type, grammar_spec = structured_output_key + result: list[tuple[SamplingParams, list[int]] | None] = [] + + for beam in beams: + # Fresh grammar per beam, replaying generated tokens. + # Backends don't support cloning grammar state, so + # replay is needed to reconstruct the FSM position. + grammar = backend.compile_grammar(request_type, grammar_spec) + assert beam.orig_prompt["type"] != "enc_dec" + prompt_len = len(beam.orig_prompt["prompt_token_ids"]) + generated_tokens = beam.tokens[prompt_len:] + + if generated_tokens: + grammar.accept_tokens("beam", generated_tokens) + + if grammar.is_terminated(): + result.append(None) + continue + + grammar.fill_bitmask(bitmask, 0) + allowed_ids = _bitmask_to_token_ids(bitmask[0], vocab_size) + + if not allowed_ids: + result.append(None) + continue + + # The engine caps the size of allowed_token_ids. While the + # grammar still allows more tokens than the cap (e.g. inside + # free-form strings), skip the engine-side constraint and rely + # on the logprobs filtering in _beam_search_step instead. + beam_params = SamplingParams( + logprobs=base_params.logprobs, + max_tokens=1, + temperature=base_params.temperature, + allowed_token_ids=( + allowed_ids + if len(allowed_ids) <= _MAX_NUM_ALLOWED_TOKEN_IDS + else None + ), + skip_clone=True, + ) + result.append((beam_params, allowed_ids)) + + return result diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index bd9dfc39311..e1e2ef72bbd 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -308,20 +308,6 @@ async def init_app_state( ) -> None: vllm_config = engine_client.vllm_config - # Propagate enable_in_reasoning to the API-server process. The engine core - # runs in a separate process, so the contextvar that backs - # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers - # call `get_enable_structured_outputs_in_reasoning()` during request - # handling and need to see the real flag, otherwise they silently fall - # back to False and mismatch the engine-side bitmask gating. - from vllm.tool_parsers.structural_tag_registry import ( - set_enable_structured_outputs_in_reasoning, - ) - - set_enable_structured_outputs_in_reasoning( - vllm_config.structured_outputs_config.enable_in_reasoning - ) - if args.tool_call_parser is not None: from vllm.parser.metrics import init_parser_metrics diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index 852a26967a0..2a0b20a3d8f 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -74,7 +74,7 @@ class OpenAIServingChatBatch(OpenAIServingChat): if error_check_ret is not None: return error_check_ret - tool_parser = render.tool_parser + parser = render.parser tool_dicts: list[dict] | None = None all_conversations: list[list[ConversationMessage]] = [] @@ -94,7 +94,7 @@ class OpenAIServingChatBatch(OpenAIServingChat): default_template_content_format=render.chat_template_content_format, default_template_kwargs=render.default_chat_template_kwargs, tool_dicts=tool_dicts, - tool_parser=tool_parser, + parser=parser, ) all_conversations.append(conversation) all_engine_prompts.append(engine_prompts[0]) diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 2da89917a8d..45b79c6a7ef 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -33,10 +33,6 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionStreamResponse, ChatMessage, ) -from vllm.entrypoints.openai.chat_completion.stream_harmony import ( - TokenState, - extract_harmony_streaming_delta, -) from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ErrorResponse, @@ -52,10 +48,6 @@ from vllm.entrypoints.openai.engine.serving import ( clamp_prompt_logprobs, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.parser.harmony_utils import ( - get_streamable_parser_for_assistant, - parse_chat_output, -) from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.tool_calls_utils import ( @@ -135,6 +127,7 @@ class OpenAIServingChat(OpenAIServing): reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, + is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) if ( is_mistral_tool_parser(self.tool_parser) @@ -155,7 +148,6 @@ class OpenAIServingChat(OpenAIServing): if mc.generation_config not in ("auto", "vllm") else getattr(mc, "override_generation_config", {}).get("max_new_tokens") ) - self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss" self.tool_call_id_type = get_tool_call_id_type(self.model_config) # NOTE(woosuk): While OpenAI's chat completion API supports browsing @@ -359,14 +351,6 @@ class OpenAIServingChat(OpenAIServing): assert len(generators) == 1 (result_generator,) = generators - parser: Parser | None = None - if self.parser_cls is not None: - parser = self.parser_cls( - tokenizer, - request.tools, - chat_template_kwargs=chat_template_kwargs, - ) - if request.stream: return self.chat_completion_stream_generator( request, @@ -387,7 +371,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, - parser, + chat_template_kwargs=chat_template_kwargs, ) def get_chat_request_role(self, request: ChatCompletionRequest) -> str: @@ -416,11 +400,6 @@ class OpenAIServingChat(OpenAIServing): finish_reason_sent = [False] * num_choices num_prompt_tokens = 0 num_cached_tokens = None - if self.use_harmony: - harmony_parsers = [ - get_streamable_parser_for_assistant() for _ in range(num_choices) - ] - harmony_tools_streamed = [False] * num_choices tools_streamed = [False] * num_choices if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): @@ -451,6 +430,7 @@ class OpenAIServingChat(OpenAIServing): ] for p in parsers: if p is not None: + # NOTE: HarmonyParser ignores _stream_state (uses its own FSM). p._stream_state.tool_call_id_type = self.tool_call_id_type p._stream_state.history_tool_call_cnt = history_tool_call_cnt else: @@ -580,32 +560,7 @@ class OpenAIServingChat(OpenAIServing): else: logprobs = None - if self.use_harmony: - harmony_parser = harmony_parsers[i] - prev_recipient = harmony_parser.current_recipient - - # Track accumulated content per token with their state - token_states: list[TokenState] = [] - for token_id in output.token_ids: - harmony_parser.process(token_id) - token_delta = harmony_parser.last_content_delta or "" - token_states.append( - TokenState( - harmony_parser.current_channel, - harmony_parser.current_recipient, - token_delta, - ) - ) - delta_text = "".join(delta for _, _, delta in token_states) - cur_channel = harmony_parser.current_channel - - # handle the case where several tokens where generated at once - # including the final token, leading to a delta in the text - # but the current channel to be empty (start state) - if not cur_channel and delta_text: - cur_channel = "final" - else: - delta_text = output.text + delta_text = output.text if ( not delta_text @@ -617,17 +572,7 @@ class OpenAIServingChat(OpenAIServing): delta_message: DeltaMessage | None - if self.use_harmony: - delta_message, tools_streamed_flag = ( - extract_harmony_streaming_delta( - harmony_parser=harmony_parser, - token_states=token_states, - prev_recipient=prev_recipient, - include_reasoning=request.include_reasoning, - ) - ) - harmony_tools_streamed[i] |= tools_streamed_flag - elif parser is not None: + if parser is not None: delta_message = parser.parse_delta( delta_text=delta_text, delta_token_ids=as_list(output.token_ids), @@ -635,8 +580,20 @@ class OpenAIServingChat(OpenAIServing): prompt_token_ids=res.prompt_token_ids, finished=output.finish_reason is not None, ) - if delta_message and delta_message.tool_calls: - tools_streamed[i] = True + if delta_message is not None: + if delta_message.tool_calls: + tools_streamed[i] = True + + if ( + delta_message.reasoning + and not request.include_reasoning + ): + delta_message.reasoning = None + if not ( + delta_message.content or delta_message.tool_calls + ): + delta_message = None + # handle streaming just a content delta (no parsers) else: delta_message = DeltaMessage(content=delta_text) @@ -714,9 +671,7 @@ class OpenAIServingChat(OpenAIServing): # finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. - if (tools_streamed[i] and not tool_choice_function_name) or ( - self.use_harmony and harmony_tools_streamed[i] - ): + if tools_streamed[i] and not tool_choice_function_name: finish_reason_ = "tool_calls" else: finish_reason_ = ( @@ -777,7 +732,7 @@ class OpenAIServingChat(OpenAIServing): completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -840,7 +795,7 @@ class OpenAIServingChat(OpenAIServing): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, - parser: Parser | None = None, + chat_template_kwargs: dict[str, Any] | None = None, ) -> ErrorResponse | ChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -871,7 +826,6 @@ class OpenAIServingChat(OpenAIServing): self._raise_if_error(output.finish_reason, request_id) token_ids = output.token_ids out_logprobs = output.logprobs - tool_call_info = None if request.logprobs and request.top_logprobs is not None: assert out_logprobs is not None, "Did not output logprobs" @@ -885,75 +839,20 @@ class OpenAIServingChat(OpenAIServing): else: logprobs = None - if self.use_harmony: - reasoning, content, _ = parse_chat_output(token_ids) - if not request.include_reasoning: - reasoning = None - - if self.tool_parser is not None: - if tokenizer is None: - raise ValueError( - "Tokenizer not available when `skip_tokenizer_init=True`" - ) - - tool_parser = self.tool_parser(tokenizer, request.tools) - # NOTE: We use token_ids for openai tool parser - tool_call_info = tool_parser.extract_tool_calls( - "", - request=request, - token_ids=token_ids, # type: ignore - ) - content = tool_call_info.content - message = ChatMessage( - role=role, - reasoning=reasoning, - content=content, - tool_calls=tool_call_info.tool_calls, - ) - else: - message = ChatMessage( - role=role, - reasoning=reasoning, - content=content, - ) - - # Encode routed_experts for transport. JSON can't carry raw - # bytes, so we write the ndarray as a ``.npy`` byte stream - # and base64-encode it. ``pybase64`` is ~3x faster than the - # stdlib ``base64`` on large payloads thanks to SIMD. - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode( - "ascii" - ) - - choice_data = ChatCompletionResponseChoice( - index=output.index, - message=message, - logprobs=logprobs, - finish_reason=( - "tool_calls" - if (tool_call_info is not None and tool_call_info.tools_called) - else output.finish_reason - if output.finish_reason - else "stop" - ), - stop_reason=output.stop_reason, - token_ids=( - as_list(output.token_ids) if request.return_token_ids else None - ), - routed_experts=routed_experts_b64, + parser: Parser | None = None + if self.parser_cls is not None: + parser = self.parser_cls( + tokenizer, + request.tools, + chat_template_kwargs=chat_template_kwargs, ) - choices.append(choice_data) - continue if parser is not None: reasoning, content, tool_calls = parser.parse( output.text, request, enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=token_ids, ) if not request.include_reasoning: reasoning = None @@ -1124,7 +1023,10 @@ class OpenAIServingChat(OpenAIServing): completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens ) diff --git a/vllm/entrypoints/openai/chat_completion/stream_harmony.py b/vllm/entrypoints/openai/chat_completion/stream_harmony.py deleted file mode 100644 index 271f8e8c85a..00000000000 --- a/vllm/entrypoints/openai/chat_completion/stream_harmony.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Harmony-specific streaming delta extraction for chat completions. - -This module handles the extraction of DeltaMessage objects from -harmony parser state during streaming chat completions. -""" - -from typing import NamedTuple - -from openai_harmony import StreamableParser - -from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, -) -from vllm.entrypoints.openai.parser.harmony_utils import ( - extract_function_from_recipient, - is_function_recipient, -) - - -class TokenState(NamedTuple): - channel: str | None - recipient: str | None - text: str - - -def extract_harmony_streaming_delta( - harmony_parser: StreamableParser, - token_states: list[TokenState], - prev_recipient: str | None, - include_reasoning: bool, -) -> tuple[DeltaMessage | None, bool]: - """ - Extract a DeltaMessage from harmony parser state during streaming. - - Args: - harmony_parser: The StreamableParser instance tracking parse state - token_states: List of TokenState tuples for each token - prev_recipient: Previous recipient for detecting tool call transitions - include_reasoning: Whether to include reasoning content - - Returns: - A tuple of (DeltaMessage or None, tools_streamed_flag) - """ - - if not token_states: - return None, False - - tools_streamed = False - - # Group consecutive tokens with same channel/recipient - groups: list[TokenState] = [] - - current_channel = token_states[0].channel - current_recipient = token_states[0].recipient - current_text = token_states[0].text - - for i in range(1, len(token_states)): - state = token_states[i] - if state.channel == current_channel and state.recipient == current_recipient: - current_text += state.text - else: - groups.append(TokenState(current_channel, current_recipient, current_text)) - current_channel = state.channel - current_recipient = state.recipient - current_text = state.text - - groups.append(TokenState(current_channel, current_recipient, current_text)) - - # Process each group and create delta messages - delta_message = None - combined_content = "" - combined_reasoning = "" - tool_messages = [] - content_encountered = False - - # Calculate base_index once before the loop - # This counts completed tool calls in messages - base_index = 0 - for msg in harmony_parser.messages: - if msg.recipient and is_function_recipient(msg.recipient): - base_index += 1 - - # If there's an ongoing tool call from previous chunk, - # the next new tool call starts at base_index + 1 - if prev_recipient and is_function_recipient(prev_recipient): - next_tool_index = base_index + 1 - # Ongoing call is at base_index - ongoing_tool_index = base_index - else: - # No ongoing call, next new call is at base_index - next_tool_index = base_index - ongoing_tool_index = None - - for group in groups: - if group.channel == "final": - combined_content += group.text - content_encountered = True - elif group.recipient and is_function_recipient(group.recipient): - opened_new_call = False - if prev_recipient != group.recipient: - # New tool call - emit the opening message - tool_name = extract_function_from_recipient(group.recipient) - tool_messages.append( - DeltaToolCall( - id=make_tool_call_id(), - type="function", - function=DeltaFunctionCall( - name=tool_name, - arguments="", - ), - index=next_tool_index, - ) - ) - opened_new_call = True - prev_recipient = group.recipient - # Increment for subsequent new tool calls - next_tool_index += 1 - - if group.text: - # Stream arguments for the ongoing tool call - if opened_new_call: - # Just opened in this group - tool_call_index = next_tool_index - 1 - else: - # Continuing from previous chunk - # If ongoing_tool_index is None here, it means - # we're continuing a call but prev_recipient - # wasn't a function. Use base_index. - tool_call_index = ( - ongoing_tool_index - if ongoing_tool_index is not None - else base_index - ) - tool_messages.append( - DeltaToolCall( - index=tool_call_index, - function=DeltaFunctionCall(arguments=group.text), - ) - ) - elif group.channel == "commentary" and group.recipient is None: - # Tool call preambles meant to be shown to the user - combined_content += group.text - content_encountered = True - elif group.channel == "analysis" and include_reasoning: - combined_reasoning += group.text - - # Combine all non-empty fields into a single message - if content_encountered or combined_reasoning or tool_messages: - delta_kwargs: dict[str, str | list[DeltaToolCall]] = {} - if content_encountered: - delta_kwargs["content"] = combined_content - if combined_reasoning: - delta_kwargs["reasoning"] = combined_reasoning - if tool_messages: - delta_kwargs["tool_calls"] = tool_messages - tools_streamed = True - delta_message = DeltaMessage(**delta_kwargs) - else: - delta_message = None - - return delta_message, tools_streamed diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index ed85323d806..bd7e26b2b16 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -443,7 +443,7 @@ class OpenAIServingCompletion(OpenAIServing): total_tokens=total_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -583,7 +583,7 @@ class OpenAIServingCompletion(OpenAIServing): if ( self.enable_prompt_tokens_details and last_final_res - and last_final_res.num_cached_tokens + and last_final_res.num_cached_tokens is not None ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=last_final_res.num_cached_tokens diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 771faabe609..82316efb86d 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import datetime -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any from openai.types.responses.tool import Tool @@ -456,65 +456,3 @@ def render_for_completion(messages: list[Message]) -> list[int]: def get_streamable_parser_for_assistant() -> StreamableParser: return StreamableParser(get_encoding(), role=Role.ASSISTANT) - - -def parse_output_into_messages(token_ids: Iterable[int]) -> StreamableParser: - parser = get_streamable_parser_for_assistant() - for token_id in token_ids: - parser.process(token_id) - return parser - - -def parse_chat_output( - token_ids: Sequence[int], -) -> tuple[str | None, str | None, bool]: - """ - Parse the output of a Harmony chat completion into reasoning and final content. - Note that when the `openai` tool parser is used, serving_chat only uses this - for the reasoning content and gets the final content from the tool call parser. - - When the `openai` tool parser is not enabled, or when `GptOssReasoningParser` is - in use,this needs to return the final content without any tool calls parsed. - - Empty reasoning or final content is returned as None instead of an empty string. - """ - parser = parse_output_into_messages(token_ids) - output_msgs = parser.messages - is_tool_call = False # TODO: update this when tool call is supported - - # Get completed messages from the parser - # - analysis channel: hidden reasoning - # - commentary channel without recipient (preambles): visible to user - # - final channel: visible to user - # - commentary with recipient (tool calls): handled separately by tool parser - reasoning_texts = [ - msg.content[0].text for msg in output_msgs if msg.channel == "analysis" - ] - final_texts = [ - msg.content[0].text - for msg in output_msgs - if msg.channel == "final" or (msg.channel == "commentary" and not msg.recipient) - ] - - # Extract partial messages from the parser - if parser.current_channel == "analysis" and parser.current_content: - reasoning_texts.append(parser.current_content) - elif parser.current_channel == "final" and parser.current_content: - final_texts.append(parser.current_content) - elif ( - parser.current_channel == "commentary" - and not parser.current_recipient - and parser.current_content - ): - # Preambles (commentary without recipient) are visible to user - final_texts.append(parser.current_content) - - # Flatten multiple messages into a single string - reasoning: str | None = "\n".join(reasoning_texts) - final_content: str | None = "\n".join(final_texts) - - # Return None instead of empty string since existing callers check for None - reasoning = reasoning or None - final_content = final_content or None - - return reasoning, final_content, is_tool_call diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 51831f60835..5b830cf6dcf 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -102,10 +102,9 @@ from vllm.logprobs import Logprob as SampleLogprob from vllm.logprobs import SampleLogprobs from vllm.lora.request import LoRARequest from vllm.outputs import CompletionOutput -from vllm.parser import ParserManager +from vllm.parser import Parser, ParserManager from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers import ToolParser from vllm.utils import random_uuid from vllm.utils.collection_utils import as_list @@ -191,6 +190,7 @@ class OpenAIServingResponses(OpenAIServing): reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, + is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage @@ -612,8 +612,7 @@ class OpenAIServingResponses(OpenAIServing): default_template_content_format=self.chat_template_content_format, default_template_kwargs=chat_template_kwargs, tool_dicts=tool_dicts, - tool_parser=self.parser.tool_parser_cls if self.parser else None, - reasoning_parser=self.parser.reasoning_parser_cls if self.parser else None, + parser=self.parser, ) return messages, engine_inputs @@ -622,7 +621,7 @@ class OpenAIServingResponses(OpenAIServing): request: ResponsesRequest, messages: list[ResponseInputOutputItem], tool_dicts: list[dict[str, Any]] | None, - tool_parser: type[ToolParser] | None, + parser: type[Parser] | None, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, ): @@ -637,8 +636,7 @@ class OpenAIServingResponses(OpenAIServing): default_template_content_format=chat_template_content_format, default_template_kwargs=chat_template_kwargs, tool_dicts=tool_dicts, - tool_parser=tool_parser, - reasoning_parser=self.parser.reasoning_parser_cls if self.parser else None, + parser=parser, ) return engine_inputs @@ -706,7 +704,7 @@ class OpenAIServingResponses(OpenAIServing): context.request, context.parser.response_messages, context.tool_dicts, - context.parser_cls.tool_parser_cls if context.parser_cls else None, + context.parser_cls, context.chat_template, context.chat_template_content_format, ) diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 72aeb843773..0bb29c68d01 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -307,7 +307,10 @@ class ServingTokens(OpenAIServing): completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): # This info is not available at the /coordinator level usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens @@ -424,7 +427,7 @@ class ServingTokens(OpenAIServing): total_tokens=num_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 9b51bc53daa..6afb26d9843 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -43,8 +43,7 @@ from vllm.inputs import ( tokens_input, ) from vllm.logger import init_logger -from vllm.parser import ParserManager -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs.preprocess import ( extract_prompt_components, @@ -52,7 +51,6 @@ from vllm.renderers.inputs.preprocess import ( parse_model_prompt, prompt_to_seq, ) -from vllm.tool_parsers import ToolParser from vllm.utils import random_uuid from vllm.utils.mistral import is_mistral_tokenizer, is_mistral_tool_parser from vllm.utils.mistral import mt as _mt @@ -89,16 +87,12 @@ class OpenAIServingRender: self.trust_request_chat_template = trust_request_chat_template self.enable_auto_tools = enable_auto_tools self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none - self.tool_parser: type[ToolParser] | None = ParserManager.get_tool_parser( + self.parser: type[Parser] | None = ParserManager.get_parser( tool_parser_name=tool_parser, + reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, model_name=model_config.model, ) - self.reasoning_parser: type[ReasoningParser] | None = ( - ParserManager.get_reasoning_parser( - reasoning_parser_name=reasoning_parser, - ) - ) self.default_chat_template_kwargs: dict[str, Any] = ( default_chat_template_kwargs or {} ) @@ -193,7 +187,7 @@ class OpenAIServingRender: """ tokenizer = self.renderer.tokenizer - tool_parser = self.tool_parser + tool_parser = self.parser.tool_parser_cls if self.parser is not None else None if is_mistral_tokenizer(tokenizer): # because of issues with pydantic we need to potentially @@ -252,9 +246,8 @@ class OpenAIServingRender: default_template_content_format=self.chat_template_content_format, default_template_kwargs=self.default_chat_template_kwargs, tool_dicts=tool_dicts, - tool_parser=tool_parser, + parser=self.parser, skip_mm_cache=skip_mm_cache, - reasoning_parser=self.reasoning_parser, ) else: # For GPT-OSS. @@ -526,8 +519,7 @@ class OpenAIServingRender: default_template_content_format: ChatTemplateContentFormatOption, default_template_kwargs: dict[str, Any] | None, tool_dicts: list[dict[str, Any]] | None = None, - tool_parser: type[ToolParser] | None = None, - reasoning_parser: type[ReasoningParser] | None = None, + parser: type[Parser] | None = None, *, skip_mm_cache: bool = False, ) -> tuple[list[ConversationMessage], list[EngineInput]]: @@ -567,14 +559,6 @@ class OpenAIServingRender: skip_mm_cache=skip_mm_cache, ) - if reasoning_parser is not None: - tokenizer = renderer.get_tokenizer() - request = reasoning_parser( - tokenizer, - model_config=self.model_config, - chat_template_kwargs=chat_params.chat_template_kwargs, - ).adjust_request(request=request) - # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser # is set, we want to prevent parsing a tool_call hallucinated by the LLM @@ -582,15 +566,22 @@ class OpenAIServingRender: # Exception: Mistral grammar-capable tokenizers always call # adjust_request — even for tool_choice="none" — so that the grammar # factory can prevent special-token leakage. - if tool_parser is not None: - tool_choice = getattr(request, "tool_choice", "none") + if parser is not None: tokenizer = renderer.get_tokenizer() + tool_parser = parser.tool_parser_cls + tool_choice = getattr(request, "tool_choice", "none") is_mistral_grammar_eligible = ( - is_mistral_tool_parser(tool_parser) + tool_parser is not None + and is_mistral_tool_parser(tool_parser) and is_mistral_tokenizer(tokenizer) and tokenizer.supports_grammar ) - if tool_choice != "none" or is_mistral_grammar_eligible: + should_adjust_request = ( + parser.reasoning_parser_cls is not None + or tool_choice != "none" + or is_mistral_grammar_eligible + ) + if should_adjust_request: if not isinstance(request, ChatCompletionRequest | ResponsesRequest): msg = ( "Tool usage is only supported " @@ -598,8 +589,13 @@ class OpenAIServingRender: f"but got {type(request).__name__}" ) raise NotImplementedError(msg) - request = tool_parser(tokenizer, request.tools).adjust_request( - request=request + request = parser( + tokenizer, + request.tools, + model_config=self.model_config, + chat_template_kwargs=chat_params.chat_template_kwargs, + ).adjust_request( + request=request, ) return conversation, [engine_input] diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index 3b6dfde447e..d24d492b61e 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -474,6 +474,13 @@ async def lifespan(app: FastAPI): finally: if task is not None: task.cancel() + for attr_name in ( + "openai_serving_transcription", + "openai_serving_translation", + ): + serving = getattr(app.state, attr_name, None) + if serving is not None and hasattr(serving, "shutdown"): + serving.shutdown() finally: # Ensure app state including engine ref is gc'd del app.state diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 1c6a0d77fe2..b60ac6ff95b 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -6,6 +6,7 @@ import math import time import zlib from collections.abc import AsyncGenerator, Callable, Set +from concurrent.futures import ThreadPoolExecutor from functools import cached_property from typing import Final, Literal, TypeAlias, TypeVar, cast @@ -37,7 +38,7 @@ from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import get_tokenizer -from vllm.utils.async_utils import merge_async_iterators +from vllm.utils.async_utils import make_async_with_semaphore, merge_async_iterators from ..transcription.protocol import ( TranscriptionResponse, @@ -63,6 +64,7 @@ T = TypeVar("T", bound=SpeechToTextResponse) V = TypeVar("V", bound=SpeechToTextResponseVerbose) S = TypeVar("S", bound=SpeechToTextSegment) + ResponseType: TypeAlias = ( TranscriptionResponse | TranslationResponse @@ -131,6 +133,19 @@ class OpenAISpeechToText(OpenAIServing): self.default_sampling_params, ) + # setup preprocess resources + # we keep separate thread pool for frontend preprocessing instead + # of reusing the one from Renderer which showed lower throughput + # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 + num_audio_preprocess_workers = envs.VLLM_MAX_AUDIO_PREPROCESS_WORKERS + self._preprocess_executor = ThreadPoolExecutor( + max_workers=num_audio_preprocess_workers, + thread_name_prefix="stt-preprocess", + ) + self._decode_and_chunk_speech_async = make_async_with_semaphore( + self._decode_and_chunk_speech, executor=self._preprocess_executor + ) + @cached_property def model_cls(self) -> type[SupportsTranscription]: from vllm.model_executor.model_loader import get_model_cls @@ -138,6 +153,51 @@ class OpenAISpeechToText(OpenAIServing): model_cls = get_model_cls(self.model_config) return cast(type[SupportsTranscription], model_cls) + def shutdown(self) -> None: + self._preprocess_executor.shutdown(wait=False) + + def _decode_and_chunk_speech( + self, + audio_data: bytes, + ) -> tuple[list[np.ndarray], float]: + # Decode audio bytes. For container formats (MP4, M4A, WebM) that + # soundfile cannot detect from a BytesIO stream, _load_audio_bytes + # transparently falls back to ffmpeg via an in-memory fd. + # NOTE resample to model SR here for efficiency. This is also a + # pre-requisite for chunking, as it assumes Whisper SR. + try: + with io.BytesIO(audio_data) as buf: + y, sr = load_audio( + buf, + sr=self.asr_config.sample_rate, + max_duration_s=self.max_audio_decode_duration_s, + ) + except ValueError: + raise + except Exception as exc: + raise ValueError("Invalid or unsupported audio file.") from exc + + duration = get_audio_duration(y=y, sr=sr) + do_split_audio = self.asr_config.allow_audio_chunking and ( + self.asr_config.max_audio_clip_s is not None + and duration > self.asr_config.max_audio_clip_s + ) + + if not do_split_audio: + chunks = [y] + else: + assert self.asr_config.max_audio_clip_s is not None + assert self.asr_config.min_energy_split_window_size is not None + chunks = split_audio( + audio_data=y, + sample_rate=int(sr), + max_clip_duration_s=self.asr_config.max_audio_clip_s, + overlap_duration_s=self.asr_config.overlap_chunk_second, + min_energy_window_size=self.asr_config.min_energy_split_window_size, + ) + + return chunks, duration + async def _detect_language( self, audio_chunk: np.ndarray, @@ -210,39 +270,8 @@ class OpenAISpeechToText(OpenAIServing): value=len(audio_data) / 1024**2, ) - # Decode audio bytes. For container formats (MP4, M4A, WebM) that - # soundfile cannot detect from a BytesIO stream, _load_audio_bytes - # transparently falls back to ffmpeg via an in-memory fd. - # NOTE resample to model SR here for efficiency. This is also a - # pre-requisite for chunking, as it assumes Whisper SR. - try: - with io.BytesIO(audio_data) as buf: - y, sr = load_audio( - buf, - sr=self.asr_config.sample_rate, - max_duration_s=self.max_audio_decode_duration_s, - ) - except Exception as exc: - raise ValueError("Invalid or unsupported audio file.") from exc - - duration = get_audio_duration(y=y, sr=sr) - do_split_audio = self.asr_config.allow_audio_chunking and ( - self.asr_config.max_audio_clip_s is not None - and duration > self.asr_config.max_audio_clip_s - ) - - if not do_split_audio: - chunks = [y] - else: - assert self.asr_config.max_audio_clip_s is not None - assert self.asr_config.min_energy_split_window_size is not None - chunks = split_audio( - audio_data=y, - sample_rate=int(sr), - max_clip_duration_s=self.asr_config.max_audio_clip_s, - overlap_duration_s=self.asr_config.overlap_chunk_second, - min_energy_window_size=self.asr_config.min_energy_split_window_size, - ) + # Run cpu intensive preprocess step in a separate thread pool executor. + chunks, duration = await self._decode_and_chunk_speech_async(audio_data) if request.language is None and getattr( self.model_cls, "supports_explicit_language_detection", False diff --git a/vllm/envs.py b/vllm/envs.py index d0133638f16..dfebcd27ae8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -78,6 +78,7 @@ if TYPE_CHECKING: VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 VLLM_MAX_AUDIO_DECODE_DURATION_S: int = 600 + VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" VLLM_MM_HASHER_ALGORITHM: str = "blake3" @@ -199,6 +200,7 @@ if TYPE_CHECKING: MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 + VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = True VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None @@ -226,7 +228,6 @@ if TYPE_CHECKING: VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False VLLM_SYSTEM_START_DATE: str | None = None VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: bool = False - VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = False VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True @@ -928,6 +929,15 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MAX_AUDIO_DECODE_DURATION_S": lambda: int( os.getenv("VLLM_MAX_AUDIO_DECODE_DURATION_S", "600") ), + # Maximum number of worker threads used for STT preprocessing. The default + # intentionally caps at 2 because that performed best in profiling. + # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS": lambda: int( + os.getenv( + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", + str(max(1, min(os.cpu_count() or 1, 2))), + ) + ), # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. @@ -1526,6 +1536,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS": lambda: int( os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1") ), + # Enforce function parameter schemas in structural-tag based tool calling. + "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: os.getenv( + "VLLM_ENFORCE_STRICT_TOOL_CALLING", "True" + ).lower() + in ("true", "1"), # Control the max chunk bytes (in MB) for the rpc message queue. # Object larger than this threshold will be broadcast to worker # processes via zmq. @@ -1649,12 +1664,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY": lambda: bool( int(os.getenv("VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY", "0")) ), - # When 1,the model structural tags will be used to enforce the model - # output conforming to the model's tool-calling format and schema. - # Default 0 (off). - "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: bool( - int(os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "0")) - ), # Add optional custom scopes for profiling, disable to avoid overheads "VLLM_CUSTOM_SCOPES_FOR_PROFILING": lambda: bool( int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0")) @@ -1997,6 +2006,7 @@ def compile_factors() -> dict[str, object]: "VLLM_MEDIA_LOADING_THREAD_COUNT", "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", "VLLM_MAX_AUDIO_DECODE_DURATION_S", + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", "VLLM_VIDEO_LOADER_BACKEND", "VLLM_MEDIA_CONNECTOR", "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", diff --git a/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json new file mode 100644 index 00000000000..23f68e88c6e --- /dev/null +++ b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json @@ -0,0 +1,1938 @@ +[ + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json new file mode 100644 index 00000000000..08a0d97ccf2 --- /dev/null +++ b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json @@ -0,0 +1,1893 @@ +[ + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 128, + "maxnreg": 256 + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py new file mode 100644 index 00000000000..8b73fac4b8e --- /dev/null +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product +from typing import Any + +import torch + +from vllm.kernels.helion.case_key import CaseKey +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion + +if not has_helion(): + raise ImportError( + "Helion kernel requires helion to be installed. " + "Install it with: pip install helion" + ) + +import helion +import helion.language as hl + +from vllm.kernels.helion.register import register_kernel + +logger = init_logger(__name__) + + +def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: + # TODO(xiaohongchen1991): it is difficult for kernel author to cover all + # input property combination. Currently, dtypes are fixed. We need + # optimization to bucket/skip some combinations + num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + hidden_size_list = [2048, 4096, 5120] + group_size_list = [128] + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + + use_ue8m0 = False + column_major = False + fp8_min, fp8_max = get_fp8_min_max() + eps = 1e-10 + + inputs = {} + + for hidden_size, group_size, num_tokens in product( + hidden_size_list, group_size_list, num_tokens_list + ): + input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) + output_q = torch.empty(input.shape, device=input.device, dtype=out_dtype) + output_s = torch.empty( + (num_tokens, hidden_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + config_key = CaseKey( + { + "hidden_size": hidden_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = ( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + use_ue8m0, + column_major, + False, + ) + + return inputs + + +_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} + + +def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: + """Pick the best pre-tuned config for the given input shape. + + Selection strategy: + 1. Find the closest hidden_size among available configs + (exact match preferred). + 2. Find the closest group_size among available configs + (exact match preferred). + 3. Among the num_tokens values tuned for that hidden_size and group_size, pick + the smallest num_tokens >= the input's num_tokens. If the input is + larger than all available num_tokens, fall back to the largest. + """ + + if not config_keys: + return None + + input, _, _, group_size, *_ = args + num_tokens, hidden_size = input.shape + + cache_key = (num_tokens, group_size, hidden_size) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, dict[int, list[int]]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["hidden_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) + best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size)) + available_num_tokens = sorted(configs[best_hidden_size][best_group_size]) + best_num_tokens = next( + (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] + ) + + result = CaseKey( + { + "hidden_size": best_hidden_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + # Unused dummy args + # Kept for consistency with existing kernel interface + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + return + + +def baseline( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + torch.ops._C.per_token_group_fp8_quant( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + dummy_is_scale_transposed, + dummy_is_tma_aligned, + ) + + +@register_kernel( + mutates_args=["output_q", "output_s"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + ), +) # type: ignore[misc] +def per_token_group_fp8_quant( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + # Unused dummy args + # Kept for consistency with existing kernel interface + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, hidden_size = input.shape + hl.specialize(hidden_size) + hl.specialize(group_size) + + groups_per_row = output_s.shape[1] + hl.specialize(groups_per_row) + assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row + assert output_s.ndim == 2 and output_s.dtype == torch.float32 + + input = input.view(num_tokens, -1, group_size) + output_q = output_q.view(num_tokens, -1, group_size) + for tile_m, tile_gn, tile_n in hl.tile( + [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] + ): + x_blk = input[tile_m, tile_gn, tile_n] + y_s_blk = torch.clamp(torch.amax(torch.abs(x_blk), dim=-1), min=eps) + y_s_blk = y_s_blk / fp8_max + + if scale_ue8m0: + y_s_blk = torch.exp2(torch.ceil(torch.log2(y_s_blk))) + + y_q_blk = torch.clamp(x_blk / y_s_blk[:, :, None], fp8_min, fp8_max).to( + output_q.dtype + ) + + output_s[tile_m, tile_gn] = y_s_blk + output_q[tile_m, tile_gn, tile_n] = y_q_blk diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index f18120da45f..764022de77d 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -260,6 +260,7 @@ class HelionKernelWrapper: op_name: str, fake_impl: Callable, config_picker: ConfigPicker, + mutates_args: list[str] | None = None, helion_settings: helion.Settings | None = None, input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, ): @@ -272,6 +273,7 @@ class HelionKernelWrapper: self.helion_settings = helion_settings self._config_picker = config_picker self._input_generator = input_generator + self._mutates_args = mutates_args self._configured_kernel: ConfiguredHelionKernel | None = None # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR, # which handles op enablement/disablement. @@ -357,7 +359,7 @@ class HelionKernelWrapper: direct_register_custom_op( op_name=self.op_name, op_func=configured_kernel._decorated_kernel, - mutates_args=None, + mutates_args=self._mutates_args, fake_impl=self._fake_impl, target_lib=vllm_helion_lib, ) @@ -402,6 +404,7 @@ def register_kernel( *, config_picker: ConfigPicker, fake_impl: Callable | None = None, + mutates_args: list[str] | None = None, helion_settings: helion.Settings | None = None, input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, ) -> Callable[[Callable], HelionKernelWrapper]: @@ -455,6 +458,7 @@ def register_kernel( op_name=final_op_name, fake_impl=final_fake_impl, config_picker=config_picker, + mutates_args=mutates_args, helion_settings=helion_settings, input_generator=input_generator, ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b04edcc513c..b067cdd00e5 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -14,7 +14,7 @@ MLA has two possible ways of computing, a data-movement friendly approach and a compute friendly approach. We generally want to use the compute friendly approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near 1) and the data-movement friendly approach for "decode" (i.e. the ratio -Sq / Skv is small). +Sq / Skv is small, often near 0). NOTE what we deem small and large is currently determined by if it is labelled prefill or decode by the scheduler, but this is something we should probably @@ -28,7 +28,7 @@ Deepseek's MLA attention works the following way: * For decode (i.e. the memory friendly approach) the attention "simulates" a multi-head attention, while the compute is similar to multi-query attention. -Below is example of both paths assuming batchsize = 1 +Below is an example of both paths assuming batch size = 1 ## More Extent Definitions: @@ -77,13 +77,13 @@ v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V) // MHA with QK headdim = P + R // V headdim = V -// spda_o shape [Sq, N, V] -spda_o = scaled_dot_product_attention( +// sdpa_o shape [Sq, N, V] +sdpa_o = scaled_dot_product_attention( torch.cat([q_nope, q_pe], dim=-1), torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1), v ) -return spda_o @ W_O +return sdpa_o @ W_O NOTE: in the actual code, `kv_b_proj` is [W_UK; W_UV] concatenated per head @@ -105,16 +105,16 @@ k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0) // MQA with QK headdim = Lkv + R // V headdim = Lkv -// spda_o shape [Sq, N, Lkv] +// sdpa_o shape [Sq, N, Lkv] // NOTE: this is less compute-friendly since Lkv > P // but is more data-movement friendly since its MQA vs MHA -spda_o = scaled_dot_product_attention( +sdpa_o = scaled_dot_product_attention( torch.cat([ql_nope, q_pe], dim=-1), torch.cat([kv_c, k_pe], dim=-1), kv_c ) -o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV) +o = einsum("snl,lnv->snv", sdpa_o.reshape(-1, N, Lkv), W_UV) return o.view(-1, N * V) @ W_O @@ -153,7 +153,7 @@ curr_o, curr_lse = scaled_dot_product_attention( torch.cat([q_nope, q_pe], dim=-1), torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1), new_v, - casual=True, + causal=True, return_softmax_lse=True ) @@ -173,7 +173,7 @@ for chunk_idx in range(cdiv(C, MCC)): cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)], dim=-1), cache_v_chunk, - casual=False, + causal=False, return_softmax_lse=True ) diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 1731cc26bc3..2ca051ad9e4 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -396,8 +396,9 @@ class MMEncoderAttention(CustomOp): if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): raise ValueError( "mm_encoder_attn_dtype='fp8' requires the FlashInfer " - "cuDNN backend with cuDNN >= 9.17.1 on a GPU with native " - "FP8 support." + "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) " + "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is " + "not available on Hopper (H100/H200) or earlier." ) self.fp8_enabled = True diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 84740fc0570..11ed775f28e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -5,7 +5,12 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm._custom_ops import CPUQuantMethod, fused_experts_cpu +from vllm._custom_ops import ( + CPUQuantAlgo, + CPUQuantMethod, + convert_weight_packed_scale_zp, + fused_experts_cpu, +) from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -17,6 +22,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, kFp8Static128BlockSym, + kInt4Static, kMxfp4Static, ) from vllm.platforms import current_platform @@ -318,3 +324,209 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): limit, True, # is_vnni ) + + +def prepare_int4_moe_layer_for_cpu( + w13_packed: torch.Tensor, + w2_packed: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + quant_algo: CPUQuantAlgo = CPUQuantAlgo.GPTQ, + w13_zeros: torch.Tensor | None = None, + w2_zeros: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Repack INT4 MoE weights via convert_weight_packed_scale_zp for CPU. + + Args: + w13_packed: [E, K//8, 2*I] int32 (packed int4) + w2_packed: [E, I//8, K] int32 (packed int4) + w13_scale: [E, num_groups, 2*I] float16/bf16 + w2_scale: [E, num_groups, K] float16/bf16 + quant_algo: CPUQuantAlgo.GPTQ or CPUQuantAlgo.AWQ + w13_zeros: optional [E, num_groups, N//8] int32 packed zeros. + If None, synthetic zeros are created for symmetric quant. + w2_zeros: optional [E, num_groups, N//8] int32 packed zeros. + If None, synthetic zeros are created for symmetric quant. + + Returns: + (blocked_w13, blocked_w2, blocked_s13, blocked_s2, + blocked_z13, blocked_z2) + """ + E = w13_packed.size(0) + + # No qzeros are available in compressed-tensors symmetric checkpoints. + # The GPTQ unpack kernel (unpack_4bit_to_32bit_signed) adds +1 to stored zeros, + # so we store 7 per nibble: 0x77777777 → +1 → 8. + if w13_zeros is None: + num_groups_w13 = w13_scale.size(1) + N_w13 = w13_scale.size(2) # 2*I + _zp = 0x77777777 + w13_zeros = torch.full( + (E, num_groups_w13, N_w13 // 8), + _zp, + dtype=torch.int32, + ) + + if w2_zeros is None: + num_groups_w2 = w2_scale.size(1) + N_w2 = w2_scale.size(2) # K + _zp = 0x77777777 + w2_zeros = torch.full( + (E, num_groups_w2, N_w2 // 8), + _zp, + dtype=torch.int32, + ) + + blocked_w13, blocked_z13, blocked_s13 = convert_weight_packed_scale_zp( + w13_packed, w13_zeros, w13_scale, quant_algo + ) + blocked_w2, blocked_z2, blocked_s2 = convert_weight_packed_scale_zp( + w2_packed, w2_zeros, w2_scale, quant_algo + ) + return (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2) + + +class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): + """CPU INT4 W4A16 group-quantized monolithic MoE experts. + + Weights are int4 (packed), activations are bf16/fp16. + Internally uses int8 compute via fused_experts_cpu with INT4_W4A8. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int | None = None, + num_dispatchers: int | None = None, + ): + super().__init__( + moe_config, + quant_config, + ) + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cpu() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kInt4Static, None), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Default, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + # grouped topk + fused topk bias parameters + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + if apply_router_weight_on_input: + raise NotImplementedError( + "CPUExpertsInt4 (W4A16) does not support " + "apply_router_weight_on_input=True. " + ) + + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import ( + select_experts, + ) + + topk_weights, topk_ids = select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + use_grouped_topk=num_expert_group is not None, + top_k=self.moe_config.experts_per_token, + renormalize=self.moe_config.routing_method + in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ), + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func="softmax", + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + e_score_correction_bias=e_score_correction_bias, + ) + + return fused_experts_cpu( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + False, # inplace + CPUQuantMethod.INT4_W4A8, + self.w1_scale, + self.w2_scale, + self.w1_zp, + self.w2_zp, + None, # block_size + None, # w1_bias + None, # w2_bias + None, # alpha + None, # limit + True, # is_vnni + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index ff259c828f4..76cd15ff5a0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -188,6 +188,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): def _supports_activation(activation: MoEActivation) -> bool: return activation in [ MoEActivation.SILU, + MoEActivation.GELU_TANH, MoEActivation.RELU2_NO_MUL, MoEActivation.SWIGLUOAI, ] @@ -267,6 +268,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): activation_str_to_value_map = { MoEActivation.SILU: ActivationType.Swiglu, # This is the default + MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 5c2aa455600..bd9b285fe74 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -351,6 +351,21 @@ def rocm_aiter_fused_experts( intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) ) + # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs + # for interleaved vs separated gate and up weights. + # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by + # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`, + # which always sets `is_guinterleave=True`. + # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling. + gate_mode = "" + if quant_config.use_mxfp4_w4a16: + try: + from aiter.ops.flydsl.moe_common import GateMode + + gate_mode = GateMode.INTERLEAVE.value + except ImportError: + pass + return rocm_aiter_ops.fused_moe( hidden_states, w1, @@ -369,6 +384,7 @@ def rocm_aiter_fused_experts( output_dtype=output_dtype, hidden_pad=hidden_pad, intermediate_pad=intermediate_pad, + gate_mode=gate_mode, bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index e90c4d6646e..e45fc77ad90 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -142,6 +142,7 @@ class TrtLlmNvFp4ExpertsBase: MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, + MoEActivation.GELU_TANH, ] @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 6ad60d62e97..8de6269e2e9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -45,6 +45,7 @@ logger = init_logger(__name__) class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" + CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" @@ -65,6 +66,12 @@ def backend_to_kernel_cls( ) return [XPUExpertsWNA16] + elif backend == WNA16MoEBackend.CPU: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsInt4, + ) + + return [CPUExpertsInt4] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -73,6 +80,8 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: """ Get available backends in priority order based on platform and config. """ + if current_platform.is_cpu(): + return [WNA16MoEBackend.CPU] if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] @@ -210,17 +219,21 @@ def make_wna16_moe_kernel( from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + CPUExpertsInt4, + ) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, ) - # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts - # and BatchedMarlinExperts + # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, + # BatchedMarlinExperts, XPUExpertsWNA16, and CPUExpertsInt4 assert experts_cls in ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, + CPUExpertsInt4, ) is_monolithic = experts_cls.is_monolithic() @@ -683,6 +696,117 @@ def _process_awq_weights_marlin( ) +def _process_weights_cpu( + quant_config: QuantizationConfig | QuantizationArgs | None, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_g_idx: torch.Tensor | None = None, + w2_g_idx: torch.Tensor | None = None, + w13_qzeros: torch.Tensor | None = None, + w2_qzeros: torch.Tensor | None = None, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor | None, # w13_g_idx + torch.Tensor | None, # w2_g_idx + torch.Tensor | None, # w13_g_idx_sort_indices + torch.Tensor | None, # w2_g_idx_sort_indices + torch.Tensor | None, # w13_qzeros + torch.Tensor | None, # w2_qzeros + torch.Tensor | None, # w13_input_global_scale + torch.Tensor | None, # w2_input_global_scale + torch.Tensor | None, # w13_bias + torch.Tensor | None, # w2_bias +]: + """CPU INT4 W4A16 weight post-processing.""" + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + prepare_int4_moe_layer_for_cpu, + ) + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + from vllm.model_executor.layers.quantization.awq_marlin import ( + AWQMarlinConfig, + ) + + # Detect packing format. + # AWQ: qweight is [E, K, 2*N//8] (packed along output/N dim). + # GPTQ: qweight is [E, K//8, 2*N] (packed along input/K dim). + # compressed-tensors: qweight is [E, K//8, 2*N] (packed along input/K dim). + if isinstance(quant_config, AWQMarlinConfig): + # AWQ: K is stored unpacked in dim 1. + cpu_quant_algo = ops.CPUQuantAlgo.AWQ + elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)): + # GPTQ / compressed-tensors: K//8 is stored packed in dim 1. + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act: + raise NotImplementedError( + "CPU WNA16 MoE backend does not support GPTQ with " + "desc_act=True. The fused MoE kernel has no g_idx " + "reordering support." + ) + cpu_quant_algo = ops.CPUQuantAlgo.GPTQ + else: + raise TypeError( + "CPU WNA16 MoE backend requires AWQMarlinConfig, AutoGPTQConfig " + f"or QuantizationArgs, got {type(quant_config).__name__}." + ) + + # Determine zero points for repacking. + w13_zeros: torch.Tensor | None = None + w2_zeros: torch.Tensor | None = None + if w13_qzeros is not None: + w13_zeros = ( + w13_qzeros.data.view(torch.int32) + if w13_qzeros.dtype != torch.int32 + else w13_qzeros.data + ) + if w2_qzeros is not None: + w2_zeros = ( + w2_qzeros.data.view(torch.int32) + if w2_qzeros.dtype != torch.int32 + else w2_qzeros.data + ) + + ( + blocked_w13, + blocked_w2, + blocked_s13, + blocked_s2, + blocked_z13, + blocked_z2, + ) = prepare_int4_moe_layer_for_cpu( + w13, + w2, + w13_scale, + w2_scale, + quant_algo=cpu_quant_algo, + w13_zeros=w13_zeros, + w2_zeros=w2_zeros, + ) + return ( + blocked_w13, + blocked_w2, + blocked_s13, + blocked_s2, + w13_g_idx, + w2_g_idx, + None, # w13_g_idx_sort_indices (unused on CPU) + None, # w2_g_idx_sort_indices (unused on CPU) + blocked_z13, + blocked_z2, + None, # w13_input_global_scale + None, # w2_input_global_scale + w13_bias.to(torch.float32) if w13_bias is not None else None, + w2_bias.to(torch.float32) if w2_bias is not None else None, + ) + + def _process_weights_xpu( layer: torch.nn.Module, quant_config: QuantizationConfig, @@ -857,6 +981,20 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) + elif backend == WNA16MoEBackend.CPU: + return _process_weights_cpu( + quant_config, + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + w13_qzeros, + w2_qzeros, + w13_bias, + w2_bias, + ) elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return _process_weights_flashinfer( w13, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 1821fd5c7f7..459a6158327 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -16,6 +16,7 @@ from vllm.model_executor.kernels.linear import ( ) from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, + FusedMoEExpertsModular, FusedMoEMethodBase, FusedMoEQuantConfig, FusedMoeWeightScaleSupported, @@ -640,8 +641,11 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - device = layer.w13_qweight.device - layer.workspace = marlin_make_workspace_new(device, 4) + if self.experts_cls is not None and issubclass( + self.experts_cls, FusedMoEExpertsModular + ): + device = layer.w13_qweight.device + layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1 @@ -660,8 +664,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2_g_idx, w13_g_idx_sort_indices, w2_g_idx_sort_indices, - _w13_qzeros, - _w2_qzeros, + w13_qzeros, + w2_qzeros, w13_input_global_scale, w2_input_global_scale, w13_bias, @@ -689,6 +693,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): replace_parameter(layer, "w2_g_idx", w2_g_idx) replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + if w13_qzeros is not None: + replace_parameter(layer, "w13_qzeros", w13_qzeros) + if w2_qzeros is not None: + replace_parameter(layer, "w2_qzeros", w2_qzeros) if w13_input_global_scale is not None: if hasattr(layer, "w13_input_global_scale"): replace_parameter( @@ -735,8 +743,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): is_k_full=self.is_k_full, w13_g_idx=layer.w13_g_idx, w2_g_idx=layer.w2_g_idx, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, + w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None), + w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None), routing_tables=layer._expert_routing_tables(), ) @@ -750,12 +758,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2_scale=layer.w2_scales, weight_bits=self.quant_config.weight_bits, group_size=self.quant_config.group_size, - w1_zp=getattr(layer, "w13_qzeros", None) - if not self.quant_config.is_sym - else None, - w2_zp=getattr(layer, "w2_qzeros", None) - if not self.quant_config.is_sym - else None, + w1_zp=getattr(layer, "w13_qzeros", None), + w2_zp=getattr(layer, "w2_qzeros", None), w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), ) @@ -794,3 +798,27 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + hidden_states=x, + w1=layer.w13_qweight, + w2=layer.w2_qweight, + router_logits=router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index c3a5bd50246..846df44a28b 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -700,8 +700,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase): is_k_full=self.is_k_full, w13_g_idx=getattr(layer, "w13_g_idx", None), w2_g_idx=getattr(layer, "w2_g_idx", None), - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, + w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None), + w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None), routing_tables=layer._expert_routing_tables(), ) @@ -757,3 +757,27 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase): shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + hidden_states=x, + w1=layer.w13_qweight, + w2=layer.w2_qweight, + router_logits=router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 2a98d444afd..a69d2a594ad 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -10,6 +10,7 @@ from compressed_tensors.quantization import ( from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( + FusedMoEExpertsModular, RoutedExperts, SharedExperts, ) @@ -414,8 +415,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) - if not self.symmetric: + if w13_qzeros is not None: replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) + if w2_qzeros is not None: replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) # Marlin-specific parameters (not needed for Flashinfer) @@ -437,9 +439,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - layer.workspace = marlin_make_workspace_new( - layer.w13_weight_g_idx.device, 4 - ) + if self.experts_cls is not None and issubclass( + self.experts_cls, FusedMoEExpertsModular + ): + layer.workspace = marlin_make_workspace_new( + layer.w13_weight_g_idx.device, 4 + ) # Alias packed weights to w13_weight/w2_weight for the modular kernel interface layer.w13_weight = layer.w13_weight_packed diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 61b52345ab8..26fea5d5244 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -34,6 +34,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" MoEActivation.GELU_NO_MUL: ActivationType.Gelu, MoEActivation.SILU: ActivationType.Swiglu, MoEActivation.GELU: ActivationType.Geglu, + MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } return ACTIVATION_TO_FI_ACTIVATION[activation] diff --git a/vllm/model_executor/model_loader/bitsandbytes_loader.py b/vllm/model_executor/model_loader/bitsandbytes_loader.py index d10f3bfcbe9..064a74023a2 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -140,8 +140,8 @@ class BitsAndBytesModelLoader(BaseModelLoader): download_safetensors_index_file_from_hf( model_name_or_path, index_file, - self.load_config.download_dir, - revision, + cache_dir=self.load_config.download_dir, + revision=revision, ) hf_weights_files = filter_duplicate_safetensors_files( hf_weights_files, hf_folder, index_file diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 47c3c99b19a..0df14227919 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -70,7 +70,10 @@ class RunaiModelStreamerLoader(BaseModelLoader): if not is_local and not is_object_storage_path: download_safetensors_index_file_from_hf( - model_name_or_path, index_file, self.load_config.download_dir, revision + model_name_or_path, + index_file, + cache_dir=self.load_config.download_dir, + revision=revision, ) if not hf_weights_files: diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index 0711fb03f84..a857769cbe1 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -252,7 +252,6 @@ class ApertusDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 317269ec3b6..66adb9a3ca7 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -397,6 +398,9 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): } # LoRA specific attributes embedding_modules = {"embed_tokens": "input_embeddings"} + # ModelOpt NVFP4 checkpoints carry raw quantizer-module state + # (e.g. "*.weight_quantizer._double_scale"); drop them before loading. See #41925. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={"_quantizer.": None}) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -453,4 +457,4 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): loader = AutoWeightsLoader( self, skip_prefixes=["lm_head", "rotary_emb.inv_freq"] ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 64d606c2890..7354771764d 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -105,6 +105,60 @@ class Gemma4Config(VerifyAndUpdateConfig): ) +class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): + @classmethod + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: + """Set up the diffusion config and defaults for DiffusionGemma. + + Auto-creates DiffusionConfig from the HF config when the user + didn't pass ``--diffusion-config``. Diffusion sampling params are + read straight from generation_config.json at sampler-build time + (see DiffusionGemma's custom_sampler), not injected here. + """ + # Inherit Gemma4's attention backend selection (FA4 on Hopper, + # TRITON_ATTN fallback for heterogeneous head dims). + Gemma4Config.verify_and_update_config(vllm_config) + + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attention_config = vllm_config.attention_config + if attention_config.backend == AttentionBackendEnum.FLASHINFER: + raise ValueError( + "FlashInfer does not support DiffusionGemma's mixed " + "causal/bidirectional attention. Use --attention-backend " + "FLASH_ATTN or TRITON_ATTN instead." + ) + if attention_config.backend is None and not attention_config.use_non_causal: + attention_config.use_non_causal = True + logger.info( + "DiffusionGemma uses mixed causal/bidirectional attention " + "within a batch; setting use_non_causal=True to exclude " + "FlashInfer from auto-selection." + ) + + # Auto-create DiffusionConfig from HF config if not provided. + if vllm_config.diffusion_config is None: + from vllm.config.diffusion import DiffusionConfig + + hf_config = vllm_config.model_config.hf_config + canvas_length = getattr(hf_config, "canvas_length", 256) + vllm_config.diffusion_config = DiffusionConfig( + canvas_length=canvas_length, + ) + + # The diffusion sampler materializes [num_seqs, canvas_length, vocab] + # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200). + # Default to 8 when the user didn't pass --max-num-seqs. + # We can't see the original None here (the engine already filled a generic + # default), so use >= DEFAULT_MAX_NUM_SEQS as a proxy, (the default is much + # larger than any deliberate value for this model) + from vllm.config.scheduler import SchedulerConfig + + sc = vllm_config.scheduler_config + if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS: + sc.max_num_seqs = 8 + + class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -591,6 +645,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "ColQwen3_5": Qwen3_5ForConditionalGenerationConfig, "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, + "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 "FalconMambaForCausalLM": MambaModelConfig, "Gemma3TextModel": Gemma3TextModelConfig, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py new file mode 100644 index 00000000000..91dd5e6b6a5 --- /dev/null +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -0,0 +1,1363 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DiffusionGemma model, ModelState, and Sampler for vLLM. + +Single Gemma4 backbone run in two modes (like YOCO): +- encoder mode: causal attention, writes KV cache +- decoder mode: bidirectional attention, reads encoder KV, doesn't write + +Same weights, same layers. The only decoder-unique component is a +self-conditioning MLP. + +Multimodal support: the model always includes a vision tower (shared with Gemma4). +Images are encoded through the vision tower and projected into the LM embedding space +via Gemma4MultimodalEmbedder. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from transformers import AutoModel + +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, +) +from vllm.model_executor.models.gemma4 import Gemma4Model +from vllm.model_executor.models.gemma4_mm import ( + Gemma4DummyInputsBuilder, + Gemma4ForConditionalGeneration, + Gemma4MultimodalEmbedder, + Gemma4MultiModalProcessor, + Gemma4ProcessingInfo, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.transformers.utils import recursive_replace_linear +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.v1.outputs import LogprobsTensors +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.penalties import use_penalty + +from .interfaces import ( + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) + +logger = init_logger(__name__) + + +class DiffusionGemmaSelfConditioning(nn.Module): + """Gated MLP that processes soft embeddings from the previous denoising step. + + Structurally identical to Gemma4MLP but with self_conditioning_size + and post_norm without learned scale. + """ + + def __init__( + self, hidden_size: int, self_conditioning_size: int, eps: float = 1e-6 + ): + super().__init__() + self.pre_norm = RMSNorm(hidden_size, eps=eps) + self.post_norm = RMSNorm(hidden_size, eps=eps, has_weight=False) + self.gate_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) + self.up_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) + self.down_proj = nn.Linear(self_conditioning_size, hidden_size, bias=False) + + def forward( + self, + inputs_embeds: torch.Tensor, + soft_embeds: torch.Tensor, + ) -> torch.Tensor: + x = self.pre_norm(soft_embeds) + sc_signal = self.down_proj( + F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x) + ) + return self.post_norm(inputs_embeds + sc_signal) + + +# --------------------------------------------------------------------------- +# Multimodal processing info (overrides Gemma4 config type check) +# --------------------------------------------------------------------------- + + +class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo): + """Processing info for DiffusionGemma. + + Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig`` + (which inherits from ``PretrainedConfig``, not ``Gemma4Config``). + Supports image and video modalities. + """ + + def get_hf_config(self): + # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we + # accept any PretrainedConfig here. + return self.ctx.get_hf_config() + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # DiffusionGemma supports image and video inputs. + return {"image": None, "video": None} + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int] | None: + return super().get_mm_max_tokens_per_item(seq_len, mm_counts) + + +@torch.compile(dynamic=True) +def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: + # fp32 before tanh for numerical stability (matches HF DiffusionGemma). + # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over + # the [num_tokens, vocab] logits instead of four separate passes. + logits = logits.float() + return torch.tanh(logits / cap) * cap + + +@MULTIMODAL_REGISTRY.register_processor( + Gemma4MultiModalProcessor, + info=DiffusionGemmaProcessingInfo, + dummy_inputs=Gemma4DummyInputsBuilder, +) +class DiffusionGemmaForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsQuant, + SupportsPP, +): + """DiffusionGemma for vLLM. + + Single Gemma4 backbone that switches between encoder and decoder mode. + The encoder path uses standard Gemma4 layers (causal attention, KV write). + The decoder path uses the same weights with bidirectional attention and + KV read-only, plus self-conditioning. + + Always includes a vision tower (same as Gemma4) for image understanding. + + In practice, the model's forward() dispatches based on the `mode` kwarg + set by DiffusionGemmaModelState.prepare_inputs(). + """ + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.decoder.": "model.", + "model.encoder.language_model.": "model.", + "model.encoder.vision_tower.": "vision_tower.", + "model.encoder.embed_vision.": "embed_vision.", + }, + orig_to_new_substr={ + ".experts.": ".moe.experts.", + }, + ) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + @staticmethod + def get_model_state_cls(): + return DiffusionGemmaModelState + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + text_config = vllm_config.model_config.hf_text_config + self.config = config + self.model_dtype = vllm_config.model_config.dtype + + # DiffusionGemma's full-attention layers have NO v_proj — V is + # computed from k_proj's output (`value_states = key_states` before + # k_norm in `DiffusionGemmaDecoderTextAttention.forward`). This is + # the "k_eq_v" variant in our Gemma4 backbone. The checkpoint has no + # v_proj weights for full-attention layers; without this flag they + # would silently load with random V projections. + text_config.attention_k_eq_v = True + + # ---- Vision tower ---- + vision_config = getattr(config, "vision_config", None) + if vision_config is not None: + quant_config = vllm_config.quant_config + if quant_config and quant_config.get_name() in [ + "bitsandbytes", + "torchao", + "compressed-tensors", + ]: + tower_quant = quant_config + else: + quantizable = ( + vision_config.hidden_size % 64 == 0 + and vision_config.intermediate_size % 64 == 0 + ) + tower_quant = quant_config if quantizable else None + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_tower = AutoModel.from_config(config=vision_config) + self.embed_vision = Gemma4MultimodalEmbedder( + vision_config, + text_config, + quant_config=tower_quant, + prefix=maybe_prefix(prefix, "embed_vision"), + ) + recursive_replace_linear( + self.vision_tower, + tower_quant, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + else: + self.vision_tower = None + self.embed_vision = None + + # ---- Language backbone (Gemma4Model) ---- + # Use maybe_prefix to ensure correct weight name prefixes for + # quantization. The quantization config uses hf_to_vllm_mapper to + # match checkpoint weight names to model parameter names. + self.model = Gemma4Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + self.lm_head = ParallelLMHead( + num_embeddings=text_config.vocab_size, + embedding_dim=text_config.hidden_size, + ) + + if text_config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + + # HF DiffusionGemma applies the final-logit softcap in fp32, before + # any other processing. Do it manually in `compute_logits` so the + # LogitsProcessor only handles the lm_head GEMM. + self.final_logit_softcapping = getattr( + text_config, "final_logit_softcapping", None + ) + self.logits_processor = LogitsProcessor( + text_config.vocab_size, + soft_cap=None, + ) + + sc_size = ( + getattr(config, "self_conditioning_size", None) + or text_config.intermediate_size + ) + self.self_conditioning = DiffusionGemmaSelfConditioning( + hidden_size=text_config.hidden_size, + self_conditioning_size=sc_size, + eps=getattr(text_config, "rms_norm_eps", 1e-6), + ) + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def compute_self_conditioning( + self, + inputs_embeds: torch.Tensor, + probs: torch.Tensor, + ) -> torch.Tensor: + embed_weight = self.model.embed_tokens.weight + soft_embeds = torch.matmul( + probs.to(embed_weight.dtype), embed_weight + ) * self.model.normalizer.to(inputs_embeds.dtype) + return self.self_conditioning(inputs_embeds, soft_embeds) + + # ------------------------------------------------------------------ # + # Multimodal: reuse Gemma4's image parsing, processing & embedding + # ------------------------------------------------------------------ # + # The vision tower, pooler, embed_vision, and their processing logic + # are architecturally identical to Gemma4. Delegate to avoid + # maintaining a duplicate copy. + + _parse_and_validate_image_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_image_input + ) + _parse_and_validate_video_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_video_input + ) + _parse_and_validate_multimodal_inputs = ( + Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs + ) + _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) + _process_image_input = Gemma4ForConditionalGeneration._process_image_input + _process_video_input = Gemma4ForConditionalGeneration._process_video_input + embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal + + def get_mm_mapping(self) -> MultiModelKeys: + """Get the module prefix mapping for multimodal models.""" + return MultiModelKeys.from_string_field( + language_model="model", + connector=["embed_vision"], + tower_model=["vision_tower"], + ) + + # ------------------------------------------------------------------ # + # Forward + # ------------------------------------------------------------------ # + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + if intermediate_tensors is not None: + inputs_embeds = None + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is not None and self.final_logit_softcapping is not None: + logits = _softcap_logits(logits, self.final_logit_softcapping) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Load weights from checkpoint. + + Checkpoint layout (HF DiffusionGemma): + model.encoder.vision_tower.* → vision tower + model.encoder.embed_vision.* → vision embedder + model.encoder.language_model.layers.* → backbone + model.decoder.layers.* → backbone (tied) + model.decoder.embed_tokens.* → embeddings + model.decoder.self_conditioning.* → self-conditioning MLP + lm_head.* → LM head (tied) + + We load encoder weights into our single ``Gemma4Model`` backbone, + skip duplicate decoder backbone weights, handle vision tower and + self-conditioning separately. + """ + + sc_params = dict( + (n, p) + for n, p in self.named_parameters() + if n.startswith("self_conditioning.") + ) + + # Collect vision tower + embedder parameters AND buffers for manual + # loading. The HF vision tower registers std_bias / std_scale as + # buffers (not parameters) when config.standardize is True, so we + # must include named_buffers() to avoid "not found in model" warnings. + vision_params: dict[str, torch.Tensor] = {} + for n, p in self.named_parameters(): + if n.startswith(("vision_tower.", "embed_vision.")): + vision_params[n] = p + for n, b in self.named_buffers(): + if n.startswith(("vision_tower.", "embed_vision.")): + vision_params[n] = b + + def _remap_weights(): + # Use full weight names (including suffixes like .weight_scale, + # .weight_packed) for dedup instead of just the base layer name. Critical + # for quantized checkpoints where each weight has multiple tensors; + # tracking only base names skips scales as duplicates. + seen_weights: set[str] = set() + for name, weight in weights: + # Self-conditioning lives under model.decoder.self_conditioning.* + # in the checkpoint but at self_conditioning.* in our model. + if "self_conditioning" in name: + sc_name = name.split("self_conditioning.", 1)[1] + sc_name = "self_conditioning." + sc_name + if sc_name in sc_params: + sc_params[sc_name].data.copy_(weight) + continue + + # Vision tower: model.encoder.vision_tower.* → vision_tower.* + # In HF, the vision tower is a sibling of language_model + # under the encoder module. + if name.startswith("model.encoder.vision_tower."): + vt_name = name[len("model.encoder.") :] + if vt_name in vision_params: + vision_params[vt_name].data.copy_(weight) + else: + logger.warning( + "Vision tower weight %s (mapped to %s) not found in model", + name, + vt_name, + ) + continue + + # Vision embedder: model.encoder.embed_vision.* → embed_vision.* + if name.startswith("model.encoder.embed_vision."): + ev_name = name[len("model.encoder.") :] + if ev_name in vision_params: + vision_params[ev_name].data.copy_(weight) + else: + logger.warning( + "Embed vision weight %s (mapped to %s) not found in model", + name, + ev_name, + ) + continue + + # Skip vestigial embed_vision.embedding weights. + if "embed_vision.embedding." in name: + continue + + # Encoder backbone → model.* + if name.startswith("model.encoder.language_model."): + name = name.replace("model.encoder.language_model.", "model.") + # Decoder backbone → model.* (skip exact duplicates) + elif name.startswith("model.decoder."): + name = name.replace("model.decoder.", "model.") + + # Skip only if we've seen the exact same weight name (including scales) + if name in seen_weights: + continue + seen_weights.add(name) + yield name, weight + + # Delegate to Gemma4ForCausalLM.load_weights for the backbone, + # which handles stacked params, MoE, k_eq_v, etc. + # Temporarily set self.config to text_config since Gemma4's + # load_weights expects it (e.g. tie_word_embeddings, layer_types). + from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM + + saved_config = self.config + self.config = self.model.config + try: + Gemma4ForCausalLM.load_weights(self, _remap_weights()) + finally: + self.config = saved_config + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "" + if modality == "video": + return "<|video|>" + raise ValueError(f"Unsupported modality: {modality}") + + +@torch.compile(dynamic=True) +def _compute_num_rejected( + num_logits: torch.Tensor, + num_sampled: torch.Tensor, + query_start_loc: torch.Tensor, +) -> torch.Tensor: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + num_rejected = num_logits - num_sampled + is_denoise = (num_logits > 0) & (num_sampled == 0) + return torch.where(is_denoise, query_lens, num_rejected) + + +@torch.compile(dynamic=True) +def _compiled_sample_step( + # Logits from the model [num_decode * CL, vocab] + logits: torch.Tensor, + # Request mapping + decode_slots: torch.Tensor, # [num_decode] int64 → slot indices + decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs + all_slots: torch.Tensor, # [num_reqs] int64 → all slot indices + valid_canvas_len: torch.Tensor, # [num_decode] int64 → real canvas length (<=CL) + # State tensors (modified in-place) + canvas: torch.Tensor, # [max_num_reqs, CL] + argmax_canvas: torch.Tensor, # [max_num_reqs, CL] + step_tensor: torch.Tensor, # [max_num_reqs] + is_encoder_phase: torch.Tensor, # [max_num_reqs] + confident_tensor: torch.Tensor, # [max_num_reqs] + sc_embeds: torch.Tensor, # [max_num_reqs, CL, hidden] + embed_weight: torch.Tensor, # [vocab, hidden] + normalizer: torch.Tensor, + history: torch.Tensor, # [max_num_reqs, ST, CL] + history_len_tensor: torch.Tensor, # [max_num_reqs] + # Output tensors (modified in-place) + sampled: torch.Tensor, # [num_reqs, CL] + num_sampled: torch.Tensor, # [num_reqs] + draft_tokens: torch.Tensor, # [max_num_reqs, >=CL] + # Scalar config + max_denoising_steps: float, + t_min: float, + t_max: float, + confidence_threshold: float, + vocab_size: int, + CL: int, + ST: int, + # Sampler config + entropy_bound: float, +) -> torch.Tensor: + """Compiled decode step: temperature → Gumbel sample → probs/confidence → + accept/renoise → convergence, all as vectorized PyTorch ops. + + Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the + caller can compute logprobs outside the compiled region.""" + num_decode = decode_slots.shape[0] + device = decode_slots.device + + # Clear outputs so prefill / non-decode slots report 0 (decode slots are + # overwritten below). + sampled.zero_() + num_sampled.zero_() + + # ---- Phase 1: Temperature schedule ---- + steps_f = step_tensor[decode_slots].float() + remaining = (max_denoising_steps - steps_f).clamp(min=1.0) + temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps) + + # ---- Phase 2: Temperature scaling + Gumbel-max sampling ---- + logits_3d = logits.reshape(num_decode, CL, -1).float() + scaled = logits_3d / temp[:, None, None].clamp(min=1e-10) + + # Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T) + u = torch.rand_like(scaled).clamp(min=1e-20) + gumbel = -torch.log(-torch.log(u)) + # Zero noise when temp==0 (greedy) + noisy = scaled + gumbel * (temp[:, None, None] > 0).float() + new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL) + argmax_tokens = ( + scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL) + ) + + # ---- Phase 3: Probs, self-conditioning, confidence ---- + log_probs = scaled.log_softmax(dim=-1) + probs = log_probs.exp() + + token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL] + # A canvas truncated near max_model_len is zero-padded up to CL by the + # caller; those padded rows are uniform (max entropy, argmax 0), so they + # never trigger early convergence and are stable, and only the real + # ``valid_canvas_len`` tokens are committed (num_sampled below). + mean_entropy = token_entropy.mean(dim=-1) # [num_decode] + confident_tensor[decode_slots] = mean_entropy < confidence_threshold + + # ---- Phase 4: Entropy-bound acceptance mask ---- + sorted_ent, sorted_idx = torch.sort(token_entropy, dim=-1) + cumsum_ent = torch.cumsum(sorted_ent, dim=-1) + cummax_ent = torch.cummax(sorted_ent, dim=-1).values + sorted_mask = (cumsum_ent - cummax_ent) <= entropy_bound + eb_mask = torch.zeros_like(sorted_mask) + eb_mask.scatter_(1, sorted_idx, sorted_mask) + + # ---- Phase 5: Post-sample ---- + is_commit = is_encoder_phase[decode_slots] # [num_decode] + is_denoise = ~is_commit + cur_step = step_tensor[decode_slots].float() + + # Step update: +1 for denoise, reset to 0 for commit + new_step_val = torch.where( + is_denoise, + (cur_step + 1).to(step_tensor.dtype), + step_tensor.new_zeros(num_decode), + ) + step_tensor[decode_slots] = new_step_val + + # Random tokens for renoise / canvas reinit + random_tokens = torch.randint( + 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype + ) + + # Compute denoise canvas (accept/renoise) + denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens) + + # Canvas: commit → random reinit, denoise → accept/renoise result + canvas[decode_slots] = torch.where( + is_commit.unsqueeze(1), random_tokens, denoise_canvas + ) + + # History: write argmax_tokens for denoise requests at circular position + hist_len = history_len_tensor[decode_slots] + write_pos = hist_len % ST + for i in range(ST): + write_here = ((write_pos == i) & is_denoise).unsqueeze(1) + history[decode_slots, i] = torch.where( + write_here, argmax_tokens, history[decode_slots, i] + ) + + # Argmax canvas: update for denoise, preserve for commit + argmax_canvas[decode_slots] = torch.where( + is_denoise.unsqueeze(1), argmax_tokens, argmax_canvas[decode_slots] + ) + + # History length: increment for denoise, reset for commit + new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode)) + history_len_tensor[decode_slots] = new_hist_len + + # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed) + sampled[decode_idx] = argmax_canvas[decode_slots].to( + sampled.dtype + ) * is_commit.unsqueeze(1).to(sampled.dtype) + # Commit only the real canvas length (== CL except for a canvas truncated + # near max_model_len); the padded tail positions are never emitted. + num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to( + num_sampled.dtype + ) + + # ---- Phase 6: Stability + convergence ---- + ref = history[decode_slots, 0] + mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32) + for h in range(1, ST): + mismatch = mismatch + (ref != history[decode_slots, h]).sum(dim=-1).int() + stable = mismatch == 0 + + step_after = step_tensor[decode_slots] + converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( + step_after >= max_denoising_steps + ) + # Commit done → denoise next (False); denoise converged → commit next (True) + is_encoder_phase[decode_slots] = torch.where( + is_commit, is_commit.new_zeros(num_decode), converged + ) + + # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's + # self-conditioning MLP consumes) only for slots that will denoise next — i.e. + # this step denoised AND it isn't about to commit (is_encoder_phase now False). + # Masking here (rather than in the consumer) lets _apply_self_conditioning read + # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full + # [.., vocab] probs avoids a giant persistent buffer. + sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] + soft_embeds = torch.matmul(probs.to(embed_weight.dtype), embed_weight) * normalizer + sc_embeds[decode_slots] = soft_embeds * sc_keep + + # Overwrite canvas with argmax for newly converged denoise requests + newly_converged = (converged & is_denoise).unsqueeze(1) + canvas[decode_slots] = torch.where( + newly_converged, argmax_canvas[decode_slots], canvas[decode_slots] + ) + + # ---- Phase 7: Copy canvas → draft_tokens for all slots ---- + draft_tokens[all_slots, :CL] = canvas[all_slots] + + return scaled + + +class DiffusionGemmaRequestStates: + """Pre-allocated GPU tensors for DiffusionGemma per-request state. + + Follows the indexed-slot pattern used by ``RequestState``. + """ + + def __init__( + self, + max_num_reqs: int, + canvas_length: int, + vocab_size: int, + max_denoising_steps: int, + device: torch.device, + hidden_size: int, + stability_threshold: int, + ): + self.max_num_reqs = max_num_reqs + self.canvas_length = canvas_length + self.vocab_size = vocab_size + self.max_denoising_steps = max_denoising_steps + self.stability_threshold = stability_threshold + self.device = device + + self.is_encoder_phase = torch.zeros( + max_num_reqs, dtype=torch.bool, device=device + ) + # Canvas tokens [max_num_reqs, canvas_length] + self.canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + # Step counter (counts up from 0 to max_denoising_steps) + self.step = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + # Accepted canvas history for stability check + self.accepted_canvas_history = torch.zeros( + max_num_reqs, + stability_threshold, + canvas_length, + dtype=torch.int64, + device=device, + ) + self.accepted_canvas_history_len = torch.zeros( + max_num_reqs, dtype=torch.int32, device=device + ) + # Latest argmax(processed_logits) per slot — what we COMMIT. + # NOT `current_canvas` (which is the post-renoise stochastic input for + # the next denoise step). We keep this separate from `canvas` because + # canvas gets renoised in-place during denoise, while argmax_canvas is + # the deterministic best-guess we ultimately emit. + self.argmax_canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + + # Per-slot prompt length (set by add_request). + self.prompt_len = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + + # Per-slot confidence flag, set by the sampler each step. + self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + + # Per-slot self-conditioning soft embedding (probs @ embed_weight) from + # the previous denoise step. Storing the [.., hidden] soft embed instead + # of the full [.., vocab] distribution shrinks this buffer by + # vocab/hidden (~170x) and moves the matmul to denoise time; the result + # is identical (SC consumes probs @ embed_weight anyway). + self.self_conditioning_embeds = torch.zeros( + max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device + ) + + def init_canvas(self, slot_indices_np: np.ndarray) -> None: + """Initialize canvas with random tokens for the given slots.""" + n = slot_indices_np.shape[0] + self.canvas[slot_indices_np] = torch.randint( + 0, + self.vocab_size, + (n, self.canvas_length), + dtype=torch.int64, + device=self.device, + ) + + def add_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = True + self.init_canvas(torch.tensor([slot_idx], device=self.device)) + self.step[slot_idx] = 0 + self.accepted_canvas_history_len[slot_idx] = 0 + self.self_conditioning_embeds[slot_idx] = 0 + + def remove_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = False + self.accepted_canvas_history_len[slot_idx] = 0 + self.self_conditioning_embeds[slot_idx] = 0 + + +class DiffusionGemmaModelState(ModelState): + """ModelState for DiffusionGemma. + + Single Gemma4 backbone in two modes: + - encoder mode (num_draft_tokens == 0): causal attention, writes KV + - decoder mode (num_draft_tokens > 0): bidirectional attention, reads KV + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: Any, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.scheduler_config = vllm_config.scheduler_config + self.model = model + self.device = device + + self.supports_mm_inputs = encoder_cache is not None + self.max_num_reqs = self.scheduler_config.max_num_seqs + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens + self.max_model_len = self.model_config.max_model_len + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() + self.dtype = self.model_config.dtype + + if self.supports_mm_inputs: + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + + assert isinstance(encoder_cache, EncoderCache) + self.encoder_cache = encoder_cache + self.encoder_runner = EncoderRunner( + model=self.model, + max_num_tokens=self.max_num_tokens, + hidden_size=self.inputs_embeds_size, + encoder_cache=encoder_cache, + dtype=self.dtype, + device=self.device, + ) + + # Per-step MM data produced by get_mm_embeddings and consumed by + # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that + # prepare_inputs can call embed_input_ids directly into the + # persistent _inputs_embeds_buf, avoiding the intermediate copy + # through encoder_runner.inputs_embeds. + self._pending_mm_embeds: tuple[list[torch.Tensor], torch.Tensor] | None = None + + diffusion_config = vllm_config.diffusion_config + canvas_length = diffusion_config.canvas_length if diffusion_config else 32 + + text_config = self.model_config.hf_text_config + self.gen_config = self.model_config.try_get_generation_config() + max_denoising_steps = ( + diffusion_config.max_denoising_steps if diffusion_config else None + ) or self.gen_config.get("max_denoising_steps", 48) + self.diffusion_states = DiffusionGemmaRequestStates( + max_num_reqs=self.max_num_reqs, + canvas_length=canvas_length, + vocab_size=self.model_config.get_vocab_size(), + max_denoising_steps=max_denoising_steps, + device=device, + hidden_size=text_config.hidden_size, + stability_threshold=self.gen_config["stability_threshold"], + ) + self._req_id_to_index: dict[str, int] = {} + + # Persistent buffer for per-request causal flags, updated in-place + # so FULL CUDA graph replay sees the latest values. + self._causal_buf = torch.zeros( + self.max_num_reqs, dtype=torch.bool, device=device + ) + + # Persistent inputs_embeds buffer — required so FULL CUDA graph + # capture and runtime point at the SAME memory address. + # `prepare_dummy_inputs` (capture path) and `prepare_inputs` (runtime + # path) both must hand the captured graph a tensor at this address. + self._inputs_embeds_buf = torch.zeros( + self.max_num_tokens, + text_config.hidden_size, + dtype=self.model_config.dtype, + device=device, + ) + + def get_supported_generation_tasks(self): + return ("generate",) + + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + diffusion_config = self.vllm_config.diffusion_config + gen = self.gen_config + sampler_cfg = gen.get("sampler_config") or {} + if "EntropyBound" not in sampler_cfg.get("_cls_name", ""): + raise ValueError("DiffusionGemma requires an EntropyBound sampler_config") + entropy_bound = sampler_cfg.get("entropy_bound") + if entropy_bound is None or entropy_bound <= 0: + raise ValueError( + f"entropy_bound must be a positive float (got {entropy_bound})" + ) + return DiffusionSampler( + sampler=sampler, + diffusion_config=diffusion_config, + vocab_size=self.model_config.get_vocab_size(), + diffusion_states=self.diffusion_states, + t_min=gen["t_min"], + t_max=gen["t_max"], + entropy_bound=entropy_bound, + confidence_threshold=gen["confidence_threshold"], + embed_weight=self.model.model.embed_tokens.weight, + normalizer=self.model.model.normalizer, + ), None + + def apply_staged_writes(self) -> None: + pass + + def add_request(self, req_index: int, new_req_data: Any) -> None: + self._req_id_to_index[new_req_data.req_id] = req_index + self.diffusion_states.add_request(req_index) + if not new_req_data.req_id.startswith("_warmup_"): + prompt_len = len(new_req_data.prompt_token_ids) + self.diffusion_states.prompt_len[req_index] = prompt_len + + def remove_request(self, req_id: str) -> None: + idx = self._req_id_to_index.pop(req_id, None) + if idx is not None: + self.diffusion_states.remove_request(idx) + + def get_mm_embeddings(self, scheduled_encoder_inputs, input_batch): + if not self.supports_mm_inputs: + return None + + mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs( + scheduled_encoder_inputs + ) + if mm_kwargs: + encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) + self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs)) + + mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings( + input_batch.req_ids, + input_batch.num_tokens, + input_batch.num_scheduled_tokens, + input_batch.query_start_loc_np, + input_batch.prefill_len_np, + input_batch.num_computed_prefill_tokens_np, + ) + + if not mm_embeds: + # No MM tokens in this batch (e.g. all-decode step). + # prepare_inputs will use embed_input_ids (text-only) directly. + self._pending_mm_embeds = None + return None + + # Stash raw MM ingredients for prepare_inputs to merge directly + # into the persistent buffer, avoiding the intermediate copy + # through encoder_runner.inputs_embeds. + self._pending_mm_embeds = (mm_embeds, is_mm_embed) + return None + + def _apply_self_conditioning( + self, + decode_slots_np: np.ndarray, + decode_idx_np: np.ndarray, + query_start_loc_np: np.ndarray, + inputs_embeds: torch.Tensor, + sc_embeds: torch.Tensor, + ) -> None: + # One self-conditioning MLP call per decode request, over that request's + # query span [start, end) = its canvas. The span is the full canvas (CL) + # or, for the final canvas truncated near max_model_len, fewer than CL + # positions. sc_embeds already holds probs @ embed_weight from the prior + # denoise step, masked to zero by the sampler for slots not denoising + # this step; only the MLP runs here. CPU metadata -> no GPU syncs. + for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()): + start = int(query_start_loc_np[idx]) + end = int(query_start_loc_np[idx + 1]) + canvas = slice(start, end) + soft = sc_embeds[slot, : end - start] + inputs_embeds[canvas] = self.model.self_conditioning( + inputs_embeds[canvas], soft.to(inputs_embeds.dtype) + ) + + def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]: + states = self.diffusion_states + num_tokens = input_batch.num_tokens + num_reqs = input_batch.num_reqs + + # Write into the PERSISTENT inputs_embeds buffer so FULL CUDA graph + # replay sees the latest values at the captured address. + num_tokens_padded = input_batch.num_tokens_after_padding + inputs_embeds = self._inputs_embeds_buf[:num_tokens_padded] + + # Populate embeddings: merge MM features when available, + # otherwise embed input_ids as text-only. + input_ids = input_batch.input_ids[:num_tokens] + if self._pending_mm_embeds is not None: + mm_embeds, is_mm_embed = self._pending_mm_embeds + self._pending_mm_embeds = None + inputs_embeds[:num_tokens].copy_( + self.model.embed_input_ids( + input_ids, + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + ) + else: + inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids)) + + # Apply self-conditioning ONLY for denoising decode requests. + if input_batch.num_draft_tokens > 0 and self._req_id_to_index: + slots_np = input_batch.idx_mapping_np[:num_reqs] + num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) + is_decode_indices_np = np.where(num_logits_np > 0)[0] + self._apply_self_conditioning( + slots_np[is_decode_indices_np], + is_decode_indices_np, + input_batch.query_start_loc_np, + inputs_embeds, + states.self_conditioning_embeds, + ) + + return {"inputs_embeds": inputs_embeds} + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + # CUDA graph capture path — return a slice of the SAME persistent + # inputs_embeds buffer that `prepare_inputs` writes to at runtime, + # so the captured graph and runtime point to identical addresses. + return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} + + def postprocess_state(self, idx_mapping, num_sampled) -> None: + return None + + def prepare_attn( + self, + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=False, + ) -> dict[str, Any]: + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + + # Per-request causal mode: encoder (commit) = causal, + # denoise = bidirectional. Pass GPU tensor so the attention + # backend can handle mixed batches. + actual_num_reqs = input_batch.num_reqs + slots = input_batch.idx_mapping[:actual_num_reqs] + # Invariant: the sampler flips is_encoder_phase to False only after a + # request's FINAL prompt chunk, so a prompt spanning multiple chunks + # (longer than the token budget) stays causal for every chunk. + self._causal_buf[:actual_num_reqs] = self.diffusion_states.is_encoder_phase[ + slots + ] + if actual_num_reqs < num_reqs: + self._causal_buf[actual_num_reqs:num_reqs] = False + causal: bool | torch.Tensor = self._causal_buf[:num_reqs] + + return build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=self.max_model_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + causal=causal, + ) + + num_new_sampled_tokens_per_step: int = 0 + + +# Penalty stub for the diffusion path: the runner reads +# penalties_state.output_bin_counts, and post_update treats None as +# "no penalty bookkeeping". +_NO_PENALTIES_STATE = SimpleNamespace(output_bin_counts=None) + + +class DiffusionSampler: + """Batched accept/renoise sampler for DiffusionGemma. + + Follows the same structure as ``vllm.v1.worker.gpu.sample.sampler.Sampler``: + decomposed into named methods, all GPU state in pre-allocated buffers, + no GPU→CPU syncs on the hot path. + """ + + def __init__( + self, + sampler: Any, + diffusion_config: Any, + vocab_size: int, + diffusion_states: DiffusionGemmaRequestStates | None = None, + *, + confidence_threshold: float, + t_min: float, + t_max: float, + entropy_bound: float, + embed_weight: torch.Tensor, + normalizer: torch.Tensor, + ): + self.sampling_states = sampler.sampling_states + self.req_states = sampler.req_states + # Self-conditioning soft embed = probs @ embed_weight * normalizer, + # computed in the sampler (see _compiled_sample_step). + self.embed_weight = embed_weight + self.normalizer = normalizer + self.canvas_length = ( + diffusion_config.canvas_length if diffusion_config is not None else 32 + ) + self.t_min = t_min + self.t_max = t_max + self.confidence_threshold = confidence_threshold + self.vocab_size = vocab_size + self.diffusion_states = diffusion_states + self.entropy_bound = entropy_bound + + max_num_reqs = diffusion_states.max_num_reqs + device = diffusion_states.device + self._sampled = torch.zeros( + max_num_reqs, + self.canvas_length, + dtype=torch.int32, + device=device, + ) + self._num_sampled = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + self._decode_slots = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self._decode_idx = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self._query_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + self._num_logits = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + + # Per-slot stash for logprobs computed on the converging denoise step. + # Populated after the post-sample kernel detects convergence; consumed + # on the subsequent commit step when num_sampled=CANVAS_LEN. + self._pending_logprobs: dict[int, LogprobsTensors] = {} + + def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None: + if use_penalty(sampling_params): + logger.warning_once( + "DiffusionGemma does not support repetition/frequency/presence " + "penalties; ignoring them for this request." + ) + # Purge any stale logprobs stashed under this slot by a prior request + # that was aborted between its converging denoise and commit steps. + self._pending_logprobs.pop(req_idx, None) + self.sampling_states.add_request(req_idx, sampling_params) + + def apply_staged_writes(self) -> None: + self.sampling_states.apply_staged_writes() + + @property + def penalties_state(self): + # Diffusion applies no penalties. The runner reads + # penalties_state.output_bin_counts, so expose a stub holding None; + # post_update treats None bin counts as "no penalty bookkeeping". + return _NO_PENALTIES_STATE + + # ------------------------------------------------------------------ + # Prefill + # ------------------------------------------------------------------ + + def _finish_prefills( + self, input_batch: Any, prefill_indices_np: np.ndarray + ) -> None: + """Transition requests whose prompt completes this step to denoising. + + Initializes their canvas, seeds draft tokens, and flips + is_encoder_phase to False. Mid-chunk requests (prompt longer than the + token budget) are left untouched so is_encoder_phase stays True and + prepare_attn keeps causal attention for their remaining chunks. + """ + states = self.diffusion_states + done_prefill_np = ( + input_batch.num_computed_prefill_tokens_np[prefill_indices_np] + + input_batch.num_scheduled_tokens[prefill_indices_np] + >= input_batch.prefill_len_np[prefill_indices_np] + ) + ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]] + if len(ps) == 0: + return + states.init_canvas(ps) + self.req_states.draft_tokens[ps, : self.canvas_length] = states.canvas[ps] + ps_gpu = async_copy_to_gpu( + ps.astype(np.int64), device=states.is_encoder_phase.device + ) + states.is_encoder_phase.index_fill_(0, ps_gpu, False) + + def _handle_prefill( + self, + input_batch: Any, + device: torch.device, + ) -> SamplerOutput: + num_reqs = input_batch.num_reqs + self._finish_prefills(input_batch, np.arange(num_reqs)) + sampled = self._sampled[:num_reqs, :1] + sampled.zero_() + num_sampled = self._num_sampled[:num_reqs] + num_sampled.zero_() + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + num_rejected=num_sampled, + ) + + # ------------------------------------------------------------------ + # Decode helpers + # ------------------------------------------------------------------ + + def _build_output( + self, + input_batch: Any, + sampled: torch.Tensor, + num_sampled: torch.Tensor, + per_req_nlogits_np: np.ndarray, + device: torch.device, + logprobs_tensors: LogprobsTensors | None = None, + ) -> SamplerOutput: + """Compute num_rejected and build SamplerOutput.""" + num_reqs = input_batch.num_reqs + + self._query_lens.np[:num_reqs] = np.diff( + input_batch.query_start_loc_np[: num_reqs + 1] + ) + self._num_logits.np[:num_reqs] = per_req_nlogits_np + self._query_lens.copy_to_uva() + self._num_logits.copy_to_uva() + + num_rejected = _compute_num_rejected( + self._num_logits.gpu[:num_reqs], + num_sampled, + input_batch.query_start_loc[: num_reqs + 1], + ) + + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=logprobs_tensors, + num_nans=None, + num_sampled=num_sampled, + num_rejected=num_rejected, + ) + + # ------------------------------------------------------------------ + # Main entry point + # ------------------------------------------------------------------ + + def __call__( + self, + logits: torch.Tensor, + input_batch: Any, + draft_logits: torch.Tensor | None = None, + ) -> SamplerOutput: + num_reqs = input_batch.num_reqs + device = logits.device + + if input_batch.num_draft_tokens == 0: + return self._handle_prefill(input_batch, device) + + # --- CPU/NumPy setup (outside compile): split decode vs prefill, init + # canvas for any new prefills, and stage decode slot indices to GPU. --- + states = self.diffusion_states + CL = self.canvas_length + slots_np = input_batch.idx_mapping_np[:num_reqs] + per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) + + decode_indices_np = np.where(per_req_nlogits_np > 0)[0] + prefill_indices_np = np.where(per_req_nlogits_np == 0)[0] + decode_slots_np = slots_np[decode_indices_np] + + if len(prefill_indices_np) > 0: + self._finish_prefills(input_batch, prefill_indices_np) + + num_decode = len(decode_indices_np) + self._decode_slots.np[:num_decode] = decode_slots_np + self._decode_idx.np[:num_decode] = decode_indices_np + self._decode_slots.copy_to_uva() + self._decode_idx.copy_to_uva() + decode_slots = self._decode_slots.gpu[:num_decode] + decode_idx = self._decode_idx.gpu[:num_decode] + + # Real canvas length per decode request. Equals CL except when a canvas + # was truncated near max_model_len, in which case the scheduler gave us + # fewer than CL logits for that request. + valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] + valid_canvas_len = async_copy_to_gpu( + valid_canvas_len_np.astype(np.int64), device=device + ) + + # Pad any truncated canvas back to CL so the uniform-CL sampler math + # holds. Phantom (padded) positions are zeroed → uniform logits → high + # entropy (no premature convergence) and argmax 0 (stable); they are + # never committed (num_sampled == real length). + if num_decode > 0 and valid_canvas_len_np.min() < CL: + ar = torch.arange(CL, device=device) + starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req + valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] + src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) + logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype) + + # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0. + sampled = self._sampled[:num_reqs] + num_sampled = self._num_sampled[:num_reqs] + + all_slots = input_batch.idx_mapping[:num_reqs] + + # Snapshot which slots are committing BEFORE the compiled step runs, + # since it mutates is_encoder_phase (commit→False, converge→True). + is_committing = states.is_encoder_phase[decode_slots].clone() + + # --- Single compiled call: temp → sample → probs → post-process --- + scaled = _compiled_sample_step( + logits, + decode_slots, + decode_idx, + all_slots, + valid_canvas_len, + # State + states.canvas, + states.argmax_canvas, + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds, + self.embed_weight, + self.normalizer, + states.accepted_canvas_history, + states.accepted_canvas_history_len, + # Output + sampled, + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=float(states.max_denoising_steps), + t_min=self.t_min, + t_max=self.t_max, + confidence_threshold=self.confidence_threshold, + vocab_size=self.vocab_size, + CL=self.canvas_length, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + ) + + # --- Logprobs: stash on convergence, return on commit --- + slots_np = input_batch.idx_mapping_np[:num_reqs] + is_decode_np = per_req_nlogits_np > 0 + + logprobs_tensors = None + max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) + if max_num_logprobs >= 0: + # Denoise steps that just converged: the compiled step flipped + # is_encoder_phase from False→True. Detect as slots where + # is_encoder_phase is now True but is_committing was False. + converged_mask = states.is_encoder_phase[decode_slots] + just_converged = converged_mask & ~is_committing + if just_converged.any(): + flat_logits = scaled.reshape(-1, scaled.shape[-1]) + argmax_tokens = scaled.argmax(dim=-1) + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = local_idx.item() + slot = decode_slots[local_idx] + # Stash only the real canvas positions (== CL unless this + # canvas was truncated near max_model_len); padded tail + # positions are never emitted. + k_i = int(valid_canvas_len_np[li]) + start = li * CL + self._pending_logprobs[slot.item()] = compute_topk_logprobs( + flat_logits[start : start + k_i], + max_num_logprobs, + argmax_tokens[local_idx][:k_i], + ) + + # Commit steps: is_committing was True at entry. Reassemble + # previously stashed logprobs and attach to SamplerOutput. + if is_committing.any() and self._pending_logprobs: + parts_ids, parts_lp, parts_ranks = [], [], [] + cu_gen: list[int] = [] + flat_offset = 0 + for i in range(num_reqs): + cu_gen.append(flat_offset) + slot = int(slots_np[i]) + if is_decode_np[i] and slot in self._pending_logprobs: + lp = self._pending_logprobs.pop(slot) + parts_ids.append(lp.logprob_token_ids) + parts_lp.append(lp.logprobs) + parts_ranks.append(lp.selected_token_ranks) + flat_offset += lp.logprobs.shape[0] + if parts_ids: + logprobs_tensors = LogprobsTensors( + logprob_token_ids=torch.cat(parts_ids), + logprobs=torch.cat(parts_lp), + selected_token_ranks=torch.cat(parts_ranks), + cu_num_generated_tokens=cu_gen, + ) + + return self._build_output( + input_batch, + sampled, + num_sampled, + per_req_nlogits_np, + device, + logprobs_tensors=logprobs_tensors, + ) diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index dca05f72c69..be45d7dfb2b 100644 --- a/vllm/model_executor/models/exaone.py +++ b/vllm/model_executor/models/exaone.py @@ -243,7 +243,6 @@ class ExaoneDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index e38dbb5ee29..a36b8e0e922 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -230,7 +230,6 @@ class Exaone4DecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index 3373983f5c9..18900557f61 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -179,7 +179,6 @@ class ExaoneMoeDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 45e82c26d95..03e67c4ada7 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -725,10 +725,8 @@ class Gemma4DecoderLayer(nn.Module): if self.enable_moe_block: hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states) - # Router and MoE experts see the residual (pre-MLP state), - # matching the HF transformers forward path - router_logits = self.router(residual) hidden_states_2 = self.pre_feedforward_layernorm_2(residual) + router_logits = self.router(residual) hidden_states_2 = self.moe(hidden_states_2, router_logits) hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 2adc29f8d25..7470e7e7381 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -199,7 +199,6 @@ class GraniteDecoderLayer(nn.Module): self.residual_multiplier = config.residual_multiplier max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/h2ovl.py b/vllm/model_executor/models/h2ovl.py index 1e3629eb42e..40240d3e4ee 100644 --- a/vllm/model_executor/models/h2ovl.py +++ b/vllm/model_executor/models/h2ovl.py @@ -157,27 +157,22 @@ class H2OVLChatModel(InternVLChatModel): config: PretrainedConfig, quant_config: QuantizationConfig | None, *, - is_mono: bool, prefix: str, ): - if not is_mono: - vision_feature_layer = config.select_layer - if vision_feature_layer < 0: - num_hidden_layers = ( - config.vision_config.num_hidden_layers + vision_feature_layer + 1 - ) - else: - num_hidden_layers = vision_feature_layer + 1 - - return InternVisionModel( - config.vision_config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers, - prefix=prefix, + vision_feature_layer = config.select_layer + if vision_feature_layer < 0: + num_hidden_layers = ( + config.vision_config.num_hidden_layers + vision_feature_layer + 1 ) else: - msg = "Monolith mode is not applicable to H2OVL" - raise NotImplementedError(msg) + num_hidden_layers = vision_feature_layer + 1 + + return InternVisionModel( + config.vision_config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers, + prefix=prefix, + ) def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: if num_image_tokens <= 0 or self.num_image_token <= 0: diff --git a/vllm/model_executor/models/internlm2_ve.py b/vllm/model_executor/models/internlm2_ve.py deleted file mode 100644 index da0dfe73e6f..00000000000 --- a/vllm/model_executor/models/internlm2_ve.py +++ /dev/null @@ -1,139 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from itertools import islice - -import torch -from torch import nn -from transformers import PretrainedConfig - -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.models.internlm2 import ( - InternLM2Attention, - InternLM2ForCausalLM, - InternLM2MLP, - InternLM2Model, -) -from vllm.sequence import IntermediateTensors - - -class InternLM2VEDecoderLayer(nn.Module): - def __init__( - self, - config: PretrainedConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - self.attention = InternLM2Attention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - rope_parameters=config.rope_parameters, - max_position_embeddings=max_position_embeddings, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attention", - ) - self.feed_forward = InternLM2MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.feed_forward", - ) - self.feed_forward_ve = InternLM2MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.feed_forward_ve", - ) - self.attention_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - visual_token_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.attention_norm(hidden_states) - else: - hidden_states, residual = self.attention_norm(hidden_states, residual) - hidden_states = self.attention( - positions=positions, - hidden_states=hidden_states, - ) - - # Fully Connected - hidden_states, residual = self.ffn_norm(hidden_states, residual) - if visual_token_mask is not None and visual_token_mask.any(): - visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool() - text_token_mask = ~visual_token_mask - hidden_states[visual_token_mask] = self.feed_forward_ve( - hidden_states[visual_token_mask].reshape(-1, self.hidden_size) - ).flatten() - if text_token_mask.any(): - hidden_states[text_token_mask] = self.feed_forward( - hidden_states[text_token_mask].reshape(-1, self.hidden_size) - ).flatten() - else: - hidden_states = self.feed_forward(hidden_states) - return hidden_states, residual - - -class InternLM2VEModel(InternLM2Model): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__( - vllm_config=vllm_config, prefix=prefix, layer_type=InternLM2VEDecoderLayer - ) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - visual_token_mask: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.tok_embeddings(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - positions, - hidden_states, - residual, - visual_token_mask=visual_token_mask, - ) - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - -class InternLM2VEForCausalLM(InternLM2ForCausalLM): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__( - vllm_config=vllm_config, prefix=prefix, model_type=InternLM2VEModel - ) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index d57614ea980..94f03a539cb 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -23,7 +23,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.awq import AWQConfig from vllm.model_executor.models.intern_vit import ( InternVisionModel, - InternVisionPatchModel, ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY @@ -582,14 +581,10 @@ class InternVLChatModel( self.downsample_ratio = config.downsample_ratio self.ps_version = config.ps_version - llm_arch_name = config.text_config.architectures[0] - self.is_mono = llm_arch_name == "InternLM2VEForCausalLM" - with self._mark_tower_model(vllm_config, {"image", "video"}): self.vision_model = self._init_vision_model( config, quant_config=quant_config, - is_mono=self.is_mono, prefix=maybe_prefix(prefix, "vision_model"), ) self.mlp1 = self._init_mlp1(config) @@ -604,7 +599,6 @@ class InternVLChatModel( self.img_context_token_id = None self.video_context_token_id = None - self.visual_token_mask = None self.make_empty_intermediate_tensors = ( self.language_model.make_empty_intermediate_tensors ) @@ -627,26 +621,22 @@ class InternVLChatModel( config: PretrainedConfig, quant_config: QuantizationConfig | None, *, - is_mono: bool, prefix: str, ): - if not is_mono: - vision_feature_layer = config.select_layer - if vision_feature_layer < 0: - num_hidden_layers = ( - config.vision_config.num_hidden_layers + vision_feature_layer + 1 - ) - else: - num_hidden_layers = vision_feature_layer + 1 - - return InternVisionModel( - config.vision_config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers, - prefix=prefix, + vision_feature_layer = config.select_layer + if vision_feature_layer < 0: + num_hidden_layers = ( + config.vision_config.num_hidden_layers + vision_feature_layer + 1 ) else: - return InternVisionPatchModel(config.vision_config) + num_hidden_layers = vision_feature_layer + 1 + + return InternVisionModel( + config.vision_config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers, + prefix=prefix, + ) def _init_mlp1(self, config: PretrainedConfig) -> nn.Module: vit_hidden_size = config.vision_config.hidden_size @@ -805,15 +795,6 @@ class InternVLChatModel( return modalities - def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None: - if self.is_mono: - assert self.img_context_token_id is not None - self.visual_token_mask = (input_ids == self.img_context_token_id).reshape( - -1, 1 - ) - else: - self.visual_token_mask = None - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: modalities = self._parse_and_validate_multimodal_inputs(**kwargs) if not modalities: @@ -844,9 +825,6 @@ class InternVLChatModel( *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - if multimodal_embeddings is not None and len(multimodal_embeddings) > 0: - self._set_visual_token_mask(input_ids) - # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) @@ -875,11 +853,6 @@ class InternVLChatModel( "inputs_embeds": inputs_embeds, } - # Only required if the model is mono-architecture - if self.visual_token_mask is not None: - forward_kwargs.update({"visual_token_mask": self.visual_token_mask}) - self.visual_token_mask = None - hidden_states = self.language_model.model(**forward_kwargs) return hidden_states diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py index dafa0f03ae9..67b0ac5033f 100644 --- a/vllm/model_executor/models/jais2.py +++ b/vllm/model_executor/models/jais2.py @@ -225,7 +225,6 @@ class Jais2DecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index 39044f5e8b4..c35896264a9 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -268,7 +268,6 @@ class LlamaDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 7c6d5363c0a..b5f618699cf 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -47,9 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType -from vllm.v1.attention.backends.flash_attn_diffkv import ( - FlashAttentionDiffKVBackend, -) +from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interfaces import MixtureOfExperts, SupportsPP from .utils import ( @@ -292,11 +290,27 @@ class MiMoV2Attention(nn.Module): sliding_window = sliding_window_size if sliding_window_size > -1 else None - # Use DiffKV backend when V has a different head dim than K + # Use DiffKV backend when V has a different head dim than K. + # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall + # back to TRITON_ATTN_DIFFKV. Users can force a choice via + # `--attention-backend `. if self.v_head_dim != self.head_dim: - FlashAttentionDiffKVBackend.set_head_size_v(self.v_head_dim) - attn_backend = FlashAttentionDiffKVBackend - logger.info_once("Using FlashAttentionDiffKVBackend for attention.") + requested = get_current_vllm_config().attention_config.backend + if requested is not None and requested.name.endswith("_DIFFKV"): + backend_enum = requested + else: + fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class() + if fa_backend.is_supported_on_current_device( + head_size=self.head_dim, + head_size_v=self.v_head_dim, + has_sinks=self.attention_sink_bias is not None, + ): + backend_enum = AttentionBackendEnum.FLASH_ATTN_DIFFKV + else: + backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV + attn_backend = backend_enum.get_class() + attn_backend.set_head_size_v(self.v_head_dim) + logger.info_once("Using %s for attention.", attn_backend.get_name()) else: attn_backend = None diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index 3fcc048f9fa..bde5bc9451f 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -75,6 +75,16 @@ class EagleMistralLarge3Model(DeepseekV2Model): ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.aux_hidden_state_layers: tuple[int, ...] = () + + # Needed by load_weights + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0) + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0) + self.use_mha = config.model_type == "deepseek" or all( + dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim) + ) + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 742dccc36f1..797826c6bf5 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -19,7 +19,7 @@ import math from collections.abc import Iterable, Mapping from itertools import tee -from typing import Annotated, Literal +from typing import Annotated, Any, Literal import torch from torch import nn @@ -78,6 +78,7 @@ from .interfaces import ( MixtureOfExperts, MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -105,7 +106,7 @@ class Llama4ImagePatchInputs(TensorSchema): patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")] """ - The number of total patches for each image in the batch. + The number of chunked image tiles for each image in the batch. This is used to split the embeddings which has the first two dimensions flattened just like `pixel_values`. @@ -731,6 +732,7 @@ class Llama4ForConditionalGeneration( SupportsMultiModal, SupportsPP, MixtureOfExperts, + SupportsEncoderCudaGraph, SupportsEagle3, SupportsLoRA, ): @@ -828,10 +830,161 @@ class Llama4ForConditionalGeneration( num_physical_experts, num_local_physical_experts ) + def get_image_patches_per_chunk(self) -> int: + return Mllama4ProcessingInfo.get_patch_per_chunk(self.config.vision_config) + + def encode_image_chunks( + self, + pixel_values: torch.Tensor, + *, + use_data_parallel: bool, + ) -> torch.Tensor: + if use_data_parallel: + vision_embeddings = run_dp_sharded_vision_model( + pixel_values, self.vision_model + ) + else: + vision_embeddings = self.vision_model(pixel_values) + + return self.multi_modal_projector(vision_embeddings) + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphConfig, + ) + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixel_values"], + out_hidden_size=self.config.text_config.hidden_size, + ) + + def get_input_modality( + self, + mm_kwargs: dict[str, Any], + ) -> str: + return "image" + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = self.get_image_patches_per_chunk() + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + patches_per_chunk = self.get_image_patches_per_chunk() + return [ + EncoderItemSpec( + input_size=num_chunks, + output_tokens=num_chunks * patches_per_chunk, + ) + for num_chunks in mm_kwargs["patches_per_image"].tolist() + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + patches_per_image = mm_kwargs["patches_per_image"] + + if len(indices) == 0: + return { + "pixel_values": pixel_values[:0], + "patches_per_image": patches_per_image[:0], + } + + cum_chunks = [0] + for num_chunks in patches_per_image.tolist(): + cum_chunks.append(cum_chunks[-1] + num_chunks) + + selected_pixel_values = torch.cat( + [pixel_values[cum_chunks[i] : cum_chunks[i + 1]] for i in indices], + dim=0, + ) + + return { + "pixel_values": selected_pixel_values, + "patches_per_image": patches_per_image[indices], + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + vision_config = self.config.vision_config + patches_per_chunk = self.get_image_patches_per_chunk() + chunks_per_capture = max( + 1, (token_budget + patches_per_chunk - 1) // patches_per_chunk + ) + dummy_pixel_values = torch.randn( + chunks_per_capture, + vision_config.num_channels, + vision_config.image_size, + vision_config.image_size, + device=device, + dtype=dtype, + ) + + return EncoderCudaGraphCaptureInputs( + values={"pixel_values": dummy_pixel_values}, + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + return EncoderCudaGraphReplayBuffers( + values={"pixel_values": mm_kwargs["pixel_values"]}, + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + ) -> torch.Tensor: + return self.encode_image_chunks( + inputs["pixel_values"], + use_data_parallel=False, + ).flatten(0, 1) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + return self.encode_image_chunks( + mm_kwargs["pixel_values"], + use_data_parallel=False, + ).flatten(0, 1) + def _parse_and_validate_image_input( self, **kwargs: object ) -> Llama4ImagePatchInputs | None: - # num_images, 1, num_chunks, channel, image_size, image_size + # total_num_chunks, channel, image_size, image_size pixel_values = kwargs.pop("pixel_values", None) if pixel_values is None: return None @@ -853,15 +1006,10 @@ class Llama4ForConditionalGeneration( pixel_values = image_input["pixel_values"] patches_per_image = image_input["patches_per_image"].tolist() - # shard image input - if self.use_data_parallel: - vision_embeddings_flat = run_dp_sharded_vision_model( - pixel_values, self.vision_model - ) - else: - vision_embeddings_flat = self.vision_model(pixel_values) - - vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat) + vision_embeddings_flat = self.encode_image_chunks( + pixel_values, + use_data_parallel=self.use_data_parallel, + ) return [ img.flatten(0, 1) diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index 7b2e6b93b27..f5c526e33ed 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -237,7 +237,6 @@ class NemotronDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py index b974a3eb085..06a2096ec69 100644 --- a/vllm/model_executor/models/nemotron_nas.py +++ b/vllm/model_executor/models/nemotron_nas.py @@ -141,7 +141,6 @@ class DeciLMDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/nvlm_d.py b/vllm/model_executor/models/nvlm_d.py index 9fd4cf0797d..2222ab09e1e 100644 --- a/vllm/model_executor/models/nvlm_d.py +++ b/vllm/model_executor/models/nvlm_d.py @@ -177,27 +177,22 @@ class NVLM_D_Model(InternVLChatModel): config: PretrainedConfig, quant_config: QuantizationConfig | None, *, - is_mono: bool, prefix: str, ): - if not is_mono: - vision_feature_layer = config.select_layer - if vision_feature_layer < 0: - num_hidden_layers = ( - config.vision_config.num_hidden_layers + vision_feature_layer + 1 - ) - else: - num_hidden_layers = vision_feature_layer + 1 - - # We added additional dummy heads to the original num of heads to - # make the number of heads divisible by 8. - return InternVisionModel( - config.vision_config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers, - num_dummy_heads=7, - prefix=prefix, + vision_feature_layer = config.select_layer + if vision_feature_layer < 0: + num_hidden_layers = ( + config.vision_config.num_hidden_layers + vision_feature_layer + 1 ) else: - msg = "Monolith mode is not applicable to NVLM_D" - raise NotImplementedError(msg) + num_hidden_layers = vision_feature_layer + 1 + + # We added additional dummy heads to the original num of heads to + # make the number of heads divisible by 8. + return InternVisionModel( + config.vision_config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers, + num_dummy_heads=7, + prefix=prefix, + ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e1ce0efae2f..ecdbe3991c9 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -140,9 +140,7 @@ _TEXT_GENERATION_MODELS = { "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"), "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"), "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"), - "InternLMForCausalLM": ("llama", "LlamaForCausalLM"), "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"), - "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"), "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), @@ -401,6 +399,10 @@ _MULTIMODAL_MODELS = { "gemma3n_mm", "Gemma3nForConditionalGeneration", ), + "DiffusionGemmaForBlockDiffusion": ( + "diffusion_gemma", + "DiffusionGemmaForConditionalGeneration", + ), "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"), "Gemma4UnifiedForConditionalGeneration": ( "gemma4_unified", @@ -713,8 +715,10 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "ErnieModel": "0.23.0", "ErnieForSequenceClassification": "0.23.0", "ErnieForTokenClassification": "0.23.0", + "InternLM2VEForCausalLM": "0.23.0", "QWenLMHeadModel": "0.23.0", "QwenVLForConditionalGeneration": "0.23.0", + "InternLMForCausalLM": "0.23.0", # encoder-decoder models except whisper # have been removed for V0 deprecation. "DonutForConditionalGeneration": "0.10.2", diff --git a/vllm/model_executor/models/skyworkr1v.py b/vllm/model_executor/models/skyworkr1v.py index a3415a20a96..685b980c3f8 100644 --- a/vllm/model_executor/models/skyworkr1v.py +++ b/vllm/model_executor/models/skyworkr1v.py @@ -22,7 +22,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.awq import AWQConfig from vllm.model_executor.models.intern_vit import ( InternVisionModel, - InternVisionPatchModel, ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.processing import BaseDummyInputsBuilder @@ -178,14 +177,10 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): self.downsample_ratio = config.downsample_ratio self.ps_version = config.ps_version - llm_arch_name = config.text_config.architectures[0] - self.is_mono = llm_arch_name == "SkyworkLM2VEForCausalLM" - with self._mark_tower_model(vllm_config, "image"): self.vision_model = self._init_vision_model( config, quant_config=quant_config, - is_mono=self.is_mono, prefix=maybe_prefix(prefix, "vision_model"), ) self.mlp1 = self._init_mlp1( @@ -223,26 +218,22 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): config: PretrainedConfig, quant_config: QuantizationConfig | None, *, - is_mono: bool, prefix: str, ): - if not is_mono: - vision_feature_layer = config.select_layer - if vision_feature_layer < 0: - num_hidden_layers = ( - config.vision_config.num_hidden_layers + vision_feature_layer + 1 - ) - else: - num_hidden_layers = vision_feature_layer + 1 - - return InternVisionModel( - config.vision_config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers, - prefix=prefix, + vision_feature_layer = config.select_layer + if vision_feature_layer < 0: + num_hidden_layers = ( + config.vision_config.num_hidden_layers + vision_feature_layer + 1 ) else: - return InternVisionPatchModel(config.vision_config) + num_hidden_layers = vision_feature_layer + 1 + + return InternVisionModel( + config.vision_config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers, + prefix=prefix, + ) def _init_mlp1( self, @@ -363,14 +354,6 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): ] return image_embeds.split(image_feature_sizes) - def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None: - if self.is_mono: - self.visual_token_mask = (input_ids == self.img_context_token_id).reshape( - -1, 1 - ) - else: - self.visual_token_mask = None - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: image_input = self._parse_and_validate_image_input(**kwargs) if image_input is None: @@ -385,9 +368,6 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP): *, is_multimodal: torch.Tensor | None = None, ) -> torch.Tensor: - if multimodal_embeddings is not None and len(multimodal_embeddings) > 0: - self._set_visual_token_mask(input_ids) - # This is to satisfy the type checker for each overload if multimodal_embeddings is None or is_multimodal is None: return super().embed_input_ids(input_ids) diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py index 454a0e97112..fcb2ae429cb 100644 --- a/vllm/model_executor/models/solar.py +++ b/vllm/model_executor/models/solar.py @@ -198,7 +198,6 @@ class SolarDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/models/deepseek_v4/xpu/mtp.py b/vllm/models/deepseek_v4/xpu/mtp.py index 8dbe40bb6ae..d4a8d293baf 100644 --- a/vllm/models/deepseek_v4/xpu/mtp.py +++ b/vllm/models/deepseek_v4/xpu/mtp.py @@ -18,7 +18,6 @@ import regex as re import torch import torch.nn as nn -from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -39,6 +38,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.deepseek_v4.common.ops import ( + fused_mtp_input_rmsnorm, + mtp_shared_head_rmsnorm, +) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -87,6 +90,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): bias=False, return_bias=False, quant_config=quant_config, + prefix=f"{prefix}.e_proj", ) self.h_proj = ReplicatedLinear( config.hidden_size, @@ -94,6 +98,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): bias=False, return_bias=False, quant_config=quant_config, + prefix=f"{prefix}.h_proj", ) self.hc_eps = config.hc_eps @@ -133,22 +138,31 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masking inputs at position 0, as not needed by MTP - inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) - inputs_embeds = self.enorm(inputs_embeds) - # Target stashes pre-hc_head residual as flat (T, hc_mult * D); - # reshape to (T, hc_mult, D) — the training-time layout. + # reshape to (T, hc_mult, D) — the training-time layout — before + # the fused norm pass so both inputs are 3D-friendly. previous_hidden_states = previous_hidden_states.view( -1, self.hc_mult, self.config.hidden_size ) - previous_hidden_states = self.hnorm(previous_hidden_states) + # Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm. + inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm( + inputs_embeds, + positions, + previous_hidden_states, + self.enorm.weight.data, + self.hnorm.weight.data, + self.enorm.variance_epsilon, + self.hc_mult, + ) hidden_states = self.h_proj(previous_hidden_states) + self.e_proj( inputs_embeds ).unsqueeze(-2) hidden_states, residual, post_mix, res_mix = self.mtp_block( positions=positions, x=hidden_states, input_ids=None ) + hidden_states = self.mtp_block.hc_post( + hidden_states, residual, post_mix, res_mix + ) # Return the flat pre-hc_head residual so it can be re-fed as the # next spec step's `previous_hidden_states` when # num_speculative_tokens > 1. hc_head is deferred to compute_logits. @@ -238,13 +252,15 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): mtp_layer.rms_norm_eps, mtp_layer.hc_eps, ) - logits = self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) + hidden_states = mtp_shared_head_rmsnorm( + hidden_states, + mtp_layer.shared_head.norm.weight.data, + mtp_layer.shared_head.norm.variance_epsilon, ) + logits = self.logits_processor(mtp_layer.shared_head.head, hidden_states) return logits -@support_torch_compile class DeepSeekV4MTP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -285,11 +301,6 @@ class DeepSeekV4MTP(nn.Module): ".emb.tok_emb.weight": ".embed_tokens.weight", ".head.weight": ".shared_head.head.weight", ".norm.weight": ".shared_head.norm.weight", - # Pre-MoE norm + gate are now owned by - # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear). - ".ffn_norm.weight": ".ffn.norm_gate.norm.weight", - ".ffn.gate.weight": ".ffn.norm_gate.gate.weight", - ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid", } def _remap_weight_name(name: str) -> str: @@ -437,11 +448,11 @@ class DeepSeekV4MTP(nn.Module): ".shared_experts.w2", ".shared_experts.down_proj" ) if name.endswith(".ffn.gate.bias"): - # ``e_score_correction_bias`` lives on - # ``norm_gate`` directly (not on the inner gate). + # ``e_score_correction_bias`` lives on the gate + # under a different attribute name. name = name.replace( ".ffn.gate.bias", - ".ffn.norm_gate.e_score_correction_bias", + ".ffn.gate.e_score_correction_bias", ) param = params_dict[name] weight_loader = getattr( diff --git a/vllm/multimodal/media/audio.py b/vllm/multimodal/media/audio.py index 1a7d6d95071..5e998be3fcb 100644 --- a/vllm/multimodal/media/audio.py +++ b/vllm/multimodal/media/audio.py @@ -92,8 +92,9 @@ def load_audio_pyav( raise ValueError( f"Audio exceeds maximum allowed duration of " f"{max_duration_s}s (metadata reports " - f"{metadata_duration_s:.1f}s). This limit " - f"prevents decompression-bomb attacks." + f"{metadata_duration_s:.1f}s). Set " + f"VLLM_MAX_AUDIO_DECODE_DURATION_S to " + f"increase this limit." ) max_samples = ( @@ -129,8 +130,9 @@ def load_audio_pyav( raise ValueError( f"Audio exceeds maximum allowed duration of " f"{max_duration_s}s (decoded {total_samples} " - f"samples at {sr}Hz). This limit prevents " - f"decompression-bomb attacks." + f"samples at {sr}Hz). Set " + f"VLLM_MAX_AUDIO_DECODE_DURATION_S to " + f"increase this limit." ) except (ValueError, ImportError): raise @@ -166,8 +168,9 @@ def load_audio_soundfile( raise ValueError( f"Audio exceeds maximum allowed duration of " f"{max_duration_s}s (file contains " - f"{file_duration_s:.1f}s at {native_sr}Hz). " - f"This limit prevents decompression-bomb attacks." + f"{file_duration_s:.1f}s at {native_sr}Hz). Set " + f"VLLM_MAX_AUDIO_DECODE_DURATION_S to " + f"increase this limit." ) y = f.read(dtype="float32", always_2d=False).T diff --git a/vllm/parser/__init__.py b/vllm/parser/__init__.py index de815b2e1fd..e13c2ece9f0 100644 --- a/vllm/parser/__init__.py +++ b/vllm/parser/__init__.py @@ -5,10 +5,12 @@ from vllm.parser.abstract_parser import ( DelegatingParser, Parser, ) +from vllm.parser.harmony import HarmonyParser from vllm.parser.parser_manager import ParserManager __all__ = [ "Parser", "DelegatingParser", + "HarmonyParser", "ParserManager", ] diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 48db01c14e0..6deba14ceaf 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -25,6 +25,7 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.parser.metrics import record_tool_parser_invocation from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.streaming import ( @@ -282,6 +283,7 @@ class Parser: model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: """Parse a complete model output, extracting reasoning and tool calls. @@ -289,6 +291,7 @@ class Parser: model_output: The complete model-generated string. request: The request object used to generate the output. enable_auto_tools: Whether to enable automatic tool call parsing. + model_output_token_ids: The generated raw output token IDs. Returns: A tuple of (reasoning, content, tool_calls). @@ -425,10 +428,50 @@ class DelegatingParser(Parser): ) -> ChatCompletionRequest | ResponsesRequest: if self._reasoning_parser is not None: request = self._reasoning_parser.adjust_request(request) + if self._tool_parser is not None: + request = self._apply_structural_tag(request) if self._tool_parser is not None: request = self._tool_parser.adjust_request(request) return request + def _apply_structural_tag( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + if ( + self._tool_parser is None + or self._tool_parser.structural_tag_model is None + or not request.tools + ): + return request + + need_tool_calling = ( + request.tool_choice == "auto" + or request.tool_choice == "required" + or isinstance( + request.tool_choice, + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ) + ) + if not need_tool_calling: + return request + + structure_tag = self._tool_parser.get_structural_tag( + request, + reasoning=False, + ) + if structure_tag is None: + return request + + structural_tag = json.dumps(structure_tag.model_dump()) + request.structured_outputs = StructuredOutputsParams( + structural_tag=structural_tag, + ) + if isinstance(request, ResponsesRequest): + request.text = None + else: + request.response_format = None + return request + def extract_reasoning_streaming( self, previous_text: str, @@ -642,6 +685,7 @@ class DelegatingParser(Parser): model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: reasoning, content = self.extract_reasoning(model_output, request) tool_calls, content = self._extract_tool_calls( diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py new file mode 100644 index 00000000000..f19d3675dab --- /dev/null +++ b/vllm/parser/harmony.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum, auto +from typing import TYPE_CHECKING, NamedTuple + +from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionCall, +) +from vllm.entrypoints.openai.parser.harmony_utils import ( + extract_function_from_recipient, + get_streamable_parser_for_assistant, + is_function_recipient, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser + +if TYPE_CHECKING: + from openai_harmony import Message, Role + from openai_harmony import StreamState as HarmonyStreamState + + +class _SegmentType(Enum): + TOOL = auto() + REASONING = auto() + CONTENT = auto() + IGNORE = auto() + + @staticmethod + def from_channel_and_recipient( + channel: str | None, recipient: str | None + ) -> _SegmentType: + if recipient and is_function_recipient(recipient): + return _SegmentType.TOOL + if channel == "analysis": + return _SegmentType.REASONING + if channel == "final" or (channel == "commentary" and recipient is None): + return _SegmentType.CONTENT + return _SegmentType.IGNORE + + +class Segment(NamedTuple): + channel: str | None + recipient: str | None + delta: str + completed_message: Message | None = None + + +@dataclass +class ChunkResult: + segments: list[Segment] + reasoning_token_count: int + + +class HarmonyParser(DelegatingParser): + def __init__(self, tokenizer, tools=None, *args, **kwargs): + super().__init__(tokenizer, tools, *args, **kwargs) + + if self._reasoning_parser and not isinstance( + self._reasoning_parser, GptOssReasoningParser + ): + raise ValueError( + "Harmony requires GptOssReasoningParser, " + f"got {self._reasoning_parser.__class__.__name__}." + ) + + if self._tool_parser and not isinstance(self._tool_parser, GptOssToolParser): + raise ValueError( + "Harmony requires GptOssToolParser, " + f"got {self._tool_parser.__class__.__name__}." + ) + + self._harmony_parser = get_streamable_parser_for_assistant() + self._next_tool_call_index = 0 + self._num_processed_messages = 0 + + @property + def state(self) -> HarmonyStreamState: + return self._harmony_parser.state + + @property + def current_role(self) -> Role | None: + return self._harmony_parser.current_role + + @property + def current_channel(self) -> str | None: + return self._harmony_parser.current_channel + + @property + def current_recipient(self) -> str | None: + return self._harmony_parser.current_recipient + + @property + def current_content(self) -> str: + return self._harmony_parser.current_content + + @property + def current_content_type(self) -> str | None: + return self._harmony_parser.current_content_type + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + """Parse Harmony output from token IDs. + + Tool calls are always extracted regardless of ``enable_auto_tools``. + Callers must decide whether to surface them. + """ + result = self.process_chunk(model_output_token_ids) + + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + tool_calls: list[FunctionCall] = [] + + def _append_parsed_message( + channel: str | None, + recipient: str | None, + text: str, + content_type: str | None = None, + ) -> None: + segment_type = _SegmentType.from_channel_and_recipient(channel, recipient) + match segment_type: + case _SegmentType.REASONING if self.reasoning_parser and text: + reasoning_parts.append(text) + case _SegmentType.CONTENT if text: + content_parts.append(text) + case _SegmentType.TOOL if self.tool_parser: + assert recipient is not None + if content_type is not None and "json" not in content_type: + arguments = text + else: + try: + arguments = json.dumps(json.loads(text)) + except json.JSONDecodeError: + arguments = text + tool_calls.append( + FunctionCall( + name=extract_function_from_recipient(recipient), + arguments=arguments, + ) + ) + + for segment in result.segments: + msg = segment.completed_message + if msg is None: + continue + if msg.author.role != "assistant" or not msg.content: + continue + _append_parsed_message( + channel=msg.channel, + recipient=msg.recipient, + text=msg.content[0].text, + content_type=msg.content_type, + ) + + if ( + self.current_channel is not None + or self.current_recipient is not None + or self.current_content + ): + _append_parsed_message( + channel=self.current_channel, + recipient=self.current_recipient, + text=self.current_content, + content_type=self.current_content_type, + ) + + reasoning = "\n".join(reasoning_parts) or None + content = "\n".join(content_parts) or None + return reasoning, content, tool_calls or None + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + prev_recipient = self.current_recipient + result = self.process_chunk(delta_token_ids) + combined_content = "" + combined_reasoning = "" + tool_messages: list[DeltaToolCall] = [] + + for segment in result.segments: + if segment.completed_message is not None: + prev_recipient = None + continue + + segment_type = _SegmentType.from_channel_and_recipient( + segment.channel, segment.recipient + ) + match segment_type: + case _SegmentType.REASONING: + combined_reasoning += segment.delta + case _SegmentType.CONTENT: + combined_content += segment.delta + case _SegmentType.TOOL: + assert segment.recipient is not None + if prev_recipient != segment.recipient: + tool_name = extract_function_from_recipient(segment.recipient) + tool_messages.append( + DeltaToolCall( + # HarmonyParser does not use _stream_state; + # "random" tool_call_id_type is always used + id=make_tool_call_id(), + type="function", + function=DeltaFunctionCall( + name=tool_name, + arguments=segment.delta, + ), + index=self._next_tool_call_index, + ) + ) + self._next_tool_call_index += 1 + prev_recipient = segment.recipient + elif segment.delta: + tool_call_index = self._next_tool_call_index - 1 + tool_messages.append( + DeltaToolCall( + index=tool_call_index, + function=DeltaFunctionCall(arguments=segment.delta), + ) + ) + + if not combined_content and not combined_reasoning and not tool_messages: + return None + + delta_message = DeltaMessage() + if combined_content: + delta_message.content = combined_content + if combined_reasoning: + delta_message.reasoning = combined_reasoning + if tool_messages: + delta_message.tool_calls = tool_messages + return delta_message + + def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: + if not token_ids: + return ChunkResult(segments=[], reasoning_token_count=0) + + segments: list[Segment] = [] + reasoning_token_count = 0 + for token_id in token_ids: + self._harmony_parser.process(token_id) + channel = self.current_channel + recipient = self.current_recipient + delta = self._harmony_parser.last_content_delta or "" + completed_message = None + _messages = self._harmony_parser.messages + if len(_messages) > self._num_processed_messages: + completed_message = _messages[self._num_processed_messages] + self._num_processed_messages += 1 + + if channel == "analysis" or ( + channel == "commentary" and recipient is not None + ): + reasoning_token_count += 1 + + segments.append( + Segment( + channel=channel, + recipient=recipient, + delta=delta, + completed_message=completed_message, + ) + ) + + # TODO: Optionally merge and suppress empty Segments + + return ChunkResult( + segments=segments, + reasoning_token_count=reasoning_token_count, + ) diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index c7f557a5a95..52f16136ee3 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -3,6 +3,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING from vllm.entrypoints.openai.engine.protocol import DeltaMessage, FunctionCall @@ -43,10 +44,14 @@ class MistralParser(DelegatingParser): model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: self._maybe_force_auto_tool_parsing(request) reasoning, content, tool_calls = super().parse( - model_output, request, enable_auto_tools + model_output, + request, + enable_auto_tools, + model_output_token_ids, ) if tool_calls: from vllm.tool_parsers.mistral_tool_parser import MistralToolCall diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 6c2fdf52dd3..1b5133f5a8f 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -79,6 +79,7 @@ class ParserManager: reasoning_parser_name: str | None = None, enable_auto_tools: bool = False, model_name: str | None = None, + is_harmony: bool = False, ) -> type[Parser] | None: """ Get a Parser that handles both reasoning and tool parsing. @@ -91,6 +92,8 @@ class ParserManager: reasoning_parser_name: The name of the reasoning parser. enable_auto_tools: Whether auto tool choice is enabled. model_name: The model name for parser-specific warnings. + is_harmony: Whether the selected model uses the Harmony format. + If True, HarmonyParser is always returned. Returns: A Parser class, or None if neither parser is specified. @@ -108,6 +111,13 @@ class ParserManager: from vllm.utils.mistral import is_mistral_tool_parser + if is_harmony: + from vllm.parser.harmony import HarmonyParser + + HarmonyParser.reasoning_parser_cls = reasoning_parser_cls + HarmonyParser.tool_parser_cls = tool_parser_cls + return HarmonyParser + if is_mistral_tool_parser(tool_parser_cls): from vllm.parser.mistral import MistralParser diff --git a/vllm/reasoning/abs_reasoning_parsers.py b/vllm/reasoning/abs_reasoning_parsers.py index 8edbc5f82ef..74b3e62abc2 100644 --- a/vllm/reasoning/abs_reasoning_parsers.py +++ b/vllm/reasoning/abs_reasoning_parsers.py @@ -181,9 +181,8 @@ class ReasoningParser: ) -> str | None: """ Instance method that is implemented for preparing the structured tag - Otherwise, None is returned """ - return None + return original_tag class ReasoningParserManager: diff --git a/vllm/reasoning/gptoss_reasoning_parser.py b/vllm/reasoning/gptoss_reasoning_parser.py index 1ba933cca31..d7bdca82912 100644 --- a/vllm/reasoning/gptoss_reasoning_parser.py +++ b/vllm/reasoning/gptoss_reasoning_parser.py @@ -8,7 +8,6 @@ from transformers import PreTrainedTokenizerBase from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.entrypoints.openai.parser.harmony_utils import parse_chat_output from vllm.logger import init_logger from vllm.reasoning import ReasoningParser @@ -132,10 +131,10 @@ class GptOssReasoningParser(ReasoningParser): return self.is_reasoning_end(input_ids[n - window :]) def extract_content_ids(self, input_ids: list[int]) -> list[int]: - _, content, _ = parse_chat_output(input_ids) - if content is None: - return [] - return self.model_tokenizer.encode(content) + raise NotImplementedError( + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." + ) def extract_reasoning_streaming( self, @@ -146,25 +145,10 @@ class GptOssReasoningParser(ReasoningParser): current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - prev_reasoning, prev_content, _ = parse_chat_output(list(previous_token_ids)) - cur_reasoning, cur_content, _ = parse_chat_output(list(current_token_ids)) - reasoning_delta = None - content_delta = None - if cur_reasoning is not None: - prev_r = prev_reasoning or "" - if cur_reasoning.startswith(prev_r): - reasoning_delta = cur_reasoning[len(prev_r) :] or None - else: - reasoning_delta = cur_reasoning - if cur_content is not None: - prev_c = prev_content or "" - if cur_content.startswith(prev_c): - content_delta = cur_content[len(prev_c) :] or None - else: - content_delta = cur_content - if reasoning_delta is None and content_delta is None: - return None - return DeltaMessage(reasoning=reasoning_delta, content=content_delta) + raise NotImplementedError( + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." + ) def extract_reasoning( self, @@ -172,7 +156,8 @@ class GptOssReasoningParser(ReasoningParser): request: "ChatCompletionRequest | ResponsesRequest", ) -> tuple[str | None, str | None]: raise NotImplementedError( - "gpt-oss has a special branch for parsing reasoning in non-streaming mode. This method shouldn't be used." # noqa: E501 + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." ) # This function prepares the structural tag to format reasoning output diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 3c1ff8ac9c3..17204093ab1 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1048,3 +1048,4 @@ class BeamSearchParams( temperature: float = 0.0 length_penalty: float = 1.0 include_stop_str_in_output: bool = False + structured_outputs: StructuredOutputsParams | None = None diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index bf832f178be..6d122b4695d 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -119,8 +119,8 @@ _TOOL_PARSERS_TO_REGISTER = { "LongcatFlashToolParser", ), "mimo": ( - "qwen3xml_tool_parser", - "Qwen3XMLToolParser", + "qwen3coder_tool_parser", + "Qwen3CoderToolParser", ), "minimax_m2": ( "minimax_m2_tool_parser", @@ -143,8 +143,8 @@ _TOOL_PARSERS_TO_REGISTER = { "Olmo3PythonicToolParser", ), "openai": ( - "openai_tool_parser", - "OpenAIToolParser", + "gptoss_tool_parser", + "GptOssToolParser", ), "phi4_mini_json": ( "phi4mini_tool_parser", @@ -159,8 +159,8 @@ _TOOL_PARSERS_TO_REGISTER = { "Qwen3CoderToolParser", ), "qwen3_xml": ( - "qwen3xml_tool_parser", - "Qwen3XMLToolParser", + "qwen3coder_tool_parser", + "Qwen3CoderToolParser", ), "seed_oss": ( "seed_oss_tool_parser", diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 94543b82350..3609bcbf457 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -6,6 +6,7 @@ import json import os from collections.abc import Callable, Sequence from functools import cached_property +from typing import Any from openai.types.responses import ( ResponseFormatTextJSONSchemaConfig, @@ -13,8 +14,8 @@ from openai.types.responses import ( ) from openai.types.responses.function_tool import FunctionTool +import vllm.envs as envs from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionToolsParam, ) @@ -25,7 +26,6 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ) -from vllm.envs import VLLM_ENFORCE_STRICT_TOOL_CALLING from vllm.logger import init_logger from vllm.sampling_params import ( StructuredOutputsParams, @@ -57,6 +57,17 @@ class ToolParser: # extract_tool_calls / extract_tool_calls_streaming methods for # required/named tool_choice, treating them the same as "auto". supports_required_and_named: bool = True + # xgrammar builtin structural tag model key. Subclasses set this when + # their parsed tool-call syntax matches a builtin xgrammar format. + structural_tag_model: str | None = None + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + if ( + cls.structural_tag_model is not None + and envs.VLLM_ENFORCE_STRICT_TOOL_CALLING + ): + cls.supports_required_and_named = False def __init__( self, @@ -112,32 +123,16 @@ class ToolParser: if not request.tools: return request - # Step 1 (highest priority for ChatCompletionRequest): apply - # vLLM-owned structural tag support for model-specific tool formats. + # Set structured output params when tool constraints are derived from + # the tool schema. Unified parsers handle model-specific structural + # tags before calling into the tool parser. + structured_outputs = getattr(request, "structured_outputs", None) if ( - isinstance(request, ChatCompletionRequest) - and VLLM_ENFORCE_STRICT_TOOL_CALLING + structured_outputs is not None + and structured_outputs.structural_tag is not None ): - need_tool_calling = ( - request.tool_choice == "auto" - or request.tool_choice == "required" - or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) - ) - if need_tool_calling: - structure_tag = self.get_structural_tag(request) - if structure_tag is not None: - if request.structured_outputs is None: - request.structured_outputs = StructuredOutputsParams( - structural_tag=json.dumps(structure_tag.model_dump()), - ) - else: - request.structured_outputs.structural_tag = json.dumps( - structure_tag.model_dump() - ) - return request + return request - # Step 2: set structured output params when tool constraints are - # derived from the tool schema. json_schema_from_tool = get_json_schema_from_tools( tool_choice=request.tool_choice, tools=request.tools ) @@ -169,8 +164,24 @@ class ToolParser: return request - def get_structural_tag(self, request: ChatCompletionRequest): - return None + def get_structural_tag( + self, + request: ChatCompletionRequest | ResponsesRequest, + *, + reasoning: bool = False, + ): + if self.structural_tag_model is None: + return None + if not envs.VLLM_ENFORCE_STRICT_TOOL_CALLING: + return None + from vllm.tool_parsers.structural_tag_registry import get_model_structural_tag + + return get_model_structural_tag( + model=self.structural_tag_model, + tools=request.tools, + tool_choice=request.tool_choice, + reasoning=reasoning, + ) def extract_tool_calls( self, model_output: str, request: ChatCompletionRequest diff --git a/vllm/tool_parsers/deepseekv31_tool_parser.py b/vllm/tool_parsers/deepseekv31_tool_parser.py index e4ade3aae98..05d33787478 100644 --- a/vllm/tool_parsers/deepseekv31_tool_parser.py +++ b/vllm/tool_parsers/deepseekv31_tool_parser.py @@ -25,6 +25,8 @@ logger = init_logger(__name__) class DeepSeekV31ToolParser(ToolParser): + structural_tag_model = "deepseek_v3_1" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/deepseekv32_tool_parser.py b/vllm/tool_parsers/deepseekv32_tool_parser.py index 7d5e299be88..c597ac61969 100644 --- a/vllm/tool_parsers/deepseekv32_tool_parser.py +++ b/vllm/tool_parsers/deepseekv32_tool_parser.py @@ -53,6 +53,7 @@ class DeepSeekV32ToolParser(ToolParser): tool_call_start_token: str = "<｜DSML｜function_calls>" tool_call_end_token: str = "" + structural_tag_model = "deepseek_v3_2" def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/deepseekv3_tool_parser.py b/vllm/tool_parsers/deepseekv3_tool_parser.py index e92af87e604..7eaa983df7e 100644 --- a/vllm/tool_parsers/deepseekv3_tool_parser.py +++ b/vllm/tool_parsers/deepseekv3_tool_parser.py @@ -28,6 +28,8 @@ logger = init_logger(__name__) class DeepSeekV3ToolParser(ToolParser): + structural_tag_model = "deepseek_r1" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/deepseekv4_tool_parser.py b/vllm/tool_parsers/deepseekv4_tool_parser.py index e32451cd8bb..2558f585f82 100644 --- a/vllm/tool_parsers/deepseekv4_tool_parser.py +++ b/vllm/tool_parsers/deepseekv4_tool_parser.py @@ -1,14 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser -from vllm.tool_parsers.structural_tag_registry import ( - get_enable_structured_outputs_in_reasoning, - get_model_structural_tag, -) class DeepSeekV4ToolParser(DeepSeekV32ToolParser): @@ -21,11 +14,4 @@ class DeepSeekV4ToolParser(DeepSeekV32ToolParser): tool_call_start_token: str = "<｜DSML｜tool_calls>" tool_call_end_token: str = "" - - def get_structural_tag(self, request: ChatCompletionRequest): - return get_model_structural_tag( - model="deepseek_v4", - tools=request.tools, - tool_choice=request.tool_choice, - reasoning=get_enable_structured_outputs_in_reasoning(), - ) + structural_tag_model = "deepseek_v4" diff --git a/vllm/tool_parsers/gemma4_tool_parser.py b/vllm/tool_parsers/gemma4_tool_parser.py index 9925284273f..a92ab9bb6cd 100644 --- a/vllm/tool_parsers/gemma4_tool_parser.py +++ b/vllm/tool_parsers/gemma4_tool_parser.py @@ -20,9 +20,11 @@ import json from collections.abc import Sequence import regex as re +from openai.types.responses import ToolChoiceFunction from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( @@ -343,6 +345,9 @@ class Gemma4ToolParser(ToolParser): tool parsers. """ + # Gemma4 emits native special-token tool calls, not generic JSON calls. + supports_required_and_named = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -390,6 +395,23 @@ class Gemma4ToolParser(ToolParser): def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: + if request.tools: + tc = request.tool_choice + if tc == "required" or isinstance( + tc, + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ): + # Do NOT call super().adjust_request() for required/named tool + # choice. The base implementation injects a JSON-array + # `structured_outputs` schema and forces xgrammar guided + # decoding, which conflicts with Gemma4's native + # `<|tool_call>call:...` (non-JSON) tool syntax and crashes + # EngineCore under MTP spec decode. The streaming/extraction + # parser already handles the native output, so guided decoding + # is skipped here (mirrors the GLM4 precedent). + if request.tool_choice != "none": + request.skip_special_tokens = False + return request request = super().adjust_request(request) if request.tools and request.tool_choice != "none": # Don't skip special tokens — <|tool_call> etc. are needed for @@ -549,22 +571,40 @@ class Gemma4ToolParser(ToolParser): return DeltaMessage(content=delta_text) return None - # Case 2: Starting a new tool call - if start_count > prev_start_count and start_count > end_count: - self.current_tool_id += 1 + # Case 2: One or more new tool calls started in this delta. + # A single delta can batch several complete calls, so advance the + # tool id once per newly-seen start token and allocate a tracking + # slot for each. + if start_count > prev_start_count: + num_new = start_count - prev_start_count + for _ in range(num_new): + self.current_tool_id += 1 + self.streamed_args_for_tool.append("") + self.prev_tool_call_arr.append({}) self.current_tool_name_sent = False - self.streamed_args_for_tool.append("") - self.prev_tool_call_arr.append({}) - logger.debug("Starting new tool call %d", self.current_tool_id) - # Don't return yet — fall through to try parsing if there's - # content after <|tool_call> in this same delta - # (but usually it's just the token itself, so return None) - if len(delta_text) <= len(self.tool_call_start_token): + logger.debug( + "Started %d new tool call(s); current_tool_id=%d", + num_new, + self.current_tool_id, + ) + # Don't return yet if this delta also contains call payload or + # the end marker; backends can batch one or more complete tool + # calls into a single streaming chunk. Only wait for more text + # when the delta is just the start token itself. + if start_count > end_count and len(delta_text) <= len( + self.tool_call_start_token + ): return None - # Case 3: Tool call just ended + # Case 3: One or more tool calls just ended (possibly several in a + # single batched delta) — drain every newly-completed call. if end_count > prev_end_count: - return self._handle_tool_call_end(current_text) + return self._handle_tool_call_end( + current_text, + prev_end_count=prev_end_count, + end_count=end_count, + start_count=start_count, + ) # Case 4: In the middle of a tool call — parse partial content if start_count > end_count: @@ -652,45 +692,111 @@ class Gemma4ToolParser(ToolParser): return None - def _handle_tool_call_end(self, current_text: str) -> DeltaMessage | None: - """Handle streaming when a tool call has just completed. + def _handle_tool_call_end( + self, + current_text: str, + prev_end_count: int, + end_count: int, + start_count: int, + ) -> DeltaMessage | None: + """Handle streaming when one or more tool calls have just completed. - Performs a final parse of the complete tool call and flushes - any remaining un-streamed argument fragments. + A single streaming delta can batch several complete tool calls + (``<|tool_call>...<|tool_call>...``). Every + call whose ```` end marker arrived in this delta — i.e. + those with index in ``[prev_end_count, end_count)`` — is drained and + emitted, with one ``DeltaToolCall`` per call in a single + ``DeltaMessage`` (this matches the OpenAI streaming wire format, and + the serving layer iterates over ``delta.tool_calls``). + + Per call: + + * If the function name was already streamed incrementally (the + token-by-token path), only the remaining argument fragment is + flushed as a diff. + * If the call is seen complete for the first time in this delta (the + batched-complete path), the id + name + full arguments JSON are + emitted exactly once. """ - if self.current_tool_id < 0 or self.current_tool_id >= len( - self.prev_tool_call_arr - ): - logger.debug( - "Tool call end detected but no active tool call (current_tool_id=%d)", - self.current_tool_id, - ) + # Parse the complete tool calls using regex for accuracy. + all_matches = self.tool_call_regex.findall(current_text) + if not all_matches: + logger.debug("Tool call end detected but no complete tool call parsed yet.") return None - # Parse the complete tool call using regex for accuracy - all_matches = self.tool_call_regex.findall(current_text) - if self.current_tool_id < len(all_matches): - _, args_str = all_matches[self.current_tool_id] + deltas: list[DeltaToolCall] = [] + for idx in range(prev_end_count, end_count): + if idx >= len(all_matches): + break + # Ensure the tracking arrays have a slot for this index (defensive; + # Case 2 normally allocates these when the start token arrives). + while len(self.prev_tool_call_arr) <= idx: + self.prev_tool_call_arr.append({}) + self.streamed_args_for_tool.append("") + + func_name, args_str = all_matches[idx] final_args = _parse_gemma4_args(args_str) final_args_json = json.dumps(final_args, ensure_ascii=False) - prev_streamed = self.streamed_args_for_tool[self.current_tool_id] - if len(final_args_json) > len(prev_streamed): - diff = final_args_json[len(prev_streamed) :] - self.streamed_args_for_tool[self.current_tool_id] = final_args_json - self.prev_tool_call_arr[self.current_tool_id]["arguments"] = final_args + # The name is sent exactly once per call. We track that via the + # per-call entry in prev_tool_call_arr (set either by the middle + # path or by the batched-complete branch below), which is robust + # even when several calls are drained in one delta. + name_already_sent = bool(self.prev_tool_call_arr[idx].get("name")) - return DeltaMessage( - tool_calls=[ + if not name_already_sent: + # Batched-complete call: emit id + name + full arguments once. + self.streamed_args_for_tool[idx] = final_args_json + self.prev_tool_call_arr[idx] = { + "name": func_name, + "arguments": final_args, + } + deltas.append( + DeltaToolCall( + index=idx, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=func_name, arguments=final_args_json + ).model_dump(exclude_none=True), + ) + ) + else: + # Incrementally-streamed call: flush the remaining argument + # tail that was withheld during the middle phase. + prev_streamed = self.streamed_args_for_tool[idx] + if len(final_args_json) > len(prev_streamed): + diff = final_args_json[len(prev_streamed) :] + self.streamed_args_for_tool[idx] = final_args_json + self.prev_tool_call_arr[idx]["arguments"] = final_args + deltas.append( DeltaToolCall( - index=self.current_tool_id, + index=idx, function=DeltaFunctionCall(arguments=diff).model_dump( exclude_none=True ), ) - ] - ) + ) + # Advance streaming state past the calls completed in this delta. If a + # further tool call is still being accumulated (start without a + # matching end), point current_tool_id at it so the middle path can + # stream its arguments next; otherwise settle on the last completed + # call. + if start_count > end_count: + self.current_tool_id = end_count + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + self.streamed_args_for_tool.append("") + self.current_tool_name_sent = bool( + self.prev_tool_call_arr[self.current_tool_id].get("name") + ) + else: + self.current_tool_id = end_count - 1 + self.current_tool_name_sent = True + + if deltas: + return DeltaMessage(tool_calls=deltas) return None def _emit_argument_diff(self, raw_args_str: str) -> DeltaMessage | None: diff --git a/vllm/tool_parsers/glm47_moe_tool_parser.py b/vllm/tool_parsers/glm47_moe_tool_parser.py index 47b6ad2f5af..80068264b70 100644 --- a/vllm/tool_parsers/glm47_moe_tool_parser.py +++ b/vllm/tool_parsers/glm47_moe_tool_parser.py @@ -24,6 +24,7 @@ logger = init_logger(__name__) class Glm47MoeModelToolParser(Glm4MoeModelToolParser): supports_required_and_named = False + structural_tag_model = "glm_4_7" def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/gptoss_tool_parser.py b/vllm/tool_parsers/gptoss_tool_parser.py new file mode 100644 index 00000000000..6857e6bbe72 --- /dev/null +++ b/vllm/tool_parsers/gptoss_tool_parser.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + ExtractedToolCallInformation, +) +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser + +if TYPE_CHECKING: + from vllm.tokenizers import TokenizerLike + + +class GptOssToolParser(ToolParser): + """ + Stub tool parser for gpt-oss/harmony models. + + All output parsing is handled by HarmonyParser. This stub exists as a + capability declaration via HarmonyParser.tool_parser_cls. + """ + + def __init__(self, tokenizer: "TokenizerLike", tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + + def extract_tool_calls( + self, model_output, request, **kwargs + ) -> ExtractedToolCallInformation: + raise NotImplementedError( + "GptOssToolParser is a stub. Use HarmonyParser for tool parsing." + ) + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request, + ) -> DeltaMessage | None: + raise NotImplementedError( + "GptOssToolParser is a stub. Use HarmonyParser for tool parsing." + ) diff --git a/vllm/tool_parsers/hermes_tool_parser.py b/vllm/tool_parsers/hermes_tool_parser.py index 546cde5cd14..3fd819297aa 100644 --- a/vllm/tool_parsers/hermes_tool_parser.py +++ b/vllm/tool_parsers/hermes_tool_parser.py @@ -32,6 +32,7 @@ logger = init_logger(__name__) class Hermes2ProToolParser(ToolParser): + structural_tag_model = "hermes" tool_call_start_token: str = "" tool_call_end_token: str = "" tool_call_regex = re.compile( diff --git a/vllm/tool_parsers/kimi_k2_tool_parser.py b/vllm/tool_parsers/kimi_k2_tool_parser.py index 7ddd8fa7a80..18f242fffe0 100644 --- a/vllm/tool_parsers/kimi_k2_tool_parser.py +++ b/vllm/tool_parsers/kimi_k2_tool_parser.py @@ -29,6 +29,8 @@ logger = init_logger(__name__) class KimiK2ToolParser(ToolParser): + structural_tag_model = "kimi" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/llama_tool_parser.py b/vllm/tool_parsers/llama_tool_parser.py index 4a041041f09..624428d992f 100644 --- a/vllm/tool_parsers/llama_tool_parser.py +++ b/vllm/tool_parsers/llama_tool_parser.py @@ -46,6 +46,7 @@ class Llama3JsonToolParser(ToolParser): """ bot_token: str = "<|python_tag|>" + structural_tag_model = "llama" # Simple regex to find opening braces - we'll use JSON decoder for parsing # This handles arbitrary nesting depth correctly tool_call_start_regex: re.Pattern = re.compile(r"\{") diff --git a/vllm/tool_parsers/minimax_m2_tool_parser.py b/vllm/tool_parsers/minimax_m2_tool_parser.py index 5a3aae81262..ba59fd77ea6 100644 --- a/vllm/tool_parsers/minimax_m2_tool_parser.py +++ b/vllm/tool_parsers/minimax_m2_tool_parser.py @@ -34,6 +34,8 @@ logger = init_logger(__name__) class MinimaxM2ToolParser(ToolParser): + structural_tag_model = "minimax" + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) diff --git a/vllm/tool_parsers/openai_tool_parser.py b/vllm/tool_parsers/openai_tool_parser.py deleted file mode 100644 index e5c37fbd3df..00000000000 --- a/vllm/tool_parsers/openai_tool_parser.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json -from collections.abc import Sequence -from typing import TYPE_CHECKING - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaMessage, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.parser.harmony_utils import ( - extract_function_from_recipient, - is_function_recipient, - parse_output_into_messages, -) -from vllm.logger import init_logger -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) - -if TYPE_CHECKING: - from vllm.tokenizers import TokenizerLike -else: - TokenizerLike = object - -logger = init_logger(__name__) - - -class OpenAIToolParser(ToolParser): - def __init__(self, tokenizer: "TokenizerLike", tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - token_ids: Sequence[int] | None = None, - ) -> ExtractedToolCallInformation: - if token_ids is None: - raise NotImplementedError( - "OpenAIToolParser requires token IDs and does not support text-based extraction." # noqa: E501 - ) - - parser = parse_output_into_messages(token_ids) - tool_calls = [] - final_content = None - commentary_content = None - - if len(parser.messages) > 0: - for msg in parser.messages: - if msg.author.role != "assistant": - continue - if len(msg.content) < 1: - continue - msg_text = msg.content[0].text - if msg.recipient and is_function_recipient(msg.recipient): - # If no content-type is given assume JSON, as that's the - # most common case with gpt-oss models. - if not msg.content_type or "json" in msg.content_type: - # load and dump the JSON text to check validity and - # remove any extra newlines or other odd formatting - try: - tool_args = json.dumps(json.loads(msg_text)) - except json.JSONDecodeError: - logger.exception( - "Error decoding JSON tool call from response." - ) - tool_args = msg_text - else: - tool_args = msg_text - tool_calls.append( - ToolCall( - type="function", - function=FunctionCall( - name=extract_function_from_recipient(msg.recipient), - arguments=tool_args, - ), - ) - ) - elif msg.channel == "final": - final_content = msg_text - elif msg.channel == "commentary" and not msg.recipient: - commentary_content = msg_text - - # Extract partial content from the parser state if the generation was truncated - if parser.current_content: - if parser.current_channel == "final": - final_content = parser.current_content - elif ( - parser.current_channel == "commentary" and not parser.current_recipient - ): - commentary_content = parser.current_content - - return ExtractedToolCallInformation( - tools_called=len(tool_calls) > 0, - tool_calls=tool_calls, - # prefer final content over commentary content if both are present - # commentary content is tool call preambles meant to be shown to the user - content=final_content or commentary_content, - ) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - raise NotImplementedError( - "Not being used, manual parsing in serving_chat.py" # noqa: E501 - ) diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 7457590c5ac..f9d777af1e9 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -18,17 +18,12 @@ from vllm.entrypoints.openai.engine.protocol import ( FunctionCall, ToolCall, ) -from vllm.envs import VLLM_ENFORCE_STRICT_TOOL_CALLING from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) -from vllm.tool_parsers.structural_tag_registry import ( - get_enable_structured_outputs_in_reasoning, - get_model_structural_tag, -) from vllm.tool_parsers.utils import ( coerce_to_schema_type, extract_types_from_schema, @@ -39,7 +34,7 @@ logger = init_logger(__name__) class Qwen3CoderToolParser(ToolParser): - supports_required_and_named: bool = not VLLM_ENFORCE_STRICT_TOOL_CALLING + structural_tag_model = "qwen_3_coder" def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -589,11 +584,3 @@ class Qwen3CoderToolParser(ToolParser): return result return None - - def get_structural_tag(self, request: ChatCompletionRequest): - return get_model_structural_tag( - model="qwen_3_5", - tools=request.tools, - tool_choice=request.tool_choice, - reasoning=get_enable_structured_outputs_in_reasoning(), - ) diff --git a/vllm/tool_parsers/qwen3xml_tool_parser.py b/vllm/tool_parsers/qwen3xml_tool_parser.py deleted file mode 100644 index e5d2b896e00..00000000000 --- a/vllm/tool_parsers/qwen3xml_tool_parser.py +++ /dev/null @@ -1,1300 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json -from collections.abc import Sequence -from typing import Any -from xml.parsers.expat import ParserCreate - -import regex as re - -from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) -from vllm.tool_parsers.utils import find_tool_properties, safe_literal_eval - -logger = init_logger(__name__) - - -class StreamingXMLToolCallParser: - """ - Simplified streaming XML tool call parser - Supports streaming input, parsing, and output - """ - - def __init__(self): - self.reset_streaming_state() - - # Tool configuration information - self.tools: list[Tool] | None = None - self.tool_call_start_token: str = "" - self.tool_call_end_token: str = "" - self.function_start_token: str = " DeltaMessage: - """ - Parse single streaming XML chunk and return Delta response - This is the actual streaming interface that receives chunks - one by one and maintains internal state - - Args: - xml_chunk: Single XML chunk string - Returns: - DeltaMessage: Contains delta information generated by this chunk, - returns empty response if no complete elements - """ - # Record delta count before processing - initial_delta_count = len(self.deltas) - - self.streaming_buffer += xml_chunk - - found_elements = self._process_complete_xml_elements() - - if found_elements: - # If complete elements found, check if end events were missed - # some tags may not have been triggered - try: - new_deltas = self.deltas[initial_delta_count:] - # If this chunk contains - # but didn't generate '}', then complete it - if ( - self.current_call_id is not None - and self.function_end_token in xml_chunk - ): - # - Added '}' (non-empty parameter ending) - # - Added '{}' (empty parameter function) - has_function_close = any( - ( - td.tool_calls - and any( - ( - tc.function - and tc.id == self.current_call_id - and isinstance(tc.function.arguments, str) - and (tc.function.arguments in ("}", "{}")) - ) - for tc in td.tool_calls - ) - ) - for td in new_deltas - ) - if not has_function_close: - # Close potentially unclosed element - if self.current_param_name: - self._end_element("parameter") - if self.current_function_name: - self._end_element("function") - # If this chunk contains - # but didn't generate final empty delta, then complete it - if ( - self.current_call_id is not None - and self.tool_call_end_token in xml_chunk - ): - has_toolcall_close = any( - ( - td.tool_calls - and any( - ( - tc.type == "function" - and tc.function - and tc.function.arguments == "" - and tc.id == self.current_call_id - ) - for tc in td.tool_calls - ) - ) - for td in new_deltas - ) - if not has_toolcall_close: - # Close potentially unclosed element - if self.current_param_name: - self._end_element("parameter") - if self.current_function_name: - self._end_element("function") - self._end_element("tool_call") - except Exception as e: - logger.warning("Error with fallback parsing: %s", e) - # Merge newly generated deltas into single response - result_delta = self._merge_new_deltas_to_single_response( - initial_delta_count - ) - return result_delta - else: - # No complete elements, check if there's unoutput text content - if self.text_content_buffer and self.tool_call_index == 0: - # Has text content but no tool_call yet, output text content - text_delta = DeltaMessage(content=self.text_content_buffer) - self._emit_delta(text_delta) - # Clear buffer to avoid duplicate output - self.text_content_buffer = "" - return text_delta - - # If this chunk contains end tags but wasn't triggered by parser, - # manually complete end events - # Only execute when still on the same call as when entered, - # to prevent accidentally closing new calls - # in multi scenarios - if self.current_call_id is not None and ( - self.function_end_token in xml_chunk - or self.tool_call_end_token in xml_chunk - ): - # Close potentially unclosed element - if self.current_param_name: - self._end_element("parameter") - if self.function_end_token in xml_chunk and self.current_function_name: - self._end_element("function") - if self.tool_call_end_token in xml_chunk: - self._end_element("tool_call") - # Return the merged delta result generated by this fallback - result_delta = self._merge_new_deltas_to_single_response( - initial_delta_count - ) - return result_delta - - # No complete elements, return empty response - return DeltaMessage(content=None) - - def _escape_xml_special_chars(self, text: str) -> str: - """ - Escape XML special characters - Args: - text: Original text - Returns: - Escaped text - """ - xml_escapes = { - "&": "&", - "<": "<", - ">": ">", - '"': """, - "'": "'", - } - - for char, escape in xml_escapes.items(): - text = text.replace(char, escape) - - return text - - def _process_complete_xml_elements(self) -> bool: - """ - Process complete XML elements in buffer - - Returns: - bool: Whether complete elements were found and processed - """ - found_any = False - - while self.last_processed_pos < len(self.streaming_buffer): - # Find next complete xml element - element, end_pos = self._find_next_complete_element(self.last_processed_pos) - if element is None: - # No complete element found, wait for more data - break - - # Check if this element should be skipped - if self._should_skip_element(element): - self.last_processed_pos = end_pos - continue - - # Found complete XML element, process it - try: - preprocessed_element = self._preprocess_xml_chunk(element) - # Check if this is the first tool_call start - if ( - ( - preprocessed_element.strip().startswith("") - or preprocessed_element.strip().startswith("") - and self.tool_call_index > 0 - and self.current_call_id - ): - # Reset parser state but preserve generated deltas - if self.current_param_name: - self._end_element("parameter") - if self.current_function_open or self.current_function_name: - self._end_element("function") - # Output final tool_call tail delta - final_delta = DeltaMessage( - role=None, - content=None, - reasoning=None, - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments=""), - ) - ], - ) - self._emit_delta(final_delta) - # Reset XML parser and current call state - self._reset_xml_parser_after_tool_call() - # Parse preprocessed element - self.parser.Parse(preprocessed_element, False) - found_any = True - - except Exception as e: - logger.warning("Error when parsing XML elements: %s", e) - - # Update processed position - self.last_processed_pos = end_pos - - return found_any - - def _should_skip_element(self, element: str) -> bool: - """ - Determine whether an element should be skipped - - Args: - element: Element to evaluate - - Returns: - bool: True means should skip, False means should process - """ - - # If it's a tool_call XML tag, don't skip - if ( - element.startswith(self.tool_call_start_token) - or element.startswith(self.function_start_token) - or element.startswith(self.parameter_start_token) - ): - return False - - # If currently not parsing tool calls and not blank, - # collect this text instead of skipping - # Only process other XML elements after tool_call appears, - # otherwise treat as plain text - if self.current_call_id is None and element: - # Collect text content to buffer - self.text_content_buffer += element - return True # Still skip, but content has been collected - - # If currently parsing tool calls, - # this might be parameter value, don't skip - if self.current_call_id is not None: - return False - - # Skip blank content - return not element - - def _find_next_complete_element(self, start_pos: int) -> tuple[str | None, int]: - """ - Find next complete XML element from specified position - - Args: - start_pos: Position to start searching - - Returns: - (Complete element string, element end position), - returns (None, start_pos) if no complete element found - """ - buffer = self.streaming_buffer[start_pos:] - - if not buffer: - return None, start_pos - - if buffer.startswith("<"): - # Need to ensure no new < appears, - # find the nearest one between < and > - tag_end = buffer.find("<", 1) - tag_end2 = buffer.find(">", 1) - if tag_end != -1 and tag_end2 != -1: - # Next nearest is < - if tag_end < tag_end2: - return buffer[:tag_end], start_pos + tag_end - # Next nearest is >, means found XML element - else: - return buffer[: tag_end2 + 1], start_pos + tag_end2 + 1 - elif tag_end != -1: - return buffer[:tag_end], start_pos + tag_end - elif tag_end2 != -1: - return buffer[: tag_end2 + 1], start_pos + tag_end2 + 1 - else: - # If currently not parsing tool calls (entering a tool_call), - # check if starts with or - if buffer == ""[: len(buffer)]: - # Might be start of , wait for more data - return None, start_pos - elif ( - buffer.startswith(" DeltaMessage: - """ - Merge newly generated deltas from this processing - into a single DeltaMessage - - Args: - initial_count: Delta count before processing - - Returns: - Merged DeltaMessage containing all newly generated delta information - """ - if len(self.deltas) <= initial_count: - return DeltaMessage(content=None) - - # Get newly generated deltas - new_deltas = self.deltas[initial_count:] - - if len(new_deltas) == 1: - # Only one new delta, return directly - return new_deltas[0] - - # Merge multiple new deltas - merged_tool_calls: list[DeltaToolCall] = [] - merged_content: str = "" - - for delta in new_deltas: - if delta.content: - merged_content += delta.content - if delta.tool_calls: - # For tool_calls, we need to intelligently merge arguments - for tool_call in delta.tool_calls: - # Find if there's already a tool_call with the same call_id - existing_call = None - for existing in merged_tool_calls: - if existing.id == tool_call.id: - existing_call = existing - break - - if existing_call and existing_call.function: - # Merge to existing tool_call - if tool_call.function and tool_call.function.name: - existing_call.function.name = tool_call.function.name - if ( - tool_call.function - and tool_call.function.arguments is not None - ): - if existing_call.function.arguments is None: - existing_call.function.arguments = "" - - # For streaming JSON parameters, - # simply concatenate in order - new_args = tool_call.function.arguments - existing_call.function.arguments += new_args - if tool_call.type: - existing_call.type = tool_call.type - else: - # Add new tool_call - merged_tool_calls.append(tool_call) - - return DeltaMessage( - content=merged_content if merged_content else None, - tool_calls=merged_tool_calls, - ) - - def _preprocess_xml_chunk(self, chunk: str) -> str: - """ - Preprocess XML chunk, handle non-standard formats, - and escape special characters - - Args: - chunk: Original XML chunk - - Returns: - Processed XML chunk - """ - - # Check if this is a tool_call related element - is_tool_call = False - if chunk.startswith(self.tool_call_start_token) or chunk.startswith( - self.tool_call_end_token - ): - is_tool_call = True - if chunk.startswith(self.function_start_token) or chunk.startswith( - self.function_end_token - ): - is_tool_call = True - if chunk.startswith(self.parameter_start_token) or chunk.startswith( - self.parameter_end_token - ): - is_tool_call = True - # Handle format -> - processed = re.sub(r"]+)>", r'', chunk) - # Handle format -> - processed = re.sub(r"]+)>", r'', processed) - - original_chunk = chunk - # If in parameter value accumulation mode - if self._pre_inside_parameter: - # Parameter end: output accumulated raw text - # safely then return - if processed.startswith(""): - body_text = self._pre_param_buffer - # Trigger deferred parsing mode - # literal_eval+json output in end_element - self.defer_current_parameter = True - self.deferred_param_raw_value = body_text - # Clean up state - self._pre_inside_parameter = False - self._pre_param_buffer = "" - self._pre_current_param_name = None - safe_text = self._escape_xml_special_chars(body_text) - return f"{safe_text}" - else: - # If this is the first block of content after entering parameter - # evaluate if deferred parsing is needed; - # If not needed, exit accumulation mode - # and pass through directly - if self._pre_param_buffer == "": - # Get current parameter type - param_type = ( - self._get_param_type(self._pre_current_param_name) - if self._pre_current_param_name - else "string" - ) - # Only these types need deferred parsing to - # handle Python literals containing single quotes - is_object_type = param_type in ["object"] - is_complex_type = ( - param_type in ["array", "arr", "sequence"] - or param_type.startswith("dict") - or param_type.startswith("list") - ) - - # Only delay when contains container symbols - # and has single quotes and is complex type - has_container_hint = ( - ("[" in original_chunk) - or ("{" in original_chunk) - or ("(" in original_chunk) - ) - - # Determine if deferred parsing is needed - need_defer = False - if is_complex_type: - # Complex type, always need deferred parsing - need_defer = True - elif ( - is_object_type - and has_container_hint - and ("'" in original_chunk) - ): - # Object type with container symbols - # and single quotes, need deferred parsing - need_defer = True - - if not need_defer: - # No need for deferred parsing, - # exit parameter mode directly - self._pre_inside_parameter = False - return self._escape_xml_special_chars(original_chunk) - self._pre_param_buffer += original_chunk - return "" - - # Parameter start: enable accumulation - if processed.startswith("', processed) - if m: - self._pre_current_param_name = m.group(1) - self._pre_inside_parameter = True - self._pre_param_buffer = "" - return processed - - # If processed doesn't contain special_token, escape processed - # This is because XML parsing encounters special characters - # and reports errors, so escaping is needed - if not is_tool_call: - processed = self._escape_xml_special_chars(processed) - return processed - - def _emit_delta(self, delta: DeltaMessage): - """Emit Delta response (streaming output)""" - self.deltas.append(delta) - - def _auto_close_open_parameter_if_needed(self, incoming_tag: str | None = None): - """Before starting to process new elements, - if there are unclosed tags from before, - automatically complete their endings to the parser. - - If there are unclosed parameters, - it's equivalent to feeding `` - - When about to start a new function or tool_call, - if there are unclosed functions, complete ``. - - When about to start a new tool_call, - if there are unclosed tool_calls, complete ``. - """ - # First close unclosed parameters - if self.current_param_name: - self._end_element("parameter") - - # If about to start new function or tool_call, - # and there are unclosed functions, close function first - if incoming_tag in ("function", "tool_call") and self.current_function_name: - self._end_element("function") - - # If about to start new tool_call, - # and there are unclosed tool_calls, close tool_call first - if incoming_tag == "tool_call" and self.current_call_id: - self._end_element("tool_call") - - def _start_element(self, name: str, attrs: dict[str, str]): - """Handle XML start element events""" - - if name == "root": - return - - if name == "tool_call": - # Before opening new tool_call, - # automatically complete previous unclosed tags - self._auto_close_open_parameter_if_needed("tool_call") - - self.parameters = {} - self.current_call_id = make_tool_call_id() - self.current_param_is_first = True - self.tool_call_index += 1 - elif name.startswith("function") or (name == "function"): - # If missing tool_call, manually complete - if not self.current_call_id: - self._start_element("tool_call", {}) - # Before opening new function, - # automatically complete previous unclosed tags (parameter/function) - self._auto_close_open_parameter_if_needed("function") - function_name = self._extract_function_name(name, attrs) - self.current_function_name = function_name - self.current_function_open = True - if function_name: - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall( - name=function_name, arguments="" - ), - ) - ] - ) - self._emit_delta(delta) - elif name.startswith("parameter") or (name == "parameter"): - # If previous parameter hasn't ended normally, - # complete its end first, then start new parameter - self._auto_close_open_parameter_if_needed("parameter") - param_name = self._extract_parameter_name(name, attrs) - self.current_param_name = param_name - self.current_param_value = "" - self.current_param_value_converted = "" - self.start_quote_emitted = False # Reset start quote flag - - # Only output parameter name and colon, - # don't output quotes - # decide after parameter value type is determined - if param_name: - if not self.parameters: - # First parameter - # start JSON, only output parameter name and colon - json_start = f'{{"{param_name}": ' - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall( - name=None, arguments=json_start - ), - ) - ] - ) - self._emit_delta(delta) - self.current_param_is_first = True - else: - # Subsequent parameters - # add comma and parameter name, no quotes - json_continue = f', "{param_name}": ' - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall( - name=None, arguments=json_continue - ), - ) - ] - ) - self._emit_delta(delta) - self.current_param_is_first = False - - def _char_data(self, data: str): - """Handle XML character data events""" - if data and self.current_param_name: - # If preprocessing stage determines deferred parsing is needed, - # only cache character data, no streaming output - if self.defer_current_parameter: - original_data = data - if self.should_emit_end_newline: - original_data = "\n" + original_data - self.should_emit_end_newline = False - if original_data.endswith("\n"): - self.should_emit_end_newline = True - original_data = original_data[:-1] - self.current_param_value += original_data - return - - param_type = self._get_param_type(self.current_param_name) - - # Check if this is the first time receiving data for this parameter - # If this is the first packet of data and starts with \n, remove \n - if not self.current_param_value and data.startswith("\n"): - data = data[1:] - - # Output start quote for string type (if not already output) - if ( - param_type in ["string", "str", "text", "varchar", "char", "enum"] - and not self.start_quote_emitted - ): - quote_delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments='"'), - ) - ] - ) - self._emit_delta(quote_delta) - self.start_quote_emitted = True - - if not data: - return - - original_data = data - # Delay output of trailing newline - if self.should_emit_end_newline: - original_data = "\n" + original_data - self.should_emit_end_newline = False - if original_data.endswith("\n"): - self.should_emit_end_newline = True - original_data = original_data[:-1] - self.current_param_value += original_data - - # convert parameter value by param_type - converted_value = self._convert_param_value( - self.current_param_value, param_type - ) - output_data = self._convert_for_json_streaming(converted_value, param_type) - - delta_data = output_data[len(self.current_param_value_converted) :] - self.current_param_value_converted = output_data - - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments=delta_data), - ) - ] - ) - self._emit_delta(delta) - - def _end_element(self, name: str): - """Handle XML end element events""" - - if name == "root": - return - - # If function or tool_call ends and there are still unclosed parameters, - # complete parameter end first - if ( - name.startswith("function") or name == "function" or name == "tool_call" - ) and self.current_param_name: - self._auto_close_open_parameter_if_needed() - - if ( - name.startswith("parameter") or name == "parameter" - ) and self.current_param_name: - # End current parameter - param_name = self.current_param_name - param_value = self.current_param_value - - # If in deferred parsing mode, - # perform overall parsing on raw content - # accumulated in preprocessing stage and output once - if self.defer_current_parameter: - raw_text = ( - self.deferred_param_raw_value - if self.deferred_param_raw_value - else param_value - ) - parsed_value = None - output_arguments = None - try: - # If previously delayed trailing newline, - # add it back before parsing - if self.should_emit_end_newline: - raw_for_parse = raw_text + "\n" - else: - raw_for_parse = raw_text - try: - parsed_value = json.loads(raw_for_parse) - except json.JSONDecodeError: - parsed_value = safe_literal_eval(raw_for_parse) - output_arguments = json.dumps(parsed_value, ensure_ascii=False) - except Exception: - # Fallback: output as string as-is - output_arguments = json.dumps(raw_text, ensure_ascii=False) - parsed_value = raw_text - - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall( - name=None, arguments=output_arguments - ), - ) - ] - ) - self._emit_delta(delta) - - # Clean up and store - self.should_emit_end_newline = False - self.parameters[param_name] = parsed_value - self.current_param_name = None - self.current_param_value = "" - self.current_param_value_converted = "" - self.start_quote_emitted = False - self.defer_current_parameter = False - self.deferred_param_raw_value = "" - return - - param_type = self._get_param_type(param_name) - - # convert complete parameter value by param_type - converted_value = self._convert_param_value(param_value, param_type) - - # Decide whether to add end quote based on parameter type - if param_type in ["string", "str", "text", "varchar", "char", "enum"]: - # For empty string parameters, need special handling - if not param_value and not self.start_quote_emitted: - # No start quote output, - # directly output complete empty string - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments='""'), - ) - ] - ) - self._emit_delta(delta) - else: - # Non-empty parameter value, output end quote - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments='"'), - ) - ] - ) - self._emit_delta(delta) - - self.should_emit_end_newline = False - # Store converted value - self.parameters[param_name] = converted_value - self.current_param_name = None - self.current_param_value = "" - self.current_param_value_converted = "" - self.start_quote_emitted = False - - elif name.startswith("function") or name == "function": - # if there are parameters, close JSON object - if self.parameters: - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments="}"), - ) - ] - ) - self._emit_delta(delta) - # return empty object - else: - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments="{}"), - ) - ] - ) - self._emit_delta(delta) - self.current_function_open = False - - elif name == "tool_call": - # Before ending tool_call, - # ensure function is closed to complete missing right brace - if self.current_function_open: - # If there are still unclosed parameters, close them first - if self.current_param_name: - self._end_element("parameter") - # Close function, ensure output '}' or '{}' - self._end_element("function") - # Final Delta - delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.tool_call_index - 1, - id=self.current_call_id, - type="function", - function=DeltaFunctionCall(name=None, arguments=""), - ) - ] - ) - self._emit_delta(delta) - - # Check if there's text content to output (between tool_calls) - if self.text_content_buffer.strip(): - text_delta = DeltaMessage(content=self.text_content_buffer) - self._emit_delta(text_delta) - - self._reset_xml_parser_after_tool_call() - - def setup_parser(self): - """Set up XML parser event handlers""" - self.parser.buffer_text = True - self.parser.StartElementHandler = self._start_element - self.parser.EndElementHandler = self._end_element - self.parser.CharacterDataHandler = self._char_data - - def set_tools(self, tools: list[Tool] | None): - """Set tool configuration information""" - self.tools = tools - - def _extract_function_name(self, name: str, attrs: dict[str, str]) -> str | None: - """Extract function name from various formats""" - if attrs and "name" in attrs: - return attrs["name"] - - if "=" in name: - parts = name.split("=", 1) - if len(parts) == 2 and parts[0] == "function": - return parts[1] - - return None - - def _extract_parameter_name(self, name: str, attrs: dict[str, str]) -> str | None: - """Extract parameter name from various formats""" - if attrs and "name" in attrs: - return attrs["name"] - - if "=" in name: - parts = name.split("=", 1) - if len(parts) == 2 and parts[0] == "parameter": - return parts[1] - - return None - - def _get_param_type(self, param_name: str) -> str: - """Get parameter type based on tool configuration, defaults to string - Args: - param_name: Parameter name - - Returns: - Parameter type - """ - if not self.tools or not self.current_function_name: - return "string" - - properties = find_tool_properties(self.tools, self.current_function_name) - if param_name in properties and isinstance(properties[param_name], dict): - return self.repair_param_type( - str(properties[param_name].get("type", "string")) - ) - return "string" - - def repair_param_type(self, param_type: str) -> str: - """Repair unknown parameter types by treating them as string - Args: - param_type: Parameter type - - Returns: - Repaired parameter type - """ - if ( - param_type in ["string", "str", "text", "varchar", "char", "enum"] - or param_type.startswith("int") - or param_type.startswith("uint") - or param_type.startswith("long") - or param_type.startswith("short") - or param_type.startswith("unsigned") - or param_type.startswith("num") - or param_type.startswith("float") - or param_type in ["boolean", "bool", "binary"] - or ( - param_type in ["object", "array", "arr", "sequence"] - or param_type.startswith("dict") - or param_type.startswith("list") - ) - ): - return param_type - else: - return "string" - - def _convert_param_value(self, param_value: str, param_type: str) -> Any: - """Convert value based on parameter type - Args: - param_value: Parameter value - param_type: Parameter type - - Returns: - Converted value - """ - if param_value.lower() == "null": - return None - - param_type = param_type.strip().lower() - if param_type in ["string", "str", "text", "varchar", "char", "enum"]: - return param_value - elif ( - param_type.startswith("int") - or param_type.startswith("uint") - or param_type.startswith("long") - or param_type.startswith("short") - or param_type.startswith("unsigned") - ): - try: - return int(param_value) - except (ValueError, TypeError): - logger.warning( - "Parsed value '%s' of parameter '%s' is not an integer " - "in tool '%s', degenerating to string.", - param_value, - ) - return param_value - elif param_type.startswith("num") or param_type.startswith("float"): - try: - float_param_value: float = float(param_value) - return ( - float_param_value - if float_param_value - int(float_param_value) != 0 - else int(float_param_value) - ) - except (ValueError, TypeError): - logger.warning( - "Parsed value '%s' of parameter '%s' is not a float " - "in tool '%s', degenerating to string.", - param_value, - ) - return param_value - elif param_type in ["boolean", "bool", "binary"]: - param_value = param_value.lower() - return param_value == "true" - else: - return param_value - - def _convert_for_json_streaming(self, converted_value: Any, param_type: str) -> str: - """Convert converted_value based on - whether it's empty and if type is string - Args: - converted_value: Converted value - param_type: Parameter type - - Returns: - Converted string for streaming output - """ - # Check if value is empty, but exclude numeric 0 - if converted_value is None or converted_value == "": - return "" - - if param_type in ["string", "str", "text", "varchar", "char", "enum"]: - # String type, remove double quotes - return json.dumps(converted_value, ensure_ascii=False)[1:-1] - else: - # Non-string type, return complete JSON string - if not isinstance(converted_value, str): - return json.dumps(converted_value, ensure_ascii=False) - else: - return converted_value - - def _reset_xml_parser_after_tool_call(self): - """ - Each tool_call is treated as a separate XML document, - so we need to reset the parser after each tool_call. - """ - - # recreate XML parser - self.parser = ParserCreate() - self.setup_parser() - - # Reset current tool_call state - if self.current_call_id: - self.last_completed_call_id = self.current_call_id - self.current_call_id = None - self.current_function_name = None - self.current_function_open = False - self.parameters = {} - self.current_param_name = None - self.current_param_value = "" - self.current_param_value_converted = "" - self.current_param_is_first = False - self.should_emit_end_newline = False - self.start_quote_emitted = False - self.text_content_buffer = "" - - # Reset preprocessing and deferred parsing state - self._pre_inside_parameter = False - self._pre_param_buffer = "" - self._pre_current_param_name = None - self.defer_current_parameter = False - self.deferred_param_raw_value = "" - - -class Qwen3XMLToolParser(ToolParser): - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - self.parser = StreamingXMLToolCallParser() - - # Add missing attributes for compatibility with serving_chat.py - self.prev_tool_call_arr: list[dict] = [] - self.streamed_args_for_tool: list[str] = [] - - logger.info( - "vLLM Successfully import tool parser %s !", self.__class__.__name__ - ) - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - self.parser.reset_streaming_state() - # Reset tool call tracking arrays for new extraction - self.prev_tool_call_arr = [] - self.streamed_args_for_tool = [] - self.parser.set_tools(self.tools) - result = self.parser.parse_single_streaming_chunks(model_output) - if not result.tool_calls: - return ExtractedToolCallInformation( - tool_calls=[], - tools_called=False, - content=result.content, - ) - else: - tool_calls = [] - for tool_call in result.tool_calls: - if tool_call.function and tool_call.function.name: - tool_calls.append( - ToolCall( - id=tool_call.id, - type=tool_call.type, - function=FunctionCall( - name=tool_call.function.name, - arguments=tool_call.function.arguments, - ), - ) - ) - - # Update tool call tracking arrays for compatibility - tool_index = ( - tool_call.index - if tool_call.index is not None - else len(self.prev_tool_call_arr) - 1 - ) - - # Ensure we have enough entries in our tracking arrays - while len(self.prev_tool_call_arr) <= tool_index: - self.prev_tool_call_arr.append({"name": "", "arguments": ""}) - while len(self.streamed_args_for_tool) <= tool_index: - self.streamed_args_for_tool.append("") - - # Update tool call information - self.prev_tool_call_arr[tool_index]["name"] = ( - tool_call.function.name - ) - self.prev_tool_call_arr[tool_index]["arguments"] = ( - tool_call.function.arguments - ) - - # Update streamed arguments - if tool_call.function.arguments: - self.streamed_args_for_tool[tool_index] = ( - tool_call.function.arguments - ) - - return ExtractedToolCallInformation( - tool_calls=tool_calls, - tools_called=len(tool_calls) > 0, - content=result.content, - ) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - if not previous_text: - self.parser.reset_streaming_state() - # Reset tool call tracking arrays for new streaming session - self.prev_tool_call_arr = [] - self.streamed_args_for_tool = [] - self.parser.set_tools(self.tools) - - # Model sometimes outputs separately causing delta_text to be empty. - # If there were tool_calls before and all current tool_calls have ended, - # return an empty tool_call for outer streaming output - # to correctly output tool_call field - if not delta_text and delta_token_ids: - open_calls = current_text.count( - self.parser.tool_call_start_token - ) - current_text.count(self.parser.tool_call_end_token) - if ( - open_calls == 0 - and self.parser.tool_call_index > 0 - or not self.parser.tool_call_index - and current_text - ): - return DeltaMessage(content="") - return None - - # Parse the delta text and get the result - delta = self.parser.parse_single_streaming_chunks(delta_text) - - # Update tool call tracking arrays based on incremental parsing results - if delta and delta.tool_calls: - for tool_call in delta.tool_calls: - if tool_call.function: - tool_index = ( - tool_call.index - if tool_call.index is not None - else len(self.prev_tool_call_arr) - 1 - ) - - # Ensure we have enough entries in our tracking arrays - while len(self.prev_tool_call_arr) <= tool_index: - self.prev_tool_call_arr.append({"name": "", "arguments": ""}) - while len(self.streamed_args_for_tool) <= tool_index: - self.streamed_args_for_tool.append("") - - # Update tool name if provided - if tool_call.function.name: - self.prev_tool_call_arr[tool_index]["name"] = ( - tool_call.function.name - ) - - # Update arguments incrementally - if tool_call.function.arguments is not None: - # Concatenate the incremental arguments - # to the existing streamed arguments - self.prev_tool_call_arr[tool_index]["arguments"] += ( - tool_call.function.arguments - ) - self.streamed_args_for_tool[tool_index] += ( - tool_call.function.arguments - ) - if delta.content is None and not delta.tool_calls and delta.reasoning is None: - # If no content and no tool calls, return None to indicate no update - return None - return delta diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 7f6638dcb94..53b3f06bb8c 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -14,7 +14,6 @@ from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, DeltaToolCall, ) -from vllm.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.mistral import is_mistral_tokenizer @@ -77,6 +76,9 @@ def extract_named_tool_call_streaming( ) else: if is_mistral_tokenizer(tokenizer): + # Import mistral_common only if we need it. + from vllm.tool_parsers.mistral_tool_parser import MistralToolCall + tool_call_id = MistralToolCall.generate_random_id() else: tool_call_id = make_tool_call_id( diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 754cc52361c..13491e95dfc 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -1,14 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Model-specific structural tag builders adapted from XGrammar's -# builtin structural tag implementations: -# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/builtin_structural_tag.py +from collections.abc import Callable, Sequence +from typing import Any, Literal, TypeAlias -from collections.abc import Callable -from typing import Any, Literal - -from xgrammar import StructuralTag +from openai.types.responses import FunctionTool +from openai.types.responses.response import ToolChoice as ResponsesToolChoice +from openai.types.responses.tool import Tool as ResponsesTool +from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed +from openai.types.responses.tool_choice_function import ToolChoiceFunction +from xgrammar import StructuralTag, normalize_tool_choice +from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag +from xgrammar.openai_tool_call_schema import ( + BuiltinToolParam, + FunctionToolParam, +) from xgrammar.structural_tag import ( AnyTextFormat, ConstStringFormat, @@ -24,23 +30,55 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionToolsParam, ) -SimplifiedToolChoice = Literal["auto", "required", "forced"] -ToolChoice = ( - Literal["none", "auto", "required"] | ChatCompletionNamedToolChoiceParam | None +ToolChoice: TypeAlias = ( + Literal["none", "auto", "required"] + | ChatCompletionNamedToolChoiceParam + | ResponsesToolChoice + | None ) -StructuralTagBuilder = Callable[ - [list[ChatCompletionToolsParam], SimplifiedToolChoice, bool], +AllowedToolRef: TypeAlias = dict[str, object] +SimplifiedToolChoice: TypeAlias = Literal["auto", "required", "forced"] +StructuralTagBuilder: TypeAlias = Callable[ + [ + list[FunctionToolParam], + list[BuiltinToolParam], + SimplifiedToolChoice, + bool, + ], StructuralTag, ] -_structural_tag_registry: dict[str, StructuralTagBuilder] = {} +# Keep this list in sync with xgrammar.builtin_structural_tag. It is used for +# vLLM-side validation and for documenting the xgrammar builtin surface that +# can be requested by tool parsers through ``structural_tag_model``. +XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( + { + "llama", + "kimi", + "deepseek_r1", + "deepseek_v3_1", + "qwen_3_5", + "qwen_3_coder", + "qwen_3", + "harmony", + "deepseek_v3_2", + "glm_4_7", + "deepseek_v4", + } +) +VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes"}) +SUPPORTED_STRUCTURAL_TAG_MODELS = ( + XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS +) + +_VLLM_STRUCTURAL_TAG_REGISTRY: dict[str, StructuralTagBuilder] = {} -def register_model_structural_tag(name: str): - """Register a vLLM-owned model-specific structural tag builder.""" +def register_vllm_structural_tag(model: str): + """Register a vLLM-owned structural tag builder.""" def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder: - _structural_tag_registry[name] = func + _VLLM_STRUCTURAL_TAG_REGISTRY[model] = func return func return decorator @@ -48,283 +86,248 @@ def register_model_structural_tag(name: str): def get_model_structural_tag( model: str, - tools: list[ChatCompletionToolsParam] | None, + tools: Sequence[ChatCompletionToolsParam | ResponsesTool] | None, tool_choice: ToolChoice, reasoning: bool, ) -> StructuralTag | None: - """Build a structural tag from vLLM-owned model-specific builders.""" + """Build a structural tag with xgrammar's builtin model templates.""" - builder = _structural_tag_registry.get(model) - if builder is None: - supported = list(_structural_tag_registry.keys()) - raise ValueError(f"Unknown format type: {model}, supported types: {supported}") - - normalized_tools, simplified_tool_choice = _normalize_tool_choice( - tools=tools, - tool_choice=tool_choice, - ) - if not normalized_tools: + if not tools or tool_choice == "none": return None - return builder(normalized_tools, simplified_tool_choice, reasoning) + dumped_tools = [_dump_tool_for_xgrammar(tool) for tool in tools] + dumped_tool_choice = _dump_tool_choice_for_xgrammar(tool_choice) + + if model in _VLLM_STRUCTURAL_TAG_REGISTRY: + function_tools, builtin_tools, simplified_tool_choice = normalize_tool_choice( + dumped_tools, + dumped_tool_choice, + ) + return _VLLM_STRUCTURAL_TAG_REGISTRY[model]( + function_tools, + builtin_tools, + simplified_tool_choice, + reasoning, + ) + + if model not in XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS: + supported = sorted(SUPPORTED_STRUCTURAL_TAG_MODELS) + raise ValueError(f"Unknown format type: {model}, supported types: {supported}") + + return get_xgrammar_model_structural_tag( + model=model, + tools=dumped_tools, + tool_choice=dumped_tool_choice, + reasoning=reasoning, + ) -def _normalize_tool_choice( - tools: list[ChatCompletionToolsParam] | None, +def _dump_tool_for_xgrammar( + tool: ChatCompletionToolsParam | ResponsesTool, +) -> dict[str, Any]: + """Convert tool objects to xgrammar's Chat Completions tool protocol.""" + + if isinstance(tool, FunctionTool): + function: dict[str, Any] = {"name": tool.name} + if tool.description is not None: + function["description"] = tool.description + if tool.parameters is not None: + function["parameters"] = tool.parameters + if tool.strict is not None: + function["strict"] = tool.strict + return {"type": "function", "function": function} + dumped_tool = tool.model_dump(mode="json", exclude_none=True) + if isinstance(tool, ChatCompletionToolsParam): + return dumped_tool + return dict(dumped_tool) + + +def _dump_tool_choice_for_xgrammar( tool_choice: ToolChoice, -) -> tuple[list[ChatCompletionToolsParam], SimplifiedToolChoice]: - """Normalize vLLM ChatCompletion tool_choice for structural tag builders.""" +) -> dict[str, Any] | str | None: + """Convert tool_choice objects to xgrammar's expected protocol.""" - if not tools: - return [], "auto" + if tool_choice is None: + return None - if tool_choice is None or tool_choice == "none": - return [], "auto" - - if tool_choice == "auto": - return tools, "auto" - - if tool_choice == "required": - return tools, "required" + if isinstance(tool_choice, str): + return tool_choice if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): - tool_name = tool_choice.function.name - filtered_tools = [tool for tool in tools if tool.function.name == tool_name] - if not filtered_tools: - raise ValueError( - f"The tool with name '{tool_name}' is not found in the tools list." - ) - return filtered_tools, "forced" + return tool_choice.model_dump(mode="json", exclude_none=True) - raise ValueError(f"Unsupported tool_choice for structural tag: {tool_choice}") + if isinstance(tool_choice, ToolChoiceFunction): + return { + "type": "function", + "function": {"name": tool_choice.name}, + } + + if isinstance(tool_choice, ToolChoiceAllowed): + return { + "type": "allowed_tools", + "allowed_tools": { + "mode": tool_choice.mode, + "tools": [ + _dump_allowed_tool_ref_for_xgrammar(tool) + for tool in tool_choice.tools + ], + }, + } + + return tool_choice.model_dump(mode="json", exclude_none=True) -def _get_function_parameters(function: Any) -> dict[str, Any] | bool: - """Return the JSON schema used for constrained tool arguments.""" +def _dump_allowed_tool_ref_for_xgrammar(tool_ref: AllowedToolRef) -> AllowedToolRef: + if ( + tool_ref.get("type") == "function" + and "function" not in tool_ref + and "name" in tool_ref + ): + return { + "type": "function", + "function": {"name": tool_ref["name"]}, + } + return tool_ref + +def _get_function_parameters(function) -> dict[str, Any] | bool: if getattr(function, "strict", None) is False: return True - if function.parameters is None: - return True - return function.parameters + return function.parameters if function.parameters is not None else True -_enable_structured_outputs_in_reasoning: bool = False +def _hermes_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: + arguments_field_prefix = '", "arguments": ' + formats = [ + # + # {"name": "t1", "arguments": {"q": "v"}} + # + ('\n{"name": "', "}\n"), + # {"name": "t1", "arguments": {"q": "v"}} + ('{"name": "', "}"), + ] - -def set_enable_structured_outputs_in_reasoning(enabled: bool) -> None: - """Publish the engine's ``enable_in_reasoning`` flag to tool parsers. - - Called once during APIServer startup so request-time parsers can read - it without going through the EngineCore-only contextvar. - """ - - global _enable_structured_outputs_in_reasoning - _enable_structured_outputs_in_reasoning = bool(enabled) - - -def get_enable_structured_outputs_in_reasoning() -> bool: - """Whether structured outputs are active during the reasoning phase. - - When ``True``, the structural tag will cover the reasoning part: - ``...`` prefix (if available); when ``False`` (default), the tag only - constrains the post-reasoning suffix. - """ - - return _enable_structured_outputs_in_reasoning - - -@register_model_structural_tag("deepseek_v4") -def get_deepseek_v4_structural_tag( - tools: list[ChatCompletionToolsParam], - tool_choice: SimplifiedToolChoice, - reasoning: bool, -) -> StructuralTag: - """Build DeepSeek V4 structural tags.""" - - invoke_begin_prefix = '<｜DSML｜invoke name="' - invoke_begin_suffix = '">\n' - invoke_end = "\n" - tool_calls_prefix = "\n\n" - function_calls_begin = "<｜DSML｜tool_calls>\n" - function_calls_end = "" - function_calls_trigger = "<｜DSML｜tool_calls>" - think_tag_end = "" - think_exclude_tokens = ["", ""] - xml_style = "deepseek_xml" - - if tool_choice == "auto": - tags = [] - for tool in tools: - function = tool.function - parameters = _get_function_parameters(function) - tags.append( - TagFormat( - begin=invoke_begin_prefix + function.name + invoke_begin_suffix, - content=JSONSchemaFormat( - json_schema=parameters, - style=xml_style, - ), - end=invoke_end, - ) - ) - - if tags: - function_calling_tags = TagsWithSeparatorFormat( - tags=tags, - separator="\n", - at_least_one=True, - ) - suffix_tag = TriggeredTagsFormat( - triggers=[function_calls_trigger], - tags=[ - TagFormat( - begin=function_calls_begin, - content=function_calling_tags, - end=function_calls_end, - ) - ], - excludes=think_exclude_tokens, - ) - else: - suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) - - elif tool_choice == "forced": - if not tools: - raise ValueError("Forced tool choice must resolve to exactly one tool.") - function = tools[0].function - suffix_tag = SequenceFormat( - elements=[ - ConstStringFormat(value=tool_calls_prefix + function_calls_begin), - TagFormat( - begin=invoke_begin_prefix + function.name + invoke_begin_suffix, - content=JSONSchemaFormat( - json_schema=_get_function_parameters(function), - style=xml_style, - ), - end=invoke_end, - ), - ConstStringFormat(value=function_calls_end), - ] - ) - - elif tool_choice == "required": - tags = [] - for tool in tools: - function = tool.function - parameters = _get_function_parameters(function) - tags.append( - TagFormat( - begin=invoke_begin_prefix + function.name + invoke_begin_suffix, - content=JSONSchemaFormat( - json_schema=parameters, - style=xml_style, - ), - end=invoke_end, - ) - ) - assert len(tags) > 0 - suffix_tag = SequenceFormat( - elements=[ - ConstStringFormat(value=tool_calls_prefix + function_calls_begin), - TagsWithSeparatorFormat( - tags=tags, - separator="\n", - at_least_one=True, - ), - ConstStringFormat(value=function_calls_end), - ] - ) - - if not reasoning: - return StructuralTag(format=suffix_tag) - - prefix_tag = TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end) - return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) - - -@register_model_structural_tag("qwen_3_5") -def get_qwen_3_5_structural_tag( - tools: list[ChatCompletionToolsParam], - tool_choice: SimplifiedToolChoice, - reasoning: bool, -) -> StructuralTag: - """Build Qwen XML structural tags. - - This format is used for Qwen3-Coder/Qwen3.5/Qwen3.6 and is compatible with - Qwen variants that use the same XML tool-call format. - """ - tool_call_begin_prefix = "\n", ""] - - if tool_choice == "auto": - tags = [] - for tool in tools: - function = tool.function - parameters = _get_function_parameters(function) - tags.append( - TagFormat( - begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", - content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), - end=tool_call_end, - ) - ) - - if tags: - suffix_tag = TriggeredTagsFormat( - triggers=[tool_call_trigger], - tags=tags, - excludes=think_exclude_tokens, - ) - else: - suffix_tag = AnyTextFormat(excludes=think_exclude_tokens) - - elif tool_choice == "forced": - if not tools: - raise ValueError("Forced tool choice must resolve to exactly one tool.") - function = tools[0].function - suffix_tag = TagFormat( - begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", + return [ + TagFormat( + begin=begin + tool.function.name + arguments_field_prefix, content=JSONSchemaFormat( - json_schema=_get_function_parameters(function), - style="qwen_xml", + json_schema=_get_function_parameters(tool.function) ), - end=tool_call_end, + end=end, ) + for tool in tools + for begin, end in formats + ] - elif tool_choice == "required": - tags = [] - for tool in tools: - function = tool.function - parameters = _get_function_parameters(function) - tags.append( - TagFormat( - begin=f"{tool_call_begin_prefix}{function.name}{tool_call_begin_suffix}", - content=JSONSchemaFormat(json_schema=parameters, style="qwen_xml"), - end=tool_call_end, - ) - ) - assert len(tags) > 0 + +@register_vllm_structural_tag("hermes") +def get_hermes_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + del builtin_tools, reasoning + + tool_call_trigger = "" + + if tool_choice == "auto": + tags = _hermes_tool_tags(tools) + suffix_tag = ( + TriggeredTagsFormat(triggers=[tool_call_trigger], tags=tags) + if tags + else AnyTextFormat() + ) + elif tool_choice == "forced": suffix_tag = TagsWithSeparatorFormat( - tags=tags, + tags=_hermes_tool_tags(tools), + separator="", + at_least_one=True, + stop_after_first=True, + ) + else: + suffix_tag = TagsWithSeparatorFormat( + tags=_hermes_tool_tags(tools), separator="", at_least_one=True, ) - if not reasoning: - result = StructuralTag(format=suffix_tag) - else: - prefix_tag = SequenceFormat( + return StructuralTag(format=suffix_tag) + + +def _minimax_tool_tags(tools: list[FunctionToolParam]) -> list[TagFormat]: + return [ + TagFormat( + begin=f'\n', + content=JSONSchemaFormat( + json_schema=_get_function_parameters(tool.function), + style="minimax_xml", + ), + end="\n", + ) + for tool in tools + ] + + +@register_vllm_structural_tag("minimax") +def get_minimax_structural_tag( + tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: SimplifiedToolChoice, + reasoning: bool, +) -> StructuralTag: + del builtin_tools, reasoning + + tool_call_begin = "\n" + tool_call_end = "" + tool_call_trigger = "" + + tags = _minimax_tool_tags(tools) + + if tool_choice == "auto": + suffix_tag = ( + TriggeredTagsFormat( + triggers=[tool_call_trigger], + tags=[ + TagFormat( + begin=tool_call_begin, + content=TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + ), + end=tool_call_end, + ) + ], + excludes=["", ""], + ) + if tags + else AnyTextFormat(excludes=["", ""]) + ) + elif tool_choice == "forced": + suffix_tag = SequenceFormat( elements=[ - TagFormat(begin="", content=AnyTextFormat(), end=think_tag_end), - ConstStringFormat(value=think_suffix), + ConstStringFormat(value="\n" + tool_call_begin), + TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + stop_after_first=True, + ), + ConstStringFormat(value=tool_call_end), + ] + ) + else: + suffix_tag = SequenceFormat( + elements=[ + ConstStringFormat(value="\n" + tool_call_begin), + TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + ), + ConstStringFormat(value=tool_call_end), ] ) - result = StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) - return result + return StructuralTag(format=suffix_tag) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 427f30b3992..3edfe932e0c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( ops_colqwen3="OpsColQwen3Config", qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig", cosmos3_omni="Cosmos3Config", + diffusion_gemma="DiffusionGemmaConfig", deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 71f7723e4c8..e91f89b2d09 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -26,6 +26,8 @@ _CLASS_TO_MODULE: dict[str, str] = { "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3", "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3", + "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma", + "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr", @@ -97,6 +99,8 @@ __all__ = [ "OpsColQwen3Config", "Qwen3VLNemotronEmbedConfig", "Cosmos3Config", + "DiffusionGemmaConfig", + "DiffusionGemmaTextConfig", "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", diff --git a/vllm/transformers_utils/configs/diffusion_gemma.py b/vllm/transformers_utils/configs/diffusion_gemma.py new file mode 100644 index 00000000000..246a25b32c6 --- /dev/null +++ b/vllm/transformers_utils/configs/diffusion_gemma.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +from transformers import PretrainedConfig +from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig + + +def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None: + PretrainedConfig.__init__(self, **kwargs) + # DiffusionGemma always uses MoE and K=V sharing for full_attention + # layers. The HF reference removed these config fields entirely. + if getattr(self, "num_experts", None): + self.enable_moe_block = True + self.attention_k_eq_v = True + + +class DiffusionGemmaTextConfig(PretrainedConfig): + model_type = "diffusion_gemma_text" + + def __init__(self, **kwargs: Any): + _init_text_config(self, **kwargs) + + +class DiffusionGemmaConfig(PretrainedConfig): + model_type = "diffusion_gemma" + + def __init__( + self, + text_config: dict[str, Any] | None = None, + canvas_length: int = 256, + self_conditioning_size: int | None = None, + **kwargs: Any, + ): + self.text_config = DiffusionGemmaTextConfig(**(text_config or {})) + self.canvas_length = canvas_length + self.self_conditioning_size = self_conditioning_size + vision_config = kwargs.pop("vision_config", None) + if isinstance(vision_config, dict): + self.vision_config = Gemma4VisionConfig(**vision_config) + else: + self.vision_config = vision_config + self.audio_config = None + PretrainedConfig.__init__(self, **kwargs) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 250aee50378..37402dcaa0b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -582,6 +582,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "cohere_asr": CohereAsrModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, + "diffusion_gemma_text": Gemma4ModelArchConfigConvertor, "ernie_mtp": ErnieMTPModelArchConfigConvertor, "falcon": FalconModelArchConfigConvertor, "falcon_mamba": MambaModelArchConfigConvertor, diff --git a/vllm/utils/async_utils.py b/vllm/utils/async_utils.py index 725868c39a3..9f368be7b2d 100644 --- a/vllm/utils/async_utils.py +++ b/vllm/utils/async_utils.py @@ -248,6 +248,32 @@ def make_async( return _async_wrapper +def make_async_with_semaphore( + func: Callable[P, T], + executor: ThreadPoolExecutor, +) -> Callable[P, Awaitable[T]]: + """ + Take a blocking function, and run it on in an executor thread. + + This function prevents the blocking function from blocking the + asyncio event loop. + The code in this function needs to be thread safe. + + The function is wrapped in a semaphore to limit the number of + concurrent executions making it easier to cancel tasks before they start. + """ + + semaphore = asyncio.Semaphore(executor._max_workers) + + async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: + loop = asyncio.get_event_loop() + p_func = partial(func, *args, **kwargs) + async with semaphore: + return await loop.run_in_executor(executor, p_func) + + return _async_wrapper + + def run_in_loop(loop: AbstractEventLoop, function: Callable, *args): if in_loop(loop): function(*args) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 95f8b4b7ec0..e0518277865 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -934,20 +934,27 @@ def should_use_flashinfer_for_blockscale_fp8_gemm( return should_use_flashinfer -_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 attention +_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 ViT attention @functools.cache def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: """Check if FP8 ViT attention is supported on this platform. - Requires native FP8 hardware support, the FlashInfer cuDNN backend, + Requires Blackwell (SM 100) or newer, the FlashInfer cuDNN backend, and cuDNN >= 9.17.1. + + cuDNN's FP8 SDPA forward path with bf16/fp16 output (used by + ``MMEncoderAttention._forward_flashinfer``) gates internally on + ``prop.major >= 10``; on Hopper it raises a misleading + ``cudnnGraphNotSupportedError: ... cuDNN version 9.13.0 and newer`` + even when the installed cuDNN is new enough. See PR #38065 for the + original Blackwell-only design intent. """ from vllm.v1.attention.backends.registry import AttentionBackendEnum - # cuDNN SDPA FP8 requires Hopper (SM 90) or newer. - if not current_platform.has_device_capability(90): + # cuDNN SDPA FP8 with bf16/fp16 output requires Blackwell (SM 100) or newer. + if not current_platform.has_device_capability(100): return False try: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 32b4b8ab9a0..152178ec2b3 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -387,7 +387,7 @@ class CommonAttentionMetadata: block_table_tensor: torch.Tensor slot_mapping: torch.Tensor - causal: bool = True + causal: bool | torch.Tensor = True # Needed by FastPrefillAttentionBuilder logits_indices_padded: torch.Tensor | None = None @@ -497,7 +497,9 @@ class CommonAttentionMetadata: max_seq_len=self.max_seq_len, block_table_tensor=self.block_table_tensor[:num_actual_reqs], slot_mapping=self.slot_mapping[:num_actual_tokens], - causal=self.causal, + causal=self.causal[:num_actual_reqs] + if isinstance(self.causal, torch.Tensor) + else self.causal, logits_indices_padded=self.logits_indices_padded, num_logits_indices=self.num_logits_indices, encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 0d6a3d298b6..474523780ff 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -131,6 +131,12 @@ def get_flash_attn_version( and head_size != head_size_v ): upgrade_reason = "Diff-KV with sinks" + elif ( + vllm_config is not None + and vllm_config.model_config is not None + and vllm_config.model_config.is_diffusion + ): + upgrade_reason = "Per-sequence causal (dynamic_causal) requires FA4" if upgrade_reason: logger.info_once( "%s: upgrading FlashAttention 3 -> 4", diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index d6774a6eb99..9e33c0d823b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -267,7 +267,7 @@ class FlashAttentionMetadata: prefix_scheduler_metadata: torch.Tensor | None = None max_num_splits: int = 0 - causal: bool = True + causal: bool | torch.Tensor = True # PrefixLM bidirectional ranges for multimodal tokens. # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range. @@ -570,6 +570,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.scheduler_metadata[n:] = 0 scheduler_metadata = self.scheduler_metadata[:n] + if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32: + causal = causal.to(torch.int32) + attn_metadata = FlashAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -824,18 +827,46 @@ class FlashAttentionImpl(AttentionImpl): if self.sliding_window is not None else None ) + + causal = attn_metadata.causal + is_dynamic_causal = isinstance(causal, torch.Tensor) + + # For non-causal (bidirectional) attention, make the + # sliding window symmetric so queries attend in both + # directions. + if ( + sliding_window_size is not None + and sliding_window_size[1] == 0 + and (is_dynamic_causal or causal is False) + ): + sliding_window_size = [ + sliding_window_size[0], + sliding_window_size[0], + ] + mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor mm_mask_mod = None mm_aux = None if ( mm_prefix_ranges is not None - and attn_metadata.causal + and not is_dynamic_causal + and causal is True and self.vllm_flash_attn_version == 4 ): max_ranges = mm_prefix_ranges.shape[1] mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges) mm_aux = [mm_prefix_ranges] + dynamic_causal = None + if isinstance(causal, torch.Tensor): + if self.vllm_flash_attn_version != 4: + raise NotImplementedError( + "Per-sequence causal requires FA4. Current version: " + f"FA{self.vllm_flash_attn_version}" + ) + dynamic_causal = causal + causal = False + flash_attn_varlen_func( q=query[:num_actual_tokens], k=key_cache, @@ -846,7 +877,7 @@ class FlashAttentionImpl(AttentionImpl): seqused_k=seqused_k, max_seqlen_k=max_seqlen_k, softmax_scale=self.scale, - causal=attn_metadata.causal, + causal=causal, alibi_slopes=self.alibi_slopes, window_size=sliding_window_size, block_table=block_table, @@ -856,6 +887,7 @@ class FlashAttentionImpl(AttentionImpl): q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, + dynamic_causal=dynamic_causal, num_splits=attn_metadata.max_num_splits, s_aux=self.sinks, mask_mod=mm_mask_mod, diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index e788b0e3496..ff8fbfc022b 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -41,6 +41,30 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): def set_head_size_v(cls, head_size_v: int) -> None: cls.head_size_v = head_size_v + @classmethod + def is_supported_on_current_device( + cls, + head_size: int, + head_size_v: int, + has_sinks: bool, + ) -> bool: + """Check whether FA3/4 with this DiffKV config is usable here. + + DiffKV (hdim_qk != hdim_v) requires FA3 or FA4 + """ + if not is_flash_attn_varlen_func_available(): + return False + try: + version = get_flash_attn_version( + requires_alibi=False, + head_size=head_size, + head_size_v=head_size_v, + has_sinks=has_sinks, + ) + except Exception: + return False + return version in (3, 4) + @staticmethod def get_name() -> str: return "FLASH_ATTN_DIFFKV" @@ -49,8 +73,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): def get_impl_cls() -> type["FlashAttentionImpl"]: return FlashAttentionDiffKVImpl - # Do not modify the interface of get_kv_cache_shape, - # but consider head_size_v when returning result. @staticmethod def get_kv_cache_shape( num_blocks: int, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 24a59f03800..2cd2bb5b986 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -46,6 +46,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend" ) TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" + TRITON_ATTN_DIFFKV = ( + "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend" + ) ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" ROCM_AITER_TRITON_MLA = ( diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 92ff08cc0f3..377e9e7ab1d 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -79,6 +79,8 @@ class TritonAttentionMetadata: softmax_segm_max: torch.Tensor softmax_segm_expsum: torch.Tensor + causal: bool | torch.Tensor + # For cascade attention. use_cascade: bool common_prefix_len: int @@ -219,6 +221,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet seq_lens=seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, + causal=common_attn_metadata.causal, use_cascade=use_cascade, common_prefix_len=common_prefix_len, cu_prefix_query_lens=cu_prefix_query_lens, @@ -271,6 +274,10 @@ class TritonAttentionBackend(AttentionBackend): forward_includes_kv_cache_update: bool = False + @classmethod + def supports_non_causal(cls) -> bool: + return True + @staticmethod def get_name() -> str: return "TRITON_ATTN" @@ -619,7 +626,7 @@ class TritonAttentionImpl(AttentionImpl): seqused_k=seqused_k, max_seqlen_k=max_seqlen_k, softmax_scale=self.scale, - causal=True, + causal=attn_metadata.causal, alibi_slopes=self.alibi_slopes, use_alibi_sqrt=self.use_alibi_sqrt, window_size=self.sliding_window, diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py new file mode 100644 index 00000000000..3420a0eba47 --- /dev/null +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton attention backend with different K/V head dimensions (DiffKV). + +The KV cache layout is identical to ``FlashAttentionDiffKVBackend`` — K +and V are packed along the last dim: + + [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v] + +so existing helpers (``triton_reshape_and_cache_flash_diffkv``) are reused. +""" + +from typing import ClassVar + +import torch + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.backend import AttentionLayer, AttentionType +from vllm.v1.attention.backends.triton_attn import ( + TritonAttentionBackend, + TritonAttentionImpl, + TritonAttentionMetadata, + TritonAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_diffkv, +) +from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( + unified_attention_diffkv, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder): + """Override the parent's softmax buffer last-dim to head_size_v. + + The parent allocates ``softmax_segm_output`` with last-dim sized to + ``next_power_of_2(head_size)`` (== Q/K head size). For DiffKV the + accumulator and per-segment partial outputs are V-shaped, so we + re-allocate with ``next_power_of_2(head_size_v)`` instead. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + head_size_v = TritonAttentionDiffKVBackend.head_size_v + head_size_v_padded = next_power_of_2(head_size_v) + self.softmax_segm_output = torch.empty( + ( + self.seq_threshold_3D, + self.num_heads_q, + self.num_par_softmax_segments, + head_size_v_padded, + ), + dtype=torch.float32, + device=device, + ) + + +class TritonAttentionDiffKVBackend(TritonAttentionBackend): + # V head dim — set per layer via ``set_head_size_v`` before instantiation. + head_size_v: int = 128 + + # No FP8 / int8 KV cache for the DiffKV path yet; require fp16/bf16/fp32. + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + ] + + @classmethod + def set_head_size_v(cls, head_size_v: int) -> None: + cls.head_size_v = head_size_v + + @staticmethod + def get_name() -> str: + return "TRITON_ATTN_DIFFKV" + + @staticmethod + def get_impl_cls() -> type["TritonAttentionDiffKVImpl"]: + return TritonAttentionDiffKVImpl + + @staticmethod + def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]: + return TritonAttentionDiffKVMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if block_size % 16 != 0: + raise ValueError("Block size must be a multiple of 16.") + return ( + num_blocks, + block_size, + num_kv_heads, + head_size + TritonAttentionDiffKVBackend.head_size_v, + ) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + cache_layout = get_kv_cache_layout() + if cache_layout == "NHD" and include_num_layers_dimension: + # (num_blocks, num_layers, block_size, + # num_kv_heads, head_size + head_size_v) + return (1, 0, 2, 3, 4) + elif cache_layout == "NHD": + return (0, 1, 2, 3) + elif cache_layout == "HND" and include_num_layers_dimension: + # (num_blocks, num_kv_heads, num_layers, + # block_size, head_size + head_size_v) + return (1, 3, 0, 2, 4) + elif cache_layout == "HND": + return (0, 2, 1, 3) + else: + raise ValueError(f"Unknown cache layout format {cache_layout}.") + + @classmethod + def supports_head_size(cls, head_size: int) -> bool: + # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed. + return head_size >= 32 + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + # DiffKV only implements decoder self-attention. Unlike the parent + # TritonAttentionBackend (which advertises all types), encoder + # attention is not supported, so gate it here at backend selection. + return attn_type == AttentionType.DECODER + + +class TritonAttentionDiffKVImpl(TritonAttentionImpl): + """Triton attention impl for the DiffKV packed KV cache layout.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if is_quantized_kv_cache(self.kv_cache_dtype): + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not yet support quantized " + f"KV cache (got kv_cache_dtype={self.kv_cache_dtype!r})." + ) + if self._is_per_token_head_quant: + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not support per-token-head " + "quantization." + ) + if self.chunk_lookback > -1: + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not support chunked " + "attention with lookback." + ) + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + # Cache is packed [..., head_size_qk + head_size_v]; the diffkv + # reshape kernel writes K to [..., :head_size_qk] and V to + # [..., head_size_qk:hqk+hv]. + triton_reshape_and_cache_flash_diffkv( + key, + value, + kv_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + def fused_rope_kvcache_supported(self): + # The fused rope+cache path assumes the standard 2-tensor layout. + return False + + def forward( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass. + + Shapes: + query: [num_tokens, num_heads, head_size_qk] + key: [num_tokens, num_kv_heads, head_size_qk] + value: [num_tokens, num_kv_heads, head_size_v] + kv_cache: [num_blocks, block_size, num_kv_heads, + head_size_qk + head_size_v] + output: [num_tokens, num_heads, head_size_v] + """ + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "fused output quantization is not supported for " + "TritonAttentionDiffKVImpl" + ) + + if attn_metadata is None: + return output.fill_(0) + + assert attn_metadata.use_cascade is False, ( + "Cascade attention not supported for TritonAttentionDiffKVImpl" + ) + + num_actual_tokens = attn_metadata.num_actual_tokens + head_size_qk = self.head_size + head_size_v = TritonAttentionDiffKVBackend.head_size_v + + # Slice the packed cache into K / V views. Strides on dims 0/1/2 + # match the original cache; dim 3 stays contiguous (stride 1). + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v] + + unified_attention_diffkv( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=attn_metadata.query_start_loc, + seqused_k=attn_metadata.seq_lens, + softmax_scale=self.scale, + causal=True, + alibi_slopes=self.alibi_slopes, + use_alibi_sqrt=self.use_alibi_sqrt, + window_size=self.sliding_window, + block_table=attn_metadata.block_table, + softcap=self.logits_soft_cap, + sinks=self.sinks, + max_seqlen_q=attn_metadata.max_query_len, + seq_threshold_3D=attn_metadata.seq_threshold_3D, + num_par_softmax_segments=attn_metadata.num_par_softmax_segments, + softmax_segm_output=attn_metadata.softmax_segm_output, + softmax_segm_max=attn_metadata.softmax_segm_max, + softmax_segm_expsum=attn_metadata.softmax_segm_expsum, + ) + return output diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 12fd3a17421..8104e808f67 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1406,6 +1406,348 @@ def _sparse_attn_decode_ragged_kernel( ) +@triton.jit +def _sparse_attn_decode_partial_kernel( + q_ptr, + main_cache_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride0, + q_stride1, + main_cache_stride0, + extra_cache_stride0, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + main_num_rows, + extra_num_rows, + main_block_size, + extra_block_size, + scale, + num_heads, + HAS_EXTRA: tl.constexpr, + NOPE_DIM: tl.constexpr, + NOPE_BLOCK: tl.constexpr, + ROPE_DIM: tl.constexpr, + IS_FNUZ: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + nope_offsets = tl.arange(0, NOPE_BLOCK) + nope_mask = nope_offsets < NOPE_DIM + rope_offsets = tl.arange(0, ROPE_DIM) + + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1 + q_nope = tl.load( + q_row_ptr + nope_offsets[None, :], + mask=head_mask[:, None] & nope_mask[None, :], + other=0.0, + ) + q_rope = tl.load( + q_row_ptr + NOPE_DIM + rope_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32) + acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32) + k_offsets = tl.arange(0, BLOCK_K) + + zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16) + zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16) + + # Each split processes a contiguous slice of this query's main (SWA) and + # extra (topk) segments. Slices are handled independently so a block never + # straddles the main/extra boundary. + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + main_lo = split_id * main_chunk + main_hi = tl.minimum(main_lo + main_chunk, main_len) + + for k_start in tl.range(main_lo, main_hi, BLOCK_K, num_stages=NUM_STAGES): + k_pos = k_start + k_offsets + in_range = k_pos < main_hi + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < main_num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size + cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + + x_uint8 = tl.load( + token_data_ptr[:, None] + nope_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + if IS_FNUZ: + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + encoded_scales = tl.load( + token_scale_ptr[:, None] + nope_offsets[None, :] // 64, + mask=valid[:, None] & nope_mask[None, :], + other=127, + ) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + k_rope = tl.load( + rope_ptr[:, None] + rope_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + k_rope = tl.where(valid[:, None], k_rope, zero_rope) + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope) + + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope)) + scores *= scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope) + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope) + m_i = m_new + l_i = l_new + + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_lo = split_id * extra_chunk + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) + + for k_start in tl.range(extra_lo, extra_hi, BLOCK_K, num_stages=NUM_STAGES): + k_pos = k_start + k_offsets + in_range = k_pos < extra_hi + slot = tl.load( + extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1 + ) + valid = in_range & (slot >= 0) & (slot < extra_num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size + cache_block_ptr = ( + extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 + ) + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = ( + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + ) + + x_uint8 = tl.load( + token_data_ptr[:, None] + nope_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + if IS_FNUZ: + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + encoded_scales = tl.load( + token_scale_ptr[:, None] + nope_offsets[None, :] // 64, + mask=valid[:, None] & nope_mask[None, :], + other=127, + ) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + k_rope = tl.load( + rope_ptr[:, None] + rope_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + k_rope = tl.where(valid[:, None], k_rope, zero_rope) + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope) + + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot( + q_rope, + tl.trans(k_rope), + ) + scores *= scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope) + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope) + m_i = m_new + l_i = l_new + + # Store raw (un-normalized) partial state for this split. Softmax sink and + # final normalization happen in the reduce kernel. + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets + tl.store(part_m_ptr + pm_base, m_i, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + acc_base = ( + part_acc_ptr + + query_idx * pa_stride0 + + split_id * pa_stride_s + + head_offsets[:, None] * pa_stride_h + ) + tl.store( + acc_base + nope_offsets[None, :], + acc_nope, + mask=head_mask[:, None] & nope_mask[None, :], + ) + tl.store( + acc_base + NOPE_DIM + rope_offsets[None, :], + acc_rope, + mask=head_mask[:, None], + ) + + +@triton.jit +def _sparse_attn_decode_reduce_kernel( + part_m_ptr, + part_l_ptr, + part_acc_ptr, + attn_sink_ptr, + out_ptr, + out_stride0, + out_stride1, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_ATTN_SINK: tl.constexpr, + COMB_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_SPLITS: tl.constexpr, + SPLITS_PAD: tl.constexpr, +): + query_idx = tl.program_id(0) + pid_h = tl.program_id(1) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + comb_offsets = tl.arange(0, COMB_DIM) + # SPLITS_PAD is NUM_SPLITS rounded up to a power of two so the parallel + # split-axis load is a legal arange for any split count; padding lanes are + # masked off. + split_offsets = tl.arange(0, SPLITS_PAD) + split_mask = split_offsets < NUM_SPLITS + + neg_large = -3.4028234663852886e38 + + # Phase 1: load every split's running max/sum at once and reduce the max + # in parallel (tl.max over the split axis) instead of walking the splits + # serially. This breaks the long online-softmax dependency chain that made + # the reduce latency-bound. + load_mask = split_mask[:, None] & head_mask[None, :] + pm_split = ( + part_m_ptr + + query_idx * pm_stride0 + + split_offsets[:, None] * pm_stride_s + + head_offsets[None, :] + ) + m_all = tl.load(pm_split, mask=load_mask, other=neg_large) # [S, H] + l_all = tl.load( + part_l_ptr + + query_idx * pm_stride0 + + split_offsets[:, None] * pm_stride_s + + head_offsets[None, :], + mask=load_mask, + other=0.0, + ) + + m_comb = tl.max(m_all, axis=0) # [H] + if HAS_ATTN_SINK: + sink = tl.load( + attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large + ).to(tl.float32) + m_final = tl.maximum(m_comb, sink) + else: + m_final = m_comb + + w_all = tl.exp(m_all - m_final[None, :]) # [S, H] + w_all = tl.where(load_mask, w_all, 0.0) + l_final = tl.sum(w_all * l_all, axis=0) # [H] + if HAS_ATTN_SINK: + l_final = l_final + tl.exp(sink - m_final) + denom = tl.maximum(l_final, 1.0e-30) + + # Phase 2: weighted sum of the per-split accumulators. The combine weight + # for each split only depends on the (already known) global max, so the + # acc loads carry no cross-split dependency and the compiler can pipeline + # them; only the cheap FMA into `acc` is loop-carried. + acc = tl.zeros((BLOCK_H, COMB_DIM), dtype=tl.float32) + for s in tl.static_range(NUM_SPLITS): + m_s = tl.load( + part_m_ptr + query_idx * pm_stride0 + s * pm_stride_s + head_offsets, + mask=head_mask, + other=neg_large, + ) + w_s = tl.exp(m_s - m_final) + acc_base = ( + part_acc_ptr + + query_idx * pa_stride0 + + s * pa_stride_s + + head_offsets[:, None] * pa_stride_h + ) + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + acc += w_s[:, None] * acc_s + + out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) + + out_row_ptr = ( + out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1 + ) + tl.store( + out_row_ptr + comb_offsets[None, :], + out, + mask=head_mask[:, None], + ) + + def _rocm_sparse_attn_prefill_ragged_triton( q: torch.Tensor, kv: torch.Tensor, @@ -1502,6 +1844,101 @@ def _rocm_sparse_attn_prefill_triton( ) +@functools.lru_cache +def _decode_cu_count() -> int: + try: + return torch.cuda.get_device_properties(0).multi_processor_count + except Exception: + return 256 # For gfx950 arch, gated behind a fallback path for other archs. + + +def _decode_partial_iters( + avg_main_len: float, avg_extra_len: float, splits: int, block_k: int +) -> int: + """BLOCK_K iterations one partial workgroup walks for ``splits`` splits. + + Each split processes ``ceil(seg_len / splits)`` tokens of a segment, walked + ``BLOCK_K`` at a time, and the main/extra segments are handled separately. + """ + main_iters = ( + math.ceil(math.ceil(avg_main_len / splits) / block_k) if avg_main_len > 0 else 0 + ) + extra_iters = ( + math.ceil(math.ceil(avg_extra_len / splits) / block_k) + if avg_extra_len > 0 + else 0 + ) + return main_iters + extra_iters + + +def _decode_num_splits( + num_queries: int, + heads_blocks: int, + avg_main_len: float = 0.0, + avg_extra_len: float = 0.0, + block_k: int = 32, +) -> int: + """Pick a flash-decode split count to keep the GPU busy across batch sizes. + + Decode launches only ``num_queries * heads_blocks`` workgroups otherwise, + which severely under-fills the device for the low-concurrency regime that + dominates latency. Splitting the KV sequence adds parallelism. + + We model the relative partial-kernel latency for a given split count ``s`` + as ``waves * (1/s + mu)`` where ``waves = ceil(base * s / CU)`` and ``mu`` + is a small per-wave overhead penalty: + + - ``waves / s`` captures the partial compute: each wave walks roughly + ``total_tokens / s`` tokens and there are ``waves`` of them, so dividing + by ``s`` makes more splits cheaper *until* they spill into extra waves. + - ``mu * waves`` charges per-wave launch/tail overhead so we do not + over-split into many mostly-idle waves (e.g. batch 224 on 256 CUs is + best left at 1 split rather than 8 splits across 7 waves). + + The minimiser naturally prefers split counts that pack the device into full + waves (``base * s`` near a multiple of ``CU``) and falls back to 1 split + once the batch already fills the device. Ties favour the smaller split + count (less reduce work). + + Finally we "snap down" the chosen split count to the smallest value that + yields the same wave count *and* the same per-workgroup BLOCK_K iteration + count. Because latency tracks iteration count (not raw token count), extra + splits that do not lower the iteration count add only reduce/HBM overhead + for no parallelism gain (e.g. batch 24: s8 and s10 both walk 4 extra iters + in one wave, so s8 is strictly better). Snapping needs the average segment + lengths, which the caller derives sync-free from the ragged index sizes. + """ + base = max(1, num_queries * heads_blocks) + # Target ~1 workgroup per CU: enough to fill the device while keeping the + # reduce cost (which grows with split count) small. Tuned on gfx950. + cu = max(1, _decode_cu_count()) + # Per-wave overhead penalty: higher values discourage split counts that + # spill into extra GPU waves. Tuned on gfx950. + mu = 0.04 + best_splits = 1 + best_cost = None + # Search up to 16 splits; beyond that the reduce/HBM overhead dominates. + for splits in range(1, 17): + waves = (base * splits + cu - 1) // cu + cost = waves * (1.0 / splits + mu) + if best_cost is None or cost < best_cost - 1e-9: + best_splits = splits + best_cost = cost + + if best_splits > 1 and (avg_main_len > 0 or avg_extra_len > 0): + target_waves = (base * best_splits + cu - 1) // cu + target_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, best_splits, block_k + ) + for splits in range(1, best_splits): + waves = (base * splits + cu - 1) // cu + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k) + if waves == target_waves and iters == target_iters: + best_splits = splits + break + return best_splits + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -1575,9 +2012,70 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32) block_h = 16 - block_k = 16 if head_dim >= 256 else 32 out = torch.empty_like(q, dtype=torch.bfloat16) - _sparse_attn_decode_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))]( + heads_blocks = triton.cdiv(num_heads, block_h) + nope_block = triton.next_power_of_2(nope_head_dim) + comb_dim = nope_head_dim + rope_head_dim + is_fnuz = current_platform.is_fp8_fnuz() + + if not _ON_GFX950: # Fallback path for un-tuned architectures. + block_k = 16 if head_dim >= 256 else 32 + _sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + attn_sink, + out, + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_ATTN_SINK=has_attn_sink, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + IS_FNUZ=is_fnuz, + BLOCK_H=block_h, + BLOCK_K=block_k, + num_warps=8, + ) + return out + + block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. + # Average per-query segment lengths, read sync-free from the ragged index + # sizes, let the split heuristic avoid over-splitting + # main_indices/extra_indices are flat [nnz] int32. + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) + + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, comb_dim), + dtype=torch.float32, + device=q.device, + ) + + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( q, main_cache, main_indices, @@ -1585,29 +2083,56 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_cache, extra_indices, extra_indptr, - attn_sink, - out, + part_m, + part_l, + part_acc, q.stride(0), q.stride(1), - out.stride(0), - out.stride(1), main_cache.stride(0), extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), main_cache.shape[0] * main_cache.shape[1], extra_cache.shape[0] * extra_cache.shape[1], main_cache.shape[1], extra_cache.shape[1], scale, num_heads, - HAS_ATTN_SINK=has_attn_sink, HAS_EXTRA=has_extra, NOPE_DIM=nope_head_dim, - NOPE_BLOCK=triton.next_power_of_2(nope_head_dim), + NOPE_BLOCK=nope_block, ROPE_DIM=rope_head_dim, - IS_FNUZ=current_platform.is_fp8_fnuz(), + IS_FNUZ=is_fnuz, BLOCK_H=block_h, BLOCK_K=block_k, - num_warps=8, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) + + _sparse_attn_decode_reduce_kernel[(num_queries, heads_blocks)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + HAS_ATTN_SINK=has_attn_sink, + COMB_DIM=comb_dim, + BLOCK_H=block_h, + NUM_SPLITS=num_splits, + SPLITS_PAD=triton.next_power_of_2(num_splits), + num_warps=4, ) return out diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index 6ed50f6a2df..ed9a38ad6cd 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -153,6 +153,8 @@ def compute_tile_loop_bounds( SLIDING_WINDOW: tl.constexpr, USE_MM_PREFIX: tl.constexpr, IS_3D: tl.constexpr, + USE_CAUSAL: tl.constexpr = True, + USE_PER_SEQ_CAUSAL: tl.constexpr = False, CHUNK_LOOKBACK: tl.constexpr = -1, CHUNK_SIZE: tl.constexpr = -1, ): @@ -163,10 +165,11 @@ def compute_tile_loop_bounds( 1. Longest prefix spanned by any query token in this q-block. Clamped to ``seq_len`` (causal) or extended to it when - mm_prefix is active (bidirectional ranges can reach past the - causal prefix). + mm_prefix is active or non-causal sequences need the full + sequence. 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to only tiles that can contain an allowed key under SWA. + For non-causal sequences, the window extends in both directions. 3. 3D scoping: when ``IS_3D`` is True, further narrows to the segment's slice via ``(segm_idx * tiles_per_segment, (segm_idx + 1) * tiles_per_segment)``. @@ -179,9 +182,10 @@ def compute_tile_loop_bounds( + (BLOCK_M - 1) // num_queries_per_kv + 1 ) - if USE_MM_PREFIX: - # image bidirectional attention ranges require a full range - # including q_block padding to make sure doc mask is correct + if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): + # Non-causal or mixed batches need the full sequence range. + # Per-element masking in compute_kv_seq_mask handles the + # actual causal/non-causal boundary per sequence. max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len) else: max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) @@ -207,12 +211,17 @@ def compute_tile_loop_bounds( # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi] q_abs = context_len + qpos_lo if CHUNK_LOOKBACK > -1: - # Chunked attention: align lower bound to the start of the - # lookback'th previous chunk. first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE else: first_allowed_key = q_abs - SLIDING_WINDOW + 1 - last_allowed_key = context_len + qpos_hi + if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): + # Non-causal: keys can be AHEAD of query within the window + last_allowed_key = tl.minimum( + context_len + qpos_hi + SLIDING_WINDOW - 1, + seq_len - 1, + ) + else: + last_allowed_key = context_len + qpos_hi # Convert to tile indices and clamp tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) @@ -262,10 +271,14 @@ def compute_kv_seq_mask( query_abs_pos, seq_offset, seq_idx, + seq_len, mm_prefix_range_ptr, SLIDING_WINDOW: tl.constexpr, USE_MM_PREFIX: tl.constexpr, MAX_MM_RANGES: tl.constexpr, + USE_CAUSAL: tl.constexpr = True, + USE_PER_SEQ_CAUSAL: tl.constexpr = False, + per_seq_causal_ptr=None, CHUNK_LOOKBACK: tl.constexpr = -1, CHUNK_SIZE: tl.constexpr = -1, ): @@ -279,9 +292,23 @@ def compute_kv_seq_mask( Chunked attention takes precedence over sliding window when both are non-default — the launcher zeros ``CHUNK_LOOKBACK`` whenever sliding window is disabled. + + When ``USE_PER_SEQ_CAUSAL`` is set, each sequence carries its own + causal flag via ``per_seq_causal_ptr``; non-causal sequences use a + simple ``key < seq_len`` bound instead. ``USE_CAUSAL=False`` + disables causal masking entirely. """ - # Compute attention mask: causal by default (key <= query) - seq_mask = seq_offset[None, :] <= query_abs_pos + if USE_PER_SEQ_CAUSAL: + is_causal = tl.load(per_seq_causal_ptr + seq_idx) + seq_mask = tl.where( + is_causal, + seq_offset[None, :] <= query_abs_pos, + seq_offset[None, :] < seq_len, + ) + elif USE_CAUSAL: + seq_mask = seq_offset[None, :] <= query_abs_pos + else: + seq_mask = seq_offset[None, :] < seq_len # Apply sliding window / chunked attention to base mask # BEFORE mm_prefix OR. @@ -293,7 +320,15 @@ def compute_kv_seq_mask( <= CHUNK_LOOKBACK ) elif SLIDING_WINDOW > 0: - seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW) + sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW + if USE_PER_SEQ_CAUSAL: + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW + seq_mask = seq_mask & tl.where(is_causal, sw_left, sw_left & sw_right) + elif not USE_CAUSAL: + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW + seq_mask = seq_mask & sw_left & sw_right + else: + seq_mask = seq_mask & sw_left # PrefixLM: extend mask with bidirectional ranges for multimodal tokens. # Applied AFTER sliding window so mm_prefix ranges override SW restriction. diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 56f1d1c1d08..f39e44286be 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -215,6 +215,9 @@ def kernel_unified_attention( USE_SOFTCAP: tl.constexpr, # bool USE_SINKS: tl.constexpr, # bool SLIDING_WINDOW: tl.constexpr, # int + USE_CAUSAL: tl.constexpr, # bool + USE_PER_SEQ_CAUSAL: tl.constexpr, # bool + per_seq_causal_ptr, # [num_seqs] bool, or None USE_MM_PREFIX: tl.constexpr, # bool MAX_MM_RANGES: tl.constexpr, # int mm_prefix_range_ptr, @@ -389,6 +392,8 @@ def kernel_unified_attention( SLIDING_WINDOW, USE_MM_PREFIX, IS_3D, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, CHUNK_LOOKBACK, CHUNK_SIZE, ) @@ -493,10 +498,14 @@ def kernel_unified_attention( query_abs_pos, seq_offset, seq_idx, + seq_len, mm_prefix_range_ptr, SLIDING_WINDOW, USE_MM_PREFIX, MAX_MM_RANGES, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, + per_seq_causal_ptr, CHUNK_LOOKBACK, CHUNK_SIZE, ) @@ -532,11 +541,19 @@ def kernel_unified_attention( if SLIDING_WINDOW: qpos_lo = q_block_local_idx * BLOCK_Q - V = tl.where( - (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, - V, - 0.0, - ) + dist = context_len + qpos_lo - seq_offset[:, None] + if USE_PER_SEQ_CAUSAL: + is_causal_seq = tl.load(per_seq_causal_ptr + seq_idx) + sw_mask_v = tl.where( + is_causal_seq, + dist < SLIDING_WINDOW, + (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW), + ) + elif USE_CAUSAL: + sw_mask_v = dist < SLIDING_WINDOW + else: + sw_mask_v = (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW) + V = tl.where(sw_mask_v, V, 0.0) if USE_PER_TOKEN_HEAD_SCALES: # Per-token-head quant: apply v_scale to P instead of V. P_v = (P * v_token_head_scales[None, :]).to(V.dtype) @@ -802,7 +819,11 @@ def unified_attention( # disabling this flag costs nothing. use_td: bool = False, ): - assert causal, "Only causal attention is supported" + # Resolve causal: bool or per-seq tensor. + use_per_seq_causal = isinstance(causal, torch.Tensor) + use_causal = bool(causal) if not use_per_seq_causal else True + per_seq_causal_ptr = causal if use_per_seq_causal else None + if sinks is not None: assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size" @@ -841,6 +862,26 @@ def unified_attention( ) BLOCK_Q = BLOCK_M // num_queries_per_kv + # Tuned launch parameters; ``None`` lets Triton pick its defaults. + launch_num_warps: int | None = None + launch_num_stages: int | None = None + + # head_size 256 with many query rows per sequence (e.g. diffusion-gemma + # bidirectional canvas passes) is prefill-shaped, but the decode-oriented + # defaults (BLOCK_Q=8, TILE=32, 4 warps) under-tile it. A wider KV tile + + # more query rows per block + 8 warps is ~2x faster on B200. + tuned_large_head = ( + head_size == 256 + and max_seqlen_q > 1 + and num_queries_per_kv <= 16 + and current_platform.is_device_capability_family(100) + ) + if tuned_large_head: + BLOCK_M = 32 + BLOCK_Q = BLOCK_M // num_queries_per_kv + launch_num_warps = 8 + launch_num_stages = 2 + # Ideally we would launch with kernel with: # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks. # However, it is slow to realize the query_lens on cpu. @@ -869,6 +910,11 @@ def unified_attention( head_size, sliding_window_val, q.element_size(), is_prefill=False ) + # Wider KV tile for the tuned large-head path (see above). Only the 2D + # path (used when max_seqlen_q > 1) reads TILE_SIZE_PREFILL. + if tuned_large_head: + TILE_SIZE_PREFILL = 128 + # USE_TD requires BLOCK_SIZE % TILE_SIZE == 0 (enforced by a # ``tl.static_assert`` in the kernel). The default prefill tile # size (32) is larger than a common ``block_size=16``, so clamp it @@ -964,6 +1010,12 @@ def unified_attention( grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) tile_size = TILE_SIZE_DECODE + launch_kwargs: dict[str, int] = {} + if launch_num_warps is not None: + launch_kwargs["num_warps"] = launch_num_warps + if launch_num_stages is not None: + launch_kwargs["num_stages"] = launch_num_stages + kernel_unified_attention[grid]( output_ptr=out, segm_output_ptr=segm_output_ptr, @@ -1002,10 +1054,13 @@ def unified_attention( USE_QQ_BIAS=use_qq_bias, USE_SOFTCAP=(softcap > 0), USE_SINKS=(sinks is not None), + SLIDING_WINDOW=(1 + window_size[0]), + USE_CAUSAL=use_causal, + USE_PER_SEQ_CAUSAL=use_per_seq_causal, + per_seq_causal_ptr=per_seq_causal_ptr, USE_MM_PREFIX=use_mm_prefix, MAX_MM_RANGES=max_mm_ranges, mm_prefix_range_ptr=mm_prefix_range, - SLIDING_WINDOW=(1 + window_size[0]), stride_k_cache_0=k.stride(0), stride_k_cache_1=k.stride(1), stride_k_cache_2=k.stride(2), @@ -1033,6 +1088,7 @@ def unified_attention( CHUNK_SIZE=chunk_size, USE_TD=use_td, USE_TD_QO=use_td_qo, + **launch_kwargs, ) if use_3d: diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py new file mode 100644 index 00000000000..eaf62b6bce6 --- /dev/null +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -0,0 +1,530 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton unified attention with different K/V head dimensions (DiffKV). + +This is a slimmed fork of ``triton_unified_attention.py`` for models like +MiMo-V2.5 where the V tensor's head dimension differs from K's. The KV cache +is the same packed layout used by ``FlashAttentionDiffKVBackend``: + + kv_cache: [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v] + +We slice ``key_cache = kv_cache[..., :head_size_qk]`` and +``value_cache = kv_cache[..., head_size_qk:]`` on the host, so the kernel +takes two cache pointers but with two distinct head sizes. + +Both 2D and 3D launches are supported: + - 2D: one program per (q-block, kv-head); tile-loop walks the full KV + sequence; final output written directly. Used for prefill and large + decode batches. + - 3D: one program per (q-block, kv-head, segm); each program covers a + KV slice and writes per-segment partials (max/expsum/output). A + follow-up ``kernel_reduce_segments_diffkv`` combines them. Selected + for decode-only batches whose 2D grid would under-fill the GPU. +""" + +from typing import Any + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_attention_helpers import ( + apply_alibi_to_score, + apply_softcap, + cdiv_fn, + compute_kv_seq_mask, + compute_tile_loop_bounds, + find_seq_idx, + init_softmax_M, + resolve_seq_and_query_len, + softmax_step, + store_segm_reduce_scalars, +) + +logger = init_logger(__name__) + +is_batch_invariant = envs.VLLM_BATCH_INVARIANT + + +@triton.jit +def kernel_unified_attention_diffkv( + # Output destinations. In 2D mode we write the final result into + # ``output_ptr``; in 3D mode we write per-segment partials into + # ``segm_*`` and ``output_ptr`` is unused (callers may pass any + # non-null pointer). + output_ptr, + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + query_ptr, + key_cache_ptr, # view of packed cache: [..., :head_size_qk] + value_cache_ptr, # view of packed cache: [..., head_size_qk:hqk+hv] + sink_ptr, + block_tables_ptr, + seq_lens_ptr, + alibi_slopes_ptr, + scale, + softcap, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, # == HEAD_SIZE_QK + output_stride_0: tl.int64, + output_stride_1: tl.int64, # == HEAD_SIZE_V + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE_QK: tl.constexpr, + HEAD_SIZE_QK_PADDED: tl.constexpr, + HEAD_SIZE_V: tl.constexpr, + HEAD_SIZE_V_PADDED: tl.constexpr, + USE_ALIBI_SLOPES: tl.constexpr, + USE_ALIBI_SQRT: tl.constexpr, + USE_SOFTCAP: tl.constexpr, + USE_SINKS: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + # Strides for both cache views (they share the same packed buffer, so + # dims 0/1/2 strides match; only the per-head extent differs). + stride_k_cache_0: tl.int64, + stride_k_cache_1: tl.int64, + stride_k_cache_2: tl.int64, + stride_k_cache_3: tl.constexpr, + stride_v_cache_0: tl.int64, + stride_v_cache_1: tl.int64, + stride_v_cache_2: tl.int64, + stride_v_cache_3: tl.constexpr, + query_start_len_ptr, + BLOCK_Q: tl.constexpr, + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + # ``IS_3D`` toggles between 2D layout (one program walks the full KV + # sequence) and 3D layout (split-KV / FlashDecoding-style: per-segm + # programs write partials, finalized by ``kernel_reduce_segments_diffkv``). + IS_3D: tl.constexpr, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) if IS_3D else 0 + + ( + seq_idx, + q_block_local_idx, + cur_batch_in_all_start_index, + cur_batch_query_len, + seq_len, + ) = resolve_seq_and_query_len( + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q + ) + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + if IS_3D: + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + else: + tiles_per_segment = 0 + + offs_m = tl.arange(0, BLOCK_M) + offs_d_qk = tl.arange(0, HEAD_SIZE_QK_PADDED) + offs_d_v = tl.arange(0, HEAD_SIZE_V_PADDED) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d_qk[None, :] + ) + + dim_mask_qk = tl.where(offs_d_qk < HEAD_SIZE_QK, 1, 0).to(tl.int1) + dim_mask_v = tl.where(offs_d_v < HEAD_SIZE_V, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Q : (BLOCK_M, HEAD_SIZE_QK_PADDED) + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask_qk[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + M = init_softmax_M( + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D + ) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + # acc : (BLOCK_M, HEAD_SIZE_V_PADDED) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_V_PADDED], dtype=tl.float32) + + context_len = seq_len - cur_batch_query_len + + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds( + context_len, + seq_len, + cur_batch_query_len, + q_block_local_idx, + segm_idx, + tiles_per_segment, + TILE_SIZE, + BLOCK_M, + BLOCK_Q, + num_queries_per_kv, + SLIDING_WINDOW, + False, # USE_MM_PREFIX + IS_3D, + ) + + for j in range(loop_lo, loop_hi): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + v_offset = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + offs_d_v[None, :] * stride_v_cache_3 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1 + ) + k_offset = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + offs_d_qk[:, None] * stride_k_cache_3 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1 + ) + # K : (HEAD_SIZE_QK_PADDED, TILE_SIZE) + K_load = tl.load( + key_cache_ptr + k_offset, + mask=dim_mask_qk[:, None] & tile_mask[None, :], + other=0.0, + ) + K = K_load.to(Q.dtype) + # V : (TILE_SIZE, HEAD_SIZE_V_PADDED) + V_load = tl.load( + value_cache_ptr + v_offset, + mask=dim_mask_v[None, :] & tile_mask[:, None], + other=0.0, + ) + V = V_load.to(Q.dtype) + + query_abs_pos = context_len + query_pos[:, None] + seq_mask = compute_kv_seq_mask( + query_abs_pos, + seq_offset, + seq_idx, + seq_len, + None, # mm_prefix_range_ptr + SLIDING_WINDOW, + False, # USE_MM_PREFIX + 0, # MAX_MM_RANGES + ) + + # S : (BLOCK_M, TILE_SIZE) + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + S += scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + S = apply_alibi_to_score( + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT + ) + + M, L, P, alpha = softmax_step(S, M, L) + acc = acc * alpha[:, None] + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, + V, + 0.0, + ) + acc += tl.dot(P.to(V.dtype), V) + + # ---- Epilogue -------------------------------------------------------- + if IS_3D: + # Store per-segment partials; finalized by reduce_segments_diffkv. + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + segm_idx * HEAD_SIZE_V_PADDED + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset, + acc, + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + store_segm_reduce_scalars( + segm_max_ptr, + segm_expsum_ptr, + query_offset_0, + query_offset_1, + segm_idx, + M, + L, + query_mask_0, + query_mask_1, + num_query_heads, + NUM_SEGMENTS_PER_SEQ, + ) + else: + acc = acc / L[:, None] + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d_v[None, :] + ) + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +@triton.jit +def kernel_reduce_segments_diffkv( + output_ptr, # [num_tokens, num_query_heads, head_size_v] + segm_output_ptr, + # [num_tokens, num_query_heads, max_num_segments, head_size_v] + segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments] + segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments] + seq_lens_ptr, # [num_seqs] + num_seqs, + num_query_heads: tl.constexpr, + output_stride_0: tl.int64, + output_stride_1: tl.int64, # == HEAD_SIZE_V + TILE_SIZE: tl.constexpr, + HEAD_SIZE_V: tl.constexpr, + HEAD_SIZE_V_PADDED: tl.constexpr, + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, +): + """Combine per-segment partials into the final softmax output. + + Mirrors ``reduce_segments`` from triton_unified_attention.py but + indexes V's head size (``HEAD_SIZE_V``) instead of the shared one. + """ + query_token_idx = tl.program_id(0) + query_head_idx = tl.program_id(1) + + seq_idx = find_seq_idx( + query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False + ) + seq_len = tl.load(seq_lens_ptr + seq_idx) + + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE) + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + ) + dim_mask = tl.where(tl.arange(0, HEAD_SIZE_V_PADDED) < HEAD_SIZE_V, 1, 0).to( + tl.int1 + ) + + segm_offset = ( + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_head_idx * NUM_SEGMENTS_PER_SEQ + + tl.arange(0, NUM_SEGMENTS_PER_SEQ) + ) + segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) + overall_max = tl.max(segm_max) + + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) + overall_expsum = tl.sum(segm_expsum) + + segm_output_offset = ( + query_token_idx.to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_V_PADDED + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :] + ) + segm_output = tl.load( + segm_output_ptr + segm_output_offset, + mask=segm_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + segm_output *= tl.exp(segm_max - overall_max)[:, None] + acc_sum = tl.sum(segm_output, axis=0) + acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) + + output_offset = ( + query_token_idx * output_stride_0 + + query_head_idx * output_stride_1 + + tl.arange(0, HEAD_SIZE_V_PADDED) + ) + tl.store(output_ptr + output_offset, acc, mask=dim_mask) + + +def unified_attention_diffkv( + q, # [num_tokens, num_query_heads, head_size_qk] + k, # view: [num_blocks, block_size, num_kv_heads, head_size_qk] + v, # view: [num_blocks, block_size, num_kv_heads, head_size_v] + out, # [num_tokens, num_query_heads, head_size_v] + cu_seqlens_q, + seqused_k, + softmax_scale, + causal, + window_size, + block_table, + softcap, + max_seqlen_q: int = 1, + alibi_slopes=None, + sinks=None, + use_alibi_sqrt=False, + # 3D / split-KV softmax buffers. When all four are provided and the + # batch is decode-only with few sequences, the 3D path is taken. + seq_threshold_3D: int | None = None, + num_par_softmax_segments: int | None = None, + softmax_segm_output: torch.Tensor | None = None, + softmax_segm_max: torch.Tensor | None = None, + softmax_segm_expsum: torch.Tensor | None = None, +): + assert causal, "Only causal attention is supported" + + if sinks is not None: + assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size" + + use_alibi_slopes = alibi_slopes is not None + + block_size = v.shape[1] + num_seqs = len(seqused_k) + num_query_heads = q.shape[1] + num_kv_heads = k.shape[2] + num_queries_per_kv = num_query_heads // num_kv_heads + head_size_qk = q.shape[2] + head_size_v = v.shape[3] + + BLOCK_M = ( + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) + ) + BLOCK_Q = BLOCK_M // num_queries_per_kv + + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0 + + # Decide between 2D and 3D launch. Mirrors the standard launcher: + # 3D requires preallocated softmax buffers, decode-only batches, and + # a small number of sequences (otherwise 2D already saturates the SM). + use_3d = not ( + seq_threshold_3D is None + or num_par_softmax_segments is None + or softmax_segm_output is None + or softmax_segm_max is None + or softmax_segm_expsum is None + or max_seqlen_q > 1 + or num_seqs > seq_threshold_3D + or is_batch_invariant + ) + + # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers + # smaller tiles to expose more parallelism along the KV dim. + tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32) + + grid: tuple[Any, ...] + if use_3d: + grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) + segm_output_ptr = softmax_segm_output + segm_max_ptr = softmax_segm_max + segm_expsum_ptr = softmax_segm_expsum + num_segments = num_par_softmax_segments + else: + grid = (total_num_q_blocks, num_kv_heads) + # 2D never touches the segm tensors but Triton wants a non-null + # pointer; reuse ``out``. + segm_output_ptr = out + segm_max_ptr = out + segm_expsum_ptr = out + num_segments = 1 + + kernel_unified_attention_diffkv[grid]( + output_ptr=out, + segm_output_ptr=segm_output_ptr, + segm_max_ptr=segm_max_ptr, + segm_expsum_ptr=segm_expsum_ptr, + query_ptr=q, + key_cache_ptr=k, + value_cache_ptr=v, + sink_ptr=sinks, + block_tables_ptr=block_table, + seq_lens_ptr=seqused_k, + alibi_slopes_ptr=alibi_slopes, + scale=softmax_scale, + softcap=softcap, + num_query_heads=num_query_heads, + num_queries_per_kv=num_queries_per_kv, + block_table_stride=block_table.stride(0), + query_stride_0=q.stride(0), + query_stride_1=q.stride(1), + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE_QK=head_size_qk, + HEAD_SIZE_QK_PADDED=triton.next_power_of_2(head_size_qk), + HEAD_SIZE_V=head_size_v, + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v), + USE_ALIBI_SLOPES=use_alibi_slopes, + USE_ALIBI_SQRT=use_alibi_sqrt, + USE_SOFTCAP=(softcap > 0), + USE_SINKS=(sinks is not None), + SLIDING_WINDOW=sliding_window_val, + stride_k_cache_0=k.stride(0), + stride_k_cache_1=k.stride(1), + stride_k_cache_2=k.stride(2), + stride_k_cache_3=k.stride(3), + stride_v_cache_0=v.stride(0), + stride_v_cache_1=v.stride(1), + stride_v_cache_2=v.stride(2), + stride_v_cache_3=v.stride(3), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + num_seqs=num_seqs, + BLOCK_M=BLOCK_M, + NUM_SEGMENTS_PER_SEQ=num_segments, + IS_3D=use_3d, + ) + + if use_3d: + kernel_reduce_segments_diffkv[(q.shape[0], num_query_heads)]( + output_ptr=out, + segm_output_ptr=softmax_segm_output, + segm_max_ptr=softmax_segm_max, + segm_expsum_ptr=softmax_segm_expsum, + seq_lens_ptr=seqused_k, + num_seqs=num_seqs, + num_query_heads=num_query_heads, + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + TILE_SIZE=tile_size, + HEAD_SIZE_V=head_size_v, + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments, + ) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9f0bfc5880c..9af54e0a249 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -17,7 +17,7 @@ from vllm.v1.kv_cache_interface import ( get_kv_cache_spec_sliding_window, ) from vllm.v1.metrics.stats import PrefixCacheStats -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @@ -122,6 +122,7 @@ class KVCacheManager: dcp_world_size: int = 1, pcp_world_size: int = 1, metrics_collector: KVCacheMetricsCollector | None = None, + watermark: float = 0.0, ) -> None: self.max_model_len = max_model_len # When unset, fall back to `max_model_len` so the recycling-aware cap @@ -155,6 +156,11 @@ class KVCacheManager: self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.kv_cache_config = kv_cache_config + + # Watermark: minimum number of KV cache blocks to keep free when + # admitting waiting/preempted requests, to avoid frequent preemptions. + assert watermark >= 0.0, "watermark must be non-negative" + self.watermark_blocks = int(watermark * kv_cache_config.num_blocks) self.kv_cache_event_metadata = tuple( ( get_kv_cache_spec_kind(group.kv_cache_spec).value, @@ -247,6 +253,7 @@ class KVCacheManager: num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. @@ -277,6 +284,8 @@ class KVCacheManager: made if it fits within (free blocks - reserved_blocks). Used to gate async KV-connector loads so their initial allocation cannot consume blocks an already in-flight (prefilling) sequence is relying on. + has_scheduled_reqs: Whether any requests are already scheduled to run + this step, controls whether watermark is applied. Blocks layout: ``` @@ -351,6 +360,15 @@ class KVCacheManager: self.max_model_len, ) + watermark_blocks = 0 + # The watermark is applied to waiting/preempted requests only, and only + # when there's at least one request already scheduled. + if has_scheduled_reqs and request.status in ( + RequestStatus.WAITING, + RequestStatus.PREEMPTED, + ): + watermark_blocks = self.watermark_blocks + if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. full_num_tokens = min(request.num_tokens, self.max_model_len) @@ -364,7 +382,8 @@ class KVCacheManager: num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) - if num_blocks_to_allocate > self.block_pool.get_num_free_blocks(): + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > self.block_pool.get_num_free_blocks(): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -392,8 +411,11 @@ class KVCacheManager: num_tokens_main_model=num_tokens_main_model, ) + # Keep `reserved_blocks` free for other in-flight sequences, and an + # additional watermark of headroom for waiting/preempted admissions. available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks - if num_blocks_to_allocate > available_blocks: + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > available_blocks: # Cannot allocate new blocks return None diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 107a89cc6b6..72ca6a2fa67 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1717,36 +1717,17 @@ def generate_scheduler_kv_cache_config( return cfg -def _report_kv_cache_config( +def get_kv_cache_capacity( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig -) -> None: +) -> tuple[int, float]: """ - Log resolved KV cache configuration. - - Args: - vllm_config: The global VllmConfig - kv_cache_config: The resolved KV cache configuration + Get the group-aware KV cache token capacity and max concurrency. """ max_model_len = vllm_config.model_config.max_model_len max_concurrency = get_max_concurrency_for_kv_cache_config( vllm_config, kv_cache_config ) - - # GPU KV cache size in tokens = max_concurrency * max_model_len: the total - # tokens of context the pool can hold at peak utilization. Sourcing this - # from the concurrency calculation handles hybrid layouts correctly: SWA / - # chunked-local groups have a per-request block count that's capped by - # their window, so a naive `num_blocks // num_groups * block_size` formula - # underestimates capacity for these models. DCP/PCP sharding is already - # accounted for in each spec's `max_memory_usage_bytes`. - num_tokens = int(max_concurrency * max_model_len) - - logger.info_once("GPU KV cache size: %s tokens", f"{num_tokens:,}") - logger.info_once( - "Maximum concurrency for %s tokens per request: %.2fx", - f"{max_model_len:,}", - max_concurrency, - ) + return int(max_concurrency * max_model_len), max_concurrency def _max_memory_usage_bytes_from_groups( @@ -2085,7 +2066,21 @@ def get_kv_cache_configs( tensor.size = tensor.size // num_blocks_old * min_num_blocks if len(kv_cache_config.kv_cache_groups) > 0: - _report_kv_cache_config(vllm_config, kv_cache_config) + max_model_len = vllm_config.model_config.max_model_len + # GPU KV cache size in tokens = max_concurrency * max_model_len: + # the total tokens of context the pool can hold at peak + # utilization. Sourcing this from the concurrency calculation + # handles hybrid layouts correctly. + num_tokens, max_concurrency = get_kv_cache_capacity( + vllm_config, kv_cache_config + ) + + logger.info_once("GPU KV cache size: %s tokens", f"{num_tokens:,}") + logger.info_once( + "Maximum concurrency for %s tokens per request: %.2fx", + f"{max_model_len:,}", + max_concurrency, + ) return kv_cache_configs diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 2fd22f4c0cb..a79e84289af 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -27,10 +27,14 @@ class AsyncScheduler(Scheduler): scheduler_output.pending_structured_output_tokens |= ( request.use_structured_output and request.num_output_placeholders > 0 ) - # The request will generate a new token plus num_spec_tokens - # in this scheduling step. + # The request will generate num_sampled_tokens_per_step new tokens + # plus num_spec_tokens in this scheduling step. Diffusion has no AR + # bonus token (num_sampled_tokens_per_step == 0) — only the canvas + # (spec) tokens. cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) - request.num_output_placeholders += 1 + cur_num_spec_tokens + request.num_output_placeholders += ( + self.num_sampled_tokens_per_step + cur_num_spec_tokens + ) # Add placeholders for the new draft/spec tokens. # We will update the actual spec token ids in the worker process. request.spec_token_ids = self._spec_token_placeholders diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 160cdb74f57..e215c698c4e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -113,6 +113,10 @@ class Scheduler(SchedulerInterface): self.kv_events_config is not None and self.kv_events_config.enable_kv_cache_events ) + # Diffusion models may not sample any tokens for a denoising step. + self.num_sampled_tokens_per_step = ( + 1 if not vllm_config.model_config.is_diffusion else 0 + ) # Create KVConnector for the Scheduler. Note that each Worker # will have a corresponding KVConnector with Role=WORKER. @@ -212,9 +216,9 @@ class Scheduler(SchedulerInterface): speculative_config = vllm_config.speculative_config self.use_eagle = False - self.num_spec_tokens = self.num_lookahead_tokens = 0 - if speculative_config: - self.num_spec_tokens = speculative_config.num_speculative_tokens + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_lookahead_tokens = 0 + if speculative_config is not None: if speculative_config.use_eagle(): self.use_eagle = True self.num_lookahead_tokens = self.num_spec_tokens @@ -242,6 +246,7 @@ class Scheduler(SchedulerInterface): scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, + watermark=self.scheduler_config.watermark, ) # Bind GPU block pool to the KV connector. This must happen after # kv_cache_manager is constructed so block_pool is available. @@ -424,7 +429,10 @@ class Scheduler(SchedulerInterface): # Make sure the input position does not exceed the max model len. # This is necessary when using spec decoding. num_new_tokens = min( - num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens + num_new_tokens, + self.max_model_len + - request.num_computed_tokens + - self.num_sampled_tokens_per_step, ) # Schedule encoder inputs. @@ -826,6 +834,7 @@ class Scheduler(SchedulerInterface): num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), ) if new_blocks is None: @@ -1403,13 +1412,6 @@ class Scheduler(SchedulerInterface): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and kv_connector_output.invalid_block_ids: @@ -1471,9 +1473,12 @@ class Scheduler(SchedulerInterface): scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) - if scheduled_spec_token_ids and generated_token_ids: + if scheduled_spec_token_ids and ( + generated_token_ids or self.num_sampled_tokens_per_step == 0 + ): num_draft_tokens = len(scheduled_spec_token_ids) - num_accepted = len(generated_token_ids) - 1 + num_sampled = self.num_sampled_tokens_per_step + num_accepted = max(len(generated_token_ids) - num_sampled, 0) num_rejected = num_draft_tokens - num_accepted # num_computed_tokens represents the number of tokens # processed in the current step, considering scheduled @@ -1653,6 +1658,23 @@ class Scheduler(SchedulerInterface): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if ( + scheduler_kv_connector_stats is not None + and not scheduler_kv_connector_stats.is_empty() + ): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # collect KV cache events from KV cache manager events = self.kv_cache_manager.take_events() @@ -1997,6 +2019,19 @@ class Scheduler(SchedulerInterface): ) return len(self.requests) > num_in_queues + def has_requests(self) -> bool: + # Override the interface default to also keep the engine alive while a + # connector still has pending push work (e.g. push-mode WRITE transfers + # in flight after all "live" requests have finished). Without this hook + # the engine would quiesce before the connector can drain completions. + # TODO: replace with a more general mechanism for connectors to keep + # the scheduler alive. + return ( + self.has_unfinished_requests() + or self.has_finished_requests() + or (self.connector is not None and self.connector.has_pending_push_work()) + ) + def reset_prefix_cache( self, reset_running_requests: bool = False, reset_connector: bool = False ) -> bool: @@ -2198,12 +2233,8 @@ class Scheduler(SchedulerInterface): ) def _inflight_prefill_reserved_blocks(self) -> int: - """Blocks in-flight prefills still need to finish (their reservation). + """Num blocks in-flight prefills still need to finish (their reservation).""" - Sums remaining full-ISL blocks over `self._inflight_prefills` (running - prefills + in-progress async loads). The candidate async load isn't yet - in the set, so it's naturally excluded. - """ return sum( self._request_remaining_blocks(req) for req in self._inflight_prefills ) diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index cf0c1d41772..6a48b6282d4 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -34,11 +34,7 @@ class CudagraphDispatcher: def __init__(self, vllm_config: VllmConfig): self.vllm_config = vllm_config self.compilation_config = vllm_config.compilation_config - self.uniform_decode_query_len = ( - 1 - if not self.vllm_config.speculative_config - else 1 + self.vllm_config.speculative_config.num_speculative_tokens - ) + self.uniform_decode_query_len = 1 + self.vllm_config.num_speculative_tokens # Dict to store valid cudagraph dispatching keys. self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = { diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 848f530ce33..fbfe1c144cc 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -78,6 +78,9 @@ class EngineCoreReadyResponse: dp_stats_address: str | None dtype: str vllm_version: str + # KV cache capacity (None for encoder-only/attention-free models). + kv_cache_size_tokens: int | None = None + kv_cache_max_concurrency: float | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 08c814ab34e..bf89f3e9d5c 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -45,6 +45,7 @@ from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, + get_kv_cache_capacity, get_kv_cache_configs, get_request_block_hasher, init_none_hash, @@ -156,6 +157,9 @@ class EngineCore: hash_block_size=hash_block_size, ) self.use_spec_decode = vllm_config.speculative_config is not None + self.check_for_draft_tokens = ( + self.use_spec_decode or vllm_config.model_config.is_diffusion + ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore @@ -283,6 +287,11 @@ class EngineCore: vllm_config.cache_config.block_size = min( g.kv_cache_spec.block_size for g in kv_cache_groups ) + num_tokens, max_concurrency = get_kv_cache_capacity( + vllm_config, scheduler_kv_cache_config + ) + vllm_config.cache_config.kv_cache_size_tokens = num_tokens + vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency vllm_config.validate_block_size() @@ -475,8 +484,7 @@ class EngineCore: # When using async scheduling we can't get draft token ids in advance, # so we update draft token ids in the worker process and don't # need to update draft token ids here. - if not self.async_scheduling and self.use_spec_decode and model_executed: - # Take the draft token ids. + if self.check_for_draft_tokens and not self.async_scheduling and model_executed: draft_token_ids = self.model_executor.take_draft_token_ids() if draft_token_ids is not None: self.scheduler.update_draft_token_ids(draft_token_ids) @@ -575,18 +583,17 @@ class EngineCore: # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. if deferred_scheduler_output: - # If we are doing speculative decoding with structured output, - # we need to get the draft token ids from the prior step before - # we can compute the grammar bitmask for the deferred request. - if self.use_spec_decode: + # When draft tokens are used with structured output, validate them + # before computing the grammar bitmask for the deferred request. + if self.check_for_draft_tokens: draft_token_ids = self.model_executor.take_draft_token_ids() - assert draft_token_ids is not None - # Update the draft token ids in the scheduler output to - # filter out the invalid spec tokens, which will be padded - # with -1 and skipped by the grammar bitmask computation. - self.scheduler.update_draft_token_ids_in_output( - draft_token_ids, deferred_scheduler_output - ) + if draft_token_ids is not None: + # Update the draft token ids in the scheduler output to + # filter out the invalid spec tokens, which will be padded + # with -1 and skipped by the grammar bitmask computation. + self.scheduler.update_draft_token_ids_in_output( + draft_token_ids, deferred_scheduler_output + ) # We now have the tokens needed to compute the bitmask for the # deferred request. Get the bitmask and call sample tokens. grammar_output = self.scheduler.get_grammar_bitmask( @@ -1493,6 +1500,12 @@ class EngineCoreProc(EngineCore): dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, + kv_cache_size_tokens=( + self.vllm_config.cache_config.kv_cache_size_tokens + ), + kv_cache_max_concurrency=( + self.vllm_config.cache_config.kv_cache_max_concurrency + ), ) ready_payload = msgspec.msgpack.encode(ready_response) for input_socket in input_sockets: diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 32f2d091eb3..195cfeecf42 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -720,14 +720,26 @@ class MPClient(EngineCoreClient): ) # Setup KV cache config with initialization state from - # engine core process. Sum values from all engines in DP case. + # engine core process. Sum num_gpu_blocks from all engines in DP case. num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks or 0 num_gpu_blocks += response.num_gpu_blocks vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks # Sync block_size: may be enlarged by _align_hybrid_block_size in the # worker for hybrid Mamba models. - vllm_config.cache_config.block_size = response.block_size + cache_config = vllm_config.cache_config + cache_config.block_size = response.block_size + # Keep these as per-engine cache_config_info values; do not sum across DP. + cache_config.kv_cache_size_tokens = ( + getattr(cache_config, "kv_cache_size_tokens", None) + if getattr(cache_config, "kv_cache_size_tokens", None) is not None + else response.kv_cache_size_tokens + ) + cache_config.kv_cache_max_concurrency = ( + getattr(cache_config, "kv_cache_max_concurrency", None) + if getattr(cache_config, "kv_cache_max_concurrency", None) is not None + else response.kv_cache_max_concurrency + ) # In external DP LB mode, the coordinator address that the # front-end procs connect to is obtained by each engine via it's diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 0052a35366a..021019dc1cd 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -110,7 +110,9 @@ class LoggingStatLogger(StatLoggerBase): self.connector_prefix_caching_metrics = CachingMetrics() self.mm_caching_metrics = CachingMetrics() - self.spec_decoding_logging = SpecDecodingLogging() + model_config = self.vllm_config.model_config + is_diffusion = model_config is not None and model_config.is_diffusion + self.spec_decoding_logging = SpecDecodingLogging(is_diffusion=is_diffusion) kv_transfer_config = self.vllm_config.kv_transfer_config self.kv_connector_logging = KVConnectorLogging(kv_transfer_config) self.cudagraph_logging = None @@ -436,7 +438,10 @@ class PrometheusStatLogger(AggregateStatLoggerBase): per_engine_labelvalues = self.per_engine_labelvalues self.spec_decoding_prom = self._spec_decoding_cls( - vllm_config.speculative_config, labelnames, per_engine_labelvalues + vllm_config.speculative_config, + labelnames, + per_engine_labelvalues, + is_diffusion=vllm_config.model_config.is_diffusion, ) self.kv_connector_prom = self._kv_connector_cls( vllm_config, labelnames, per_engine_labelvalues diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 38135b9b158..3336fca606a 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -396,6 +396,20 @@ class AttentionQuantizationConfigParser(Parser): return args +class AttentionDetectionParser(Parser): + """ + Prevents standard AttentionMetrics from being instantiated for MLA models. + MLA models should use MLAAttentionMetrics instead. + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + if vllm_config.model_config.is_deepseek_mla: + raise InvalidComponent( + "Model uses MLA attention; use MLAAttentionMetrics instead" + ) + return args + + class AttentionMetrics(ComponentMetrics): # From BaseConfigParser num_hidden_layers: int = Field(..., gt=0) @@ -423,6 +437,7 @@ class AttentionMetrics(ComponentMetrics): @classmethod def get_parser(cls) -> ParserChain: return ParserChain( + AttentionDetectionParser(), BaseConfigParser(), BaseAttentionConfigParser(), AttentionQuantizationConfigParser(), @@ -525,6 +540,276 @@ class AttentionMetrics(ComponentMetrics): } +#### MLA Attention #### + + +class MLADetectionParser(Parser): + """ + Validates that the model uses MLA attention. + Raises InvalidComponent if the model does not use MLA, + so MLAAttentionMetrics is silently skipped for non-MLA models. + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + if not vllm_config.model_config.is_deepseek_mla: + raise InvalidComponent("Model does not use MLA attention") + return args + + +class MLAConfigParser(Parser): + """ + Parses MLA-specific configuration fields. + Provides: kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, + v_head_dim, q_lora_rank + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + model_config = vllm_config.model_config + cfg = model_config.hf_text_config + + args.kv_lora_rank = get_required(cfg, "kv_lora_rank") + args.qk_nope_head_dim = get_required(cfg, "qk_nope_head_dim") + args.qk_rope_head_dim = get_required(cfg, "qk_rope_head_dim") + args.v_head_dim = get_required(cfg, "v_head_dim") + args.q_lora_rank = getattr(cfg, "q_lora_rank", None) + + model_dtype = vllm_config.model_config.dtype + cache_dtype = vllm_config.cache_config.cache_dtype + kv_cache_torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype) + args.cache_byte_size = get_dtype_size(kv_cache_torch_dtype) + + return args + + +class MLAAttentionMetrics(ComponentMetrics): + """ + Performance metrics for Multi-Latent Attention (MLA) layers. + + MLA uses a compressed latent representation for KV cache: + - KV cache stores a single compressed vector of size + (kv_lora_rank + qk_rope_head_dim) per token per layer, + instead of 2 * num_kv_heads * head_dim as in standard MHA/GQA. + - Q path uses optional low-rank compression: + h -> q_lora_rank -> num_heads * qk_head_dim + - KV path: h -> (kv_lora_rank + qk_rope_head_dim), + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim) + + Used by DeepSeek-V2, DeepSeek-V3, DeepSeek-R1, and similar models. + """ + + # From BaseConfigParser + num_hidden_layers: int = Field(..., gt=0) + hidden_size: int = Field(..., gt=0) + num_attention_heads: int = Field(..., gt=0) + activation_byte_size: int = Field(..., gt=0) + tp_size: int = Field(..., gt=0) + pp_size: int = Field(..., gt=0) + + # From BaseConfigParser, can be overridden by AttentionQuantizationConfigParser + weight_byte_size: int | float = Field(..., gt=0) + + # From MLAConfigParser + kv_lora_rank: int = Field(..., gt=0) + qk_nope_head_dim: int = Field(..., gt=0) + qk_rope_head_dim: int = Field(..., gt=0) + v_head_dim: int = Field(..., gt=0) + q_lora_rank: int | None = Field(None) + cache_byte_size: int = Field(..., gt=0) + + @classmethod + def component_type(cls) -> str: + return "mla_attn" + + @classmethod + def get_parser(cls) -> ParserChain: + return ParserChain( + MLADetectionParser(), + BaseConfigParser(), + MLAConfigParser(), + AttentionQuantizationConfigParser(), + ) + + def get_num_flops_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate flops breakdown for MLA attention layers. + + MLA projection structure: + - Q path: h -> q_lora_rank -> num_heads * qk_head_dim + (or h -> num_heads * qk_head_dim if q_lora_rank is None) + - KV path: h -> (kv_lora_rank + qk_rope_head_dim), + then kv_lora_rank -> num_heads * (qk_nope_head_dim + v_head_dim) + - Attention: Q @ K^T and attn @ V + - Output: num_heads * v_head_dim -> h + """ + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + TC = ctx.total_token_context_product() + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + flops: dict[str, int] = {} + + # Q projection + if q_rank is not None: + # Two-stage: h -> q_lora_rank -> num_heads * qk_head_dim + flops["q_a_proj"] = 2 * T * D * q_rank * L + flops["q_b_proj"] = 2 * T * q_rank * q * qk_head_dim * L + else: + # Direct: h -> num_heads * qk_head_dim + flops["q_proj"] = 2 * T * D * q * qk_head_dim * L + + # KV projection (always compressed, shared across heads) + # kv_a: h -> (kv_lora_rank + qk_rope_head_dim) [replicated] + flops["kv_a_proj"] = 2 * T * D * (c + r) * L + # kv_b: kv_lora_rank -> num_heads * (qk_nope + v_head_dim) + flops["kv_b_proj"] = 2 * T * c * q * (self.qk_nope_head_dim + v_d) * L + + # Attention core + flops["attn_qk"] = 2 * q * TC * qk_head_dim * L + flops["attn_av"] = 2 * q * TC * v_d * L + + # Output projection: num_heads * v_head_dim -> h + flops["out_proj"] = 2 * T * q * v_d * D * L + + return flops + + def get_read_bytes_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate read memory traffic for MLA attention layers.""" + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + # Compressed KV cache size per token + kv_compressed_dim = c + r + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + read_bytes: dict[str, int] = {} + + # Q projection weight + input reads + if q_rank is not None: + read_bytes["q_a_input"] = T * D * self.activation_byte_size * L + read_bytes["q_a_weight"] = int(D * q_rank * self.weight_byte_size * L) + read_bytes["q_b_input"] = T * q_rank * self.activation_byte_size * L + read_bytes["q_b_weight"] = int( + q_rank * q * qk_head_dim * self.weight_byte_size * L + ) + else: + read_bytes["q_input"] = T * D * self.activation_byte_size * L + read_bytes["q_weight"] = int( + D * q * qk_head_dim * self.weight_byte_size * L + ) + + # KV projection weight + input reads + # kv_a is replicated (not TP-sharded) + read_bytes["kv_a_input"] = T * D * self.activation_byte_size * L + read_bytes["kv_a_weight"] = int( + D * kv_compressed_dim * self.weight_byte_size * L + ) + # kv_b is TP-sharded along heads + read_bytes["kv_b_input"] = T * c * self.activation_byte_size * L + read_bytes["kv_b_weight"] = int( + c * q * (self.qk_nope_head_dim + v_d) * self.weight_byte_size * L + ) + + # Attention input reads + # Prefill: read Q activations + K,V from kv_b_proj output + if ctx.prefill_num_tokens > 0: + read_bytes["attn_input"] = ( + ctx.prefill_num_tokens * q * qk_head_dim * self.activation_byte_size * L + + ctx.prefill_context_len + * q + * (qk_head_dim + v_d) + * self.activation_byte_size + * L + ) + + # Decode: read Q activations + read compressed KV from cache + if ctx.decode_num_tokens > 0: + read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + ( + ctx.decode_num_tokens * q * qk_head_dim * self.activation_byte_size * L + + ctx.decode_context_len * kv_compressed_dim * self.cache_byte_size * L + ) + + # Output projection reads + read_bytes["out_input"] = T * q * v_d * self.activation_byte_size * L + read_bytes["out_weight"] = int(q * v_d * D * self.weight_byte_size * L) + + return read_bytes + + def get_write_bytes_breakdown( + self, ctx: ExecutionContext, per_gpu: bool = True + ) -> dict[str, int]: + """Calculate write memory traffic for MLA attention layers.""" + L = self.num_hidden_layers + D = self.hidden_size + q = self.num_attention_heads + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + v_d = self.v_head_dim + c = self.kv_lora_rank + r = self.qk_rope_head_dim + q_rank = self.q_lora_rank + + T = ctx.total_num_tokens() + kv_compressed_dim = c + r + + if per_gpu: + L //= self.pp_size + q = max(1, q // self.tp_size) + + write_bytes: dict[str, int] = {} + + # Q projection outputs + if q_rank is not None: + write_bytes["q_a_output"] = T * q_rank * self.activation_byte_size * L + write_bytes["q_b_output"] = ( + T * q * qk_head_dim * self.activation_byte_size * L + ) + else: + write_bytes["q_output"] = ( + T * q * qk_head_dim * self.activation_byte_size * L + ) + + # KV projection outputs + write_bytes["kv_a_output"] = ( + T * kv_compressed_dim * self.activation_byte_size * L + ) + write_bytes["kv_b_output"] = ( + T * q * (self.qk_nope_head_dim + v_d) * self.activation_byte_size * L + ) + + # KV cache write: one compressed vector per token + # (kv_lora_rank + qk_rope_head_dim) instead of + # 2 * num_kv_heads * head_dim in standard MHA + write_bytes["kv_cache"] = T * kv_compressed_dim * self.cache_byte_size * L + + # Output projection + write_bytes["out_output"] = T * D * self.activation_byte_size * L + + return write_bytes + + #### Ffn #### diff --git a/vllm/v1/metrics/prometheus.py b/vllm/v1/metrics/prometheus.py index 1eacb785aa8..c8740276713 100644 --- a/vllm/v1/metrics/prometheus.py +++ b/vllm/v1/metrics/prometheus.py @@ -64,7 +64,7 @@ def unregister_vllm_metrics(): registry = REGISTRY # Unregister any existing vLLM collectors for collector in list(registry._collector_to_names): - if hasattr(collector, "_name") and "vllm" in collector._name: + if hasattr(collector, "_name") and collector._name.startswith("vllm:"): registry.unregister(collector) diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py index 9a41ff5c818..5da41510b4d 100644 --- a/vllm/v1/spec_decode/metrics.py +++ b/vllm/v1/spec_decode/metrics.py @@ -53,7 +53,11 @@ class SpecDecodingLogging: before resetting to zero. """ - def __init__(self): + def __init__(self, is_diffusion: bool = False): + # Diffusion (dLLM) models reuse the spec-decode data path with + # overloaded semantics, so the raw spec-decode framing (drafts, bonus + # token, per-position vector) is logged with diffusion-native terms. + self.is_diffusion = is_diffusion self.reset() def reset(self): @@ -85,6 +89,17 @@ class SpecDecodingLogging: draft_throughput = num_draft_tokens / elapsed_time accepted_throughput = num_accepted_tokens / elapsed_time + if self.is_diffusion: + self._log_diffusion( + log_fn, + num_denoising_steps=num_drafts, + num_canvas_tokens=num_draft_tokens, + num_committed_tokens=num_accepted_tokens, + committed_throughput=accepted_throughput, + ) + self.reset() + return + draft_acceptance_rate = ( num_accepted_tokens / num_draft_tokens * 100 if num_draft_tokens > 0 @@ -117,6 +132,43 @@ class SpecDecodingLogging: ) self.reset() + def _log_diffusion( + self, + log_fn, + num_denoising_steps: int, + num_canvas_tokens: int, + num_committed_tokens: int, + committed_throughput: float, + ): + # Each "draft" is one denoising step that re-evaluates the canvas block + # and finalizes some of its positions. + mean_committed_per_step = ( + num_committed_tokens / num_denoising_steps + if num_denoising_steps > 0 + else float("nan") + ) + mean_steps_per_canvas = ( + num_canvas_tokens / num_committed_tokens + if num_committed_tokens > 0 + else float("nan") + ) + + log_fn( + "DiffusionDecoding metrics: " + "Committed token throughput: %.2f tokens/s, " + "Mean denoising steps per canvas: %.2f, " + "Mean tokens committed per denoising step: %.2f, " + "Committed: %d tokens, " + "Denoising steps: %d, " + "Canvas positions evaluated: %d", + committed_throughput, + mean_steps_per_canvas, + mean_committed_per_step, + num_committed_tokens, + num_denoising_steps, + num_canvas_tokens, + ) + class SpecDecodingProm: """Record spec decoding metrics in Prometheus. @@ -146,56 +198,66 @@ class SpecDecodingProm: speculative_config: SpeculativeConfig | None, labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], + is_diffusion: bool = False, ): - self.spec_decoding_enabled = speculative_config is not None + # Diffusion (dLLM) models reuse the spec-decode counters but expose them + # under diffusion-native names; the per-position acceptance vector does + # not apply, so it is omitted. + self.is_diffusion = is_diffusion + self.spec_decoding_enabled = speculative_config is not None or is_diffusion if not self.spec_decoding_enabled: return - counter_drafts = self._counter_cls( - name="vllm:spec_decode_num_drafts", - documentation="Number of spec decoding drafts.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_drafts = create_metric_per_engine( - counter_drafts, per_engine_labelvalues - ) + if is_diffusion: + counter_specs = [ + ("vllm:diffusion_num_denoising_steps", "Number of denoising steps."), + ( + "vllm:diffusion_num_canvas_positions", + "Number of canvas positions evaluated.", + ), + ( + "vllm:diffusion_num_committed_tokens", + "Number of committed (finalized) tokens.", + ), + ] + else: + counter_specs = [ + ("vllm:spec_decode_num_drafts", "Number of spec decoding drafts."), + ("vllm:spec_decode_num_draft_tokens", "Number of draft tokens."), + ("vllm:spec_decode_num_accepted_tokens", "Number of accepted tokens."), + ] - counter_draft_tokens = self._counter_cls( - name="vllm:spec_decode_num_draft_tokens", - documentation="Number of draft tokens.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_draft_tokens = create_metric_per_engine( - counter_draft_tokens, per_engine_labelvalues - ) + counters = [ + create_metric_per_engine( + self._counter_cls(name=name, documentation=doc, labelnames=labelnames), + per_engine_labelvalues, + ) + for name, doc in counter_specs + ] + # num_drafts/num_draft_tokens/num_accepted_tokens map onto denoising + # steps/canvas positions/committed tokens in the diffusion path. + self.counter_spec_decode_num_drafts = counters[0] + self.counter_spec_decode_num_draft_tokens = counters[1] + self.counter_spec_decode_num_accepted_tokens = counters[2] - counter_accepted_tokens = self._counter_cls( - name="vllm:spec_decode_num_accepted_tokens", - documentation="Number of accepted tokens.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_accepted_tokens = create_metric_per_engine( - counter_accepted_tokens, per_engine_labelvalues - ) - - assert speculative_config is not None - num_spec_tokens = ( - speculative_config.num_speculative_tokens - if self.spec_decoding_enabled - else 0 - ) - pos_labelnames = labelnames + ["position"] - base_counter = self._counter_cls( - name="vllm:spec_decode_num_accepted_tokens_per_pos", - documentation="Accepted tokens per draft position.", - labelnames=pos_labelnames, - ) self.counter_spec_decode_num_accepted_tokens_per_pos: dict[ int, list[prometheus_client.Counter] - ] = { - idx: [base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens)] - for idx, lv in per_engine_labelvalues.items() - } + ] = {} + if not is_diffusion: + assert speculative_config is not None + num_spec_tokens = speculative_config.num_speculative_tokens + pos_labelnames = labelnames + ["position"] + base_counter = self._counter_cls( + name="vllm:spec_decode_num_accepted_tokens_per_pos", + documentation="Accepted tokens per draft position.", + labelnames=pos_labelnames, + ) + self.counter_spec_decode_num_accepted_tokens_per_pos = { + idx: [ + base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens) + ] + for idx, lv in per_engine_labelvalues.items() + } def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0): if not self.spec_decoding_enabled: @@ -210,6 +272,6 @@ class SpecDecodingProm: spec_decoding_stats.num_accepted_tokens ) for pos, counter in enumerate( - self.counter_spec_decode_num_accepted_tokens_per_pos[engine_idx] + self.counter_spec_decode_num_accepted_tokens_per_pos.get(engine_idx, []) ): counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos]) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 6a4fcbb629f..30921f3d74a 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -211,11 +211,8 @@ class StructuredOutputManager: if not structured_output_request_ids: return None - max_num_spec_tokens = 0 - if self.vllm_config.speculative_config is not None: - max_num_spec_tokens = ( - self.vllm_config.speculative_config.num_speculative_tokens - ) + # Covers both speculative decoding and diffusion LLMs (canvas_length). + max_num_spec_tokens = self.vllm_config.num_speculative_tokens if self._grammar_bitmask is None: assert self.backend is not None @@ -277,7 +274,13 @@ class StructuredOutputManager: state_advancements = 0 req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - for token in itertools.chain(req_tokens, (-1,)): + if self.vllm_config.model_config.is_diffusion and req_tokens: + # Diffusion LLMs don't sample a bonus token after the + # scheduled positions, so don't append the -1 placeholder. + token_iter: Iterable[int] = req_tokens + else: + token_iter = itertools.chain(req_tokens, (-1,)) + for token in token_iter: self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) if token == -1: # Stop advancing the grammar once we hit a padding token. diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index f905d09e45f..6b750fe7ebf 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -302,6 +302,7 @@ def _combine_sampled_and_draft_tokens_kernel( cu_num_logits_ptr, logits_indices_ptr, BLOCK_SIZE: tl.constexpr, + NUM_NEW_SAMPLED_TOKENS: tl.constexpr = 1, ): batch_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + batch_idx) @@ -310,7 +311,7 @@ def _combine_sampled_and_draft_tokens_kernel( cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx) cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1) num_logits = cu_num_logits_end - cu_num_logits_start - num_draft_tokens = num_logits - 1 + num_draft_tokens = num_logits - NUM_NEW_SAMPLED_TOKENS # Compute the logits indices. block = tl.arange(0, BLOCK_SIZE) @@ -328,9 +329,10 @@ def _combine_sampled_and_draft_tokens_kernel( # Handling prefill tokens. No sampled or draft tokens. return - # Write the last sampled token ID to input_ids. - last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx) - tl.store(input_ids_ptr + query_end - num_logits, last_token_id) + if NUM_NEW_SAMPLED_TOKENS > 0: + # Write the last sampled token ID to input_ids. + last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx) + tl.store(input_ids_ptr + query_end - num_logits, last_token_id) # Write the draft tokens (if any) to input_ids. if num_draft_tokens > 0: @@ -356,7 +358,11 @@ def combine_sampled_and_draft_tokens( draft_tokens: torch.Tensor, cu_num_logits: torch.Tensor, num_logits: int, + num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens ) -> torch.Tensor: + assert num_new_sampled_tokens in (0, 1), ( + f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" + ) # use idx_mapping.shape[0] for actual request count num_reqs = idx_mapping.shape[0] num_speculative_steps = draft_tokens.shape[-1] @@ -377,9 +383,12 @@ def combine_sampled_and_draft_tokens( draft_tokens.stride(0), cu_num_logits, logits_indices, - # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token - # in addition to all draft tokens. - BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1), + NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens, + # NOTE(woosuk): Add num_new_sampled_tokens to ensure the block covers the + # last sampled token in addition to all draft tokens. + BLOCK_SIZE=triton.next_power_of_2( + num_speculative_steps + num_new_sampled_tokens + ), ) return logits_indices diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7cd1e6c5c86..328b521bfc8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,7 +78,6 @@ from vllm.v1.worker.gpu.input_batch import ( InputBuffers, combine_sampled_and_draft_tokens, expand_idx_mapping, - get_num_sampled_and_rejected, post_update, post_update_num_computed_tokens, prepare_pos_seq_lens, @@ -185,11 +184,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Speculative decoding. self.speculator = None - self.num_speculative_steps = 0 self.use_aux_hidden_state_outputs = False + self.num_speculative_steps = vllm_config.num_speculative_tokens if self.speculative_config is not None: - self.num_speculative_steps = self.speculative_config.num_speculative_tokens - if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) @@ -204,7 +201,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) - self.uniform_decode_query_len = 1 + self.num_speculative_steps # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" @@ -232,38 +228,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): device=self.device, ) + # Samplers and decode_query_len created in load_model() after + # model_state exists (num_new_sampled_tokens_per_step from ModelState). self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None self.prompt_logprobs_worker: PromptLogprobsWorker | None = None self.structured_outputs_worker: StructuredOutputsWorker | None = None - if self.is_last_pp_rank and not self.is_pooling_model: - # Initialize sampling-related workers. - # These components are only set up on the last PP rank and - # for generative (non-pooling) models. - self.sampler = Sampler( - max_num_reqs=self.max_num_reqs, - vocab_size=self.vocab_size, - device=self.device, - req_states=self.req_states, - logprobs_mode=self.model_config.logprobs_mode, - num_speculative_tokens=self.num_speculative_steps + 1, - use_fp64_gumbel=self.model_config.use_fp64_gumbel, - ) - if self.speculative_config is not None: - self.rejection_sampler = RejectionSampler( - self.sampler, - self.speculative_config, - self.device, - ) - self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs) - self.structured_outputs_worker = StructuredOutputsWorker( - max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1), - vocab_size=self.vocab_size, - device=self.device, - ) - - # For CUDA graphs, and will init cudagraph_manager after init_attn_backend. - self.decode_query_len = self.num_speculative_steps + 1 self.cudagraph_manager: ModelCudaGraphManager | None = None # LoRA-related workers. self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) @@ -335,6 +305,40 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.model_state = init_model_state( self.vllm_config, self.model, self.encoder_cache, self.device ) + + self.decode_query_len = ( + self.num_speculative_steps + + self.model_state.num_new_sampled_tokens_per_step + ) + + # Initialize samplers. Model states may override via custom_sampler(). + if self.is_last_pp_rank and not self.is_pooling_model: + self.sampler = Sampler( + max_num_reqs=self.max_num_reqs, + vocab_size=self.vocab_size, + device=self.device, + req_states=self.req_states, + logprobs_mode=self.model_config.logprobs_mode, + num_speculative_tokens=self.decode_query_len, + use_fp64_gumbel=self.model_config.use_fp64_gumbel, + ) + custom = self.model_state.custom_sampler(self.sampler) + + if custom: + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + ) + self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs) + self.structured_outputs_worker = StructuredOutputsWorker( + max_num_logits=self.max_num_reqs * self.decode_query_len, + vocab_size=self.vocab_size, + device=self.device, + ) + if self.is_pooling_model and self.is_last_pp_rank: self.pooling_runner = PoolingRunner(self.model) eplb_models_added |= self.eplb.maybe_register_model( @@ -363,8 +367,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 GPUModelRunnerV1.reload_weights(self, *args, **kwargs) # type: ignore[arg-type] - self.reset_encoder_cache() - self.reset_mm_cache() def apply_sparse_weight_patches(self, *args, **kwargs) -> None: # TODO: Use full version instead of import when fully migrated to v2 @@ -447,7 +449,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, - self.uniform_decode_query_len, + self.decode_query_len, self.parallel_config.tensor_parallel_size, self.kv_cache_config, self.max_num_reqs, @@ -710,6 +712,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): return cuda_graph_size def _remove_request(self, req_id: str) -> bool: + # Call model_state.remove_request *before* req_states.remove_request + # so the model_state can still look up the slot index. + self.model_state.remove_request(req_id) req_idx = self.req_states.remove_request(req_id) if req_idx is None: return False @@ -857,16 +862,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): dtype=np.int32, count=num_reqs, ) + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) - total_num_logits = num_reqs + total_num_draft_tokens - - num_logits = num_draft_tokens_per_req + 1 + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + num_logits = num_draft_tokens_per_req + num_bonus_tokens cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) - max_expand_len = self.num_speculative_steps + 1 + max_expand_len = self.decode_query_len expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, max_expand_len ) @@ -935,6 +940,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.draft_tokens, cu_num_logits, total_num_logits, + self.model_state.num_new_sampled_tokens_per_step, ) # CPU upper bound on seq_lens; padded entries left at zero. @@ -1027,8 +1033,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): grammar_output.grammar_bitmask, ) - if input_batch.num_draft_tokens == 0: - # No draft tokens (common case). + if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: assert self.sampler is not None sampler_output = self.sampler(logits, input_batch) else: @@ -1042,16 +1047,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.speculator.draft_logits, ) - # Get the number of sampled and rejected tokens. - # For chunked prefills, num_sampled and num_rejected are both 0. - num_sampled, num_rejected = get_num_sampled_and_rejected( - sampler_output.num_sampled, - input_batch.seq_lens, - input_batch.cu_num_logits, - input_batch.idx_mapping, - self.req_states.prefill_len.gpu, - ) - return sampler_output, num_sampled, num_rejected + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected def postprocess_sampled( self, @@ -1448,7 +1444,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) + + if self.num_speculative_steps > 0: + # Spec-decode and diffusion LLMs both use draft tokens but the latter does + # not have a speculator (i.e. self.speculator is None) + self.draft_tokens_handler.set_draft_tokens( + input_batch, + self.req_states.draft_tokens[input_batch.idx_mapping], + ) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index b096fcaf5e6..e24c7e9b1cb 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -13,6 +13,11 @@ def init_model_state( encoder_cache: EncoderCache | None, device: torch.device, ): + # Let the model provide its own ModelState if it defines one. + if hasattr(model, "get_model_state_cls"): + cls = model.get_model_state_cls() + return cls(vllm_config, model, encoder_cache, device) + if ( "WhisperForConditionalGeneration" in vllm_config.model_config.architectures or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 55bf8d473cc..86f28e08ea9 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -53,6 +53,9 @@ class ModelState(ABC): def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: return None + def remove_request(self, req_id: str) -> None: + return None + def apply_staged_writes(self) -> None: return None @@ -89,3 +92,16 @@ class ModelState(ABC): for_capture: bool = False, ) -> dict[str, Any]: raise NotImplementedError + + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + """Wrap or replace the default sampler. + + Called after model loading with the already-constructed base + ``Sampler``. Return ``None`` to keep the defaults, or + ``(sampler, rejection_sampler | None)`` to override. + """ + return None + + num_new_sampled_tokens_per_step: int = 1 + """New tokens sampled on each decode step + (excluding accepted draft tokens, a.k.a num bonus tokens).""" diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py index f38ac8affd8..130f4ddbf8a 100644 --- a/vllm/v1/worker/gpu/sample/output.py +++ b/vllm/v1/worker/gpu/sample/output.py @@ -13,3 +13,4 @@ class SamplerOutput: logprobs_tensors: LogprobsTensors | None num_nans: torch.Tensor | None num_sampled: torch.Tensor | None + num_rejected: torch.Tensor | None = None diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 6b545aef3a2..b269de9eaed 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -12,7 +12,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import ( flashinfer_sample, flashinfer_sampler_supported, ) -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected from vllm.v1.worker.gpu.metrics.logits import get_num_nans from vllm.v1.worker.gpu.sample.bad_words import BadWordsState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @@ -44,6 +44,7 @@ class Sampler: self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default. self.use_fp64_gumbel = use_fp64_gumbel + self.req_states = req_states self.sampling_states = SamplingStates(max_num_reqs, vocab_size) self.penalties_state = PenaltiesState(req_states) self.logit_bias_state = LogitBiasState(max_num_reqs, device) @@ -118,6 +119,17 @@ class Sampler: else: logprobs_tensors = None + # 1 sampled token per request, except chunked-prefill requests + # (seq_len < prefill_len) which aren't done prefilling and produce no + # output token. num_rejected is always 0 here (one logit per request). + num_sampled, num_rejected = get_num_sampled_and_rejected( + input_batch.seq_lens.new_ones(input_batch.num_reqs), + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping, + self.req_states.prefill_len.gpu, + ) + # These are GPU tensors. sampler_output = SamplerOutput( # The sampled tokens are expanded to 2D tensor with shape @@ -126,7 +138,8 @@ class Sampler: sampled_token_ids=sampled.view(-1, 1), logprobs_tensors=logprobs_tensors, num_nans=num_nans, - num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs), + num_sampled=num_sampled, + num_rejected=num_rejected, ) return sampler_output diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 1fe079a43e7..3868604d3ae 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -6,7 +6,10 @@ from vllm.config import SpeculativeConfig from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.input_batch import ( + InputBatch, + get_num_sampled_and_rejected, +) from vllm.v1.worker.gpu.metrics.logits import get_num_nans from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -136,9 +139,18 @@ class RejectionSampler: else logits, ) + num_sampled, num_rejected = get_num_sampled_and_rejected( + num_sampled, + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping, + self.sampler.req_states.prefill_len.gpu, + ) + return SamplerOutput( sampled_token_ids=sampled, logprobs_tensors=logprobs_tensors, num_nans=num_nans, num_sampled=num_sampled, + num_rejected=num_rejected, ) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 7bfd981ee0c..4ab45b2ae27 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -35,6 +35,10 @@ class DraftTokensHandler: self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): self.draft_tokens_np = async_copy_to_np(draft_tokens) + # draft_tokens is a temporary allocation on the main stream and read here on + # copy_stream; without record_stream, the caching allocator may reuse its + # memory before the async copy executes. + draft_tokens.record_stream(self.copy_stream) self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 83d87c74a4a..0da845a0673 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -30,17 +30,18 @@ def warmup_kernels( pipeline parallel coordination. The first iteration simulates a prefill with requests of - 2 + num_spec_steps prompt tokens each. The second iteration simulates - a decode step with all requests generating 1 + num_spec_steps tokens. + decode_query_len + 1 prompt tokens each. The second iteration simulates + a decode step with all requests generating decode_query_len tokens. """ num_spec_steps = model_runner.num_speculative_steps - # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request - # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing - # it from being misclassified as a uniform decode batch. - prompt_len = 2 + num_spec_steps + decode_query_len = model_runner.decode_query_len + # Use decode_query_len + 1 tokens so the prefill batch's per-request query + # length exceeds decode_query_len, preventing it from being misclassified as + # a uniform decode batch. + prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates 1 verified + num_spec_steps draft tokens. - decode_len = prompt_len + 1 + num_spec_steps + # After prefill, decode generates decode_query_len tokens. + decode_len = prompt_len + decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) @@ -57,7 +58,7 @@ def warmup_kernels( num_reqs = min( model_runner.scheduler_config.max_num_seqs, model_runner.scheduler_config.max_num_batched_tokens - // max(prompt_len, 1 + num_spec_steps), + // max(prompt_len, decode_query_len), # Reserve block 0 (null block) and ensure we have enough blocks. max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), ) @@ -79,7 +80,7 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) - # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each. + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params), @@ -117,7 +118,7 @@ def warmup_kernels( worker_sample_tokens(grammar_output) - # Step 2: Decode all requests with 1 + num_spec_steps tokens each. + # Step 2: Decode all requests with decode_query_len tokens each. cached_req_data = CachedRequestData.make_empty() cached_req_data.req_ids = list(req_ids) cached_req_data.num_computed_tokens = [prompt_len] * num_reqs @@ -131,7 +132,7 @@ def warmup_kernels( decode_output = SchedulerOutput.make_empty() decode_output.scheduled_cached_reqs = cached_req_data decode_output.num_scheduled_tokens = { - req_id: 1 + num_spec_steps for req_id in req_ids + req_id: decode_query_len for req_id in req_ids } if num_spec_steps > 0: decode_output.scheduled_spec_decode_tokens = { diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f3f52c75d8b..cb607c0b7b0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5391,6 +5391,9 @@ class GPUModelRunner( weights_not_loaded, ) + self.reset_encoder_cache() + self.reset_mm_cache() + def _get_prompt_logprobs_dict( self, hidden_states: torch.Tensor, diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 5004ba9c8f2..276b9b4250f 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -209,6 +209,7 @@ def flash_attn_varlen_func( # FA4 only mask_mod=None, aux_tensors=None, + dynamic_causal: "torch.Tensor | None" = None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -392,6 +393,7 @@ def flash_attn_varlen_func( page_table=block_table, softmax_scale=softmax_scale, causal=causal, + dynamic_causal=dynamic_causal, softcap=softcap, window_size_left=real_window_size[0] if real_window_size[0] >= 0 else None, window_size_right=real_window_size[1] if real_window_size[1] >= 0 else None,