diff --git a/.buildkite/ci_config_rocm.yaml b/.buildkite/ci_config_rocm.yaml index 23f32340071..408ccee58ed 100644 --- a/.buildkite/ci_config_rocm.yaml +++ b/.buildkite/ci_config_rocm.yaml @@ -8,6 +8,7 @@ run_all_patterns: - "docker/docker-bake-rocm.hcl" - ".buildkite/hardware_tests/amd.yaml" - ".buildkite/scripts/ci-bake-rocm.sh" + - ".buildkite/scripts/rocm/" - ".buildkite/scripts/hardware_ci/run-amd-test.py" - ".buildkite/scripts/hardware_ci/run-amd-test.sh" - "CMakeLists.txt" diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index a18241cf18b..d47a7e394a4 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -1,5 +1,28 @@ group: Hardware - AMD Build + +# ROCm image flow: +# 1. Refresh the long-lived ROCm base image only when Dockerfile.rocm_base changes. +# 2. Build ci_base from either the stable base or the freshly refreshed base. +# 3. Build the per-commit ROCm CI image and smoke-test it before GPU jobs run. steps: + - label: "AMD: :docker: refresh ROCm base" + key: refresh-rocm-base-amd + depends_on: [] + device: amd_cpu + no_plugin: true + commands: + - bash .buildkite/scripts/rocm/refresh-base-image.sh + env: + DOCKER_BUILDKIT: "1" + BUILDKIT_PROGRESS: "tty" + TERM: "xterm-256color" + retry: + automatic: + - exit_status: -1 # Agent was lost + limit: 1 + - exit_status: -10 # Agent was lost + limit: 1 + # Ensure ci_base is up-to-date before building the test image. # Compares a content hash of ci_base-affecting files against the remote # image label. If hashes match the build is skipped (< 30 s); if they @@ -7,13 +30,16 @@ steps: - label: "AMD: :docker: ensure ci_base" key: ensure-ci-base-amd soft_fail: false - depends_on: [] + depends_on: + - refresh-rocm-base-amd device: amd_cpu no_plugin: true commands: - - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps + - bash .buildkite/scripts/rocm/build-ci-base.sh env: DOCKER_BUILDKIT: "1" + BUILDKIT_PROGRESS: "tty" + TERM: "xterm-256color" VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl" PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950" REMOTE_VLLM: "1" @@ -33,35 +59,12 @@ steps: device: amd_cpu no_plugin: true commands: - - | - if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then - echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only" - IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts - else - bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel - fi - - | - docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec ' - if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi - if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi - if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi - if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi - command -v python3 - command -v uv - command -v pytest - if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then - echo No ROCm CLI found in image >&2 - exit 1 - fi - python3 - <- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'export VLLM_XPU_FUSED_MOE_USE_REF=1 && + cd tests && + pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]' + - label: XPU CPU Offload timeout_in_minutes: 60 device: intel_gpu @@ -110,7 +134,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/models_distributed_intel.yaml b/.buildkite/intel_jobs/models_distributed_intel.yaml index 7b574f2a8e3..604f7744c10 100644 --- a/.buildkite/intel_jobs/models_distributed_intel.yaml +++ b/.buildkite/intel_jobs/models_distributed_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 2+ - mem: 24+ + mem: 16+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 8dae59f4ea2..e12c2658f50 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -9,7 +9,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -81,7 +81,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: @@ -105,7 +105,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 16+ + mem: 24+ no_plugin: true working_dir: "." env: diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 203e63f25e2..d20ae36a2e3 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -81,7 +81,7 @@ steps: agent_tags: label: production gpu: 1+ - mem: 16+ + mem: 24+ no_plugin: true env: REGISTRY: "public.ecr.aws/q9t5s3a7" diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 4ccbbb352d9..90078fc0c1d 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -15,7 +15,7 @@ set -euo pipefail DEFAULT_REPO_SLUG="vllm-project/vllm" DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" -DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh" +DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh" DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base" DEFAULT_CI_BASE_METADATA_VERSION="1" @@ -393,6 +393,16 @@ should_upload_wheel_artifacts() { || "${TARGET}" == *"artifact"* ]] } +set_buildkite_metadata() { + local key="$1" + local value="$2" + + [[ -n "${value}" ]] || return 0 + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data set "${key}" "${value}" || true + fi +} + get_remote_image_label() { local image_ref="$1" local label_key="$2" @@ -733,6 +743,8 @@ configure_ci_base_image_refs() { if is_ci_base_target; then IMAGE_TAG="${primary_tag}" + CI_BASE_IMAGE="${primary_tag}" + export CI_BASE_IMAGE export IMAGE_TAG echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}" @@ -750,6 +762,10 @@ configure_ci_base_image_refs() { echo "ci_base stable alias will not be pushed for this build" echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}" fi + set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}" + set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}" + set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}" + set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}" return 0 fi @@ -1779,7 +1795,10 @@ seed_dependency_caches_if_needed() { echo "--- :docker: Seeding ${target}" echo "Expected cache ref: ${cache_ref}" - docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}" + docker buildx bake \ + "${BAKE_FILES[@]}" \ + --progress "${BUILDKIT_PROGRESS:-plain}" \ + "${target}" verify_dependency_cache_ref "${cache_ref}" done } @@ -1807,7 +1826,10 @@ run_bake() { local build_rc=0 echo "--- :docker: Building ${TARGET}" - docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$? + docker buildx bake \ + "${BAKE_FILES[@]}" \ + --progress "${BUILDKIT_PROGRESS:-plain}" \ + "${BAKE_TARGETS[@]}" || build_rc=$? if [[ ${build_rc} -eq 0 ]]; then echo "--- :white_check_mark: Build complete" diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 9db805c4585..fee9ab04f4b 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -28,6 +28,17 @@ ############################################################################### set -o pipefail +: "${BUILDKIT_PROGRESS:=plain}" +: "${TERM:=xterm-256color}" +: "${FORCE_COLOR:=1}" +: "${CLICOLOR_FORCE:=1}" +: "${PY_COLORS:=1}" +: "${ROCM_DOCKER_TTY:=1}" +if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then + PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes" +fi +export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY + # Export Python path for commands that run directly on the host. Containerized # tests set this to /vllm-workspace below so spawned Python processes do not # depend on their current working directory. @@ -149,6 +160,7 @@ EOF echo "--- Building local ROCm test image" docker build \ --pull=false \ + --progress "${BUILDKIT_PROGRESS}" \ --build-arg "BASE_IMAGE=${base_image}" \ -t "${artifact_image}" \ "${context_dir}" || return 1 @@ -535,6 +547,13 @@ if is_multi_node "$commands"; then else echo "--- Single-node job" echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" + docker_run_terminal_args=(-i) + if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then + docker_run_terminal_args+=(-t) + echo "Docker interactive stdin: enabled; TTY allocation: enabled" + else + echo "Docker interactive stdin: enabled; TTY allocation: disabled" + fi ulimit_core_hard=$(ulimit -H -c) if [[ "$ulimit_core_hard" == "unlimited" ]]; then @@ -551,6 +570,7 @@ else fi docker run \ + "${docker_run_terminal_args[@]}" \ --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ $RDMA_FLAGS \ --network=host \ @@ -565,6 +585,11 @@ else -e AWS_SECRET_ACCESS_KEY \ -e BUILDKITE_PARALLEL_JOB \ -e BUILDKITE_PARALLEL_JOB_COUNT \ + -e TERM \ + -e FORCE_COLOR \ + -e CLICOLOR_FORCE \ + -e PY_COLORS \ + -e PYTEST_ADDOPTS \ -v "${HF_CACHE}:${HF_MOUNT}" \ -e "HF_HOME=${HF_MOUNT}" \ -e "PYTHONPATH=${MYPYTHONPATH}" \ diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 252eeeef8ce..2d11dd477ea 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -38,7 +38,9 @@ function cpu_tests() { pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/core/test_cpu_activation.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits @@ -62,7 +64,6 @@ function cpu_tests() { set -e pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs" - # basic online serving docker exec cpu-test bash -c ' set -e diff --git a/.buildkite/scripts/rocm/build-ci-base.sh b/.buildkite/scripts/rocm/build-ci-base.sh new file mode 100755 index 00000000000..23d17e17b4d --- /dev/null +++ b/.buildkite/scripts/rocm/build-ci-base.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base. + +set -euo pipefail + +metadata_get() { + local key="$1" + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data get "${key}" 2>/dev/null || true + fi +} + +main() { + local base_refreshed="" + + base_refreshed="$(metadata_get rocm-base-refresh)" + if [[ "${base_refreshed}" == "1" ]]; then + export BASE_IMAGE + export CI_BASE_PUSH_STABLE_TAG + + BASE_IMAGE="$(metadata_get rocm-base-image)" + CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)" + CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}" + + echo "Using refreshed ROCm base image for ci_base: ${BASE_IMAGE}" + echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}" + fi + + bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps +} + +main "$@" diff --git a/.buildkite/scripts/rocm/build-test-image.sh b/.buildkite/scripts/rocm/build-test-image.sh new file mode 100755 index 00000000000..9803e20d02e --- /dev/null +++ b/.buildkite/scripts/rocm/build-test-image.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# Build the ROCm CI test image or wheel artifact. +# +# When Dockerfile.rocm_base changes, always build the full image so downstream +# ROCm tests can validate the freshly rebuilt base -> ci_base -> ci image chain. + +set -euo pipefail + +metadata_get() { + local key="$1" + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data get "${key}" 2>/dev/null || true + fi +} + +use_refreshed_base_if_present() { + local base_refreshed="" + + base_refreshed="$(metadata_get rocm-base-refresh)" + if [[ "${base_refreshed}" != "1" ]]; then + return 1 + fi + + export BASE_IMAGE + export CI_BASE_IMAGE + export IMAGE_TAG_LATEST + + BASE_IMAGE="$(metadata_get rocm-base-image)" + CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)" + IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)" + + echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}" + echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}" + if [[ -n "${IMAGE_TAG_LATEST}" ]]; then + echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}" + fi + + return 0 +} + +main() { + local base_refreshed=0 + + if use_refreshed_base_if_present; then + base_refreshed=1 + fi + + if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then + echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only" + IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts + return + fi + + bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel +} + +main "$@" diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh new file mode 100755 index 00000000000..06e3e80c967 --- /dev/null +++ b/.buildkite/scripts/rocm/refresh-base-image.sh @@ -0,0 +1,513 @@ +#!/usr/bin/env bash +# Build and publish a fresh ROCm base image when Dockerfile.rocm_base changes. +# +# Normal AMD CI builds should not pay for this path. The script no-ops unless +# docker/Dockerfile.rocm_base changed relative to the branch base, the previous +# main commit, or ROCM_BASE_REFRESH_FORCE=1 is set. + +set -euo pipefail + +DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}" +BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}" +CI_IMAGE_REPO="${ROCM_CI_IMAGE_REPO:-rocm/vllm-ci}" +BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}" +DEFAULT_ROCM_BASE_METADATA_VERSION="1" +DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}" +DEFAULT_ROCM_BASE_CONTENT_ARGS="BASE_IMAGE TRITON_BRANCH TRITON_REPO PYTORCH_BRANCH PYTORCH_REPO PYTORCH_VISION_BRANCH PYTORCH_VISION_REPO PYTORCH_AUDIO_BRANCH PYTORCH_AUDIO_REPO FA_BRANCH FA_REPO AITER_BRANCH AITER_REPO MORI_BRANCH MORI_REPO PYTORCH_ROCM_ARCH PYTHON_VERSION USE_SCCACHE" + +metadata_set() { + local key="$1" + local value="$2" + + [[ -n "${value}" ]] || return 0 + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent meta-data set "${key}" "${value}" || true + fi +} + +compute_content_hash() { + local path="" + local file="" + + for path in "$@"; do + if [[ -d "${path}" ]]; then + while IFS= read -r -d '' file; do + printf 'file:%s\n' "${file}" + sha256sum "${file}" + done < <(find "${path}" -type f -print0 | sort -z) + elif [[ -f "${path}" ]]; then + printf 'file:%s\n' "${path}" + sha256sum "${path}" + else + printf 'missing:%s\n' "${path}" + fi + done | sha256sum | cut -d' ' -f1 +} + +clean_docker_tag() { + local input="$1" + echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128 +} + +tag_component() { + local input="$1" + local max_chars="${2:-24}" + + clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}" +} + +extract_arg_default() { + local arg_name="$1" + + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + "${DOCKERFILE}" | head -1 +} + +resolve_image_digest() { + local image_ref="$1" + + docker buildx imagetools inspect "${image_ref}" 2>/dev/null \ + | sed -n -E 's/^Digest:[[:space:]]+//p' \ + | head -1 || true +} + +resolve_rocm_base_arg_value() { + local arg_name="$1" + local use_sccache="$2" + + case "${arg_name}" in + USE_SCCACHE) + printf '%s\n' "${use_sccache}" + ;; + *) + extract_arg_default "${arg_name}" + ;; + esac +} + +hash_rocm_base_arg_values() { + local use_sccache="$1" + local base_image_digest="$2" + local arg_name="" + local arg_value="" + shift 2 || true + + for arg_name in "$@"; do + [[ -n "${arg_name}" ]] || continue + arg_value=$(resolve_rocm_base_arg_value "${arg_name}" "${use_sccache}") + printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-}" + if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then + printf 'arg:%s.digest=%s\n' "${arg_name}" "${base_image_digest:-unknown}" + fi + done +} + +rocm_version_from_base_image() { + local base_image="$1" + local version="" + + version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")" + tag_component "${version:-${base_image}}" 16 +} + +git_diff_changed_base() { + local range="$1" + [[ -n "$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null)" ]] +} + +short_git_ref() { + local ref="$1" + + git rev-parse --short "${ref}" 2>/dev/null || printf '%s\n' "${ref}" +} + +extract_arg_default_from_ref() { + local ref="$1" + local arg_name="$2" + local content="" + + content="$(git show "${ref}:${DOCKERFILE}" 2>/dev/null || true)" + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + <<<"${content}" | head -1 +} + +log_arg_default_changes() { + local old_ref="$1" + local new_ref="$2" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local arg_name="" + local old_value="" + local new_value="" + local changed=0 + + echo "Changed ROCm base ARG defaults:" + for arg_name in ${content_args}; do + old_value="$(extract_arg_default_from_ref "${old_ref}" "${arg_name}")" + new_value="$(extract_arg_default_from_ref "${new_ref}" "${arg_name}")" + if [[ "${old_value}" != "${new_value}" ]]; then + echo " - ${arg_name}: ${old_value:-} -> ${new_value:-}" + changed=1 + fi + done + + if [[ "${changed}" == "0" ]]; then + echo " - none detected; Dockerfile instructions changed outside tracked ARG defaults" + fi +} + +log_arg_line_diff() { + local range="$1" + local arg_diff="" + + arg_diff="$( + git diff --unified=0 "${range}" -- "${DOCKERFILE}" 2>/dev/null \ + | awk '/^[+-][[:space:]]*ARG[[:space:]]/ && $0 !~ /^(---|\+\+\+)/ { print " " $0 }' \ + || true + )" + + if [[ -n "${arg_diff}" ]]; then + echo "Changed Dockerfile ARG lines:" + printf '%s\n' "${arg_diff}" + fi +} + +log_rocm_base_change_check() { + local context="$1" + local range="$2" + local old_ref="$3" + local old_short="" + local head_short="" + + old_short="$(short_git_ref "${old_ref}")" + head_short="$(short_git_ref HEAD)" + + echo "--- :mag: ROCm base refresh check" + echo "Context: ${context}" + echo "Dockerfile: ${DOCKERFILE}" + echo "Base revision: ${old_short}" + echo "Head revision: ${head_short}" + echo "Git diff range: ${range}" +} + +log_rocm_base_rebuild_reason() { + local context="$1" + local range="$2" + local old_ref="$3" + local changed_files="" + + log_rocm_base_change_check "${context}" "${range}" "${old_ref}" + + changed_files="$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null || true)" + echo "Changed files:" + if [[ -n "${changed_files}" ]]; then + sed 's/^/ - /' <<<"${changed_files}" + else + echo " - ${DOCKERFILE}" + fi + log_arg_default_changes "${old_ref}" HEAD + log_arg_line_diff "${range}" + echo "Decision: rebuilding ROCm base image because ${DOCKERFILE} changed." +} + +rocm_base_changed_in_range() { + local context="$1" + local range="$2" + local old_ref="$3" + + if git_diff_changed_base "${range}"; then + log_rocm_base_rebuild_reason "${context}" "${range}" "${old_ref}" + return 0 + fi + + log_rocm_base_change_check "${context}" "${range}" "${old_ref}" + echo "Decision: ROCm base refresh not required; ${DOCKERFILE} is unchanged." + return 1 +} + +rocm_base_changed() { + local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" + local base_ref="refs/remotes/origin/${base_branch}" + local merge_base="" + + if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then + echo "ROCM_BASE_REFRESH_SKIP=1 set; skipping ROCm base refresh" + return 1 + fi + + if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" == "1" ]]; then + echo "ROCM_BASE_REFRESH_FORCE=1 set; refreshing ROCm base image" + return 0 + fi + + if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "Not in a git checkout; skipping ROCm base refresh unless forced" + return 1 + fi + + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + git fetch --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true) + if [[ -z "${merge_base}" ]]; then + echo "Unable to determine merge base with PR base ${base_ref}; skipping ROCm base refresh unless forced" + return 1 + fi + if rocm_base_changed_in_range \ + "pull request build against ${base_ref}" \ + "${merge_base}...HEAD" \ + "${merge_base}"; then + return 0 + fi + elif [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] \ + && git rev-parse --verify HEAD~1 >/dev/null 2>&1; then + if rocm_base_changed_in_range \ + "stable branch build; comparing against previous ${ROCM_BASE_STABLE_BRANCH:-main} commit" \ + "HEAD~1..HEAD" \ + "HEAD~1"; then + return 0 + fi + else + git fetch --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true) + if [[ -z "${merge_base}" ]]; then + echo "Unable to determine merge base with branch base ${base_ref}; skipping ROCm base refresh unless forced" + return 1 + fi + if rocm_base_changed_in_range \ + "branch build against ${base_ref}" \ + "${merge_base}...HEAD" \ + "${merge_base}"; then + return 0 + fi + fi + + return 1 +} + +should_push_stable_tag() { + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + return 1 + fi + + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then + return 0 + fi + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then + return 1 + fi + + [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" \ + && "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] +} + +setup_builder() { + echo "--- :buildkite: Setting up buildx builder for ROCm base" + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then + docker buildx use "${BUILDER_NAME}" + else + docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use + fi + docker buildx inspect --bootstrap +} + +compute_base_content_hash() { + local use_sccache="$1" + local base_image_digest="$2" + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local -a content_paths=() + local -a content_arg_names=() + + read -r -a content_paths <<< "${content_files}" + read -r -a content_arg_names <<< "${content_args}" + + { + printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")" + printf 'dockerfile:%s\n' "${DOCKERFILE}" + printf 'resolved-build-args:\n' + hash_rocm_base_arg_values \ + "${use_sccache}" "${base_image_digest}" "${content_arg_names[@]}" + } | sha256sum | cut -d' ' -f1 +} + +build_base_image() { + local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}" + local base_hash="" + local build_date="" + local build_suffix="" + local base_image_arg="" + local base_image_digest="" + local rocm_version="" + local triton_arg="" + local pytorch_arg="" + local pytorch_vision_arg="" + local pytorch_audio_arg="" + local fa_arg="" + local aiter_arg="" + local mori_arg="" + local python_version_arg="" + local pytorch_rocm_arch_arg="" + local pytorch_branch="" + local aiter_branch="" + local dependency_summary="" + local descriptor="" + local ci_descriptor="" + local descriptive_tag="" + local stable_tag="${BASE_REPO}:base" + local ci_descriptive_tag="" + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}" + local content_files_hash="" + local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}" + local -a tags=() + local -a no_cache_args=() + local -a sccache_args=() + local -a content_paths=() + + if [[ ! -f "${DOCKERFILE}" ]]; then + echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2 + exit 1 + fi + + build_date="${ROCM_BASE_TAG_DATE:-$(date -u +%Y%m%d)}" + if [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then + build_suffix="_bk_${BUILDKITE_BUILD_NUMBER}" + fi + base_image_arg="$(extract_arg_default BASE_IMAGE)" + base_image_digest="$(resolve_image_digest "${base_image_arg}")" + read -r -a content_paths <<< "${content_files}" + content_files_hash="$(compute_content_hash "${content_paths[@]}")" + base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}") + rocm_version="$(rocm_version_from_base_image "${base_image_arg}")" + triton_arg="$(extract_arg_default TRITON_BRANCH)" + pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)" + pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)" + pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)" + fa_arg="$(extract_arg_default FA_BRANCH)" + aiter_arg="$(extract_arg_default AITER_BRANCH)" + mori_arg="$(extract_arg_default MORI_BRANCH)" + python_version_arg="$(extract_arg_default PYTHON_VERSION)" + pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)" + pytorch_branch="$(tag_component "${pytorch_arg}" 16)" + aiter_branch="$(tag_component "${aiter_arg}" 24)" + dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}" + descriptor="$(clean_docker_tag "base_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")" + ci_descriptor="$(clean_docker_tag "ci_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")" + + descriptive_tag="${BASE_REPO}:${descriptor}" + ci_descriptive_tag="${CI_IMAGE_REPO}:${ci_descriptor}" + + tags=(-t "${descriptive_tag}") + if should_push_stable_tag; then + tags+=(-t "${stable_tag}") + metadata_set "rocm-base-push-stable-tag" "1" + else + metadata_set "rocm-base-push-stable-tag" "0" + fi + + if [[ "${ROCM_BASE_NO_CACHE:-1}" == "1" ]]; then + no_cache_args=(--no-cache) + fi + + for env_name in \ + SCCACHE_DOWNLOAD_URL \ + SCCACHE_ENDPOINT \ + SCCACHE_BUCKET_NAME \ + SCCACHE_REGION_NAME \ + SCCACHE_S3_NO_CREDENTIALS; do + if [[ -n "${!env_name:-}" ]]; then + sccache_args+=(--build-arg "${env_name}=${!env_name}") + fi + done + + echo "--- :docker: Building ROCm base image" + echo "Dockerfile: ${DOCKERFILE}" + echo "Descriptive tag: ${descriptive_tag}" + echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))" + echo "Content hash: ${base_hash}" + echo "Dependency summary: ${dependency_summary}" + echo "USE_SCCACHE: ${use_sccache}" + + docker buildx build \ + "${no_cache_args[@]}" \ + --pull \ + --progress "${BUILDKIT_PROGRESS:-plain}" \ + --file "${DOCKERFILE}" \ + --build-arg "USE_SCCACHE=${use_sccache}" \ + "${sccache_args[@]}" \ + --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \ + --label "org.opencontainers.image.vendor=vLLM" \ + --label "org.opencontainers.image.title=vLLM ROCm base" \ + --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT:-}" \ + --label "vllm.rocm_base.metadata_version=${metadata_version}" \ + --label "vllm.rocm_base.content_hash=${base_hash}" \ + --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \ + --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \ + --label "vllm.rocm_base.image.descriptive=${descriptive_tag}" \ + --label "vllm.rocm_base.image.stable=${stable_tag}" \ + --label "vllm.rocm_base.git_commit=${BUILDKITE_COMMIT:-}" \ + --label "vllm.rocm_base.stable_branch=${ROCM_BASE_STABLE_BRANCH:-main}" \ + --label "vllm.rocm_base.descriptor=${descriptor}" \ + --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \ + --label "vllm.rocm_base.base_image=${base_image_arg}" \ + --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \ + --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \ + --label "vllm.rocm_base.dependency.python=${python_version_arg}" \ + --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \ + --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \ + --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \ + --label "vllm.rocm_base.dependency.triton=${triton_arg}" \ + --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \ + --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \ + --label "vllm.rocm_base.dependency.mori=${mori_arg}" \ + --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \ + "${tags[@]}" \ + --push \ + . + + docker buildx imagetools inspect "${descriptive_tag}" >/dev/null + + metadata_set "rocm-base-refresh" "1" + metadata_set "rocm-base-image" "${descriptive_tag}" + metadata_set "rocm-base-image-descriptive" "${descriptive_tag}" + metadata_set "rocm-base-image-stable" "${stable_tag}" + metadata_set "rocm-base-image-ci-descriptive" "${ci_descriptive_tag}" + metadata_set "rocm-base-metadata-version" "${metadata_version}" + metadata_set "rocm-base-content-hash" "${base_hash}" + metadata_set "rocm-base-content-files-hash" "${content_files_hash}" + metadata_set "rocm-base-content-files" "${content_files}" + metadata_set "rocm-base-content-args" "${content_args}" + metadata_set "rocm-base-base-image-digest" "${base_image_digest}" + metadata_set "rocm-base-dockerfile" "${DOCKERFILE}" + metadata_set "rocm-base-descriptor" "${descriptor}" + metadata_set "rocm-base-dependency-summary" "${dependency_summary}" + metadata_set "rocm-base-dependency-rocm" "${rocm_version}" + metadata_set "rocm-base-dependency-python" "${python_version_arg}" + metadata_set "rocm-base-dependency-pytorch" "${pytorch_arg}" + metadata_set "rocm-base-dependency-torchvision" "${pytorch_vision_arg}" + metadata_set "rocm-base-dependency-torchaudio" "${pytorch_audio_arg}" + metadata_set "rocm-base-dependency-triton" "${triton_arg}" + metadata_set "rocm-base-dependency-flash-attention" "${fa_arg}" + metadata_set "rocm-base-dependency-aiter" "${aiter_arg}" + metadata_set "rocm-base-dependency-mori" "${mori_arg}" + metadata_set "rocm-base-pytorch-rocm-arch" "${pytorch_rocm_arch_arg}" + metadata_set "rocm-ci-image-descriptive" "${ci_descriptive_tag}" + + echo "--- :white_check_mark: ROCm base image published" + echo "Use BASE_IMAGE=${descriptive_tag} for downstream ROCm CI builds" +} + +main() { + metadata_set "rocm-base-refresh" "0" + + if ! rocm_base_changed; then + echo "ROCm base Dockerfile did not change; skipping base image refresh" + return 0 + fi + + setup_builder + build_base_image +} + +main "$@" diff --git a/.buildkite/scripts/rocm/smoke-test-image.sh b/.buildkite/scripts/rocm/smoke-test-image.sh new file mode 100755 index 00000000000..ed511c9b77a --- /dev/null +++ b/.buildkite/scripts/rocm/smoke-test-image.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Fast structural smoke test for the full ROCm CI image. + +set -euo pipefail + +image_ref="${VLLM_CI_SMOKE_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}" + +docker run --rm --network=none --entrypoint /bin/bash "${image_ref}" -ec ' + if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi + if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi + if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi + if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi + + command -v python3 + command -v uv + command -v pytest + + if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then + echo No ROCm CLI found in image >&2 + exit 1 + fi + + python3 - <= 0.0.2" - pytest -v -s plugins_tests/gguf +#------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# + +- label: Rust Frontend OpenAI Coverage # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/benchmarks/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/sample/ + - tests/utils.py + - tests/benchmarks/test_serve_cli.py + - tests/entrypoints/openai/chat_completion/test_chat_completion.py + - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py + - tests/entrypoints/openai/completion/test_shutdown.py + - tests/entrypoints/openai/test_return_token_ids.py + - tests/entrypoints/openai/test_uds.py + - tests/v1/sample/test_logprobs_e2e.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" + - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" + - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" + - pytest -v -s entrypoints/openai/test_uds.py + - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" + +- label: Rust Frontend Serve Admin Coverage # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - vllm/entrypoints/serve/ + - vllm/v1/engine/ + - tests/utils.py + - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py + - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py + - tests/entrypoints/serve/instrumentator/test_basic.py + - tests/entrypoints/serve/instrumentator/test_metrics.py + - tests/entrypoints/serve/tokenize/test_tokenization.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" + - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" + +- label: Rust Frontend Core Correctness # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - tests/utils.py + - tests/entrypoints/openai/correctness/test_lmeval.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine + +- label: Rust Frontend Tool Use # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/entrypoints/openai/ + - vllm/tool_parsers/ + - tests/utils.py + - tests/tool_use/ + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" + +- label: Rust Frontend Distributed # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - rust/ + - vllm/distributed/ + - vllm/engine/ + - vllm/executor/ + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/utils.py + - tests/v1/distributed/test_external_lb_dp.py + - tests/v1/distributed/test_hybrid_lb_dp.py + - tests/v1/distributed/test_internal_lb_dp.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RUST_FRONTEND=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info" + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info" + #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# - label: Quantization # TBD @@ -2248,7 +2407,7 @@ steps: - pytest -v -s v1/worker - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api # - export HSA_NO_SCRATCH_RECLAIM=1 - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine @@ -2439,6 +2598,59 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh + +- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh + +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + - label: V1 e2e (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2729,15 +2941,19 @@ steps: - vllm/envs.py - examples/offline_inference/data_parallel.py - tests/distributed/test_context_parallel.py + - tests/distributed/test_rocm_aiter_custom_ar.py - tests/distributed/test_rocm_quick_reduce.py - tests/distributed/test_quick_all_reduce.py + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - tests/v1/distributed/test_dbo.py - tests/utils.py commands: - pytest -v -s tests/distributed/test_context_parallel.py - - pytest -v -s tests/v1/distributed/test_dbo.py + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py - pytest -v -s tests/distributed/test_rocm_quick_reduce.py - pytest -v -s tests/distributed/test_quick_all_reduce.py + - pytest -v -s tests/v1/distributed/test_dbo.py #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# @@ -3318,7 +3534,7 @@ steps: - pytest -v -s v1/worker - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - label: V1 Sample + Logits # TBD diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 8f47e7b5c21..4cf7774ebd4 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -150,9 +150,9 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: LM Eval Humming (A100 - TEMPORARY) - key: lm-eval-humming-a100 - timeout_in_minutes: 30 +- label: LM Eval Humming f16 (A100 - TEMPORARY) + key: lm-eval-humming-f16-a100 + timeout_in_minutes: 120 device: a100 optional: true num_devices: 1 @@ -160,13 +160,29 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt -- label: LM Eval Humming (H100 - TEMPORARY) - key: lm-eval-humming-h100 - timeout_in_minutes: 30 +- label: LM Eval Humming Act int8 (A100 - TEMPORARY) + key: lm-eval-humming-act-a100 + timeout_in_minutes: 120 + device: a100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (H100 - TEMPORARY) + key: lm-eval-humming-f16-h100 + timeout_in_minutes: 120 device: h100 optional: true num_devices: 1 @@ -174,14 +190,30 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt -- label: LM Eval Humming (B200 - TEMPORARY) - key: lm-eval-humming-b200 - timeout_in_minutes: 30 +- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) + key: lm-eval-humming-act-h100 + timeout_in_minutes: 120 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + +- label: LM Eval Humming f16 (B200 - TEMPORARY) + key: lm-eval-humming-f16-b200 + timeout_in_minutes: 120 device: b200-k8s optional: true num_devices: 1 @@ -189,10 +221,26 @@ steps: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + +- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY) + key: lm-eval-humming-act-b200 + timeout_in_minutes: 120 + device: b200-k8s + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt - label: LM Eval TurboQuant KV Cache key: lm-eval-turboquant-kv-cache diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 1ad04c28970..fd6ef2e61ba 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -103,7 +103,7 @@ steps: - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics # Integration test for streaming correctness (requires special branch). - - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api + - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine mirror: amd: @@ -351,6 +351,7 @@ steps: - tests/test_outputs.py - tests/test_pooling_params.py - tests/test_ray_env.py + - tests/test_sampling_params.py - tests/multimodal - tests/renderers - tests/standalone_tests/lazy_imports.py @@ -368,6 +369,7 @@ steps: - pytest -v -s test_outputs.py - pytest -v -s test_pooling_params.py - pytest -v -s test_ray_env.py + - pytest -v -s test_sampling_params.py - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s renderers - pytest -v -s reasoning diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 5227bbc1f3b..3a113f1982a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -36,10 +36,10 @@ steps: source_file_dependencies: - vllm/ - tests/models/test_terratorch.py - - tests/models/test_transformers.py + - tests/models/transformers/test_backend.py - tests/models/test_registry.py commands: - - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py mirror: amd: device: mi325_1 @@ -55,6 +55,7 @@ steps: - vllm/ - tests/models/test_utils.py - tests/models/test_vision.py + - tests/models/transformers/fusers/ device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py + - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index b5758c55aff..a3ee7666ed0 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -17,7 +17,7 @@ steps: - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' # Avoid importing model tests that cause CUDA reinitialization error - - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 9ecfe01d400..a721efa6067 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -27,6 +27,7 @@ steps: - tests/models/multimodal commands: - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma" + - pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model mirror: amd: @@ -58,7 +59,7 @@ steps: - vllm/ - tests/models/multimodal commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 310c2a8fd2a..5effe17513d 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -37,6 +37,11 @@ steps: - pytest -v -s plugins_tests/test_stats_logger_plugins.py - pip uninstall dummy_stat_logger -y # end stat_logger plugins test + # begin endpoint plugins test + - pip install -e ./plugins/vllm_add_dummy_endpoint_plugin + - pytest -v -s plugins_tests/test_endpoint_plugins.py + - pip uninstall vllm_add_dummy_endpoint_plugin -y + # end endpoint plugins test # other tests continue here: - pytest -v -s plugins_tests/test_scheduler_plugins.py - pip install -e ./plugins/vllm_add_dummy_model diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index adb27c4a049..1dfe912aa89 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -15,24 +15,26 @@ steps: - tests/utils.py - tests/benchmarks/test_serve_cli.py - tests/entrypoints/openai/chat_completion/test_chat_completion.py - # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py + - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py # - tests/entrypoints/openai/completion/test_prompt_validation.py - tests/entrypoints/openai/completion/test_shutdown.py - # - tests/entrypoints/openai/test_return_token_ids.py - # - tests/entrypoints/openai/test_uds.py + - tests/entrypoints/openai/test_return_token_ids.py + - tests/entrypoints/openai/test_uds.py - tests/v1/sample/test_logprobs_e2e.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex" - # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid" + - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple" # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds" - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" - # - pytest -v -s entrypoints/openai/test_return_token_ids.py - # - pytest -v -s entrypoints/openai/test_uds.py + # test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and + # finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk. + - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison" + - pytest -v -s entrypoints/openai/test_uds.py - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server" - label: Rust Frontend Serve/Admin Coverage @@ -45,19 +47,24 @@ steps: - vllm/entrypoints/serve/ - vllm/v1/engine/ - tests/utils.py - # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py + - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py - tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py # - tests/entrypoints/serve/dev/test_sleep.py + - tests/entrypoints/serve/tokenize/test_tokenization.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py + # server_load can be flaky under the Rust frontend; keep it excluded for now. - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" # - pytest -v -s entrypoints/serve/dev/test_sleep.py + # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op). + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" - label: Rust Frontend Core Correctness timeout_in_minutes: 30 diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 8ca6fc22d64..57166d9d9b7 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -119,7 +119,7 @@ # Transformers modeling backend /vllm/model_executor/models/transformers @hmellor -/tests/models/test_transformers.py @hmellor +/tests/models/transformers @hmellor # Docs /docs/mkdocs @hmellor diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index edfb6179d85..143fc427a49 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -28,7 +28,8 @@ jobs: pull_number: context.payload.pull_request.number, }); - const hasReadyLabel = pr.labels.some(l => l.name === 'ready'); + const readyLabels = ['ready', 'ready-run-all-tests']; + const hasReadyLabel = pr.labels.some(l => readyLabels.includes(l.name)); const hasVerifiedLabel = pr.labels.some(l => l.name === 'verified'); const { data: mergedPRs } = await github.rest.search.issuesAndPullRequests({ @@ -40,7 +41,7 @@ jobs: if (hasReadyLabel || hasVerifiedLabel || mergedCount >= 4) { core.info(`Check passed: verified label=${hasVerifiedLabel}, ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`); } else { - core.setFailed(`PR must have the 'verified' or 'ready' (which also triggers tests) label or the author must have at least 4 merged PRs (found ${mergedCount}).`); + core.setFailed(`PR must have the 'verified', 'ready', or 'ready-run-all-tests' label (the ready labels also trigger tests) or the author must have at least 4 merged PRs (found ${mergedCount}).`); } pre-commit: diff --git a/AGENTS.md b/AGENTS.md index 241eab38818..a53b81873cf 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -29,6 +29,7 @@ Do not open one-off PRs for tiny edits (single typo, isolated style change, one - PR descriptions for AI-assisted work **must** include: - Why this is not duplicating an existing PR. - Test commands run and results. + - Model evaluation results when the change affects output, accuracy, or serving. - Clear statement that AI assistance was used. ### Fail-closed behavior @@ -66,23 +67,38 @@ VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto uv pip install -e . --torch-backend=auto ``` -### Running tests +### Tests > Requires [Environment setup](#environment-setup) and [Installing dependencies](#installing-dependencies). ```bash -# Install test dependencies. -# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the -# unpinned source file instead: -uv pip install -r requirements/test/cuda.in # resolves for current platform -# Or on x86_64: -uv pip install -r requirements/test/cuda.txt +# Install test dependencies (use cuda.in on non-x86_64): +uv pip install -r requirements/test/cuda.in -# Run a specific test file (use .venv/bin/python directly; -# `source activate` does not persist in non-interactive shells): +# Run a specific test file: .venv/bin/python -m pytest tests/path/to/test_file.py -v ``` +When adding tests: + +- **Design before you write.** Answer four questions first: what is the module + for, what is its I/O contract, what failure am I guarding against, and what is + the cheapest level that catches it (unit over integration over e2e)? +- **Reuse before create.** Extend existing test files, `conftest.py` fixtures, and + helpers; add a new file only when no nearby suite fits. +- **Test behavior with intent.** Assert observable outcomes through public APIs; + state why in the name or docstring. Skip trivial wiring; flaky tests are worse + than no tests. +- **Keep it minimal.** One behavior per test and the smallest setup that + triggers it; if the test diff dwarfs the code change, cut scope. +- **No one-off kernel benchmarks in `tests/`.** Put kernel perf work in + `benchmarks/kernels/`; prove correctness in existing pytest suites. +- **Run model evals for model-affecting changes.** Search `tests/evals/` or use + `vllm bench` and include results in the PR — do not wait for reviewers to ask. + +For model-specific requirements, see +[`docs/contributing/model/tests.md`](docs/contributing/model/tests.md). + ### Running linters > Requires [Environment setup](#environment-setup). @@ -107,23 +123,18 @@ Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#3 ### 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. +- Match existing code style +- Minimize use of comments. Eliminate comments which are redundant, preferring legible and self-documenting code. When used, keep docstrings and comments brief and direct. - Assume the reader is familiar with vLLM. ### Commit messages -Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example: +Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`): ```text Your commit message here -Co-authored-by: GitHub Copilot -Co-authored-by: Claude -Co-authored-by: gemini-code-assist +Co-authored-by: Agent Name Here Signed-off-by: Your Name ``` diff --git a/CMakeLists.txt b/CMakeLists.txt index 901f2be6bbb..48c0270e2c6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -400,7 +400,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/topk.cu" "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" "csrc/libtorch_stable/custom_all_reduce.cu" "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py index a69637bfc43..6349095ad72 100644 --- a/benchmarks/backend_request_func.py +++ b/benchmarks/backend_request_func.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field import aiohttp import huggingface_hub.constants from tqdm.asyncio import tqdm -from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast +from transformers import AutoTokenizer, PythonBackend, TokenizersBackend # NOTE(simon): do not import vLLM here so the benchmark script # can run without vLLM installed. @@ -609,7 +609,7 @@ def get_tokenizer( tokenizer_mode: str = "auto", trust_remote_code: bool = False, **kwargs, -) -> PreTrainedTokenizer | PreTrainedTokenizerFast: +) -> PythonBackend | TokenizersBackend: if pretrained_model_name_or_path is not None and not os.path.exists( pretrained_model_name_or_path ): diff --git a/benchmarks/benchmark_topk_topp.py b/benchmarks/benchmark_topk_topp.py index 27b6dd8d6be..f727f16ea29 100644 --- a/benchmarks/benchmark_topk_topp.py +++ b/benchmarks/benchmark_topk_topp.py @@ -132,8 +132,10 @@ def benchmark_function( reset_memory_stats() # Benchmark - start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)] - end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)] + start_events = [ + torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters) + ] + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)] for i in range(benchmark_iters): logits_copy = logits.clone() diff --git a/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py index 9e4f4157a8a..09a01b301be 100644 --- a/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py +++ b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py @@ -17,7 +17,7 @@ from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_m from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 compressed_tensors_moe_w4a16_flydsl, ) -from vllm.platforms import current_platform +from vllm.utils.platform_utils import get_device_name_as_file_name RoutingBuffers = tuple[ torch.Tensor, # sorted_token_ids @@ -259,7 +259,7 @@ def tune_flydsl_moe_w4a16( ) us_best = us tuned_config[str(num_tokens)] = tile_config - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() tuned_config_file_name = ( f"E={num_experts},N={inter_dim},device_name={device_name}," f"dtype=int4_w4a16,backend=flydsl.json" diff --git a/benchmarks/kernels/benchmark_moe_defaults.py b/benchmarks/kernels/benchmark_moe_defaults.py index 7f000e01137..f6ad59366dc 100644 --- a/benchmarks/kernels/benchmark_moe_defaults.py +++ b/benchmarks/kernels/benchmark_moe_defaults.py @@ -134,8 +134,8 @@ def benchmark_config( torch.accelerator.synchronize() # Benchmark - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(num_iters): with override_config(config): diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index 5a3a6e88a63..a8b73da2aa9 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -170,8 +170,8 @@ def benchmark_config( graph.replay() torch.accelerator.synchronize() - start = torch.Event(enable_timing=True) - end = torch.Event(enable_timing=True) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) latencies: list[float] = [] for _ in range(num_iters): start.record() diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py index 36dce1b6388..590d4cfdc6d 100644 --- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py +++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( from vllm.platforms import current_platform from vllm.triton_utils import triton from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.platform_utils import get_device_name_as_file_name mp.set_start_method("spawn", force=True) @@ -264,7 +265,7 @@ def save_configs( input_type="fp8", ) -> None: os.makedirs(save_path, exist_ok=True) - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() json_file_name = ( f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8," f"block_shape=[{block_n},{block_k}].json" diff --git a/build_vllm_ppc64le.sh b/build_vllm_ppc64le.sh new file mode 100644 index 00000000000..3c0b74cc74d --- /dev/null +++ b/build_vllm_ppc64le.sh @@ -0,0 +1,241 @@ +#!/bin/bash +set -eoux pipefail + +######################################## +# Resolve repo root (IMPORTANT) +######################################## +REPO_ROOT="$(pwd)" + +cd "$REPO_ROOT" + +######################################## +# DevPI configuration +######################################## + +IBM_DEVPI_URL=${IBM_DEVPI_URL:-"https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/"} +RHOAI_INDEX_URL=${RHOAI_INDEX_URL:-"https://console.redhat.com/api/pypi/public-rhai/rhoai/3.4/cpu-ubi9/simple/"} + +######################################## +# wheel dir +######################################## + +WHEEL_DIR=${WHEEL_DIR:-"/tmp/wheels"} +mkdir -p "$WHEEL_DIR" + +######################################## +# Helpers +######################################## +try_install_from_devpi() { + local pkg=$1 + uv pip install \ + --extra-index-url "${IBM_DEVPI_URL}" \ + --index-strategy unsafe-best-match \ + --no-build-isolation \ + "${pkg}" +} + +######################################## +# Package Versions +######################################## +cd "$REPO_ROOT" +TORCH_VERSION=${TORCH_VERSION:-$(grep -E '^torch==.+==\s*"ppc64le"' requirements/cpu.txt | grep -Eo '\b[0-9\.]+\b' || true)} +TORCH_VERSION=${TORCH_VERSION:-2.11.0} + +TORCHVISION_VERSION=${TORCHVISION_VERSION:-0.26.0} +TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION:-${TORCH_VERSION}} + +export TORCH_VERSION +export TORCHVISION_VERSION +export TORCHAUDIO_VERSION +export OPENCV_VERSION=${OPENCV_VERSION:-4.13.0.92} +export XGRAMMAR_VERSION=${XGRAMMAR_VERSION:-0.2.1} + +######################################## +# install system dependencies +######################################## + +rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm || true + +microdnf install -y \ + python3.12 python3.12-devel python3.12-pip gcc \ + git jq gcc-toolset-14 gcc-toolset-14-libatomic-devel \ + automake libtool clang-devel openssl-devel \ + harfbuzz-devel kmod lcms2-devel libimagequant-devel libjpeg-turbo-devel \ + llvm15-devel libraqm-devel libtiff-devel libwebp-devel libxcb-devel \ + ninja-build openjpeg2-devel pkgconfig \ + tcl-devel tk-devel xsimd-devel zeromq-devel zlib-devel patchelf file openblas openblas-devel protobuf numactl numactl-devel openmpi openmpi-devel + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-17.el9.ppc64le.rpm + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-devel-3.14.0-17.el9.ppc64le.rpm + +rpm -ivh --nodeps \ + https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-compiler-3.14.0-17.el9.ppc64le.rpm + +######################################## +# Python 3.12 virtual environment +######################################## + +python3.12 -m venv /opt/vllm +source /opt/vllm/bin/activate + +export PATH=/opt/vllm/bin:$PATH + +python --version + +######################################## +# install build tools (stable uv) +######################################## + +pip install -U pip setuptools-rust +pip install uv +pip install "setuptools<70" build wheel cmake auditwheel +uv pip install "setuptools<70" cython meson-python pybind11 "sympy>=1.13.3" --no-build-isolation + +######################################## +# Rust +######################################## + +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +source /root/.cargo/env + +######################################## +# Compiler env +######################################## + +source /opt/rh/gcc-toolset-14/enable + +export PATH=/usr/lib64/llvm15/bin:$PATH +export LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config +export CMAKE_ARGS="-DPython3_EXECUTABLE=python" + +export MAX_JOBS=${MAX_JOBS:-$(nproc)} +export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 + +######################################## +# Install Packages From Devpi +######################################## +uv pip install numpy==2.3.5 pillow==12.2.0 --extra-index-url "$IBM_DEVPI_URL" +try_install_from_devpi "opencv-python-headless==${OPENCV_VERSION}" +try_install_from_devpi "torch==${TORCH_VERSION}" +try_install_from_devpi "torchvision==${TORCHVISION_VERSION}" + +######################################## +# torch audio +######################################## + +TEMP_BUILD_DIR=$(mktemp -d) +cd "${TEMP_BUILD_DIR}" +export BUILD_SOX=1 BUILD_KALDI=1 BUILD_RNNT=1 USE_FFMPEG=0 USE_ROCM=0 USE_CUDA=0 +export TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 +git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} +cd audio +#patching +sed -i ' +s|_CSRC_DIR / "_torchaudio.cpp"|str(_CSRC_DIR / "_torchaudio.cpp")|; +s|_CSRC_DIR / "utils.cpp"|str(_CSRC_DIR / "utils.cpp")|; +s|sources=\[_CSRC_DIR / s for s in sources\]|sources=[str(_CSRC_DIR / s) for s in sources]|; +' tools/setup_helpers/extension.py +MAX_JOBS=${MAX_JOBS:-$(nproc)} \ +BUILD_VERSION=${TORCHAUDIO_VERSION} \ +uv build --wheel --out-dir "${WHEEL_DIR}" --no-build-isolation +uv pip install "${WHEEL_DIR}"/torchaudio*.whl +cd "${REPO_ROOT}" +rm -rf "${TEMP_BUILD_DIR}" + +######################################## +# Xgrammar +######################################## +uv pip install \ + "scikit-build-core==0.11.6" \ + "pyproject-metadata<0.8" \ + pathspec \ + packaging \ + distro \ + "setuptools<70" \ + setuptools_scm \ + cmake \ + ninja \ + pybind11 \ + nanobind +uv pip install apache-tvm-ffi==0.1.12 \ + --no-build-isolation \ + --no-cache + +TEMP_BUILD_DIR=$(mktemp -d) + +pushd "${TEMP_BUILD_DIR}" + +export CFLAGS="-fno-lto -mcpu=power9" +export CXXFLAGS="-fno-lto -mcpu=power9" +export LDFLAGS="-fno-lto" +export PATH=/opt/vllm/bin:$PATH + +export Python_EXECUTABLE=/opt/vllm/bin/python3 +export Python3_EXECUTABLE=/opt/vllm/bin/python3 +export PYTHON_EXECUTABLE=/opt/vllm/bin/python3 + +export Python_ROOT_DIR=/opt/vllm +export Python3_ROOT_DIR=/opt/vllm + +git clone \ + --recursive \ + https://github.com/mlc-ai/xgrammar \ + -b "v${XGRAMMAR_VERSION}" + +cd xgrammar + +cp cmake/config.cmake . +export PYTHONPATH=/opt/vllm/lib64/python3.12/site-packages:/opt/vllm/lib/python3.12/site-packages:${PYTHONPATH:-} + +uv build \ + --wheel \ + --out-dir "${WHEEL_DIR}" \ + --no-build-isolation + +uv pip install "${WHEEL_DIR}"/xgrammar*.whl -v + +popd + +rm -rf "${TEMP_BUILD_DIR}" +cd "${REPO_ROOT}" + +######################################## +# RHOAI Binary Downloads +######################################## +pip download \ + --index-url "${RHOAI_INDEX_URL}" \ + --only-binary=:all: \ + --no-deps \ + llvmlite==0.47.0 \ + -d "${WHEEL_DIR}" + +pip download \ + --index-url "${RHOAI_INDEX_URL}" \ + --only-binary=:all: \ + --no-deps \ + Numba==0.65.0 \ + -d "${WHEEL_DIR}" + +######################################## +# install built wheels +######################################## +uv pip install setuptools_scm maturin setuptools-rust ninja scikit-build-core pybind11 nanobind \ + --no-build-isolation +uv pip install "${WHEEL_DIR}"/*.whl + +######################################## +# install remaining deps +######################################## + +sed -i.bak -e 's/.*torch.*//g' pyproject.toml requirements/*.txt + +uv pip install "setuptools>=78.1.1" --no-build-isolation + +export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/local/lib64/pkgconfig:/usr/lib64/pkgconfig + +uv pip install -r requirements/common.txt \ + -r requirements/cpu.txt \ + -r requirements/build/cpu.txt --index-strategy unsafe-best-match diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 177c420776f..3aca9bcea91 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -15,6 +15,7 @@ endif() # set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86}) set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16}) +set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16}) include_directories("${CMAKE_SOURCE_DIR}/csrc") @@ -110,6 +111,13 @@ else() set(ARM_BF16_FOUND ON) message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable") endif() + # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin + # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1 + # overrides the detection result. + if (ENABLE_RVV_BF16) + set(RVV_BF16_FOUND ON) + message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable") + endif() endif() if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) @@ -178,7 +186,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV. if(NOT DEFINED VLLM_RVV_VLEN) # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. - if(EXISTS /proc/cpuinfo) + # Skip when cross-compiling — /proc/cpuinfo describes the build host. + if(CMAKE_CROSSCOMPILING) + message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo") + elseif(EXISTS /proc/cpuinfo) file(READ /proc/cpuinfo _cpuinfo) set(_best 0) foreach(_n IN ITEMS 128 256 512 1024) @@ -186,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") set(_best ${_n}) endif() endforeach() + # Only VLEN=128 and VLEN=256 are supported by the RVV kernels. + if(_best GREATER 256) + message(WARNING + "Detected VLEN=${_best} but only 128/256 are supported; " + "clamping to 256") + set(_best 256) + endif() if(_best GREATER 0) set(VLLM_RVV_VLEN ${_best}) endif() @@ -195,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) message(FATAL_ERROR "RISC-V RVV is available but VLEN could not be auto-detected. " - "Please specify VLEN explicitly:\n" - " -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n" - " -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)") + "Please specify VLEN explicitly via CMAKE_ARGS:\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n" + " CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)") endif() endif() if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0) @@ -209,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "BF16 extension detected") set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) elseif(RVV_FP16_FOUND) - message(WARNING "BF16 functionality is not available") + message(WARNING "BF16 functionality is not available.") set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d) else() message(STATUS "compile riscv with scalar (no FP16/BF16)") diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 3897f8e1b3f..052966b2755 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5 + GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73 GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index ec1a2b162de..fa22861157e 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( bool cpu_attn_has_isa(const std::string& isa) { if (isa == "rvv") { -#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128 +#if defined(__riscv) && defined(__riscv_v_min_vlen) && \ + (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) return true; #else return false; diff --git a/csrc/cpu/cpu_attn_vsx.hpp b/csrc/cpu/cpu_attn_vsx.hpp index dd24b95ba02..562a5312571 100644 --- a/csrc/cpu/cpu_attn_vsx.hpp +++ b/csrc/cpu/cpu_attn_vsx.hpp @@ -323,8 +323,6 @@ class AttentionImpl { const int64_t num_blocks_stride, const int64_t cache_head_num_stride, const int64_t block_size, const int64_t block_size_stride, const float k_inv = 0.0f, const float v_inv = 0.0f) { - // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on - // PowerPC. The parameters are present to match the common interface. #pragma omp parallel for collapse(2) for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { diff --git a/csrc/cpu/cpu_types.hpp b/csrc/cpu/cpu_types.hpp index 7b2c3d3b74c..50ef5e10b83 100644 --- a/csrc/cpu/cpu_types.hpp +++ b/csrc/cpu/cpu_types.hpp @@ -4,7 +4,7 @@ #if defined(__x86_64__) // x86 implementation #include "cpu_types_x86.hpp" -#elif defined(__POWER9_VECTOR__) +#elif defined(__powerpc__) // ppc implementation #include "cpu_types_vsx.hpp" #elif defined(__s390x__) @@ -41,4 +41,4 @@ inline int get_max_threads() { } } // namespace cpu_utils -#endif \ No newline at end of file +#endif diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index d0ce67a5afe..70cb0ab52de 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec { explicit BF16Vec32(const BF16Vec8& v) { fixed_u16x8_t u16_val = bf16_to_u16(v.reg); - fixed_u16x32_t u16_combined = - RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)( - u16_val, u16_val, u16_val, u16_val); - reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, - LMUL_512)(u16_combined); + // Widen LMUL_128 → LMUL_256 so vslideup operands share a type. + // At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2. + fixed_u16x16_t ext = + RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val); + // Build 16-element half: place the 8 elements at offsets 0 and 8. + fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8); + half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16); + // Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128). + fixed_u16x32_t dst = + RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half); + reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst); }; void save(void* ptr) const { @@ -623,17 +630,29 @@ struct FP32Vec16 : public Vec { data.reg, data.reg)) {}; explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { - const uint64_t q_values = static_cast(value); - auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM); - auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM); - auto shifts = - RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM); - auto shifted = - RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM); - auto idx64 = - RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM); - auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM); - reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM); + // Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on + // VLEN=128 / m4 on VLEN=256), which causes heavy register spilling. + constexpr int HALF = VEC_ELEM_NUM / 2; + const auto q = static_cast(value); + const uint32_t lo = static_cast(q); + const uint32_t hi = static_cast(q >> 32); + + auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF); + auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF); + + auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF); + auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF, + HALF); + + auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF); + auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)( + RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF, + HALF); + + auto idx = + RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi); + reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM); } explicit FP32Vec16(const FP16Vec16& v); diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 2031e4c14a8..250c870dbe4 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -344,53 +344,133 @@ struct FP32Vec8 : public Vec { return result; } - FP32Vec8 exp() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::exp(ar.values[0]); - ret.val[0][1] = std::exp(ar.values[1]); - ret.val[0][2] = std::exp(ar.values[2]); - ret.val[0][3] = std::exp(ar.values[3]); - ret.val[1][0] = std::exp(ar.values[4]); - ret.val[1][1] = std::exp(ar.values[5]); - ret.val[1][2] = std::exp(ar.values[6]); - ret.val[1][3] = std::exp(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + f32x4x2_t out; + const __vector float log2e = vec_splats(1.44269504088896341f); + const __vector float one = vec_splats(1.0f); + const __vector float min_x = vec_splats(-87.3f); + const __vector float max_x = vec_splats(88.7f); + + // 5th-degree minimax polynomial for 2^r (r in [0,1)) + const __vector float c1 = vec_splats(0.6931471805599453f); + const __vector float c2 = vec_splats(0.240226506959101f); + const __vector float c3 = vec_splats(0.05550410866482158f); + const __vector float c4 = vec_splats(0.009618129107628477f); + const __vector float c5 = vec_splats(0.0013333558146428443f); + + for (int i = 0; i < 2; i++) { + __vector float x = reg.val[i]; + x = vec_max(x, min_x); + x = vec_min(x, max_x); + + __vector float y = vec_mul(x, log2e); + + __vector float kf = vec_floor(y); + __vector float r = vec_sub(y, kf); + + // Convert float to signed integer. Use vec_cts for PowerPC AltiVec + // compatibility. + __vector signed int k = vec_cts(kf, 0); + const __vector signed int min_k = vec_splats((signed int)-126); + const __vector signed int max_k = vec_splats((signed int)127); + k = vec_min(vec_max(k, min_k), max_k); + + // Build 2^k from exponent bits + __vector signed int exp_int = vec_add(k, vec_splats((signed int)127)); + __vector unsigned int bits = (__vector unsigned int)exp_int; + bits = vec_sl(bits, vec_splats((unsigned int)23)); + __vector float pow2k = (__vector float)bits; + + // Improved minimax polynomial + __vector float poly = vec_madd(c5, r, c4); + poly = vec_madd(poly, r, c3); + poly = vec_madd(poly, r, c2); + poly = vec_madd(poly, r, c1); + poly = vec_madd(poly, r, one); + + out.val[i] = vec_mul(pow2k, poly); + } + return FP32Vec8(out); } FP32Vec8 tanh() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::tanh(ar.values[0]); - ret.val[0][1] = std::tanh(ar.values[1]); - ret.val[0][2] = std::tanh(ar.values[2]); - ret.val[0][3] = std::tanh(ar.values[3]); - ret.val[1][0] = std::tanh(ar.values[4]); - ret.val[1][1] = std::tanh(ar.values[5]); - ret.val[1][2] = std::tanh(ar.values[6]); - ret.val[1][3] = std::tanh(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + const __vector float one = vec_splats(1.0f); + const __vector float two = vec_splats(2.0f); + const __vector float zero = vec_splats(0.0f); + const __vector float sat = vec_splats(9.0f); + + f32x4x2_t out; + + for (int i = 0; i < 2; i++) { + __vector float x = reg.val[i]; + __vector float ax = vec_abs(x); + + __vector bool int mask = vec_cmpge(x, zero); + __vector float sign = vec_sel(vec_splats(-1.0f), one, mask); + + __vector bool int saturated = vec_cmpge(ax, sat); + + __vector float two_x = vec_mul(x, two); + f32x4x2_t tmp; + tmp.val[0] = two_x; + tmp.val[1] = two_x; + FP32Vec8 temp_vec(tmp); + vector float e = temp_vec.exp().reg.val[0]; + + vector float num = vec_sub(e, one); + vector float den = vec_add(e, one); + vector float t = vec_div(num, den); + + out.val[i] = vec_sel(t, sign, saturated); + } + return FP32Vec8(out); } FP32Vec8 er() const { - // TODO: Vectorize this - AliasReg ar; - ar.reg = reg; - f32x4x4_t ret; - ret.val[0][0] = std::erf(ar.values[0]); - ret.val[0][1] = std::erf(ar.values[1]); - ret.val[0][2] = std::erf(ar.values[2]); - ret.val[0][3] = std::erf(ar.values[3]); - ret.val[1][0] = std::erf(ar.values[4]); - ret.val[1][1] = std::erf(ar.values[5]); - ret.val[1][2] = std::erf(ar.values[6]); - ret.val[1][3] = std::erf(ar.values[7]); - return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]})); + const vector float a1 = vec_splats(0.254829592f); + const vector float a2 = vec_splats(-0.284496736f); + const vector float a3 = vec_splats(1.421413741f); + const vector float a4 = vec_splats(-1.453152027f); + const vector float a5 = vec_splats(1.061405429f); + const vector float p = vec_splats(0.3275911f); + const vector float one = vec_splats(1.0f); + const vector float zero = vec_splats(0.0f); + const vector float sat = vec_splats(6.0f); + + f32x4x2_t ret; + + for (int i = 0; i < 2; i++) { + vector float x = reg.val[i]; + vector float ax = vec_abs(x); + + vector bool int mask = vec_cmpge(x, zero); + vector float sign = vec_sel(vec_splats(-1.0f), one, mask); + + vector bool int saturated = vec_cmpge(ax, sat); + + vector float t = vec_div(one, vec_madd(p, ax, one)); + + vector float poly = a5; + poly = vec_madd(poly, t, a4); + poly = vec_madd(poly, t, a3); + poly = vec_madd(poly, t, a2); + poly = vec_madd(poly, t, a1); + poly = vec_mul(poly, t); + + vector float x_squared = vec_mul(x, x); + vector float neg_x_squared = vec_mul(vec_splats(-1.0f), x_squared); + f32x4x2_t tmp; + tmp.val[0] = neg_x_squared; + tmp.val[1] = neg_x_squared; + FP32Vec8 exp_input(tmp); + vector float exp_term = exp_input.exp().reg.val[0]; + + vector float y = vec_nmsub(poly, exp_term, one); + vector float erf_val = vec_mul(sign, y); + + ret.val[i] = vec_sel(erf_val, sign, saturated); + } + return FP32Vec8(ret); } FP32Vec8 operator*(const FP32Vec8& b) const { diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index e17c9ab3a7e..cfa296e73b6 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "dynamic_4bit_int_moe(" "Tensor x, Tensor topk_ids, Tensor topk_weights," - "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2," + "Tensor w13_packed, Tensor w2_packed," + "int hidden_size, int intermediate_size," "int group_size, bool apply_router_weight_on_input, int activation_kind" ") -> Tensor"); diff --git a/csrc/cpu/utils.hpp b/csrc/cpu/utils.hpp index ec10a0f3524..78ee7081b24 100644 --- a/csrc/cpu/utils.hpp +++ b/csrc/cpu/utils.hpp @@ -76,14 +76,14 @@ inline int64_t get_available_l2_size() { if (l2_cache_size == 0) { l2_cache_size = 256 * 1024; } - return static_cast(l2_cache_size) >> 1; // use 50% of L2 cache + return static_cast(l2_cache_size) >> 1; }(); return size; #else static int64_t size = []() { auto caps = at::cpu::get_cpu_capabilities(); const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt(); - return l2_cache_size >> 1; // use 50% of L2 cache + return l2_cache_size >> 1; }(); return size; #endif diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 43cbb7f86d3..b60d2d548f5 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -15,7 +15,8 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, torch::stable::Tensor& gating_output, bool renormalize, - std::optional bias); + std::optional bias, + double routed_scaling_factor); void topk_softplus_sqrt( torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 9f2e3640e23..b4bcd9479e9 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -173,7 +173,8 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int start_expert, const int end_expert, const bool renormalize, - const float* bias) + const float* bias, + const double routed_scaling_factor) { using cub_kvp = cub::KeyValuePair; @@ -241,14 +242,16 @@ __launch_bounds__(TPB) __global__ void moeTopK( __syncthreads(); } - // Renormalize the k weights for this row to sum to 1, if requested. - if (renormalize) { - if (threadIdx.x == 0) { + // Apply renormalization and routed scaling factor to final weights. + if (threadIdx.x == 0) { + float scale = static_cast(routed_scaling_factor); + if (renormalize) { const float denom = selected_sum > 0.f ? selected_sum : 1.f; - for (int k_idx = 0; k_idx < k; ++k_idx) { - const int idx = k * block_row + k_idx; - output[idx] = output[idx] / denom; - } + scale /= denom; + } + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] = output[idx] * scale; } } } @@ -274,7 +277,7 @@ template || std::is_same_v || std::is_same_v, @@ -570,17 +573,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ } } - // Renormalize the k weights for this row to sum to 1, if requested. - if (renormalize) { - if (thread_group_idx == 0) - { - const float denom = selected_sum > 0.f ? selected_sum : 1.f; - for (int k_idx = 0; k_idx < k; ++k_idx) - { - const int idx = k * thread_row + k_idx; - output[idx] = output[idx] / denom; - } - } + // Apply renormalization and routed scaling factor to final weights. + if (thread_group_idx == 0) { + float scale = static_cast(routed_scaling_factor); + if (renormalize) { + const float denom = selected_sum > 0.f ? selected_sum : 1.f; + scale /= denom; + } + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * thread_row + k_idx; + output[idx] = output[idx] * scale; + } } } @@ -602,7 +605,7 @@ struct TopkConstants template void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, - const float* bias, cudaStream_t stream) + const float* bias, const double routed_scaling_factor, cudaStream_t stream) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = detail::TopkConstants; @@ -613,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); topkGating<<>>( - input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias); + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor); } #ifndef USE_ROCM @@ -624,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); + bias, routed_scaling_factor, stream); #else #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -632,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); \ + bias, routed_scaling_factor, stream); \ } else if (WARP_SIZE == 32) { \ topkGatingLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, stream); \ + bias, routed_scaling_factor, stream); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -657,6 +660,7 @@ void topkGatingKernelLauncher( const int topk, const bool renormalize, const float* bias, + const double routed_scaling_factor, cudaStream_t stream) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; @@ -732,7 +736,7 @@ void topkGatingKernelLauncher( } moeTopK<<>>( workspace, nullptr, topk_weights, topk_indices, token_expert_indices, - num_experts, topk, 0, num_experts, renormalize, bias); + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor); } } } @@ -750,6 +754,7 @@ void dispatch_topk_launch( torch::stable::Tensor& softmax_workspace, int num_tokens, int num_experts, int topk, bool renormalize, std::optional bias, + double routed_scaling_factor, cudaStream_t stream) { const float* bias_ptr = nullptr; @@ -772,7 +777,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); + bias_ptr, routed_scaling_factor, stream); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { vllm::moe::topkGatingKernelLauncher( reinterpret_cast(gating_output.const_data_ptr()), @@ -781,7 +786,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); + bias_ptr, routed_scaling_factor, stream); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); vllm::moe::topkGatingKernelLauncher( @@ -791,7 +796,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, stream); + bias_ptr, routed_scaling_factor, stream); } } @@ -820,15 +825,15 @@ void topk_softmax( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, 1.0, stream); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, 1.0, stream); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, 1.0, stream); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } @@ -840,7 +845,8 @@ void topk_sigmoid( torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, - std::optional bias) + std::optional bias, + double routed_scaling_factor) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -859,15 +865,15 @@ void topk_sigmoid( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, routed_scaling_factor, stream); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, routed_scaling_factor, stream); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, stream); + bias, routed_scaling_factor, stream); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index bfcb0074e5b..ba5b9b896f1 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -13,8 +13,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { // Apply topk sigmoid to the gating outputs. m.def( "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! " - "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? " - "bias) -> ()"); + "token_expert_indices, Tensor gating_output, bool renormalize, " + "Tensor? bias, float routed_scaling_factor) -> ()"); m.def( "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " diff --git a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp index 58dc4020168..1b071d334ff 100644 --- a/csrc/moe/dynamic_4bit_int_moe_cpu.cpp +++ b/csrc/moe/dynamic_4bit_int_moe_cpu.cpp @@ -29,25 +29,37 @@ enum ActivationKind : int64_t { torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind) { + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind) { TORCH_CHECK(x.dim() == 2, "x must be 2D"); TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2, "topk tensors must be [T, K]"); TORCH_CHECK( w13_packed.size(0) == w2_packed.size(0), "w13_packed and w2_packed must have same number of experts in dim 0"); - TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I"); const int64_t T = x.size(0); const int64_t K = topk_ids.size(1); const int64_t E = w13_packed.size(0); const int64_t N = T * K; + const int64_t w13_out_features = 2 * intermediate_size; auto x_c = x.contiguous(); + // _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation + // dtypes: + // - fp32: with channelwise and groupwise + // - bf16: with channelwise -> upcast to fp32 for groupwise + // - fp16: not supported -> upcast to fp32 for groupwise & channelwise + const auto output_dtype = x_c.scalar_type(); + const bool should_cast_input = + ((group_size != -1) && output_dtype == at::kBFloat16) || + output_dtype == at::kHalf; + if (should_cast_input) { + x_c = x_c.to(at::kFloat); + } auto ids_c = topk_ids.contiguous(); - auto gates_c = topk_weights.to(at::kFloat).contiguous(); + auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous(); // bucketing tokens -> experts c10::SmallVector counts( @@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu( c10::SmallVector offsets(E + 1, 0); // ( E +1 ) for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e]; + // expert_tokens = [tokens indices for expert 0, ...] + // expert_gates = [router weights for tokens assigned to expert 0, ...] auto expert_tokens = at::empty({offsets[E]}, ids_c.options()); auto expert_gates = at::empty({offsets[E]}, gates_c.options()); { c10::SmallVector cursor(E, 0); - const auto* ids_ptr = ids_c.data_ptr(); - const auto* gts_ptr = gates_c.data_ptr(); - auto* tok_ptr = expert_tokens.data_ptr(); - auto* gate_ptr = expert_gates.data_ptr(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(), + "bucket_expert_tokens_and_gates", [&] { + const auto* ids_ptr = ids_c.data_ptr(); + const auto* gts_ptr = gates_c.data_ptr(); + auto* tok_ptr = expert_tokens.data_ptr(); + auto* gate_ptr = expert_gates.data_ptr(); - for (int64_t t = 0; t < T; ++t) { - const int64_t base = t * K; - for (int64_t k = 0; k < K; ++k) { - const int64_t idx = base + k; - const int64_t e = ids_ptr[idx]; - const int64_t p = offsets[e] + (cursor[e]++); - tok_ptr[p] = t; - gate_ptr[p] = gts_ptr[idx]; - } - } + for (int64_t t = 0; t < T; ++t) { + const int64_t base = t * K; + for (int64_t k = 0; k < K; ++k) { + const int64_t idx = base + k; + const int64_t e = ids_ptr[idx]; + const int64_t p = offsets[e] + (cursor[e]++); + tok_ptr[p] = t; + gate_ptr[p] = gts_ptr[idx]; + } + } + }); } - const int64_t g_eff_13 = (group_size != -1) ? group_size : H; - const int64_t g_eff_2 = (group_size != -1) ? group_size : I; + const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size; + const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size; + // X_all [num_tokens * K, hidden_size] auto X_all = x_c.index_select(/*dim=*/0, expert_tokens); if (apply_router_weight_on_input) { X_all = X_all.mul(expert_gates.unsqueeze(1)); } - auto Y_all = at::empty({offsets[E], H}, x_c.options()); + auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options()); at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) { c10::InferenceMode guard; @@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu( auto w2_e = w2_packed.select(/*dim=*/0, e); // W13 - auto y13 = - mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2); + auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size, + /*out_features=*/w13_out_features); - auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I); - auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I); + auto g_part = + y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size); + auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size, + /*length=*/intermediate_size); torch::Tensor act; if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI @@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu( } // W2 - auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H); + auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size, + /*out_features=*/hidden_size); // Store per-expert result Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y); @@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu( if (!apply_router_weight_on_input) { Y_all = Y_all.mul(expert_gates.unsqueeze(1)); } + if (Y_all.scalar_type() != output_dtype) { + Y_all = Y_all.to(output_dtype); + } - auto out = at::zeros({T, H}, x.options()); + auto out = at::zeros({T, hidden_size}, x.options()); out = at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all); diff --git a/csrc/ops.h b/csrc/ops.h index cd18b1e5e0d..274cd52bea4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, - torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, - int64_t I2, int64_t group_size, bool apply_router_weight_on_input, - int64_t activation_kind); + torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size, + int64_t intermediate_size, int64_t group_size, + bool apply_router_weight_on_input, int64_t activation_kind); using fptr_t = int64_t; #ifdef USE_ROCM diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index ee5d5daf649..a528ffbd8d1 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common WORKDIR /workspace ARG PYTHON_VERSION=3.12 -ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" ARG max_jobs=32 ENV MAX_JOBS=${max_jobs} @@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH" ENV UV_HTTP_TIMEOUT=500 # Install Python dependencies -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" @@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --upgrade pip && \ - uv pip install -r requirements/cpu.txt + uv pip install -r requirements/cpu.txt --torch-backend cpu ARG TARGETARCH ENV TARGETARCH=${TARGETARCH} @@ -149,7 +146,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \ COPY requirements/build/cpu.txt requirements/build/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/build/cpu.txt + uv pip install -r requirements/build/cpu.txt --torch-backend cpu COPY . . @@ -205,7 +202,7 @@ RUN case "$(uname -m)" in \ esac RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/test/cpu.txt + uv pip install -r requirements/test/cpu.txt --torch-backend cpu ######################### DEV IMAGE ######################### FROM vllm-build AS vllm-dev @@ -231,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements/lint.txt && \ - uv pip install -r requirements/test/cpu.txt && \ + uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \ pre-commit install --hook-type pre-commit --hook-type commit-msg ENTRYPOINT ["bash"] diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 845d900c39c..f0363d43be2 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -1,275 +1,80 @@ +# Base UBI image ARG BASE_UBI_IMAGE_TAG=9.6-1754584681 ############################################################### -# Stage to build openblas +# BUILDER STAGE # ############################################################### -FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS openblas-builder - -ARG MAX_JOBS -ARG OPENBLAS_VERSION=0.3.30 -RUN microdnf install -y dnf && dnf install -y gcc-toolset-14 make wget unzip \ - && source /opt/rh/gcc-toolset-14/enable \ - && wget https://github.com/OpenMathLib/OpenBLAS/releases/download/v$OPENBLAS_VERSION/OpenBLAS-$OPENBLAS_VERSION.zip \ - && unzip OpenBLAS-$OPENBLAS_VERSION.zip \ - && cd OpenBLAS-$OPENBLAS_VERSION \ - && make -j${MAX_JOBS} TARGET=POWER9 BINARY=64 USE_OPENMP=1 USE_THREAD=1 NUM_THREADS=120 DYNAMIC_ARCH=1 INTERFACE64=0 \ - && cd /tmp && touch control - - -############################################################### -# base stage with dependencies coming from centos mirrors -############################################################### -FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS centos-deps-builder -RUN microdnf install -y dnf && \ - dnf install -y https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-gpg-keys-9.0-26.el9.noarch.rpm \ - https://mirror.stream.centos.org/9-stream/BaseOS/`arch`/os/Packages/centos-stream-repos-9.0-26.el9.noarch.rpm \ - https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ - dnf config-manager --set-enabled crb - -RUN dnf install -y openjpeg2-devel lcms2-devel tcl-devel tk-devel fribidi-devel yajl-devel && \ - dnf remove -y centos-gpg-keys-9.0-24.el9.noarch centos-stream-repos-9.0-26.el9.noarch - - -############################################################### -# base stage with basic dependencies -############################################################### - -FROM centos-deps-builder AS base-builder +FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS builder-base +ARG VLLM_VERSION="0.22.1" ARG PYTHON_VERSION=3.12 -ARG OPENBLAS_VERSION=0.3.30 - -# Set Environment Variables for venv, cargo & openblas -ENV VIRTUAL_ENV=/opt/vllm -ENV PATH=${VIRTUAL_ENV}/bin:/root/.cargo/bin:$PATH -ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig/ -ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib64:/usr/local/lib:/usr/lib64:/usr/lib -ENV UV_LINK_MODE=copy - -# install gcc-13, python, rust, openblas -# Note: A symlink for libatomic.so is created for gcc-13 (linker fails to find libatomic otherwise - reqd. for sentencepiece) -# Note: A dummy file 'control' is created in /tmp/ to artificially create dependencies between stages when building stages in parallel -# when `--jobs=` is passed with podman build command - -COPY --from=openblas-builder /tmp/control /dev/null - -RUN --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \ - dnf install -y openssl-devel \ - && dnf install -y \ - git tar gcc-toolset-14 automake libtool \ - pkgconfig xsimd zeromq-devel kmod findutils protobuf* \ - libtiff-devel libjpeg-devel zlib-devel freetype-devel libwebp-devel \ - harfbuzz-devel libraqm-devel libimagequant-devel libxcb-devel \ - python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip clang-devel \ - && dnf clean all \ - && PREFIX=/usr/local make -C /openblas install \ - && ln -sf /usr/lib64/libatomic.so.1 /usr/lib64/libatomic.so \ - && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ - && python -m pip install -U pip uv \ - && uv pip install wheel build "setuptools<70" setuptools_scm setuptools_rust meson-python 'cmake<4' ninja cython scikit_build_core scikit_build \ - && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ - && cd /tmp && touch control - - -############################################################### -# Stage to build torch family -############################################################### - -FROM base-builder AS torch-builder - -ARG MAX_JOBS -ARG TORCH_VERSION=2.7.0 -ARG _GLIBCXX_USE_CXX11_ABI=1 -ARG OPENBLAS_VERSION=0.3.30 - -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/pytorch.git -b v${TORCH_VERSION} && \ - cd pytorch && \ - uv pip install -r requirements.txt && \ - python setup.py develop && \ - rm -f dist/torch*+git*whl && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir /torchwheels/ - -ARG TORCHVISION_VERSION=0.22.0 -ARG TORCHVISION_USE_NVJPEG=0 -ARG TORCHVISION_USE_FFMPEG=0 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/vision.git -b v${TORCHVISION_VERSION} && \ - cd vision && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - BUILD_VERSION=${TORCHVISION_VERSION} \ - uv build --wheel --out-dir /torchwheels/ --no-build-isolation - -ARG TORCHAUDIO_VERSION=2.7.0 -ARG BUILD_SOX=1 -ARG BUILD_KALDI=1 -ARG BUILD_RNNT=1 -ARG USE_FFMPEG=0 -ARG USE_ROCM=0 -ARG USE_CUDA=0 -ARG TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} && \ - cd audio && \ - MAX_JOBS=${MAX_JOBS:-$(nproc)} \ - BUILD_VERSION=${TORCHAUDIO_VERSION} \ - uv build --wheel --out-dir /torchwheels/ --no-build-isolation - -############################################################### -# Stage to build pyarrow -############################################################### - -FROM base-builder AS arrow-builder - -ARG MAX_JOBS -ARG PYARROW_PARALLEL -ARG PYARROW_VERSION=21.0.0 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/apache/arrow.git -b apache-arrow-${PYARROW_VERSION} && \ - cd arrow/cpp && \ - mkdir build && cd build && \ - cmake -DCMAKE_BUILD_TYPE=release \ - -DCMAKE_INSTALL_PREFIX=/usr/local \ - -DARROW_PYTHON=ON \ - -DARROW_BUILD_TESTS=OFF \ - -DARROW_JEMALLOC=ON \ - -DARROW_BUILD_STATIC="OFF" \ - -DARROW_PARQUET=ON \ - .. && \ - make install -j ${MAX_JOBS:-$(nproc)} && \ - cd ../../python/ && \ - uv pip install -v -r requirements-build.txt && uv pip install numpy==2.1.3 && \ - PYARROW_PARALLEL=${PYARROW_PARALLEL:-$(nproc)} \ - python setup.py build_ext \ - --build-type=release --bundle-arrow-cpp \ - bdist_wheel --dist-dir /arrowwheels/ - -############################################################### -# Stage to build opencv -############################################################### - -FROM base-builder AS cv-builder - -ARG MAX_JOBS -ARG OPENCV_VERSION=86 -# patch for version 4.11.0.86 -ARG OPENCV_PATCH=97f3f39 -ARG ENABLE_HEADLESS=1 -RUN --mount=type=cache,target=/root/.cache/uv \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone --recursive https://github.com/opencv/opencv-python.git -b ${OPENCV_VERSION} && \ - cd opencv-python && \ - sed -i -E -e 's/"setuptools.+",/"setuptools",/g' pyproject.toml && \ - cd opencv && git cherry-pick --no-commit $OPENCV_PATCH && cd .. && \ - uv pip install scikit-build && \ - python -m build --wheel --installer=uv --outdir /opencvwheels/ - -############################################################### -# Stage to build numactl -############################################################### - -FROM base-builder AS numa-builder - -# Note: Building numactl with gcc-11. Compiling with gcc-13 in this builder stage will -# trigger recompilation with gcc-11 (and require libtool) in the final stage where we do not have gcc-13 -ARG MAX_JOBS -ARG NUMACTL_VERSION=2.0.19 -RUN git clone --recursive https://github.com/numactl/numactl.git -b v${NUMACTL_VERSION} \ - && cd numactl \ - && autoreconf -i && ./configure \ - && make -j ${MAX_JOBS:-$(nproc)} - - -############################################################### -# Stage to build numba -############################################################### - -FROM base-builder AS numba-builder - -ARG MAX_JOBS -ARG NUMBA_VERSION=0.61.2 - -# Clone all required dependencies -RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-14/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \ - git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \ - cd ./numba && \ - if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \ - sed -i '/#include "internal\/pycore_atomic.h"/i\#include "dynamic_annotations.h"' numba/_dispatcher.cpp; \ - fi && python -m build --wheel --installer=uv --outdir /numbawheels/ - -############################################################### -# Stage to build vllm - this stage builds and installs -# vllm, tensorizer and vllm-tgis-adapter and builds uv cache -# for transitive dependencies - eg. grpcio -############################################################### - -FROM base-builder AS vllmcache-builder - -ENV LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config -ENV PATH=/usr/lib64/llvm15/bin:$PATH - -COPY --from=torch-builder /tmp/control /dev/null -COPY --from=arrow-builder /tmp/control /dev/null -COPY --from=cv-builder /tmp/control /dev/null -COPY --from=numa-builder /tmp/control /dev/null -COPY --from=numba-builder /tmp/control /dev/null - ARG VLLM_TARGET_DEVICE=cpu -ARG GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 -# this step installs vllm and populates uv cache -# with all the transitive dependencies +USER root +WORKDIR /root + +ENV HOME=/root \ + WHEEL_DIR=/wheelsdir \ + VIRTUAL_ENV=/opt/vllm \ + GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1 \ + CARGO_HOME=/root/.cargo \ + RUSTUP_HOME=/root/.rustup \ + UV_CACHE_DIR=$HOME/.cache/uv \ + PATH=/root/.cargo/bin:/root/.rustup/bin:${VIRTUAL_ENV}/bin:$PATH + +RUN echo "DEBUG: VLLM_VERSION=${VLLM_VERSION}" +RUN --mount=type=cache,target=/var/cache/dnf \ + microdnf install -y \ + python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip \ + && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ + && python${PYTHON_VERSION} -m pip install -U pip uv --no-cache + +# Important: Copy only bare minimum required for the script to run +COPY requirements/ requirements/ +COPY pyproject.toml ./ + +# The script is expected to install whatever python dependencies are missing +# as well as whatever system libraries need to be installed from source +COPY build_vllm_*.sh ./ + RUN --mount=type=cache,target=/root/.cache/uv \ - dnf install llvm15 llvm15-devel -y && \ - rpm -ivh --nodeps https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-16.el9.ppc64le.rpm && \ - source /opt/rh/gcc-toolset-14/enable && \ - git clone https://github.com/huggingface/xet-core.git && cd xet-core/hf_xet/ && \ - uv pip install maturin && \ - uv build --wheel --out-dir /hf_wheels/ + sh ./build_vllm_$(uname -m).sh + +# copy vllm source code to build cache +COPY . . -ENV CXXFLAGS="-fno-lto -Wno-error=free-nonheap-object" \ - CFLAGS="-fno-lto" RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \ - --mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \ - --mount=type=bind,from=cv-builder,source=/opencvwheels/,target=/opencvwheels/,ro \ - --mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \ - --mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \ - --mount=type=bind,src=.,dst=/src/,rw \ source /opt/rh/gcc-toolset-14/enable && \ - export PATH=$PATH:/usr/lib64/llvm15/bin && \ - uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl && \ - sed -i -e 's/.*torch.*//g' /src/pyproject.toml /src/requirements/*.txt && \ - sed -i -e 's/.*sentencepiece.*//g' /src/pyproject.toml /src/requirements/*.txt && \ - uv pip install sentencepiece==0.2.0 pandas pythran nanobind pybind11 /hf_wheels/*.whl && \ - make -C /numactl install && \ - # sentencepiece.pc is in some pkgconfig inside uv cache - export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':') && \ - nanobind_DIR=$(uv pip show nanobind | grep Location | sed 's/^Location: //;s/$/\/nanobind\/cmake/') && uv pip install -r /src/requirements/common.txt -r /src/requirements/cpu.txt -r /src/requirements/build/cuda.txt --no-build-isolation && \ - cd /src/ && \ - uv build --wheel --out-dir /vllmwheel/ --no-build-isolation && \ - uv pip install /vllmwheel/*.whl + pip install -U uv +# build & install vLLM so that all transitive dependencies are build/downloaded into the uv cache +RUN --mount=type=cache,target=/root/.cache/uv \ + source /opt/rh/gcc-toolset-14/enable && \ + export PATH=/opt/rh/gcc-toolset-14/root/usr/bin:$PATH && \ + export CC=/opt/rh/gcc-toolset-14/root/usr/bin/gcc && \ + export CXX=/opt/rh/gcc-toolset-14/root/usr/bin/g++ && \ + export PKG_CONFIG_PATH=/usr/local/lib/pkgconfig:/usr/lib64/pkgconfig:$PKG_CONFIG_PATH && \ + export CMAKE_PREFIX_PATH=/usr/local:/usr:$CMAKE_PREFIX_PATH && \ + export Protobuf_PROTOC_EXECUTABLE=/usr/bin/protoc && \ + export CFLAGS="-mcpu=power10 -mtune=power10" && \ + export CXXFLAGS="-mcpu=power10 -mtune=power10" && \ + export DNNL_ARCH_OPT_FLAGS="-mcpu=power10 -mtune=power10" && \ + export C_INCLUDE_PATH=/usr/local/include:$C_INCLUDE_PATH && \ + export CPLUS_INCLUDE_PATH=/usr/local/include:$CPLUS_INCLUDE_PATH && \ + uv pip install 'setuptools>=78.1.1' && \ + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/OpenBLAS/lib/:/usr/local/lib64:/usr/local/lib && \ + export LIBGOMP=/opt/rh/gcc-toolset-14/root/usr/lib/gcc/ppc64le-redhat-linux/14/libgomp.so && \ -############################################################### -# Stage to build lapack -############################################################### - -FROM base-builder AS lapack-builder - -ARG MAX_JOBS -ARG LAPACK_VERSION=3.12.1 -RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${LAPACK_VERSION} \ - && cd lapack && source /opt/rh/gcc-toolset-14/enable \ - && cmake -B build -S . \ - && cmake --build build -j ${MAX_JOBS:-$(nproc)} + export CMAKE_LIBRARY_PATH=$(dirname $LIBGOMP):${CMAKE_LIBRARY_PATH} && \ + export LIBRARY_PATH=$(dirname $LIBGOMP):${LIBRARY_PATH} && \ + export LD_LIBRARY_PATH=$(dirname $LIBGOMP):${LD_LIBRARY_PATH} && \ + echo "LIBGOMP=${LIBGOMP}" && \ + find /root/.cache/uv -name "*.whl" && \ + SETUPTOOLS_SCM_PRETEND_VERSION="$VLLM_VERSION" uv build \ + --wheel --out-dir ${WHEEL_DIR} --no-build-isolation && \ + uv pip install "$(echo ${WHEEL_DIR}/vllm*.whl)[tensorizer]" --refresh ############################################################### # FINAL VLLM IMAGE STAGE # @@ -278,72 +83,74 @@ RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${L FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS vllm-openai ARG PYTHON_VERSION=3.12 -ARG OPENBLAS_VERSION=0.3.30 +ENV VLLM_NO_USAGE_STATS=1 # Set Environment Variables for venv & openblas ENV VIRTUAL_ENV=/opt/vllm -ENV PATH=${VIRTUAL_ENV}/bin:$PATH -ENV PKG_CONFIG_PATH=/usr/local/lib/pkgconfig/ -ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib64:/usr/local/lib:/usr/lib64:/usr/lib +ENV PCP_DIR=/opt/rh/gcc-toolset-14/root +ENV PATH=${VIRTUAL_ENV}/bin:${PCP_DIR}/usr/bin:/usr/local/bin:$PATH +ENV PKG_CONFIG_PATH=${PCP_DIR}/usr/lib64/pkgconfig:/usr/local/lib/pkgconfig/ +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH" +ENV LD_LIBRARY_PATH=${PCP_DIR}/usr/lib64:${PCP_DIR}/usr/lib:${VIRTUAL_ENV}/lib64/python${PYTHON_VERSION}/site-packages/torch/lib:/usr/local/lib:$LD_LIBRARY_PATH:/usr/local/lib64:/usr/lib64:/usr/lib ENV UV_LINK_MODE=copy -ENV OMP_NUM_THREADS=16 +ARG VLLM_VERSION="0.22.1" +ARG UV_EXTRA_INDEX_URL="https://wheels.developerfirst.ibm.com/ppc64le/linux/+simple/" +ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL} +ENV UV_INDEX_STRATEGY=first-match -# create artificial dependencies between stages for independent stages to build in parallel -COPY --from=torch-builder /tmp/control /dev/null -COPY --from=arrow-builder /tmp/control /dev/null -COPY --from=cv-builder /tmp/control /dev/null -COPY --from=vllmcache-builder /tmp/control /dev/null -COPY --from=numa-builder /tmp/control /dev/null -COPY --from=lapack-builder /tmp/control /dev/null -COPY --from=openblas-builder /tmp/control /dev/null -COPY --from=numba-builder /tmp/control /dev/null -# install gcc-11, python, openblas, numactl, lapack RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \ - --mount=type=bind,from=lapack-builder,source=/lapack/,target=/lapack/,rw \ - --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \ rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \ microdnf install --nodocs -y \ - libomp libicu tar findutils openssl llvm15 llvm15-devel \ - pkgconfig xsimd g++ gcc-fortran libsndfile \ + libomp libicu tar autoconf automake libtool findutils openssl numactl numactl-devel \ + pkgconfig xsimd gcc-toolset-14 libsndfile \ libtiff libjpeg openjpeg2 zlib zeromq \ freetype lcms2 libwebp tcl tk utf8proc \ - harfbuzz fribidi libraqm libimagequant libxcb util-linux \ + harfbuzz fribidi libraqm libimagequant libxcb util-linux gperftools-libs \ python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip \ - && export PATH=$PATH:/usr/lib64/llvm15/bin && microdnf clean all \ - && python${PYTHON_VERSION} -m venv ${VIRTUAL_ENV} \ - && python -m pip install -U pip uv --no-cache \ - && make -C /numactl install \ - && PREFIX=/usr/local make -C /openblas install \ - && uv pip install 'cmake<4' \ - && cmake --install /lapack/build \ - && uv pip uninstall cmake + && source /opt/rh/gcc-toolset-14/enable \ + && microdnf update -y \ + && microdnf clean all -# consume previously built wheels (including vllm) -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \ - --mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \ - --mount=type=bind,from=cv-builder,source=/opencvwheels/,target=/opencvwheels/,ro \ - --mount=type=bind,from=vllmcache-builder,source=/hf_wheels/,target=/hf_wheels/,ro \ - --mount=type=bind,from=vllmcache-builder,source=/vllmwheel/,target=/vllmwheel/,ro \ - --mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \ - export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':') && uv pip install sentencepiece==0.2.0 && \ - HOME=/root uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl /hf_wheels/*.whl /vllmwheel/*.whl +# The `lscpu` command was added as a requirement in part of https://github.com/vllm-project/vllm/pull/21032, so installing it. +RUN microdnf install --nodocs -y util-linux && \ + microdnf clean all +COPY --from=builder-base /usr/lib64/libprotobuf.so.25 /usr/lib64/ +COPY --from=builder-base /usr/lib64/libprotobuf.so.25.0.0 /usr/lib64/ -COPY ./ /workspace/vllm -WORKDIR /workspace/vllm -ARG GIT_REPO_CHECK=0 -RUN --mount=type=bind,source=.git,target=.git \ - if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi +# Use builder venv in final stage instead of wheel reinstallation +COPY --from=builder-base /opt/vllm /opt/vllm -# install development dependencies (for testing) -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -e tests/vllm_test_utils +ENV LD_PRELOAD=/usr/lib64/libtcmalloc.so.4 -WORKDIR /workspace/ +WORKDIR /home/vllm -RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks +# setup non-root user for OpenShift +RUN umask 002 && \ + useradd --uid 2000 --gid 0 vllm && \ + mkdir -p /home/vllm && \ + chmod g+rwx /home/vllm + +ENV HOME=/home/vllm + +# Add labels to document build configuration +LABEL org.opencontainers.image.title="vLLM CPU" +LABEL org.opencontainers.image.description="vLLM inference engine for CPU platforms" +LABEL org.opencontainers.image.vendor="vLLM Project" +LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" + +# Build configuration labels +ARG TARGETARCH +ARG VLLM_CPU_PPC64LE +ARG PYTHON_VERSION + +LABEL ai.vllm.build.target-arch="${TARGETARCH}" +LABEL ai.vllm.build.cpu-ppc64le="${VLLM_CPU_PPC64LE:-false}" +LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}" + +USER 2000 ENTRYPOINT ["vllm", "serve"] + + diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 3f307a5fa0f..02e4086d625 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -252,6 +252,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/doc COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -543,6 +545,8 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.h COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 # ----------------------- @@ -576,6 +580,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + unzip \ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index fbd5e1e60e3..2faaf774cf6 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.16.post2" +ARG AITER_BRANCH="v0.1.16.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index 554a7257c23..6d1c0c3452f 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \ OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \ GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \ - uv pip install -v \ + uv pip install -v \ $ARROW_WHL_FILE \ $VISION_WHL_FILE \ $HF_XET_WHL_FILE \ @@ -257,6 +257,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ $NUMBA_WHL_FILE \ $OPENCV_WHL_FILE \ $GUIDANCE_WHL_FILE \ + --torch-backend cpu \ --index-strategy unsafe-best-match \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt diff --git a/docs/cli/README.md b/docs/cli/README.md index 43857704522..123f3f109a9 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -50,6 +50,21 @@ vllm serve --help=max-num-seqs vllm serve --help=max ``` +!!! tip "Human-readable integer arguments" + Many integer arguments accept human-readable suffixes for convenience. For example: + + - `1k` = 1,000 (decimal kilo) + - `1K` = 1,024 (binary kibibyte) + - `1m` = 1,000,000 (decimal mega) + - `1M` = 1,048,576 (binary mebibyte) + - `1g` / `1G` = 1 billion / 1 gibibyte + - `1t` / `1T` = 1 trillion / 1 tebibyte + + Decimal suffixes (`k`, `m`, `g`, `t`) also accept floating point: `25.6k` = 25,600. + Binary suffixes (`K`, `M`, `G`, `T`) require integers: `32K` = 32,768. + + Supported arguments include: `--max-model-len`, `--max-num-batched-tokens`, `--max-num-scheduled-tokens`, `--kv-cache-memory-bytes`, `--safetensors-prefetch-block-size`. + See [vllm serve](./serve.md) for the full reference of all available arguments. ## launch diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index c6d64b25035..efa0f8b9046 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -16,6 +16,14 @@ vLLM provides 4 optimization levels (`-O0`, `-O1`, `-O2`, `-O3`) that allow user For more information, see the [optimization level documentation](../design/optimization_levels.md). +## Faster Startup + +Beyond the optimization levels, three mechanisms reduce time-to-first-token on repeated boots of the same (model, config, hardware) combination: + +- **Reuse the compile cache.** vLLM persists `torch.compile` artifacts under `VLLM_CACHE_ROOT` (default `~/.cache/vllm`), and the cache directory can be copied between machines or baked into a container image; see the [torch.compile design doc](../design/torch_compile.md). Set `VLLM_FORCE_AOT_LOAD=1` to fail loudly instead of silently recompiling when the cache misses (any change to the model, config, relevant `VLLM_*` environment variables, torch build, or GPU model invalidates it). +- **Skip memory profiling with `--kv-cache-memory`.** On startup, vLLM logs the exact `--kv-cache-memory` value that reproduces the current allocation. Passing it back on the next boot skips the memory-profiling measurement and the CUDA-graph memory estimation pass. Note that this has performance implications: the KV cache is sized to exactly the given value instead of being measured, so a conservative value caps batch concurrency (and therefore throughput), while an optimistic one fails at allocation time. The value is only valid on the same GPU with the same initial free memory; if a boot OOMs after hardware or co-tenant changes, remove the flag to re-profile. +- **Serve without CUDA graphs using `--enforce-eager`.** Skips both compilation and CUDA-graph capture for the fastest possible startup, at the cost of steady-state decode performance. Useful for development loops and for measuring how much of a boot is compile/capture. + ## Preemption Due to the autoregressive nature of transformer architecture, there are times when KV cache space is insufficient to handle all batched requests. @@ -276,8 +284,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.) you can switch to the [fastokens](https://github.com/crusoecloud/fastokens) Rust backend, a drop-in replacement that's substantially faster on -encode/decode and on streaming detokenization. Enable it by setting -`VLLM_USE_FASTOKENS=1`: +encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is +available in vLLM v0.23.0 and later. If your installed vLLM version does not +recognize the environment variable, upgrade vLLM before enabling the override: ```console VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B diff --git a/docs/contributing/ci/nightly_builds.md b/docs/contributing/ci/nightly_builds.md index 8f3512db3d4..10c4a437240 100644 --- a/docs/contributing/ci/nightly_builds.md +++ b/docs/contributing/ci/nightly_builds.md @@ -14,7 +14,7 @@ Wheels are built in the `Release` pipeline (`.buildkite/release-pipeline.yaml`) Each build step: 1. Builds the wheel in a Docker container. -2. Renames the wheel filename to use the correct manylinux tag (currently `manylinux_2_31`) for PEP 600 compliance. +2. Renames the wheel filename to use the correct manylinux tag (currently `manylinux_2_28`) for PEP 600 compliance. 3. Uploads the wheel to S3 bucket `vllm-wheels` under `/{commit_hash}/`. ### Index Generation diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index d28de8fc36a..23df00672fb 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -168,9 +168,9 @@ Priority is **1 = highest** (tried first). | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | +| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | | `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` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %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`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any | | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | diff --git a/docs/design/endpoint_plugins.md b/docs/design/endpoint_plugins.md new file mode 100644 index 00000000000..9f38fe5da18 --- /dev/null +++ b/docs/design/endpoint_plugins.md @@ -0,0 +1,136 @@ +# Endpoint Plugins + +Endpoint plugins let out-of-tree packages add HTTP routes to the OpenAI compatible API server without editing `vllm/entrypoints/openai/api_server.py`. Their scope is +the **HTTP surface only** registering routes and optionally per app state used by those routes. A plugin reaches the engine the same way an in-tree serving handler does, through the `EngineClient` it is handed at startup (e.g. `engine_client.collective_rpc(...)`). No new engine access path is introduced. + +!!! warning "Security" + Endpoint plugins are **not loaded by default** and must be explicitly allowlisted. Read [Endpoint Plugins security posture](../usage/security.md#endpoint-plugins) before enabling one, especially the route shadowing warning. + +## The `EndpointPlugin` protocol + +Endpoint plugins implement the [`EndpointPlugin`][vllm.plugins.endpoint_plugins.interface.EndpointPlugin] runtime checkable `Protocol`: + +```python +class EndpointPlugin(Protocol): + name: str + required_tasks: tuple[SupportedTask, ...] | None + + def attach_router(self, app: FastAPI) -> None: ... + + async def init_state( + self, engine_client: EngineClient | None, state: State, args: Namespace + ) -> None: ... +``` + +- `name`: a unique identifier used in logs and for `VLLM_PLUGINS` allowlisting +- `required_tasks`: the tasks the server must support for this plugin to load. `None` means the plugin has no task requirement +- `attach_router`: registers routes on `app` +- `init_state`: initializes per app state the routes read at request time + +## The two phase lifecycle + +Routes are registered before the engine exists. This means the interface has to expose two hooks that run at two different points in server startup: + +| Phase | Called from | `engine_client` available? | Work | +| --- | --- | --- | --- | +| A. Route registration | `build_app()` | No | `attach_router(app)` add routes. Do not touch the engine here. | +| B. State init | `init_app_state()` | Usually but `None` on the CPU only render server | `init_state(engine_client, state, args)` build a serving handler holding `engine_client` and store it on `state`. | + +Because `app.state` *is* the `state` object passed to `init_app_state()`, an object stored during phase A is visible in phase B and an object stored in phase B is visible to route handlers at request time via `request.app.state`. This is the same pattern in-tree endpoints already use. + +### Engine less servers (the render server) + +The CPU only render server (`init_render_app_state()`) has no `EngineClient`. It still runs both phases for any plugin eligible for the `render` task (`required_tasks` is `None` or includes `"render"`). `attach_router` is called as usual but `init_state` is called with `engine_client=None`. + +A plugin that needs an engine to function has two options: + +- Exclude `"render"` from `required_tasks` so it is never loaded on the render server in the first place +- Accept being loaded on `render` and check for `None` in `init_state` or in the route handler returning an error response (e.g. HTTP 503) instead of dereferencing a client that doesn't exist + +`tests/plugins/vllm_add_dummy_endpoint_plugin` demonstrates the second option. Its route handler returns a 503 when `state.dummy_engine_client` is `None`. + +### Reaching the engine from a route handler + +`init_state` is where a plugin captures `engine_client` into a small serving handler and stashes it on `state`. The route added in `attach_router` reads that handler off `request.app.state` at request time and calls the engine through it, typically via `engine_client.collective_rpc(...)`. + +This minimal example omits the `None` check from the previous section for brevity since `required_tasks` is `None` here. It is in fact eligible for `render` and should handle `engine_client=None` the way `tests/plugins/vllm_add_dummy_endpoint_plugin` does before shipping it: + +```python +from fastapi import FastAPI, Request + + +class MyAdminEndpointPlugin: + name = "my_admin_endpoint_plugin" + required_tasks: tuple[str, ...] | None = None + + def attach_router(self, app: FastAPI) -> None: + @app.get("/plugins/my_admin_endpoint_plugin/scheduler_config") + async def scheduler_config(raw_request: Request): + engine_client = raw_request.app.state.my_engine_client + results = await engine_client.collective_rpc("get_scheduler_config") + return {"scheduler_config": results} + + async def init_state(self, engine_client, state, args) -> None: + state.my_engine_client = engine_client +``` + +A complete and tested version of this example is in-repo as `tests/plugins/vllm_add_dummy_endpoint_plugin` and is exercised e2e (including a real HTTP request) in `tests/plugins_tests/test_endpoint_plugins.py`. + +## Registering the entry point + +Register a zero argument factory (a class or function) under the `vllm.endpoint_plugins` group. The factory must return an object satisfying `EndpointPlugin`: + +```toml +# pyproject.toml +[project.entry-points."vllm.endpoint_plugins"] +my_admin_api = "my_pkg.endpoints:MyAdminEndpointPlugin" +``` + +```python +# setup.py equivalent +setup( + name="my_pkg", + entry_points={ + "vllm.endpoint_plugins": [ + "my_admin_api = my_pkg.endpoints:MyAdminEndpointPlugin" + ] + }, +) +``` + +The entry point name (`my_admin_api` above) is independent of the plugin's `name` attribute. `VLLM_PLUGINS` allowlisting matches on the **entry point name** following the same convention as `vllm.general_plugins` (see [Plugin System](plugin_system.md)). + +## Gating: `VLLM_PLUGINS` and `required_tasks` + +Endpoint plugins are discovered and gated by [`load_endpoint_plugins`][vllm.plugins.load_endpoint_plugins] which is stricter than the loader used for other plugin groups: + +- **Nothing loads unless `VLLM_PLUGINS` is set and names the plugin.** Other plugin groups load everything unless `VLLM_PLUGINS` narrows the set. Endpoint plugins invert that default because they add network exposed surface. See [Security](../usage/security.md#endpoint-plugins). +- **`required_tasks` must intersect the server's supported tasks** unless it is `None`. Use this to keep a plugin from attaching routes on a server that can't service them (e.g. a pooling only deployment). +- A factory that raises an issue during instantiation is logged and skipped. It does not abort server startup. + +Only the front end API server process loads endpoint plugins. There is no need to guard for worker or engine core processes. + +## Pairing with `vllm.general_plugins` + +Endpoint plugins cover the HTTP surface only. If a plugin also needs new engine side behavior (a new worker-side RPC method, a custom stat) that half ships separately through the existing `vllm.general_plugins` group which loads in worker processes (see [Plugin System](plugin_system.md)). The two entry points are registered and loaded **independently**. Neither implies the other. The recommended distribution shape is a single package exposing both: + +```toml +[project.entry-points."vllm.general_plugins"] +my_admin_engine = "my_pkg.engine:register" # adds the worker side method + +[project.entry-points."vllm.endpoint_plugins"] +my_admin_api = "my_pkg.endpoints:MyAdminEndpointPlugin" # adds the HTTP route +``` + +Do not expect a single endpoint plugin to also mutate engine/worker state. If your route needs a worker side method that doesn't already exist then add it via a paired `general_plugins` entry point. + +## Path-prefix convention + +There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)). A plugin's `attach_router` can register a path that collides with a core route and routes attached later win. To avoid surprising operators: + +- Namespace your routes under a distinct prefix, e.g. `/plugins//...`, rather than reusing `/v1/...` or other core prefixes +- Only register routes under a core prefix (like the worked example's `/v1/admin/scheduler_config`) if you specifically intend to override or extend existing behavior and document that clearly for operators allowlisting your plugin + +## Compatibility + +`state`/serving handler internals (e.g. the shape of in-tree `OpenAIServing*` classes) are not a stable public contract yet. Treat them as use-at-your-own-risk and expect them to change between vLLM versions. `FastAPI`, `EngineClient` and the `EndpointPlugin` protocol itself are the supported surface. diff --git a/docs/design/plugin_system.md b/docs/design/plugin_system.md index e5c9cea17c2..dd49df0ef2f 100644 --- a/docs/design/plugin_system.md +++ b/docs/design/plugin_system.md @@ -53,6 +53,8 @@ Every plugin has three parts: - **Stat logger plugins** (with group name `vllm.stat_logger_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree loggers into vLLM. The entry point should be a class that subclasses StatLoggerBase. +- **Endpoint plugins** (with group name `vllm.endpoint_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree HTTP routes on the OpenAI compatible API server. Unlike the other plugin groups above, endpoint plugins are loaded only in the API server front end process and are **not loaded by default**. See [Endpoint Plugins](endpoint_plugins.md) for the interface and [Security](../usage/security.md#endpoint-plugins) for the opt-in and trust model. + ## Guidelines for Writing Plugins - **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes. diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 847743dfff1..e44596626f8 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -816,6 +816,44 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu export VLLM_VIDEO_FETCH_TIMEOUT= ``` +#### Video Decoding Backend + +vLLM decodes video bytes into frames using a selectable decoding backend. Three +backends are supported: + +- `opencv` (default): OpenCV-based decoder. +- `pyav`: PyAV decoder. +- `torchcodec`: TorchCodec (PyTorch-native) decoder. + +All three backends are ultimately backed by FFmpeg. `torchcodec` lets +you choose which FFmpeg version is used while `opencv` and `pyav` rely on +whichever FFmpeg build they were linked against. + +Select the backend by passing the `backend` parameter via `--media-io-kwargs`: + +```bash +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec"}}' +``` + +**TorchCodec-specific parameters:** + +The following parameters only apply to the `torchcodec` backend: + +- `num_ffmpeg_threads`: Number of FFmpeg decoding threads. `0` (default) relies + on the FFmpeg default, which is `min(cpu_count + 1, 16)`. This allows you to + control thread over-subscription. +- `seek_mode`: Seek mode for the decoder. `"exact"` (default) guarantees + frame-accurate sampling by scanning the file when the decoder is created. + `"approximate"` skips that scan for faster decoder creation, at the cost of + relying on the file's metadata (which may yield less accurate seeking). + +```bash +# Example: TorchCodec with approximate seek mode and 4 FFmpeg threads +vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \ + --media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}' +``` + #### Video Frame Recovery For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place. diff --git a/docs/features/per_request_metrics.md b/docs/features/per_request_metrics.md new file mode 100644 index 00000000000..9bc64d2b86d --- /dev/null +++ b/docs/features/per_request_metrics.md @@ -0,0 +1,127 @@ +# Per-Request Metrics + +vLLM can return per-request timing metrics directly in API responses. +This is useful for billing, SLA monitoring, and latency analysis at the +individual request level, as a complement to the server-aggregated Prometheus +metrics exposed at `/metrics`. + +## Enabling + +Start the server with `--enable-per-request-metrics`: + +```bash +vllm serve meta-llama/Llama-3.1-8B-Instruct --enable-per-request-metrics +``` + +When this flag is set, supported API responses include metrics for each +attributable request. + +!!! note + At high concurrency, enabling per-request metrics computation may introduce + non-negligible CPU overhead. Benchmark your specific workload to evaluate the + impact before enabling in production. + +## Response Format + +When per-request metrics are enabled, the response includes a `metrics` object: + +```json +{ + "id": "chatcmpl-abc123", + "object": "chat.completion", + "model": "meta-llama/Llama-3.1-8B-Instruct", + "choices": [ ... ], + "usage": { + "prompt_tokens": 42, + "completion_tokens": 128, + "total_tokens": 170 + }, + "metrics": { + "time_to_first_token_ms": 85.2, + "generation_time_ms": 1240.5, + "queue_time_ms": 12.3, + "mean_itl_ms": 9.1, + "tokens_per_second": 103.2 + } +} +``` + +| Field | Description | +| --- | --- | +| `time_to_first_token_ms` | Time from when the request was scheduled until the first output token was generated (TTFT). | +| `generation_time_ms` | Decode time: time from the first output token to the last output token. Excludes both queue wait and prefill/TTFT. | +| `queue_time_ms` | Time the request spent waiting in the scheduler queue before processing began. | +| `mean_itl_ms` | Mean inter-token latency (average time between successive output tokens) during the decode phase. `null` for single-token responses. | +| `tokens_per_second` | Overall output token throughput: all generated tokens over the inference interval (scheduling to last output token). Unlike `generation_time_ms`, this includes the prefill phase, so it reflects end-to-end generation speed rather than pure decode speed. | + +All fields are `null` if the underlying timing data is not available for that +request. + +!!! note + Timing metrics describe a single generation stream, so they are only + returned when the request maps to exactly one. They are suppressed (the + `metrics` object is `null`) for requests with `n > 1`, because the + underlying timing data reflects only one of the `n` sequences and cannot be + accurately attributed to the request as a whole. Token usage + (`prompt_tokens`, `completion_tokens`) remains accurate in these cases. + Per-request metrics also require server-side statistics logging, which is + on by default. vLLM rejects `--enable-per-request-metrics` when + `--disable-log-stats` is also set. + +## Example Request + +=== "Non-streaming" + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + response = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + ) + + print(response.usage) + print(response.model_extra.get("metrics")) + ``` + +=== "Streaming" + + In streaming responses, metrics are attached to the final usage chunk (the + chunk sent after all content chunks). That chunk is only emitted when usage + reporting is enabled with `stream_options.include_usage: true` or forced + server-side with `--enable-force-include-usage`. Without forced usage, a + streaming client must set `stream_options.include_usage: true` to receive + metrics. + + ```python + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="token") + + stream = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[{"role": "user", "content": "What is the capital of France?"}], + stream=True, + stream_options={"include_usage": True}, + ) + + for chunk in stream: + if chunk.usage: + print("Usage:", chunk.usage) + print("Metrics:", chunk.model_extra.get("metrics")) + ``` + +## Completions API + +Per-request metrics are also available on the `/v1/completions` endpoint using +the same `metrics` response field. As with `n > 1`, metrics are omitted for +requests with multiple prompts, because the timing data cannot be attributed to +a single prompt's generation. + +## Relationship to Prometheus Metrics + +The `metrics` response field provides per-request values for a single request. +The `/metrics` Prometheus endpoint exposes server-level histograms (e.g. +`vllm:time_to_first_token_seconds`) that aggregate across all requests. diff --git a/docs/features/quantization/inc.md b/docs/features/quantization/inc.md index adb6b3ae8e2..ffb90cec8c1 100644 --- a/docs/features/quantization/inc.md +++ b/docs/features/quantization/inc.md @@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \ --max-model-len 4096 ``` -!!! note - To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now. - ## Evaluating the Quantized Model with vLLM ```bash lm_eval --model vllm \ - --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \ + --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \ --tasks gsm8k \ --num_fewshot 5 \ --batch_size 128 diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 92563a8b4bb..50a58b8b3e3 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -439,7 +439,7 @@ Additionally, to enable structured output, you'll need to create a new `Reasoner end_token: str = "" @classmethod - def from_tokenizer(cls, tokenizer: PreTrainedTokenizer) -> Reasoner: + def from_tokenizer(cls, tokenizer: PythonBackend) -> Reasoner: return cls( start_token_id=tokenizer.encode("", add_special_tokens=False)[0], end_token_id=tokenizer.encode("", add_special_tokens=False)[0], diff --git a/docs/features/speculative_decoding/dynamic_speculative_decoding.md b/docs/features/speculative_decoding/dynamic_speculative_decoding.md index aa55c5f08d4..682eaafd29d 100644 --- a/docs/features/speculative_decoding/dynamic_speculative_decoding.md +++ b/docs/features/speculative_decoding/dynamic_speculative_decoding.md @@ -73,3 +73,4 @@ VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \ * Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box * Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature +* Not compatible with data parallelism (`--data-parallel-size > 1`). Each DP rank schedules independently, so ranks can pick different K values, causing DP collective divergence and deadlocks. When DP is enabled, vLLM automatically disables `num_speculative_tokens_per_batch_size` and falls back to the static `num_speculative_tokens` value. diff --git a/docs/getting_started/installation/cpu.apple.inc.md b/docs/getting_started/installation/cpu.apple.inc.md index e312964ec8a..479b6d2c011 100644 --- a/docs/getting_started/installation/cpu.apple.inc.md +++ b/docs/getting_started/installation/cpu.apple.inc.md @@ -35,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match +uv pip install -r requirements/cpu.txt uv pip install -e . ``` -!!! tip - The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts. - - The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation. - !!! note On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device. diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index 7a783b53c65..3950adc0251 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -20,12 +20,12 @@ Pre-built vLLM wheels for Arm are available since version 0.11.2. These wheels c ```bash export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl --torch-backend cpu +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --torch-backend cpu ``` ??? console "pip" ```bash - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" @@ -63,7 +63,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/nightly/cpu --index If you insist on using `pip`, you have to specify the full URL (link address) of the wheel file (which can be obtained from https://wheels.vllm.ai/nightly/cpu/vllm). ```bash - pip install https://wheels.vllm.ai/4fa7ce46f31cbd97b4651694caf9991cc395a259/vllm-0.13.0rc2.dev104%2Bg4fa7ce46f.cpu-cp38-abi3-manylinux_2_35_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu # current nightly build (the filename will change!) + pip install https://wheels.vllm.ai/2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7/vllm-0.23.1rc1.dev901%2Bg2f3f441f8.cpu-cp38-abi3-manylinux_2_34_aarch64.whl --extra-index-url https://download.pytorch.org/whl/cpu # current nightly build (the filename will change!) ``` #### Install specific revisions diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 1e36b431764..15baa487c2a 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source. ```bash uv pip install -v \ - --extra-index-url https://download.pytorch.org/whl/cpu \ - --torch-backend auto \ -r requirements/build/cpu.txt \ -r requirements/cpu.txt \ + --torch-backend cpu \ + --index-strategy unsafe-best-match && \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv pip install dist/*.whl ``` diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index 32295d63879..6ded3b50832 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -24,13 +24,13 @@ Pre-built vLLM wheels for x86 with AVX512/AVX2 are available since version 0.17. export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') # use uv -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl --torch-backend cpu +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --torch-backend cpu ``` ??? console "pip" ```bash # use pip - pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu + pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cpu-cp38-abi3-manylinux_2_34_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cpu ``` !!! warning "set `LD_PRELOAD`" Before use vLLM CPU installed via wheels, make sure TCMalloc and Intel OpenMP are installed and added to `LD_PRELOAD`: diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 5f774952d59..0e86c0e6049 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -43,7 +43,7 @@ As of now, vLLM's binaries are compiled with CUDA 12.9 and public PyTorch releas export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//') export CUDA_VERSION=130 # or other export CPU_ARCH=$(uname -m) # x86_64 or aarch64 -uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_35_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION} +uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION} ``` #### Install the latest code @@ -68,8 +68,8 @@ uv pip install -U vllm \ If you insist on using `pip`, you have to specify the full URL of the wheel file (which can be obtained from the web page). ```bash - pip install -U https://wheels.vllm.ai/nightly/vllm-0.11.2.dev399%2Bg3c7461c18-cp38-abi3-manylinux_2_31_x86_64.whl # current nightly build (the filename will change!) - pip install -U https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-0.11.2.dev399%2Bg3c7461c18-cp38-abi3-manylinux_2_31_x86_64.whl # from specific commit + pip install -U https://wheels.vllm.ai/2f3f441f84bd5b35ec8aa9fcfffb540f107da8a7/vllm-0.23.1rc1.dev901%2Bg2f3f441f8-cp38-abi3-manylinux_2_28_x86_64.whl # current nightly build (the filename will change!) + pip install -U https://wheels.vllm.ai/${VLLM_COMMIT}/vllm-0.23.1rc1.dev901%2Bg2f3f441f8-cp38-abi3-manylinux_2_28_x86_64.whl # from specific commit ``` ##### Install specific revisions diff --git a/docs/models/pooling_models/README.md b/docs/models/pooling_models/README.md index 37fca366eba..f8de9d437ad 100644 --- a/docs/models/pooling_models/README.md +++ b/docs/models/pooling_models/README.md @@ -2,8 +2,7 @@ !!! note We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance -improvements over using Hugging Face Transformers or Sentence Transformers directly. - + improvements over using Hugging Face Transformers or Sentence Transformers directly. We plan to optimize pooling models in vLLM. Please comment on if you have any suggestions! ## What are pooling models? @@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). !!! note Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models -are a subset of classification models that accept two prompts as input and output num_labels equal to 1. + are a subset of classification models that accept two prompts as input and output num_labels equal to 1. ### Pooling Types diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 562e38109ff..a191bb2d3c5 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang ### Transformers -vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend". +vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation. Currently, the Transformers modeling backend works for the following: @@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded: That's it! -For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class: +For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
configuration_my_model.py @@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported). + - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated. - `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s: - You only need to do this for layers which are not present on all pipeline stages - vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages + - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage - The `list` in the first element of the `tuple` contains the names of the input arguments - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code @@ -240,50 +242,24 @@ Use the Hugging Face CLI to [manage models](https://huggingface.co/docs/huggingf ```bash # List cached models -hf scan-cache +hf cache list -q # Show detailed (verbose) output -hf scan-cache -v +hf cache list # Specify a custom cache directory -hf scan-cache --dir ~/.cache/huggingface/hub +hf cache list --dir ~/.cache/huggingface/hub ``` #### Delete a cached model -Use the Hugging Face CLI to interactively [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache: +Use the Hugging Face CLI to [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache: -
-Commands - -```console -# The `delete-cache` command requires extra dependencies to work with the TUI. -# Please run `pip install huggingface_hub[cli]` to install them. - -# Launch the interactive TUI to select models to delete -$ hf delete-cache -? Select revisions to delete: 1 revisions selected counting for 438.9M. - ○ None of the following (if selected, nothing will be deleted). -Model BAAI/bge-base-en-v1.5 (438.9M, used 1 week ago) -❯ ◉ a5beb1e3: main # modified 1 week ago - -Model BAAI/bge-large-en-v1.5 (1.3G, used 1 week ago) - ○ d4aa6901: main # modified 1 week ago - -Model BAAI/bge-reranker-base (1.1G, used 4 weeks ago) - ○ 2cfc18c9: main # modified 4 weeks ago - -Press to select, to validate and to quit without modification. - -# Need to confirm after selected -? Select revisions to delete: 1 revision(s) selected. -? 1 revisions selected counting for 438.9M. Confirm deletion ? Yes -Start deletion. -Done. Deleted 1 repo(s) and 0 revision(s) for a total of 438.9M. +```bash +# delete all the cached objects +hf cache rm $(hf cache list -q) ``` -
- #### Using a proxy Here are some tips for loading/downloading models from Hugging Face using a proxy: @@ -593,6 +569,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ | | `Molmo2ForConditionalGeneration` | Molmo2 | T + I+ / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`^, `allenai/MolmoWeb-8B`^ | ✅︎ | ✅︎ | | `MossAudioModel` | MOSS-Audio | T + A+ | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ | +| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ | @@ -695,6 +672,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ | +| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | diff --git a/docs/usage/security.md b/docs/usage/security.md index ab4b5d6ded3..d222155b770 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -326,6 +326,19 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the **Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md). +## Endpoint Plugins + +vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` entry point group (see [Endpoint Plugins](../design/endpoint_plugins.md) for how to write one). An endpoint plugin can register arbitrary FastAPI routes, including routes that reach the engine via `EngineClient.collective_rpc`, so it must be treated as part of the server's trusted code base and not as sandboxed or reviewed input. + +**Endpoint plugins are not loaded by default.** Unlike other vLLM plugin groups (`vllm.general_plugins`, `vllm.platform_plugins`, etc.), which load every discovered plugin unless `VLLM_PLUGINS` narrows the set, endpoint plugins load **none** unless `VLLM_PLUGINS` is set and explicitly names them. This mirrors the "off by default in production" posture used for development endpoints gated behind `VLLM_SERVER_DEV_MODE`. Both surfaces are only present when an operator has explicitly opted in. + +### Recommended Security Practices + +1. **Only allowlist plugins you trust.** Set `VLLM_PLUGINS` to the exact plugin names you intend to run and never wildcard or copy an allowlist between deployments without reviewing what each named plugin does. +2. **Audit routes before deploying.** A plugin's `attach_router` can add routes under any path, including ones that duplicate existing `/v1/*` paths. There is currently no route conflict enforcement (tracked as a follow-up to RFC [#46565](https://github.com/vllm-project/vllm/issues/46565)), so a malicious or buggy plugin can **shadow a core route** and silently replace its behavior. Prefer plugins that namespace their routes under a distinct prefix (e.g. `/plugins//...`) instead of reusing `/v1/...` and review `app.routes` after startup if you need certainty about what is actually being served. +3. **Treat plugin routes like any other unauthenticated by default surface.** `--api-key` only protects the `/v1`, `/v2`, and `/inference` path prefixes (see [API Key Authentication Limitations](#api-key-authentication-limitations)). A plugin route outside those prefixes is unauthenticated unless the plugin implements its own authentication. Deploy behind a reverse proxy that allowlists only the plugin routes you intend to expose externally. +4. **Remember the `vllm.general_plugins` pairing.** A plugin that also needs new engine side behavior ships that half separately via `vllm.general_plugins` which loads in every worker process under the default (load all unless restricted) posture. Allowlisting the endpoint plugin does not by itself restrict its paired engine side plugin. Need to review both. + ## gRPC Interface vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server. diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index b8561c73980..cfe131f000d 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -99,7 +99,7 @@ def main(): replay.send((last_seq + 1).to_bytes(8, "big")) while poller.poll(timeout=200): - seq_bytes, replay_payload = replay.recv_multipart() + _, seq_bytes, replay_payload = replay.recv_multipart() if not replay_payload: # End of replay marker is sent as an empty frame # for the payload diff --git a/examples/features/prompt_embed/prompt_embed_offline.py b/examples/features/prompt_embed/prompt_embed_offline.py index 29853bce967..9e90aa46b5b 100644 --- a/examples/features/prompt_embed/prompt_embed_offline.py +++ b/examples/features/prompt_embed/prompt_embed_offline.py @@ -19,7 +19,7 @@ Run: """ import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer +from transformers import AutoModelForCausalLM, AutoTokenizer, PythonBackend from vllm import LLM @@ -34,7 +34,7 @@ def init_tokenizer_and_llm(model_name: str): def get_prompt_embeds( chat: list[dict[str, str]], - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, embedding_layer: torch.nn.Module, ): token_ids = tokenizer.apply_chat_template( @@ -45,7 +45,7 @@ def get_prompt_embeds( def single_prompt_inference( - llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module + llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module ): chat = [{"role": "user", "content": "Please tell me about the capital of France."}] prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer) @@ -64,7 +64,7 @@ def single_prompt_inference( def batch_prompt_inference( - llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module + llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module ): chats = [ [{"role": "user", "content": "Please tell me about the capital of France."}], diff --git a/examples/pooling/token_embed/jina_reranker_v3_online.py b/examples/pooling/token_embed/jina_reranker_v3_online.py new file mode 100644 index 00000000000..8350aee2f14 --- /dev/null +++ b/examples/pooling/token_embed/jina_reranker_v3_online.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +""" +Example online usage of the Jina Reranker v3 score and rerank APIs with a task +instruction. + +Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the +server in vLLM. +""" + +import argparse +import json + +import requests + + +def post_http_request(prompt: dict, api_url: str) -> requests.Response: + headers = {"User-Agent": "Test Client"} + response = requests.post(api_url, headers=headers, json=prompt) + return response + + +def print_response(name: str, prompt: dict, response: requests.Response) -> None: + print(f"\n{name} request:") + print(json.dumps(prompt, indent=2)) + print(f"\n{name} response:") + print(json.dumps(response.json(), indent=2)) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="localhost") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3") + return parser.parse_args() + + +def main(args): + score_url = f"http://{args.host}:{args.port}/score" + rerank_url = f"http://{args.host}:{args.port}/rerank" + model_name = args.model + + query = "Which passage is about sports?" + documents = [ + "Basketball is played by two teams on a court.", + "Green tea contains antioxidants and may support metabolism.", + ] + instruction = "Rank passages about sports higher than passages about nutrition." + + score_prompt = { + "model": model_name, + "queries": query, + "documents": documents, + "instruction": instruction, + } + score_response = post_http_request(prompt=score_prompt, api_url=score_url) + print_response("Score", score_prompt, score_response) + + rerank_prompt = { + "model": model_name, + "query": query, + "documents": documents, + "instruction": instruction, + } + rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url) + print_response("Rerank", rerank_prompt, rerank_response) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/requirements/build/cpu.txt b/requirements/build/cpu.txt index 640432ddd8c..27a3ac65c98 100644 --- a/requirements/build/cpu.txt +++ b/requirements/build/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu cmake>=3.26.1 ninja packaging>=24.2 diff --git a/requirements/cpu.txt b/requirements/cpu.txt index 5ec338af736..c0b98d22c9b 100644 --- a/requirements/cpu.txt +++ b/requirements/cpu.txt @@ -1,4 +1,3 @@ ---extra-index-url https://download.pytorch.org/whl/cpu # Common dependencies -r common.txt @@ -16,6 +15,9 @@ torchaudio; platform_machine != "s390x" and platform_machine != "riscv64" # required for the image processor of phi3v, this must be updated alongside torch torchvision; platform_machine != "s390x" and platform_machine != "riscv64" +# required for the torchcodec video decoding backend +torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64" and platform_machine != "ppc64le" + # Intel Extension for PyTorch, only for x86_64 CPUs intel-openmp==2024.2.1; platform_machine == "x86_64" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 5545d3344f0..1d90c7ef404 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -8,6 +8,7 @@ torch==2.11.0 torchaudio==2.11.0 # These must be updated alongside torch torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version +torchcodec >= 0.14 PyNvVideoCodec==2.1.0 # FlashInfer should be updated together with the Dockerfile flashinfer-python==0.6.13 @@ -25,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode -tokenspeed-mla==0.1.2 +tokenspeed-mla==0.1.2; platform_system == "Linux" # Humming kernels for quantization gemm -humming-kernels[cu13]==0.1.6 +humming-kernels[cu13]==0.1.10 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index e953419242d..da2c7b77210 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -112,9 +112,10 @@ charset-normalizer==3.4.0 # via requests chz==0.3.0 # via gpt-oss -click==8.1.7 +click==8.4.2 # via # black + # huggingface-hub # jiwer # nltk # ray @@ -309,7 +310,7 @@ h2==4.3.0 # via httpx harfile==0.5.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub hiredis==3.0.0 # via tensorizer @@ -335,7 +336,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.10.2 +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -671,7 +672,7 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.18.1 +peft==0.19.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in @@ -1133,6 +1134,8 @@ torchaudio==2.11.0+cpu # -r requirements/test/cuda.in # encodec # vocos +torchcodec==0.14.0+cpu + # via -r requirements/test/cuda.in torchvision==0.26.0+cpu # via # -r requirements/test/cuda.in @@ -1156,7 +1159,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.10.4 # via # -r requirements/test/../common.txt # -r requirements/test/cuda.in @@ -1182,7 +1185,6 @@ typer==0.26.8 # fastapi-cli # fastapi-cloud-cli # fastsafetensors - # huggingface-hub # perceptron # transformers typing-extensions==4.15.0 diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 9a6e46712cb..a0061fbf78c 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -13,13 +13,14 @@ pytest-cov # testing utils albumentations # required for Nemotron Parse in test_common.py av # required for audio_in_video tests +torchcodec >= 0.14 # required for torchcodec video backend tests backoff # required for phi4mm test blobfile # required for kimi-vl test httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.18.1 # required for phi-4-mm test +peft>=0.19.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -38,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.10.4 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index e8d600ba632..2f7b942e911 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -117,9 +117,10 @@ charset-normalizer==3.4.0 # via requests chz==0.3.0 # via gpt-oss -click==8.1.7 +click==8.4.2 # via # black + # huggingface-hub # jiwer # nltk # ray @@ -330,7 +331,7 @@ h2==4.3.0 # via httpx harfile==0.5.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub hiredis==3.0.0 # via tensorizer @@ -356,7 +357,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.10.2 +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -754,7 +755,7 @@ pathvalidate==3.2.1 # via pytablewriter patsy==1.0.1 # via statsmodels -peft==0.18.1 +peft==0.19.1 # via -r requirements/test/cuda.in perceptron==0.1.4 # via -r requirements/test/cuda.in @@ -1232,6 +1233,10 @@ torchaudio==2.11.0+cu130 # -r requirements/test/cuda.in # encodec # vocos +torchcodec==0.14.0+cu130 + # via + # -c requirements/cuda.txt + # -r requirements/test/cuda.in torchvision==0.26.0+cu130 # via # -c requirements/cuda.txt @@ -1256,7 +1261,7 @@ tqdm==4.67.3 # segmentation-models-pytorch # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.10.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1285,7 +1290,6 @@ typer==0.26.8 # fastapi-cli # fastapi-cloud-cli # fastsafetensors - # huggingface-hub # perceptron # transformers typing-extensions==4.15.0 diff --git a/requirements/test/nightly-torch.txt b/requirements/test/nightly-torch.txt index 08f721771c8..826473db1b2 100644 --- a/requirements/test/nightly-torch.txt +++ b/requirements/test/nightly-torch.txt @@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.10.4 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test. # quantization diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 5afa6fcec92..b1c9a473e2f 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -19,7 +19,7 @@ httpx librosa # required for audio tests vector_quantize_pytorch # required for minicpmo_26 test vocos # required for minicpmo_26 test -peft>=0.15.0 # required for phi-4-mm test +peft>=0.19.1 # required for phi-4-mm test pqdm ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests sentence-transformers>=5.2.0 # required for embedding tests @@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_ datamodel_code_generator # required for minicpm3 test lm-eval[api]>=0.4.12 # required for model evaluation test mteb[bm25s]>=2, <3 # required for mteb test -transformers==5.5.3 +transformers==5.10.4 tokenizers==0.22.2 schemathesis>=4.0.0 # Required for openai schema test # quantization diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 55cac6f5243..52a0ad88c9c 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -116,9 +116,10 @@ choreographer==1.2.1 # via kaleido chz==0.4.0 # via gpt-oss -click==8.3.1 +click==8.4.2 # via # black + # huggingface-hub # jiwer # nltk # ray @@ -323,7 +324,7 @@ h2==4.3.0 # via httpx harfile==0.5.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub hiredis==3.3.1 # via tensorizer @@ -349,7 +350,7 @@ httpx==0.27.2 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.10.2 +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -730,7 +731,7 @@ pathvalidate==3.3.1 # via pytablewriter patsy==1.0.2 # via statsmodels -peft==0.18.1 +peft==0.19.1 # via -r requirements/test/rocm.in perceptron==0.1.4 # via -r requirements/test/rocm.in @@ -1217,7 +1218,7 @@ tqdm==4.67.3 # sentence-transformers # tilelang # transformers -transformers==5.5.3 +transformers==5.10.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt @@ -1244,7 +1245,6 @@ typer==0.24.1 # fastapi-cli # fastapi-cloud-cli # fastsafetensors - # huggingface-hub # perceptron # transformers typing-extensions==4.15.0 diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index b17d921d392..1172553c4ac 100644 --- a/requirements/test/xpu.in +++ b/requirements/test/xpu.in @@ -17,6 +17,7 @@ accelerate arctic-inference lm_eval[api]>=0.4.12 modelscope<1.38 +transformers==5.10.4 # --- Audio Processing --- librosa diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 16169b99863..6335fc90cff 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -83,8 +83,9 @@ charset-normalizer==3.4.6 # via requests chz==0.4.0 # via gpt-oss -click==8.3.1 +click==8.4.2 # via + # huggingface-hub # jiwer # nltk # rich-toolkit @@ -206,7 +207,7 @@ h11==0.16.0 # uvicorn harfile==0.4.0 # via schemathesis -hf-xet==1.4.3 +hf-xet==1.5.1 # via huggingface-hub html2text==2025.4.15 # via gpt-oss @@ -227,7 +228,7 @@ httpx==0.28.1 # schemathesis httpx-sse==0.4.3 # via mcp -huggingface-hub==1.10.2 +huggingface-hub==1.22.0 # via # accelerate # datasets @@ -939,10 +940,11 @@ tqdm==4.67.3 # pqdm # sentence-transformers # transformers -transformers==5.5.3 +transformers==5.10.4 # via # -c requirements/common.txt # -r requirements/test/../common.txt + # -r requirements/test/xpu.in # compressed-tensors # sentence-transformers # xgrammar @@ -959,7 +961,6 @@ typer==0.24.1 # via # fastapi-cli # fastapi-cloud-cli - # huggingface-hub # transformers typing-extensions==4.15.0 # via diff --git a/requirements/tpu.txt b/requirements/tpu.txt index f0e23f89276..da477a68461 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.23.0 +tpu-inference==0.24.0 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 684d0ef30f0..68b8eb13010 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -15,6 +15,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding torch==2.12.0 torchaudio torchvision +torchcodec >= 0.14 # Required for the torchcodec video decoding backend -auto_round_lib>=0.13.3 +auto_round_lib>=0.14.0 vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 8284ddd1285..8e7ed02a403 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -172,6 +172,9 @@ impl ChatLlm { pub async fn chat(&self, mut request: ChatRequest) -> Result { request.validate()?; + // Stamp before rendering so render and tokenize count toward TTFT/e2e. + let arrival_time = vllm_llm::current_unix_timestamp_secs(); + let output_processor = self.backend.new_chat_output_processor( &mut request, NewChatOutputProcessorOptions { @@ -210,6 +213,7 @@ impl ChatLlm { data_parallel_rank: request.data_parallel_rank, reasoning_parser_kwargs, lora_request: request.lora_request, + arrival_time: Some(arrival_time), }; let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 3ad6a5c7c75..d9031d73a4b 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -155,11 +155,23 @@ impl HfChatRenderer { effective_template: &CompiledChatTemplate, request: &ChatRequest, ) -> Result { - let messages = to_template_messages( + let mut messages = to_template_messages( &request.messages, effective_template.content_format(), self.multimodal.as_ref(), )?; + + // Handling of `continue_final_message`: + // Append a sentinel tag to the final message content, render as usual, then + // truncate the rendered prompt at the tag so any template suffix after the + // final message content (e.g. the end-of-turn marker) is dropped. + let final_message_text = if request.chat_options.continue_final_message() { + let final_message = messages.last_mut().ok_or(Error::EmptyMessages)?; + Some(append_continue_final_message_tag(final_message)?) + } else { + None + }; + let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools)); trace!( message_count = messages.len(), @@ -183,6 +195,13 @@ impl HfChatRenderer { }) .map_err(|error| Error::ChatTemplate(error.to_report_string()))?; + let prompt = match &final_message_text { + Some(final_message_text) => { + truncate_prompt_at_continue_final_message_tag(prompt, final_message_text)? + } + None => prompt, + }; + trace!( prompt_len = prompt.len(), prompt, "rendered chat template prompt" @@ -429,6 +448,74 @@ fn to_template_string_content( } } +/// Sentinel appended to the final message content when `continue_final_message` +/// is requested, used to locate the truncation point in the rendered prompt. +/// +/// Same literal as `transformers`. Occurrences of this string earlier in the +/// prompt are harmless because truncation uses the rightmost match, and the +/// appended sentinel ends up last as long as the template renders messages in +/// order. +const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG "; + +/// Append [`CONTINUE_FINAL_MESSAGE_TAG`] to the trailing text of the final +/// message, returning the original text for post-render validation. +// TODO: transformers v5 also allows continuing a non-`content` field (e.g. +// `reasoning_content`) by passing a field name; only the boolean form is +// supported here. +fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result { + let text = match &mut message.content { + TemplateContent::String(text) => Some(text), + // Pick the last text part in the message. + TemplateContent::OpenAi(parts) => parts.iter_mut().rev().find_map(|part| match part { + TemplateContentPart::Text { text } => Some(text), + TemplateContentPart::Image => None, + }), + }; + let text = text.ok_or_else(|| { + Error::ChatTemplate( + "continue_final_message is set but there is no text to continue \ + in the final message" + .to_string(), + ) + })?; + + let original = text.clone(); + text.push_str(CONTINUE_FINAL_MESSAGE_TAG); + Ok(original) +} + +/// Truncate the rendered prompt at [`CONTINUE_FINAL_MESSAGE_TAG`] so that it +/// ends exactly with the final message content, dropping any template suffix +/// such as end-of-turn markers. +fn truncate_prompt_at_continue_final_message_tag( + mut rendered: String, + final_message_text: &str, +) -> Result { + let tag_loc = rendered + .rfind(CONTINUE_FINAL_MESSAGE_TAG.trim_end()) + .filter(|_| rendered.contains(final_message_text.trim())); + let Some(tag_loc) = tag_loc else { + return Err(Error::ChatTemplate(format!( + "continue_final_message is set but the final message does not appear \ + in the prompt after applying the chat template! This can happen if \ + the chat template deletes portions of the final message. Final \ + message to continue: {}", + final_message_text.trim(), + ))); + }; + + if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) { + // The template preserved spacing, so a plain cut at the tag suffices. + rendered.truncate(tag_loc); + } else { + // The template trimmed the trailing spacing of the message content, so + // apply the same trimming to the retained prefix. + rendered.truncate(tag_loc); + rendered.truncate(rendered.trim_end().len()); + } + Ok(rendered) +} + fn to_template_tools(tools: &[ChatTool]) -> Vec { tools .iter() @@ -548,28 +635,124 @@ mod tests { ChatRole::Assistant, "The capital of", )]); + let template = + "{% if continue_final_message %}continue:{% endif %}{{ messages[0].content }}"; - assert_eq!( - render( - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"), - &request, - ) - .unwrap(), - "new" - ); + assert_eq!(render(Some(template), &request).unwrap(), "The capital of"); request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; assert_eq!( - render( - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"), - &request, - ) - .unwrap(), - "continue" + render(Some(template), &request).unwrap(), + "continue:The capital of" ); } + #[test] + fn continue_final_message_truncates_template_suffix() { + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "What is the capital of France?"), + ChatMessage::text(ChatRole::Assistant, "The capital of"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + // The Qwen3 template is unaware of `continue_final_message`; the + // end-of-turn marker it appends must still be stripped. + let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap(); + + expect![[r#" + <|im_start|>user + What is the capital of France?<|im_end|> + <|im_start|>assistant + + + + + The capital of"#]] + .assert_eq(&rendered); + } + + #[test] + fn continue_final_message_trims_like_the_template_does() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::Assistant, "Sure, ")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + // The template trims the trailing spacing of the message content, so + // the truncated prompt must be trimmed the same way. + let rendered = render( + Some("{{ messages[0].content.strip() }}<|im_end|>"), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "Sure,"); + } + + #[test] + fn continue_final_message_appends_to_last_text_part() { + // The renderer itself is role-agnostic like transformers (the + // assistant-final restriction is enforced by request validation + // upstream), so a multimodal user message exercises the part + // selection: the sentinel must attach to the last *text* part, + // skipping the trailing image. + let mut request = sample_request(vec![ChatMessage::user(vec![ + ChatContentPart::text("Sure,"), + ChatContentPart::image_url("data:image/png;base64,test"), + ])]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let rendered = render_mm( + "{% for item in messages[0].content %}{% if item.type == 'image' %}{% else %}{{ item.text }}{% endif %}{% endfor %}<|im_end|>", + &request, + ChatTemplateContentFormatOption::OpenAi, + ) + .unwrap() + .prompt; + + // Anything rendered after the continued text (here the image + // placeholder and the end marker) is truncated away, matching + // transformers. + assert_eq!(rendered, Prompt::Text("Sure,".to_string())); + } + + #[test] + fn continue_final_message_composes_with_aware_templates() { + // A template that reads `continue_final_message` and skips its own + // end-of-turn marker must produce the same prompt as an unaware one: + // the sentinel truncation degenerates to a cut at the very end. + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "hi"), + ChatMessage::text(ChatRole::Assistant, "Sure,"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let aware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}{% if not (loop.last and continue_final_message) %}<|im_end|>\n{% endif %}{% endfor %}"; + let unaware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}"; + + let expected = "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\nSure,"; + assert_eq!(render(Some(aware), &request).unwrap(), expected); + assert_eq!(render(Some(unaware), &request).unwrap(), expected); + } + + #[test] + fn continue_final_message_errors_when_template_drops_final_message() { + let mut request = sample_request(vec![ + ChatMessage::text(ChatRole::User, "hi"), + ChatMessage::text(ChatRole::Assistant, "Sure,"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant; + + let error = render( + Some( + "{% for m in messages %}{% if m.role == 'user' %}{{ m.content }}{% endif %}{% endfor %}", + ), + &request, + ) + .unwrap_err(); + + assert!(matches!(error, Error::ChatTemplate(_))); + } + #[test] fn chat_template_flattens_text_parts_for_string_templates() { let request = sample_request(vec![ChatMessage::user(vec![ diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index e6044614aa9..b1670814ef6 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -130,6 +130,19 @@ impl RoundtripCase { } } + /// DeepSeek V3.2 DSML tool-call format. + fn deepseek_v32() -> Self { + Self { + model_id: "deepseek-ai/DeepSeek-V3.2-Exp", + assistant_stop_suffix: "<｜end▁of▁sentence｜>", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: false }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// GLM-4.7 XML-like argument format with `` reasoning tags. fn glm47() -> Self { Self { @@ -235,6 +248,7 @@ roundtrip_tests! { qwen35 => [reasoning_and_content, tool_call_mix], minimax_m25 => [reasoning_and_content, tool_call_mix], deepseek_v4 => [reasoning_and_content, tool_call_mix], + deepseek_v32 => [tool_call_mix], glm47 => [reasoning_and_content, tool_call_mix], seed_oss => [reasoning_and_content], step3p5 => [reasoning_and_content], diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 521a188d6ac..548001412cb 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -456,13 +456,17 @@ pub struct ServerUnsupportedArgs { /// Enable the `/tokenizer_info` endpoint. May expose chat /// templates and other tokenizer configuration. + /// + /// Accepted as a no-op: the Rust frontend serves `/tokenize` and + /// `/detokenize`, but does not implement `/tokenizer_info` yet. #[arg( long, visible_alias = "no-enable-tokenizer-info-endpoint", default_missing_value = "true", - num_args = 0..=1 + num_args = 0..=1, + hide = true )] - pub enable_tokenizer_info_endpoint: Option, + pub enable_tokenizer_info_endpoint: Option, /// If set to True, log model outputs (generations). /// Requires `--enable-log-requests`. As with `--enable-log-requests`, diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index babc8179329..eb899e3e784 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -512,6 +512,7 @@ impl EngineCoreClient { Ok(EngineCoreOutputStream::new( request_id, + engine_id.engine_index().unwrap_or(0), self.abort_tx.clone(), rx, )) diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 3465817a453..c91a93d2714 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -15,7 +15,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{LoraInfoExporter, record_scheduler_stats}; +use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::EngineCoreRequestType; @@ -29,6 +29,7 @@ pub(crate) struct ClientInner { /// The runtime handle used for sending messages to the engine. handle: Handle, model_name: String, + scheduler_stats_recorder: SchedulerStatsRecorder, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, @@ -43,10 +44,13 @@ impl ClientInner { model_name: String, engines: &[ConnectedEngine], ) -> Self { + let scheduler_stats_recorder = + SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines); Self { input_send, handle, model_name, + scheduler_stats_recorder, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), @@ -389,12 +393,7 @@ pub(crate) async fn run_output_dispatcher_loop( "dropping scheduler stats for unknown engine" ); } - record_scheduler_stats( - &METRICS.scheduler, - inner.model_name(), - batch.engine_index, - scheduler_stats, - ); + inner.scheduler_stats_recorder.record(batch.engine_index, scheduler_stats); } // The engine's scheduler stats never carry adapter names; diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index 56c6a7cb663..b0ea180795a 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -45,6 +45,7 @@ impl Deref for EngineCoreStreamOutput { /// `finish_reason` is non-`None`. pub struct EngineCoreOutputStream { request_id: String, + engine_index: u32, abort_tx: mpsc::UnboundedSender, state: State, rx: OutputReceiver, @@ -53,11 +54,13 @@ pub struct EngineCoreOutputStream { impl EngineCoreOutputStream { pub(crate) fn new( request_id: String, + engine_index: u32, abort_tx: mpsc::UnboundedSender, rx: OutputReceiver, ) -> Self { Self { request_id, + engine_index, abort_tx, state: State::Running, rx, @@ -68,6 +71,11 @@ impl EngineCoreOutputStream { pub fn request_id(&self) -> &str { &self.request_id } + + /// Return the index of the engine that owns this request. + pub fn engine_index(&self) -> u32 { + self.engine_index + } } impl Stream for EngineCoreOutputStream { diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e129fdd6a27..05744db42df 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -1,90 +1,190 @@ +use std::collections::BTreeMap; use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ - EngineLabels, EnginePositionLabels, LoraAdapterNames, LoraInfoLabels, SchedulerMetrics, + EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, + LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::SchedulerStats; +use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; -/// Record the scheduler-stats-backed metrics for one engine at one point in -/// time. -pub(crate) fn record_scheduler_stats( - metrics: &SchedulerMetrics, - model_name: impl Into, - engine: u32, - stats: &SchedulerStats, -) { - let model_name = model_name.into(); - let labels = EngineLabels { - model_name: model_name.clone(), - engine, - }; +/// Cached scheduler-stats metric handles for all engines connected to one +/// frontend client. +pub(crate) struct SchedulerStatsRecorder { + engines: BTreeMap, +} + +/// Per-engine cached metric handles used while recording `SchedulerStats`. +struct SchedulerStatsHandles { + // Base labels reused for dynamic child labels. + labels: EngineLabels, // Scheduler state gauges. - metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs); - metrics - .scheduler_waiting - .get_or_create(&labels) - .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs); - metrics - .scheduler_waiting_by_reason - .get_or_create(&WaitingReasonLabels { - model_name: model_name.clone(), - engine, - reason: WAITING_REASON_CAPACITY, - }) - .set(stats.num_waiting_reqs); - metrics - .scheduler_waiting_by_reason - .get_or_create(&WaitingReasonLabels { - model_name: model_name.clone(), - engine, - reason: WAITING_REASON_DEFERRED, - }) - .set(stats.num_skipped_waiting_reqs); - metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage); + scheduler_running: U64Gauge, + scheduler_waiting: U64Gauge, + scheduler_waiting_capacity: U64Gauge, + scheduler_waiting_deferred: U64Gauge, + kv_cache_usage: F64Gauge, // Prefix-cache counters, including the connector-backed external cache path. - metrics - .prefix_cache_queries - .get_or_create(&labels) - .inc_by(stats.prefix_cache_stats.base.queries); - metrics - .prefix_cache_hits - .get_or_create(&labels) - .inc_by(stats.prefix_cache_stats.base.hits); + prefix_cache_queries: U64Counter, + prefix_cache_hits: U64Counter, + external_prefix_cache_queries: U64Counter, + external_prefix_cache_hits: U64Counter, + + // Speculative decoding counters. + spec_decode_num_drafts: U64Counter, + spec_decode_num_draft_tokens: U64Counter, + spec_decode_num_accepted_tokens: U64Counter, + spec_decode_num_accepted_tokens_per_pos: Family, + + // Per-engine performance / MFU counters. + estimated_flops_per_gpu: U64Counter, + estimated_read_bytes_per_gpu: U64Counter, + estimated_write_bytes_per_gpu: U64Counter, + + // Sampled KV-cache residency histograms. + kv_block_lifetime_seconds: HistogramMetric, + kv_block_idle_before_evict_seconds: HistogramMetric, + kv_block_reuse_gap_seconds: HistogramMetric, + + // Non-Prometheus interval accumulator for periodic text-log helpers. + log_stats: SchedulerLogStatsAccumulator, +} + +impl SchedulerStatsRecorder { + /// Resolve the fixed-label metric handles for the connected engines. + pub(crate) fn new( + metrics: &SchedulerMetrics, + model_name: &str, + engines: &[ConnectedEngine], + ) -> Self { + let engines = engines + .iter() + .filter_map(|engine| { + let engine = engine.engine_id.engine_index()?; + Some(( + engine, + resolve_scheduler_stats_handles(metrics, model_name, engine), + )) + }) + .collect(); + + Self { engines } + } + + /// Record one scheduler-stats payload for the given engine index. + pub(crate) fn record(&self, engine_index: u32, stats: &SchedulerStats) { + if let Some(handles) = self.engines.get(&engine_index) { + record_scheduler_stats_with_handles(handles, stats); + } + } +} + +/// Resolve all fixed-label scheduler metrics for one engine. +fn resolve_scheduler_stats_handles( + metrics: &SchedulerMetrics, + model_name: &str, + engine: u32, +) -> SchedulerStatsHandles { + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; + let capacity = WaitingReasonLabels { + model_name: model_name.to_string(), + engine, + reason: WAITING_REASON_CAPACITY, + }; + let deferred = WaitingReasonLabels { + model_name: model_name.to_string(), + engine, + reason: WAITING_REASON_DEFERRED, + }; + + SchedulerStatsHandles { + scheduler_running: metrics.scheduler_running.get_or_create_owned(&labels), + scheduler_waiting: metrics.scheduler_waiting.get_or_create_owned(&labels), + scheduler_waiting_capacity: metrics + .scheduler_waiting_by_reason + .get_or_create_owned(&capacity), + scheduler_waiting_deferred: metrics + .scheduler_waiting_by_reason + .get_or_create_owned(&deferred), + kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels), + prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels), + prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels), + external_prefix_cache_queries: metrics + .external_prefix_cache_queries + .get_or_create_owned(&labels), + external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels), + spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels), + spec_decode_num_draft_tokens: metrics + .spec_decode_num_draft_tokens + .get_or_create_owned(&labels), + spec_decode_num_accepted_tokens: metrics + .spec_decode_num_accepted_tokens + .get_or_create_owned(&labels), + spec_decode_num_accepted_tokens_per_pos: metrics + .spec_decode_num_accepted_tokens_per_pos + .clone(), + log_stats: metrics.log_stats.get_or_create_owned(&labels), + estimated_flops_per_gpu: metrics.estimated_flops_per_gpu.get_or_create_owned(&labels), + estimated_read_bytes_per_gpu: metrics + .estimated_read_bytes_per_gpu + .get_or_create_owned(&labels), + estimated_write_bytes_per_gpu: metrics + .estimated_write_bytes_per_gpu + .get_or_create_owned(&labels), + kv_block_lifetime_seconds: metrics.kv_block_lifetime_seconds.get_or_create_owned(&labels), + kv_block_idle_before_evict_seconds: metrics + .kv_block_idle_before_evict_seconds + .get_or_create_owned(&labels), + kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels), + labels, + } +} + +/// Record scheduler-stats values through pre-resolved metric handles. +fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &SchedulerStats) { + // Scheduler state gauges. + handles.scheduler_running.set(stats.num_running_reqs); + handles + .scheduler_waiting + .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs); + handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs); + handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs); + handles.kv_cache_usage.set(stats.kv_cache_usage); + + // Prefix-cache counters, including the connector-backed external cache path. + handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries); + handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits); if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats { - metrics + handles .external_prefix_cache_queries - .get_or_create(&labels) .inc_by(connector_prefix_cache_stats.base.queries); - metrics + handles .external_prefix_cache_hits - .get_or_create(&labels) .inc_by(connector_prefix_cache_stats.base.hits); } // Speculative decoding counters. if let Some(spec_decoding_stats) = &stats.spec_decoding_stats { - metrics - .spec_decode_num_drafts - .get_or_create(&labels) - .inc_by(spec_decoding_stats.num_drafts); - metrics + handles.spec_decode_num_drafts.inc_by(spec_decoding_stats.num_drafts); + handles .spec_decode_num_draft_tokens - .get_or_create(&labels) .inc_by(spec_decoding_stats.num_draft_tokens); - metrics + handles .spec_decode_num_accepted_tokens - .get_or_create(&labels) .inc_by(spec_decoding_stats.num_accepted_tokens); - metrics.log_stats.get_or_create(&labels).observe_spec_decode( + handles.log_stats.observe_spec_decode( spec_decoding_stats.num_drafts, &spec_decoding_stats.num_accepted_tokens_per_pos, ); @@ -92,11 +192,11 @@ pub(crate) fn record_scheduler_stats( for (position, accepted_tokens) in spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate() { - metrics + handles .spec_decode_num_accepted_tokens_per_pos .get_or_create(&EnginePositionLabels { - model_name: model_name.clone(), - engine, + model_name: handles.labels.model_name.clone(), + engine: handles.labels.engine, position: position as u32, }) .inc_by(accepted_tokens); @@ -109,22 +209,13 @@ pub(crate) fn record_scheduler_stats( || perf_stats.num_read_bytes_per_gpu != 0 || perf_stats.num_write_bytes_per_gpu != 0) { - metrics - .estimated_flops_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_flops_per_gpu); - metrics - .estimated_read_bytes_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_read_bytes_per_gpu); - metrics - .estimated_write_bytes_per_gpu - .get_or_create(&labels) - .inc_by(perf_stats.num_write_bytes_per_gpu); + handles.estimated_flops_per_gpu.inc_by(perf_stats.num_flops_per_gpu); + handles.estimated_read_bytes_per_gpu.inc_by(perf_stats.num_read_bytes_per_gpu); + handles.estimated_write_bytes_per_gpu.inc_by(perf_stats.num_write_bytes_per_gpu); } if let Some(cudagraph_stats) = &stats.cudagraph_stats { - metrics.log_stats.get_or_create(&labels).observe_cudagraph( + handles.log_stats.observe_cudagraph( cudagraph_stats.num_unpadded_tokens, cudagraph_stats.num_padded_tokens, cudagraph_stats.num_paddings, @@ -134,16 +225,11 @@ pub(crate) fn record_scheduler_stats( // Sampled KV-cache residency histograms. if !stats.kv_cache_eviction_events.is_empty() { - let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels); - let kv_block_idle_before_evict_seconds = - metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels); - let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels); - for event in &stats.kv_cache_eviction_events { - kv_block_lifetime_seconds.observe(event.lifetime_seconds); - kv_block_idle_before_evict_seconds.observe(event.idle_seconds); + handles.kv_block_lifetime_seconds.observe(event.lifetime_seconds); + handles.kv_block_idle_before_evict_seconds.observe(event.idle_seconds); for reuse_gap_seconds in &event.reuse_gaps_seconds { - kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds); + handles.kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds); } } } diff --git a/rust/src/engine-core-client/src/protocol/tensor.rs b/rust/src/engine-core-client/src/protocol/tensor.rs index b6711215481..b80472129b7 100644 --- a/rust/src/engine-core-client/src/protocol/tensor.rs +++ b/rust/src/engine-core-client/src/protocol/tensor.rs @@ -11,6 +11,21 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; /// const CUSTOM_TYPE_RAW_VIEW: i8 = 3; +#[derive(Serialize)] +#[serde(rename = "_ExtStruct")] +struct MsgpackExtRef<'a>((i8, ByteSlice<'a>)); + +struct ByteSlice<'a>(&'a [u8]); + +impl Serialize for ByteSlice<'_> { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + serializer.serialize_bytes(self.0) + } +} + #[easy_ext::ext(ShapeExt)] impl [usize] { /// Returned the total number of elements implied by this shape, or `None` @@ -184,7 +199,7 @@ impl Serialize for WireArrayData { match self { Self::AuxIndex(index) => serializer.serialize_u64(*index as u64), Self::RawView(bytes) => { - Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer) + MsgpackExtRef((CUSTOM_TYPE_RAW_VIEW, ByteSlice(bytes))).serialize(serializer) } } } @@ -194,6 +209,21 @@ impl Serialize for WireArrayData { mod tests { use super::*; + #[test] + fn raw_view_serializes_as_msgpack_ext() { + let bytes = vec![1, 2, 3, 4]; + let encoded = + rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode"); + let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone())) + .expect("encode expected"); + + assert_eq!(encoded, expected); + assert_eq!( + rmpv::decode::read_value(&mut std::io::Cursor::new(encoded)).expect("decode"), + Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes) + ); + } + #[test] fn constructors_build_raw_view_tensors() { let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap(); diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index ce15a970c63..942bf55c288 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -14,6 +14,7 @@ pub use output::{ GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage, }; pub use request::GenerateRequest; +pub use request_metrics::current_unix_timestamp_secs; pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob}; use crate::inflight::InflightRequests; @@ -88,14 +89,21 @@ impl Llm { // Record internal engine-core request ID in the current tracing span. Span::current().record("engine_request_id", &internal_request_id); + let arrival_time = prepared.engine_request.arrival_time; + let max_tokens_param = + (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens); + let prompt_len = prepared.prompt_token_ids().len() as u32; + + let stream = self.client.call(prepared.engine_request).await?; + let request_metrics = RequestMetricsTracker::new( self.client.model_name().to_string(), - prepared.engine_request.arrival_time, - prepared.prompt_token_ids().len() as u32, - (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens), + stream.engine_index(), + arrival_time, + prompt_len, + max_tokens_param, 1, ); - let stream = self.client.call(prepared.engine_request).await?; let guard = self.inflight.track(external_request_id, internal_request_id); Ok(GenerateOutputStream::new( diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 609088ec6b2..d1d7e5f46e5 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -248,12 +248,7 @@ impl Stream for GenerateOutputStream { }; let received_at = current_unix_timestamp_secs(); - self.request_metrics.observe_output( - raw.engine_index, - raw.timestamp, - received_at, - &raw.output, - ); + self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); let raw = raw.output; diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs index 159cf823de4..45e5bd1ca64 100644 --- a/rust/src/llm/src/request.rs +++ b/rust/src/llm/src/request.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::time::{SystemTime, UNIX_EPOCH}; use uuid::Uuid; use vllm_engine_core_client::protocol::lora::LoraRequest; @@ -8,6 +7,7 @@ use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningPar use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use crate::error::{Error, Result}; +use crate::request_metrics::current_unix_timestamp_secs; /// Tokenized decoder-only generate request accepted by [`crate::Llm`]. /// @@ -30,8 +30,9 @@ pub struct GenerateRequest { pub mm_features: Option, /// Unix timestamp, in seconds, when this request arrived at the frontend. /// - /// When omitted, the Rust frontend fills it immediately before sending the - /// request to engine-core, matching Python's default arrival-time behavior. + /// Stamped at the frontend entry, before render and tokenization, to match + /// Python's renderer-entry arrival_time. When omitted, it is filled as a + /// fallback before the request is sent to engine-core. pub arrival_time: Option, /// Optional salt used to partition prefix-cache entries for this request. pub cache_salt: Option, @@ -122,13 +123,6 @@ impl PreparedGenerateRequest { } } -fn current_unix_timestamp_secs() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("system clock is before unix epoch") - .as_secs_f64() -} - #[cfg(test)] mod tests { use std::collections::BTreeMap; diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 38795a70928..4f1673db154 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -5,15 +5,12 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_metrics::{ - EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics, + EngineLabels, Family, FinishedReasonLabels, HistogramMetric, METRICS, PromptTokenSourceLabels, + U64Counter, }; use crate::FinishReason; -fn metrics() -> &'static RequestMetrics { - &METRICS.request -} - const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute"; const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit"; const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; @@ -29,9 +26,11 @@ const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer"; /// /// Original Python update flow: /// -#[derive(Debug, Clone)] +#[derive(Clone)] pub(crate) struct RequestMetricsTracker { - model_name: String, + /// Cached request metric handles for this request's model and engine index. + handles: RequestMetricHandles, + arrival_time: f64, prompt_len: u32, max_tokens_param: Option, @@ -44,7 +43,38 @@ pub(crate) struct RequestMetricsTracker { first_token_latency: f64, num_generation_tokens: u32, latest_num_cached_tokens: u32, - last_seen_engine_index: u32, +} + +/// Cached request metric handles for one model and engine index. +#[derive(Clone)] +struct RequestMetricHandles { + labels: EngineLabels, + + // Request-derived counters. + num_preemptions: U64Counter, + prompt_tokens: U64Counter, + prompt_tokens_local_compute: U64Counter, + prompt_tokens_local_cache_hit: U64Counter, + prompt_tokens_external_kv_transfer: U64Counter, + prompt_tokens_cached: U64Counter, + generation_tokens: U64Counter, + + // Request lifecycle counters and histograms. + request_success: Family, + request_prompt_tokens: HistogramMetric, + request_generation_tokens: HistogramMetric, + request_max_num_generation_tokens: HistogramMetric, + request_params_max_tokens: HistogramMetric, + request_params_n: HistogramMetric, + request_prefill_kv_computed_tokens: HistogramMetric, + time_to_first_token_seconds: HistogramMetric, + inter_token_latency_seconds: HistogramMetric, + e2e_request_latency_seconds: HistogramMetric, + request_queue_time_seconds: HistogramMetric, + request_prefill_time_seconds: HistogramMetric, + request_decode_time_seconds: HistogramMetric, + request_inference_time_seconds: HistogramMetric, + request_time_per_output_token_seconds: HistogramMetric, } impl RequestMetricsTracker { @@ -52,13 +82,14 @@ impl RequestMetricsTracker { /// context. pub(crate) fn new( model_name: String, + engine_index: u32, arrival_time: f64, prompt_len: u32, max_tokens_param: Option, n_param: u32, ) -> Self { Self { - model_name, + handles: resolve_request_metric_handles(&model_name, engine_index), arrival_time, prompt_len, max_tokens_param, @@ -71,7 +102,6 @@ impl RequestMetricsTracker { first_token_latency: 0.0, num_generation_tokens: 0, latest_num_cached_tokens: 0, - last_seen_engine_index: 0, } } @@ -81,23 +111,18 @@ impl RequestMetricsTracker { /// pub(crate) fn observe_output( &mut self, - engine_index: u32, batch_timestamp: f64, received_at: f64, output: &EngineCoreOutput, ) { - self.last_seen_engine_index = engine_index; if let Some(prefill_stats) = &output.prefill_stats { self.latest_num_cached_tokens = prefill_stats.num_cached_tokens; } self.num_generation_tokens += output.new_token_ids.len() as u32; - metrics() - .generation_tokens - .get_or_create(&engine_labels(&self.model_name, engine_index)) - .inc_by(output.new_token_ids.len() as u64); + self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64); if let Some(events) = &output.events { - self.observe_events(engine_index, events); + self.observe_events(events); } // Only outputs that actually carry tokens drive token-timing metrics. @@ -107,22 +132,16 @@ impl RequestMetricsTracker { if !output.new_token_ids.is_empty() { if self.is_prefilling { if let Some(prefill_stats) = &output.prefill_stats { - record_prompt_tokens(&self.model_name, engine_index, prefill_stats); + self.record_prompt_tokens(prefill_stats); } self.first_token_latency = received_at - self.arrival_time; - observe_time_to_first_token_seconds( - &self.model_name, - engine_index, - self.first_token_latency, - ); + self.handles.time_to_first_token_seconds.observe(self.first_token_latency); self.first_token_ts = batch_timestamp; self.is_prefilling = false; } else if self.last_token_ts > 0.0 { - observe_inter_token_latency_seconds( - &self.model_name, - engine_index, - batch_timestamp - self.last_token_ts, - ); + self.handles + .inter_token_latency_seconds + .observe(batch_timestamp - self.last_token_ts); } self.last_token_ts = batch_timestamp; @@ -135,7 +154,6 @@ impl RequestMetricsTracker { /// Original Python finished-request stats: /// pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) { - let labels = engine_labels(&self.model_name, self.last_seen_engine_index); let prefill_kv_computed_tokens = self.prompt_len.saturating_sub(self.latest_num_cached_tokens); let e2e_latency_seconds = received_at - self.arrival_time; @@ -150,57 +168,47 @@ impl RequestMetricsTracker { 0.0 }; - record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason); - metrics() - .request_prompt_tokens - .get_or_create(&labels) - .observe(self.prompt_len as f64); - metrics() + self.record_request_success(finish_reason); + + self.handles.request_prompt_tokens.observe(self.prompt_len as f64); + self.handles .request_generation_tokens - .get_or_create(&labels) .observe(self.num_generation_tokens as f64); - metrics() + self.handles .request_max_num_generation_tokens - .get_or_create(&labels) .observe(self.num_generation_tokens as f64); if let Some(max_tokens_param) = self.max_tokens_param { - metrics() - .request_params_max_tokens - .get_or_create(&labels) - .observe(max_tokens_param as f64); + self.handles.request_params_max_tokens.observe(max_tokens_param as f64); } - metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64); - metrics() + self.handles.request_params_n.observe(self.n_param as f64); + self.handles .request_prefill_kv_computed_tokens - .get_or_create(&labels) .observe(prefill_kv_computed_tokens as f64); - metrics() - .e2e_request_latency_seconds - .get_or_create(&labels) - .observe(e2e_latency_seconds); - metrics() - .request_queue_time_seconds - .get_or_create(&labels) - .observe(queue_time_seconds); - metrics() - .request_prefill_time_seconds - .get_or_create(&labels) - .observe(prefill_time_seconds); - metrics() - .request_decode_time_seconds - .get_or_create(&labels) - .observe(decode_time_seconds); - metrics() - .request_inference_time_seconds - .get_or_create(&labels) - .observe(inference_time_seconds); - metrics() + self.handles.e2e_request_latency_seconds.observe(e2e_latency_seconds); + self.handles.request_queue_time_seconds.observe(queue_time_seconds); + self.handles.request_prefill_time_seconds.observe(prefill_time_seconds); + self.handles.request_decode_time_seconds.observe(decode_time_seconds); + self.handles.request_inference_time_seconds.observe(inference_time_seconds); + self.handles .request_time_per_output_token_seconds - .get_or_create(&labels) .observe(time_per_output_token_seconds); } - fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) { + /// Record prompt token counters through cached metric handles. + fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) { + let computed = prefill_stats.num_computed_tokens as u64; + let local_cache_hit = prefill_stats.num_local_cached_tokens as u64; + let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64; + + self.handles.prompt_tokens.inc_by(prefill_stats.num_prompt_tokens as u64); + self.handles.prompt_tokens_local_compute.inc_by(computed); + self.handles.prompt_tokens_local_cache_hit.inc_by(local_cache_hit); + self.handles.prompt_tokens_external_kv_transfer.inc_by(external_kv_transfer); + self.handles.prompt_tokens_cached.inc_by(prefill_stats.num_cached_tokens as u64); + } + + /// Record request event counters through cached metric handles. + fn observe_events(&mut self, events: &[EngineCoreEvent]) { for event in events { match event.r#type { EngineCoreEventType::Queued => { @@ -212,46 +220,86 @@ impl RequestMetricsTracker { } } EngineCoreEventType::Preempted => { - metrics() - .num_preemptions - .get_or_create(&engine_labels(&self.model_name, engine_index)) - .inc(); + self.handles.num_preemptions.inc(); } } } } -} -fn engine_labels(model_name: &str, engine: u32) -> EngineLabels { - EngineLabels { - model_name: model_name.to_string(), - engine, + /// Increment the request-success counter for the terminal finish reason. + fn record_request_success(&self, finish_reason: FinishReason) { + self.handles + .request_success + .get_or_create(&FinishedReasonLabels { + model_name: self.handles.labels.model_name.clone(), + engine: self.handles.labels.engine, + finished_reason: finish_reason.as_str(), + }) + .inc(); } } -fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) { - metrics() - .time_to_first_token_seconds - .get_or_create(&engine_labels(model_name, engine)) - .observe(seconds); -} +/// Resolve fixed request metric handles for one model and engine index. +fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetricHandles { + let metrics = &METRICS.request; + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; -fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) { - metrics() - .inter_token_latency_seconds - .get_or_create(&engine_labels(model_name, engine)) - .observe(seconds); -} - -fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) { - metrics() - .request_success - .get_or_create(&FinishedReasonLabels { - model_name: model_name.to_string(), - engine, - finished_reason: finish_reason.as_str(), - }) - .inc(); + RequestMetricHandles { + num_preemptions: metrics.num_preemptions.get_or_create_owned(&labels), + prompt_tokens: metrics.prompt_tokens.get_or_create_owned(&labels), + prompt_tokens_local_compute: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE), + ), + prompt_tokens_local_cache_hit: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT), + ), + prompt_tokens_external_kv_transfer: metrics.prompt_tokens_by_source.get_or_create_owned( + &prompt_token_source_labels( + model_name, + engine, + PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER, + ), + ), + prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels), + generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels), + request_success: metrics.request_success.clone(), + request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels), + request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels), + request_max_num_generation_tokens: metrics + .request_max_num_generation_tokens + .get_or_create_owned(&labels), + request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels), + request_params_n: metrics.request_params_n.get_or_create_owned(&labels), + request_prefill_kv_computed_tokens: metrics + .request_prefill_kv_computed_tokens + .get_or_create_owned(&labels), + time_to_first_token_seconds: metrics + .time_to_first_token_seconds + .get_or_create_owned(&labels), + inter_token_latency_seconds: metrics + .inter_token_latency_seconds + .get_or_create_owned(&labels), + e2e_request_latency_seconds: metrics + .e2e_request_latency_seconds + .get_or_create_owned(&labels), + request_queue_time_seconds: metrics.request_queue_time_seconds.get_or_create_owned(&labels), + request_prefill_time_seconds: metrics + .request_prefill_time_seconds + .get_or_create_owned(&labels), + request_decode_time_seconds: metrics + .request_decode_time_seconds + .get_or_create_owned(&labels), + request_inference_time_seconds: metrics + .request_inference_time_seconds + .get_or_create_owned(&labels), + request_time_per_output_token_seconds: metrics + .request_time_per_output_token_seconds + .get_or_create_owned(&labels), + labels, + } } fn prompt_token_source_labels( @@ -266,45 +314,6 @@ fn prompt_token_source_labels( } } -fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) { - let computed = prefill_stats.num_computed_tokens as u64; - let local_cache_hit = prefill_stats.num_local_cached_tokens as u64; - let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64; - - metrics() - .prompt_tokens - .get_or_create(&engine_labels(model_name, engine)) - .inc_by(prefill_stats.num_prompt_tokens as u64); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE, - )) - .inc_by(computed); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT, - )) - .inc_by(local_cache_hit); - metrics() - .prompt_tokens_by_source - .get_or_create(&prompt_token_source_labels( - model_name, - engine, - PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER, - )) - .inc_by(external_kv_transfer); - metrics() - .prompt_tokens_cached - .get_or_create(&engine_labels(model_name, engine)) - .inc_by(prefill_stats.num_cached_tokens as u64); -} - fn diff_or_zero(end: f64, start: f64) -> f64 { if end > 0.0 && start > 0.0 && end >= start { end - start @@ -321,7 +330,7 @@ fn diff_or_zero(end: f64, start: f64) -> f64 { /// /// Original Python request timestamp source: /// -pub(crate) fn current_unix_timestamp_secs() -> f64 { +pub fn current_unix_timestamp_secs() -> f64 { SystemTime::now() .duration_since(UNIX_EPOCH) .expect("system clock is before unix epoch") @@ -337,10 +346,10 @@ mod tests { #[test] fn tracker_updates_timing_state_across_prefill_decode_and_finish() { - let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1); + let mut tracker = + RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1); tracker.observe_output( - 2, 10.0, 100.2, &vllm_engine_core_client::protocol::output::EngineCoreOutput { @@ -368,7 +377,6 @@ mod tests { }, ); tracker.observe_output( - 2, 11.5, 100.4, &vllm_engine_core_client::protocol::output::EngineCoreOutput { @@ -384,7 +392,7 @@ mod tests { ); assert!(!tracker.is_prefilling); - assert_eq!(tracker.last_seen_engine_index, 2); + assert_eq!(tracker.handles.labels.engine, 2); assert_eq!(tracker.num_generation_tokens, 3); assert_eq!(tracker.queued_ts, 8.0); assert_eq!(tracker.scheduled_ts, 9.0); diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 9911841eade..8581b1ac08f 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -17,7 +17,7 @@ use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task}; -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig}; +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId}; use vllm_llm::{ Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm, }; @@ -699,7 +699,7 @@ async fn abort_by_external_id_aborts_all_internal_requests() { async fn generate_records_request_metrics_in_prometheus_output() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-metrics".to_vec(); + let engine_id = EngineId::from_engine_index(4); let model_name = request_metrics_model_name("metrics-model"); let (shutdown_tx, engine_task) = spawn_mock_engine_task( @@ -832,7 +832,7 @@ async fn generate_records_request_metrics_in_prometheus_output() { async fn dropping_stream_records_abort_terminal_request_metrics() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-metrics-drop".to_vec(); + let engine_id = EngineId::from_engine_index(5); let model_name = request_metrics_model_name("metrics-drop-model"); let (shutdown_tx, engine_task) = spawn_mock_engine_task( diff --git a/rust/src/metrics/src/lib.rs b/rust/src/metrics/src/lib.rs index 8f0db53d3ff..ca650fbadae 100644 --- a/rust/src/metrics/src/lib.rs +++ b/rust/src/metrics/src/lib.rs @@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64; use prometheus_client::encoding::text::encode; use prometheus_client::metrics::counter::Counter; -use prometheus_client::metrics::family::Family; +pub use prometheus_client::metrics::family::Family; use prometheus_client::metrics::gauge::Gauge; use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; @@ -23,6 +23,8 @@ pub use scheduler::*; pub type U64Counter = Counter; pub type U64Gauge = Gauge; pub type F64Gauge = Gauge; +/// Histogram metric handle cloned out of a Prometheus family. +pub type HistogramMetric = Histogram; pub(crate) type HistogramFamily = Family Histogram>; /// Shared Prometheus registry for frontend metrics. diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 9de89727205..4327221221d 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -94,6 +94,7 @@ pub fn to_text_request( data_parallel_rank: None, reasoning_parser_kwargs: None, lora_request: None, + arrival_time: None, }) } diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 7a6dcdfc45c..965155b5825 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -70,6 +70,7 @@ pub(super) fn prepare_generate_request( data_parallel_rank: ctx.data_parallel_rank, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), + arrival_time: None, }; Ok(PreparedRequest { diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index d4567c44aa6..28855968df0 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -29,7 +29,9 @@ pub struct GenerateRequest { impl Normalizable for GenerateRequest {} /// Mirrors the Python vLLM `GenerateResponseChoice` class. -#[serde_with::skip_serializing_none] +/// +/// Do not skip serializing `None` fields here: non-streaming response types +/// should serialize `None` as explicit `null`. #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateResponseChoice { pub index: u32, @@ -58,7 +60,6 @@ pub(super) struct GenerateStreamResponse { } /// Mirrors the Python vLLM `GenerateResponse` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateResponse { pub request_id: String, @@ -68,7 +69,6 @@ pub(super) struct GenerateResponse { } /// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub(super) struct GenerateLogprob { pub logprob: f32, diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 0429d9ac8b5..f0368c5614a 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -211,7 +211,7 @@ async fn collect_chat_completion( Some(prefix) => Some(format!("{prefix}{}", message.text())), None => Some(message.text()).filter(|t| !t.is_empty()), }, - tool_calls: Some(tool_calls).filter(|calls| !calls.is_empty()), + tool_calls, reasoning: if include_reasoning { reasoning } else { None }, }, logprobs, 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 26a76271495..a284b2197fc 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -830,7 +830,7 @@ mod tests { let message = ChatCompletionMessage { role: AssistantRole, content: Some("answer".to_string()), - tool_calls: None, + tool_calls: Vec::new(), reasoning: Some("inner".to_string()), }; let message_json = serde_json::to_value(message).expect("message serializes"); diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index b5681d0abdb..f3a4dd24131 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -331,7 +331,9 @@ impl Normalizable for ChatCompletionRequest { } /// Mirrors the Python vLLM `ChatCompletionResponse` class. -#[serde_with::skip_serializing_none] +/// +/// Do not skip serializing `None` fields here: non-streaming response types +/// should serialize `None` as explicit `null`. #[derive(Debug, Clone, Serialize)] pub(super) struct ChatCompletionResponse { pub id: String, @@ -347,7 +349,6 @@ pub(super) struct ChatCompletionResponse { } /// Mirrors the Python vLLM `ChatCompletionResponseChoice` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub(super) struct ChatCompletionChoice { pub index: u32, @@ -370,12 +371,12 @@ impl fmt::Display for AssistantRole { } /// Mirrors the Python vLLM response `ChatMessage` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub(super) struct ChatCompletionMessage { pub role: AssistantRole, pub content: Option, - pub tool_calls: Option>, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub tool_calls: Vec, pub reasoning: Option, } diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index fa80e16e0f9..1355481b49b 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request( data_parallel_rank: ctx.data_parallel_rank, reasoning_parser_kwargs: None, lora_request: lora_resolution.lora_request.clone(), + arrival_time: None, }; Ok(PreparedRequest { diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index fe637d82dea..32542b8b351 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -196,7 +196,9 @@ impl Normalizable for CompletionRequest { } /// Mirrors the Python vLLM `CompletionResponse` class. -#[serde_with::skip_serializing_none] +/// +/// Do not skip serializing `None` fields here: non-streaming response types +/// should serialize `None` as explicit `null`. #[derive(Debug, Clone, Serialize)] pub(super) struct CompletionResponse { pub id: String, @@ -210,7 +212,6 @@ pub(super) struct CompletionResponse { } /// Mirrors the Python vLLM `CompletionResponseChoice` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub(super) struct CompletionChoice { pub index: u32, diff --git a/rust/src/server/src/routes/openai/utils/types.rs b/rust/src/server/src/routes/openai/utils/types.rs index 8b079bbcc13..98eb10f937a 100644 --- a/rust/src/server/src/routes/openai/utils/types.rs +++ b/rust/src/server/src/routes/openai/utils/types.rs @@ -311,7 +311,9 @@ pub enum MessageContent { // ============================================================================ /// Mirrors the Python vLLM `UsageInfo` class. -#[serde_with::skip_serializing_none] +/// +/// Do not skip serializing `None` fields here: non-streaming response types +/// should serialize `None` as explicit `null`. #[derive(Debug, Clone, Serialize)] pub struct Usage { pub prompt_tokens: usize, @@ -402,14 +404,12 @@ pub struct LogProbs { } /// Mirrors the Python vLLM `ChatCompletionLogProbs` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub struct ChatLogProbs { pub content: Option>, } /// Mirrors the Python vLLM `ChatCompletionLogProbsContent` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub struct ChatLogProbsContent { pub token: String, @@ -419,7 +419,6 @@ pub struct ChatLogProbsContent { } /// Mirrors the Python vLLM `ChatCompletionLogProb` class. -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] pub struct TopLogProb { pub token: String, @@ -436,7 +435,6 @@ pub struct ErrorResponse { pub error: ErrorDetail, } -#[serde_with::skip_serializing_none] #[derive(Debug, Clone, Deserialize, Serialize)] pub struct ErrorDetail { pub message: String, diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 02348151ea7..520d5043832 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -2191,6 +2191,28 @@ async fn non_stream_chat_returns_json_response() { assert_eq!(json["usage"]["prompt_tokens"], 22); assert_eq!(json["usage"]["completion_tokens"], 3); assert_eq!(json["usage"]["total_tokens"], 25); + + // Unset optional fields are serialized as explicit `null` on + // non-streaming responses... + let response_object = json.as_object().expect("response object"); + let choice = json["choices"][0].as_object().expect("choice object"); + let message = choice["message"].as_object().expect("message object"); + for (object, key) in [ + (response_object, "system_fingerprint"), + (response_object, "prompt_token_ids"), + (response_object, "kv_transfer_params"), + (choice, "logprobs"), + (choice, "stop_reason"), + (choice, "token_ids"), + (message, "reasoning"), + ] { + assert!( + object.contains_key(key) && object[key].is_null(), + "expected explicit null `{key}`: {json}" + ); + } + // ...except `tool_calls`, which Python pops from the payload when empty. + assert!(!message.contains_key("tool_calls"), "{json}"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -2956,6 +2978,25 @@ async fn non_stream_completions_return_json_response() { assert_eq!(json["choices"][0]["text"], "hi"); assert_eq!(json["choices"][0]["finish_reason"], "stop"); assert_eq!(json["usage"]["completion_tokens"], 3); + + // Unset optional fields are serialized as explicit `null` on + // non-streaming responses. + let response_object = json.as_object().expect("response object"); + let choice = json["choices"][0].as_object().expect("choice object"); + for (object, key) in [ + (response_object, "system_fingerprint"), + (response_object, "kv_transfer_params"), + (choice, "logprobs"), + (choice, "stop_reason"), + (choice, "prompt_logprobs"), + (choice, "token_ids"), + (choice, "prompt_token_ids"), + ] { + assert!( + object.contains_key(key) && object[key].is_null(), + "expected explicit null `{key}`: {json}" + ); + } } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -4371,10 +4412,10 @@ async fn include_reasoning_false_suppresses_reasoning_in_non_stream_chat() { let json: serde_json::Value = serde_json::from_str(&text).expect("decode json"); assert_eq!(json["choices"][0]["message"]["content"], "answer"); + // Suppressed fields are serialized as explicit `null` on non-streaming + // responses. assert!( - json["choices"][0]["message"] - .as_object() - .is_some_and(|message| !message.contains_key("reasoning")), + json["choices"][0]["message"]["reasoning"].is_null(), "{text}" ); } @@ -4476,14 +4517,14 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() { let choice = json["choices"][0].as_object().expect("choice object"); assert_eq!(json["choices"][0]["message"]["content"], "answer"); + // Suppressed fields are serialized as explicit `null` on non-streaming + // responses. assert!( - json["choices"][0]["message"] - .as_object() - .is_some_and(|message| !message.contains_key("reasoning")), + json["choices"][0]["message"]["reasoning"].is_null(), "{text}" ); - assert!(!choice.contains_key("logprobs"), "{text}"); - assert!(!choice.contains_key("token_ids"), "{text}"); + assert!(choice["logprobs"].is_null(), "{text}"); + assert!(choice["token_ids"].is_null(), "{text}"); assert!(json["prompt_token_ids"].is_array(), "{text}"); } diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 987e0e23f39..f8c3e1cee85 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -102,12 +102,13 @@ pub struct DetokenizeRequest { pub tokens: Vec, } +/// Do not skip serializing `None` fields here: non-streaming response types +/// should serialize `None` as explicit `null`. #[derive(Debug, Clone, Serialize)] pub struct TokenizeResponse { pub count: usize, pub max_model_len: u32, pub tokens: Vec, - #[serde(skip_serializing_if = "Option::is_none")] pub token_strs: Option>, } diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index 7f84c90f39c..d8ddb0139c9 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -378,7 +378,7 @@ mod tests { fs::write(dir.path().join("tokenizer.json"), "{}").expect("write tokenizer"); fs::write( dir.path().join("tokenizer_config.json"), - r#"{"tokenizer_class":"PreTrainedTokenizerFast"}"#, + r#"{"tokenizer_class":"TokenizersBackend"}"#, ) .expect("write tokenizer config"); fs::write(dir.path().join("config.json"), "{}").expect("write config"); diff --git a/rust/src/text/src/lib.rs b/rust/src/text/src/lib.rs index 2987ef93e57..130eaec7f42 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -132,6 +132,10 @@ impl TextLlm { ) -> Result<(TextRequest, GenerateOutputStream)> { request.validate()?; + if request.arrival_time.is_none() { + request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs()); + } + let tokenizer = self.backend.tokenizer(); let prompt_token_ids = match take(&mut request.prompt) { Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index eabece37f09..164d2a3db06 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -18,7 +18,7 @@ use crate::request::{SamplingParams, TextRequest}; /// One text request after it has been lowered into the raw generate boundary. #[derive(Debug)] pub struct PreparedTextRequest { - /// The original high-level request, preserved for response-side metadata + /// The high-level request fields still needed for response-side metadata /// and decoding options. pub text_request: TextRequest, /// The southbound request ready to be sent to `vllm-llm`. @@ -28,7 +28,7 @@ pub struct PreparedTextRequest { /// Convert a high-level [`TextRequest`] into one lower-level /// [`GenerateRequest`] ready for the `llm` crate. pub fn lower_text_request( - request: TextRequest, + mut request: TextRequest, prompt_token_ids: Vec, sampling_hints: SamplingHints, sampling_limits: SamplingLimits, @@ -40,7 +40,10 @@ pub fn lower_text_request( let generate_request = GenerateRequest { request_id: request.request_id.clone(), prompt_token_ids, - mm_features: request.mm_features.clone(), + // Align with Python's response path: decoded output state does not retain + // `mm_features`; move them to the engine request to avoid cloning large + // multimodal tensor payloads. + mm_features: request.mm_features.take(), sampling_params: lower_sampling_params( request.sampling_params.clone(), sampling_hints, @@ -53,7 +56,7 @@ pub fn lower_text_request( data_parallel_rank: request.data_parallel_rank, reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(), lora_request: request.lora_request.clone(), - arrival_time: None, + arrival_time: request.arrival_time, trace_headers: None, }; @@ -307,6 +310,7 @@ mod tests { use std::collections::{BTreeSet, HashMap}; use serial_test::file_serial; + use vllm_engine_core_client::protocol::multimodal::{MmFeatureSpec, PlaceholderRange}; use vllm_tokenizer::test_utils::TestTokenizer; use super::*; @@ -574,6 +578,35 @@ mod tests { .assert_debug_eq(¶ms); } + #[test] + fn lower_text_request_moves_multimodal_features_to_generate_request() { + let features = vec![MmFeatureSpec { + data: None, + modality: "image".to_string(), + identifier: "image-1".to_string(), + mm_position: PlaceholderRange { + offset: 2, + length: 4, + is_embed: None, + }, + mm_hash: Some("hash-1".to_string()), + }]; + let mut request = sample_request(); + request.mm_features = Some(features.clone()); + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.mm_features, Some(features)); + assert_eq!(prepared.text_request.mm_features, None); + } + #[test] fn lower_text_request_uses_union_vocab_for_prompt_token_ids() { lower_text_request( @@ -1110,6 +1143,44 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_passes_arrival_time_through() { + let request = TextRequest { + arrival_time: Some(42.5), + ..sample_request() + }; + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.arrival_time, Some(42.5)); + } + + #[test] + fn lower_text_request_leaves_arrival_time_unset_when_absent() { + let request = TextRequest { + arrival_time: None, + ..sample_request() + }; + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!(prepared.generate_request.arrival_time, None); + } + #[test] fn resolve_max_tokens_user_smaller_than_model_limit() { let result = resolve_max_tokens(Some(50), None, 200, 100); diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index c2371858837..ad37d864904 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -10,7 +10,7 @@ pub enum TokenIdsError { #[error("allowed_token_ids should not be empty")] EmptyAllowedTokenIds, #[error( - "token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \ + "token_id(s) {token_ids:?} in {parameter} are out of vocabulary. \ Vocabulary size: {vocab_size}" )] OutOfVocab { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index c64e9ca05e6..09522868872 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -187,6 +187,12 @@ pub struct TextRequest { /// LoRA adapter selected for this request. #[serde(default)] pub lora_request: Option, + /// Wall-clock unix timestamp (seconds) when this request arrived at the + /// frontend, stamped before render/tokenize to match Python's + /// renderer-entry arrival_time. When unset, it is stamped before + /// tokenization. + #[serde(default)] + pub arrival_time: Option, } impl TextRequest { @@ -205,6 +211,7 @@ impl TextRequest { data_parallel_rank: None, reasoning_parser_kwargs: None, lora_request: None, + arrival_time: None, } } diff --git a/tests/benchmarks/test_bfcl_dataset.py b/tests/benchmarks/test_bfcl_dataset.py index e5110c50985..d1919223197 100644 --- a/tests/benchmarks/test_bfcl_dataset.py +++ b/tests/benchmarks/test_bfcl_dataset.py @@ -21,7 +21,7 @@ def _patch_hf_api(side_effect): @pytest.fixture(scope="session") def hf_tokenizer() -> PreTrainedTokenizerBase: - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") _FAKE_ROWS = { diff --git a/tests/benchmarks/test_custom_dataset_seed.py b/tests/benchmarks/test_custom_dataset_seed.py index dac87e6e6d9..d23ce40b53e 100644 --- a/tests/benchmarks/test_custom_dataset_seed.py +++ b/tests/benchmarks/test_custom_dataset_seed.py @@ -12,7 +12,7 @@ from vllm.benchmarks.datasets import get_samples @pytest.fixture(scope="session") def hf_tokenizer() -> PreTrainedTokenizerBase: - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") def _write_jsonl(path: Path, n_rows: int) -> None: diff --git a/tests/benchmarks/test_random_dataset.py b/tests/benchmarks/test_random_dataset.py index 57f68930618..ff691ae15d0 100644 --- a/tests/benchmarks/test_random_dataset.py +++ b/tests/benchmarks/test_random_dataset.py @@ -17,7 +17,7 @@ from vllm.benchmarks.datasets import ( @pytest.fixture(scope="session") def hf_tokenizer() -> PreTrainedTokenizerBase: # Use a small, commonly available tokenizer - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") class Params(NamedTuple): diff --git a/tests/benchmarks/test_random_multimodal_dataset_video.py b/tests/benchmarks/test_random_multimodal_dataset_video.py index bd37a520d01..b394ea2c0d7 100644 --- a/tests/benchmarks/test_random_multimodal_dataset_video.py +++ b/tests/benchmarks/test_random_multimodal_dataset_video.py @@ -16,7 +16,7 @@ from vllm.benchmarks.datasets import RandomMultiModalDataset, SampleRequest @pytest.fixture(scope="session") def hf_tokenizer() -> PreTrainedTokenizerBase: """Use a small, commonly available tokenizer.""" - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") @pytest.fixture diff --git a/tests/benchmarks/test_txt_slices_dataset.py b/tests/benchmarks/test_txt_slices_dataset.py index 7821e9a925a..8741805d0d5 100644 --- a/tests/benchmarks/test_txt_slices_dataset.py +++ b/tests/benchmarks/test_txt_slices_dataset.py @@ -13,7 +13,7 @@ from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices @pytest.fixture(scope="session") def hf_tokenizer() -> PreTrainedTokenizerBase: # Use a small, commonly available tokenizer - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") text_content = """ @@ -39,7 +39,7 @@ def test_create_txt_slices_jsonl( create_txt_slices_jsonl( input_path=str(txt_path), output_path=str(jsonl_path), - tokenizer_name="gpt2", + tokenizer_name="openai-community/gpt2", num_prompts=10, input_len=10, output_len=10, diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 9f34d25c46d..a4ed63ffe7b 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -79,6 +79,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): ): monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0") monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0") from vllm._aiter_ops import rocm_aiter_ops rocm_aiter_ops.refresh_env_variables() diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 83ce458aafd..b86018a7555 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -13,6 +13,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant from vllm.compilation.passes.fusion.allreduce_rms_fusion import ( AllReduceFusionPass, RocmAiterAllReduceFusionPass, + _select_flashinfer_allreduce_use_oneshot, ) from vllm.compilation.passes.fx_utils import find_op_nodes from vllm.compilation.passes.utility.fix_functionalization import ( @@ -30,6 +31,9 @@ from vllm.config import ( set_current_vllm_config, ) from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( + AiterCustomAllreduce, +) from vllm.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, @@ -45,6 +49,35 @@ from vllm.utils.torch_utils import set_random_seed DEVICE_TYPE = current_platform.device_type +@pytest.mark.parametrize( + ("workspace_backend", "device_capability", "world_size", "tensor_size", "expected"), + [ + ("mnnvl", 103, 8, 2 * 1024 * 1024, None), + ("trtllm", 103, 8, 2 * 1024 * 1024, True), + ("trtllm", 103, 8, 2 * 1024 * 1024 + 1, False), + ("trtllm", 100, 4, 4 * 1024 * 1024, True), + ("trtllm", 100, 4, 4 * 1024 * 1024 + 1, False), + ("trtllm", None, 8, 128 * 1024 * 1024, True), + ], +) +def test_select_flashinfer_allreduce_use_oneshot( + workspace_backend: str, + device_capability: int | None, + world_size: int, + tensor_size: int, + expected: bool | None, +): + assert ( + _select_flashinfer_allreduce_use_oneshot( + workspace_backend, + device_capability, + world_size, + tensor_size, + ) + is expected + ) + + class TestAllReduceRMSNormModel(torch.nn.Module): def __init__( self, @@ -504,8 +537,12 @@ def all_reduce_fusion_pass_on_test_model( "MASTER_ADDR": "localhost", "MASTER_PORT": "12345", "VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend, + "VLLM_ROCM_USE_AITER": str(int(use_aiter)), + "VLLM_ROCM_USE_AITER_CUSTOM_AR": str(int(use_aiter)), } ) + if use_aiter: + rocm_aiter_ops.refresh_env_variables() init_distributed_environment() @@ -616,7 +653,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( m.setenv("VLLM_ROCM_USE_AITER", "1") rocm_aiter_ops.refresh_env_variables() - if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group(): + if not AiterCustomAllreduce.build_supports_per_group_quant(): pytest.skip( "aiter build is missing 'fused_ar_rms_per_group_quant' (needs " "ROCm/aiter PR #2823); the new patterns aren't registered." @@ -671,6 +708,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( "MASTER_ADDR": "localhost", "MASTER_PORT": "12345", "VLLM_ROCM_USE_AITER": "1", + "VLLM_ROCM_USE_AITER_CUSTOM_AR": "1", } ) rocm_aiter_ops.refresh_env_variables() diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 531d26e008a..b776f6af98a 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -306,11 +306,6 @@ def test_attention_quant_pattern( torch.manual_seed(42) backend_cls = backend.get_class() - - # TODO: drop once AITER reenables fp16 unified attention. - if dtype not in backend_cls.supported_dtypes: - pytest.skip(f"{backend.name} does not support dtype {dtype}") - block_size = backend_cls.get_preferred_block_size(16) model_config = ModelConfig( diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index 5ff0fac6c82..a7f32483a70 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -502,7 +502,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch): m.setenv("VLLM_USE_AOT_COMPILE", "1") # First compilation - initialize model and generate llm_model = LLM( - model="gpt2", + model="openai-community/gpt2", compilation_config=CompilationConfig( mode=CompilationMode.VLLM_COMPILE, ), @@ -519,7 +519,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch): # Second compilation - should hit cache m.setenv("VLLM_FORCE_AOT_LOAD", "1") llm_model = LLM( - model="gpt2", + model="openai-community/gpt2", compilation_config=CompilationConfig( mode=CompilationMode.VLLM_COMPILE, ), diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index b8c18fa6cdc..96c3f49aba3 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -24,7 +24,7 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer def get_test_models(): """Get list of models to test based on PyTorch version""" models = [ - "gpt2", + "openai-community/gpt2", "Qwen/Qwen2-7B-Instruct", "meta-llama/Llama-3.1-8B", ] diff --git a/tests/config/test_bailing_mtp_config.py b/tests/config/test_bailing_mtp_config.py new file mode 100644 index 00000000000..8fae29959f2 --- /dev/null +++ b/tests/config/test_bailing_mtp_config.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers import PretrainedConfig + +from vllm.config.speculative import MTPModelTypes, SpeculativeConfig +from vllm.transformers_utils.model_arch_config_convertor import ( + BailingHybridMTPModelArchConfigConvertor, +) + + +def _bailing_config() -> PretrainedConfig: + config = PretrainedConfig( + architectures=["BailingMoeV2_5ForCausalLM"], + hidden_size=4096, + kv_lora_rank=512, + num_attention_heads=32, + num_experts=256, + num_hidden_layers=32, + num_key_value_heads=32, + num_nextn_predict_layers=1, + qk_rope_head_dim=64, + vocab_size=157184, + ) + config.model_type = "bailing_hybrid" + return config + + +def test_bailing_hybrid_mtp_hf_config_override(): + config = _bailing_config() + + overridden = SpeculativeConfig.hf_config_override(config) + + assert overridden.model_type == "bailing_hybrid_mtp" + assert overridden.architectures == ["BailingMoeV25MTPModel"] + assert overridden.n_predict == 1 + assert "bailing_hybrid_mtp" in MTPModelTypes.__args__ + + +def test_bailing_hybrid_mtp_model_arch_config(): + config = _bailing_config() + config.model_type = "bailing_hybrid_mtp" + config.architectures = ["BailingMoeV25MTPModel"] + + model_arch_config = BailingHybridMTPModelArchConfigConvertor( + config, config + ).convert() + + assert model_arch_config.model_type == "bailing_hybrid_mtp" + assert model_arch_config.architectures == ["BailingMoeV25MTPModel"] + assert model_arch_config.total_num_hidden_layers == 1 + assert model_arch_config.is_deepseek_mla diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 23451c475ea..3cc26e6e476 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -6,6 +6,7 @@ from enum import Enum import pytest +from vllm.config.cache import CacheConfig from vllm.config.utils import get_hash_factors, hash_factors, normalize_value # Helpers @@ -201,3 +202,15 @@ print(hash_factors(envs.compile_factors())) "compile_factors hash differs between fresh initializations - " "dynamic env vars may not be properly ignored" ) + + +def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): + """kv_cache_memory_bytes only sizes the KV cache allocation (like + gpu_memory_utilization, which is already ignored); it does not affect + the compiled computation graph. If it leaks into the hash, setting the + documented fast-boot knob silently invalidates the torch.compile cache + and forces a full recompile. + """ + base_hash = CacheConfig().compute_hash() + assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash + assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py new file mode 100644 index 00000000000..ddb8752a80d --- /dev/null +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for SpeculativeConfig.compose_draft_hf_overrides. + +Callable ``hf_overrides`` on the target model config (e.g. the +``dummy_hf_overrides`` shrink used by ``tests/models/test_initialization.py``) +must also be applied when building the draft ``ModelConfig``. Otherwise a +draft belonging to a large target model is instantiated at full size even +when the target itself is shrunk — which is what kept spec-decode archs like +``EagleMistralLarge3ForCausalLM`` stuck at ``is_available_online=False`` +("TODO: revert once figuring out OOM in CI"). +""" + +import functools + +import pytest +from transformers import PretrainedConfig + +from vllm.config.speculative import SpeculativeConfig + + +def _make_hf_config(**kwargs) -> PretrainedConfig: + defaults = dict( + architectures=["LlamaForCausalLM"], + model_type="llama", + num_hidden_layers=64, + ) + defaults.update(kwargs) + return PretrainedConfig(**defaults) + + +@pytest.mark.cpu_test +def test_dict_overrides_are_not_forwarded_to_draft(): + """Dict overrides are target-specific key patches; the draft must get + only the architecture-mapping override.""" + composed = SpeculativeConfig.compose_draft_hf_overrides( + {"max_position_embeddings": 1234} + ) + assert composed is SpeculativeConfig.hf_config_override + + +@pytest.mark.cpu_test +def test_none_overrides_fall_back_to_arch_mapping(): + composed = SpeculativeConfig.compose_draft_hf_overrides(None) + assert composed is SpeculativeConfig.hf_config_override + + +@pytest.mark.cpu_test +def test_callable_overrides_reach_the_draft_config(): + """A callable override (config-to-config transform) composes with the + architecture-mapping override and is applied to the draft config.""" + + def shrink(hf_config: PretrainedConfig) -> PretrainedConfig: + hf_config.num_hidden_layers = 1 + return hf_config + + composed = SpeculativeConfig.compose_draft_hf_overrides(shrink) + assert composed is not SpeculativeConfig.hf_config_override + + out = composed(_make_hf_config()) + # The shrink transform must have been applied to the draft config. + assert out.num_hidden_layers == 1 + + +@pytest.mark.cpu_test +def test_arch_mapping_applies_before_callable_override(): + """The static arch-mapping override runs first, so the user callable + observes (and may adjust) the post-mapping config.""" + seen_architectures: list[str] = [] + + def record(hf_config: PretrainedConfig) -> PretrainedConfig: + seen_architectures.append(hf_config.architectures[0]) + return hf_config + + composed = SpeculativeConfig.compose_draft_hf_overrides(record) + + # MiMo is one of the arch-mapped model types: hf_config_override + # rewrites architectures to ["MiMoMTPModel"]. + mimo = _make_hf_config( + architectures=["MiMoForCausalLM"], + model_type="mimo", + num_nextn_predict_layers=1, + ) + composed(mimo) + assert seen_architectures == ["MiMoMTPModel"] + + +def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig: + hf_config.num_hidden_layers = 1 + return hf_config + + +@pytest.mark.cpu_test +def test_composed_override_is_picklable(): + """The draft ``ModelConfig`` is sent to spawned engine-core processes, so + the composed override must be picklable. A nested local closure is not + (it raised ``Can't get local object`` on DFlashDraftModel); a + ``functools.partial`` over a module-referenceable static method is. + Guard against regressing to a closure.""" + composed = SpeculativeConfig.compose_draft_hf_overrides(_module_level_shrink) + + assert isinstance(composed, functools.partial) + assert composed.func is SpeculativeConfig._apply_composed_hf_override + + out = composed(_make_hf_config()) + assert out.num_hidden_layers == 1 diff --git a/tests/conftest.py b/tests/conftest.py index 6f9c8fa120f..94f7a83dd24 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -74,7 +74,7 @@ from torch._inductor.utils import fresh_cache if TYPE_CHECKING: - from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast + from transformers import PythonBackend, TokenizersBackend from transformers.generation.utils import GenerateOutput @@ -499,7 +499,7 @@ class HfRunner: self.model = model if not skip_tokenizer_init: - self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = ( + self.tokenizer: "PythonBackend | TokenizersBackend" = ( AutoTokenizer.from_pretrained( tokenizer_name or model_name, trust_remote_code=trust_remote_code, diff --git a/tests/cuda/test_cuda_context.py b/tests/cuda/test_cuda_context.py index 6336f2112c6..16d2f16c2d8 100644 --- a/tests/cuda/test_cuda_context.py +++ b/tests/cuda/test_cuda_context.py @@ -77,5 +77,43 @@ class TestSetCudaContext: current_platform.set_device(torch.device("cpu")) +def test_get_device_capability_uses_visible_device_ordinal(monkeypatch): + import vllm.platforms.interface as platform_interface + from vllm.platforms.cuda import NvmlCudaPlatform, pynvml + + seen_indices: list[int] = [] + + def record_handle(index: int) -> str: + seen_indices.append(index) + return f"handle-{index}" + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [1]) + monkeypatch.setenv(NvmlCudaPlatform.device_control_env_var, "0,1") + monkeypatch.setattr( + NvmlCudaPlatform, + "device_control_id_to_physical_device_id", + classmethod(lambda _cls, device_id: int(device_id)), + ) + monkeypatch.setattr(pynvml, "nvmlInit", lambda: None) + monkeypatch.setattr(pynvml, "nvmlShutdown", lambda: None) + monkeypatch.setattr( + pynvml, + "nvmlDeviceGetHandleByIndex", + record_handle, + ) + monkeypatch.setattr( + pynvml, + "nvmlDeviceGetCudaComputeCapability", + lambda _handle: (9, 0), + ) + NvmlCudaPlatform.get_device_capability.cache_clear() + + capability = NvmlCudaPlatform.get_device_capability(device_id=1) + + assert capability is not None + assert capability.to_int() == 90 + assert seen_indices == [1] + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index da661c5e13b..48df856f5e7 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -97,11 +97,12 @@ class MockSubscriber: for endpoint in pub_endpoints: self.sub.connect(endpoint) - # Set up replay sockets if provided + # Set up replay sockets if provided. + # DEALER allows receiving multiple replies per request. self.replay_sockets = [] if replay_endpoints: for replay_endpoint in replay_endpoints: - replay = self.ctx.socket(zmq.REQ) + replay = self.ctx.socket(zmq.DEALER) replay.connect(replay_endpoint) self.replay_sockets.append(replay) @@ -132,7 +133,9 @@ class MockSubscriber: if socket_idx >= len(self.replay_sockets): raise ValueError(f"Invalid socket index {socket_idx}") - self.replay_sockets[socket_idx].send(start_seq.to_bytes(8, "big")) + self.replay_sockets[socket_idx].send_multipart( + [b"", start_seq.to_bytes(8, "big")] + ) def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]: """Receive replayed messages from a specific replay socket""" @@ -148,12 +151,16 @@ class MockSubscriber: if not replay_socket.poll(1000): break + # DEALER receives [empty_delim, topic, seq, payload] frames = replay_socket.recv_multipart() - if not frames or not frames[-1]: + if frames and frames[0] == b"": + frames = frames[1:] + if len(frames) != 3 or not frames[-1]: # End of replay marker break - seq_bytes, payload = frames + topic, seq_bytes, payload = frames + assert topic == self.topic_bytes seq = int.from_bytes(seq_bytes, "big") data = self.decoder.decode(payload) replayed.append((seq, data)) diff --git a/tests/distributed/test_events.py b/tests/distributed/test_events.py index f17b7997c58..9b5601ad1d9 100644 --- a/tests/distributed/test_events.py +++ b/tests/distributed/test_events.py @@ -80,20 +80,38 @@ def test_replay_mechanism(publisher, subscriber): batch = create_test_events(1) publisher.publish(batch) - time.sleep(0.5) # Need publisher to process above requests - subscriber.request_replay(10) + # Drain live events to ensure publisher has buffered them. + for _ in range(19): + assert subscriber.receive_one(timeout=1000) is not None - batch = create_test_events(1) - publisher.publish(batch) # 20th message + subscriber.request_replay(10) replayed = subscriber.receive_replay() - assert len(replayed) > 0, "No replayed messages received" - seqs = [seq for seq, _ in replayed] - assert all(seq >= 10 for seq in seqs), "Replayed messages not in order" - assert seqs == list(range(min(seqs), max(seqs) + 1)), ( - "Replayed messages not consecutive" + assert len(replayed) == 9, ( + f"Expected 9 replayed messages (seq 10-18), got {len(replayed)}" ) + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(10, 19)), "Replayed sequences should be 10-18" + + +def test_replay_includes_topic(publisher, subscriber, publisher_config): + """Test that replay responses include the topic, matching PUB format""" + for _ in range(5): + publisher.publish(create_test_events(1)) + + # Drain live events to ensure publisher has processed them. + for _ in range(5): + assert subscriber.receive_one(timeout=1000) is not None + + subscriber.request_replay(0) + + # receive_replay unpacks (topic, seq, payload) and asserts + # topic == publisher topic for each message. + replayed = subscriber.receive_replay() + assert len(replayed) == 5, f"Expected 5 replayed messages, got {len(replayed)}" + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(5)), "Replayed sequences should be 0-4" def test_buffer_limit(publisher, subscriber, publisher_config): @@ -108,15 +126,16 @@ def test_buffer_limit(publisher, subscriber, publisher_config): time.sleep(0.5) # Need publisher to process above requests subscriber.request_replay(0) - batch = create_test_events(1) - publisher.publish(batch) - replayed = subscriber.receive_replay() - assert len(replayed) <= buffer_size, "Can't replay more than buffer size" + assert len(replayed) == buffer_size, ( + f"Expected {buffer_size} replayed messages, got {len(replayed)}" + ) - oldest_seq = min(seq for seq, _ in replayed) - assert oldest_seq >= 10, "The oldest sequence should be at least 10" + seqs = [seq for seq, _ in replayed] + assert seqs == list(range(10, buffer_size + 10)), ( + "Should replay seq 11 through buffer_size+10" + ) def test_topic_filtering(publisher_config): diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index e773c7d826a..762a95fb987 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -114,7 +114,7 @@ TEXT_GENERATION_MODELS = { "tiiuae/falcon-7b": PPTestSettings.fast(), "google/gemma-1.1-2b-it": PPTestSettings.fast(), "google/gemma-2-9b": PPTestSettings.fast(), - "gpt2": PPTestSettings.fast(), + "openai-community/gpt2": PPTestSettings.fast(), "EleutherAI/gpt-j-6b": PPTestSettings.fast(), "EleutherAI/pythia-1.4b": PPTestSettings.fast(), "ibm/PowerLM-3b": PPTestSettings.fast(), diff --git a/tests/distributed/test_rocm_aiter_custom_ar.py b/tests/distributed/test_rocm_aiter_custom_ar.py new file mode 100644 index 00000000000..0b85f36410d --- /dev/null +++ b/tests/distributed/test_rocm_aiter_custom_ar.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import ray +import torch +import torch.distributed as dist + +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa +from vllm.distributed.parallel_state import get_tp_group, graph_capture +from vllm.envs import disable_envs_cache +from vllm.platforms import current_platform + +from ..utils import ( + assert_rocm_custom_allreduce_backend_state, + ensure_model_parallel_initialized, + init_test_distributed_environment, + multi_gpu_test, + multi_process_parallel, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm-only AITER custom allreduce tests", +) + +test_cases = [ + ((2, 7168), torch.float16), + ((2, 7168), torch.bfloat16), + ((128, 8192), torch.float16), + ((128, 8192), torch.bfloat16), +] + + +def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1") + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE") + disable_envs_cache() + rocm_aiter_ops.refresh_env_variables() + + +def _assert_aiter_handles_input(inp: torch.Tensor) -> None: + aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm + assert aiter_ar_comm is not None + assert aiter_ar_comm.should_custom_ar(inp), ( + f"AITER CustomAllreduce does not support input shape {inp.shape}." + ) + + +@ray.remote(num_gpus=1, max_calls=1) +def graph_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + assert_rocm_custom_allreduce_backend_state(True, "NONE") + group = get_tp_group().device_group + + # A small all_reduce for warmup. + # this is needed because device communicators might be created lazily + # (e.g. NCCL). This will ensure that the communicator is initialized + # before any communication happens, so that this group can be used for + # graph capture immediately. + data = torch.zeros(1) + data = data.to(device=device) + dist.all_reduce(data, group=group) + torch.accelerator.synchronize() + del data + + for shape, dtype in test_cases: + with graph_capture(device=device) as graph_capture_context: + inp = torch.ones(shape, dtype=dtype, device=device) + _assert_aiter_handles_input(inp) + expected = inp * tp_size + + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=graph_capture_context.stream): + out = tensor_model_parallel_all_reduce(inp) + + graph.replay() + torch.testing.assert_close(out, expected) + + +@ray.remote(num_gpus=1, max_calls=1) +def eager_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pp_size, + rank, + distributed_init_port, +) -> None: + with monkeypatch.context() as m: + _configure_aiter_custom_ar_env(m) + + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + assert_rocm_custom_allreduce_backend_state(True, "NONE") + + for shape, dtype in test_cases: + inp = torch.ones(shape, dtype=dtype, device=device) + _assert_aiter_handles_input(inp) + expected = inp * tp_size + out = tensor_model_parallel_all_reduce(inp) + torch.testing.assert_close(out, expected) + + +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("tp_size", [2]) +@pytest.mark.parametrize("pipeline_parallel_size", [1]) +@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce]) +def test_rocm_aiter_custom_allreduce( + monkeypatch: pytest.MonkeyPatch, + tp_size, + pipeline_parallel_size, + test_target, +): + multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target) diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 8a13b24dc52..f3423745ca5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -191,7 +191,7 @@ class TestNCCLEngineParsing: return NCCLWeightTransferEngine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -240,21 +240,30 @@ class TestEngineRegistry: def test_create_engine_nccl(self): config = WeightTransferConfig(backend="nccl") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, NCCLWeightTransferEngine) def test_create_engine_ipc(self): config = WeightTransferConfig(backend="ipc") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, IPCWeightTransferEngine) def test_create_engine_sparse_nccl(self): config = WeightTransferConfig(backend="sparse_nccl") engine = WeightTransferEngineFactory.create_engine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) assert isinstance(engine, SparseNCCLWeightTransferEngine) @@ -264,7 +273,7 @@ class TestEngineRegistry: WeightTransferEngineFactory.create_engine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -284,7 +293,7 @@ class TestSparseNCCLPatchApplication: def _make_engine(self, model): config = WeightTransferConfig(backend="sparse_nccl") return SparseNCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cpu", model + config, create_mock_vllm_config(), torch.device("cpu"), model ) def _make_model(self, numel: int = 8): @@ -382,7 +391,10 @@ def test_nccl_receive_weights_without_init_raises(): config = WeightTransferConfig(backend="nccl") engine = NCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) update_info = NCCLWeightTransferUpdateInfo( @@ -400,7 +412,10 @@ def test_sparse_nccl_receive_weights_without_init_raises(): config = WeightTransferConfig(backend="sparse_nccl") engine = SparseNCCLWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda"), + MagicMock(spec=torch.nn.Module), ) update_info = SparseNCCLWeightTransferUpdateInfo( @@ -495,7 +510,9 @@ def inference_receive_tensor( vllm_config.model_config = MagicMock() recorder = Recorder() - engine = NCCLWeightTransferEngine(config, vllm_config, "cuda", recorder) + engine = NCCLWeightTransferEngine( + config, vllm_config, torch.device("cuda"), recorder + ) # Transport-only test: bypass the set_current_vllm_config context that # receive_weights enters, since vllm_config here is a mock. import vllm.config as _vllm_config_mod @@ -664,7 +681,9 @@ def inference_receive_sparse_tensor( num_updates_list=[3], ) - engine = SparseNCCLWeightTransferEngine(config, vllm_config, "cuda", model) + engine = SparseNCCLWeightTransferEngine( + config, vllm_config, torch.device("cuda"), model + ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, ) @@ -879,7 +898,7 @@ class TestIPCEngineParsing: return IPCWeightTransferEngine( config, create_mock_vllm_config(), - "cuda", + torch.device("cuda"), MagicMock(spec=torch.nn.Module), ) @@ -1068,7 +1087,9 @@ def inference_receive_ipc_tensor( vllm_config.model_config = MagicMock() recorder = Recorder() - engine = IPCWeightTransferEngine(config, vllm_config, "cuda", recorder) + engine = IPCWeightTransferEngine( + config, vllm_config, _get_ray_assigned_device(), recorder + ) # Transport-only test: bypass the set_current_vllm_config context that # receive_weights enters, since vllm_config here is a mock. import vllm.config as _vllm_config_mod @@ -1173,7 +1194,10 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises(): config = WeightTransferConfig(backend="ipc") engine = IPCWeightTransferEngine( - config, create_mock_vllm_config(), "cuda", MagicMock(spec=torch.nn.Module) + config, + create_mock_vllm_config(), + torch.device("cuda:0"), + MagicMock(spec=torch.nn.Module), ) dummy_tensor = torch.ones(10, 10, device="cuda:0") diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 7da3aa66c9c..2feb9f7a039 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -565,6 +565,40 @@ def test_human_readable_model_len(): parser.parse_args(["--max-model-len", invalid]) +def test_human_readable_other_args(): + # Test human-readable parsing for other integer args + # that were added to use human_readable_int parser + parser = EngineArgs.add_cli_args(FlexibleArgumentParser(exit_on_error=False)) + + # Test max_num_scheduled_tokens + args = parser.parse_args(["--max-num-scheduled-tokens", "1024"]) + assert args.max_num_scheduled_tokens == 1024 + args = parser.parse_args(["--max-num-scheduled-tokens", "2k"]) + assert args.max_num_scheduled_tokens == 2_000 + args = parser.parse_args(["--max-num-scheduled-tokens", "4K"]) + assert args.max_num_scheduled_tokens == 2**10 * 4 + args = parser.parse_args(["--max-num-scheduled-tokens", "10.5k"]) + assert args.max_num_scheduled_tokens == 10500 + + # Test kv_cache_memory_bytes (existing human-readable arg) + args = parser.parse_args(["--kv-cache-memory-bytes", "100000"]) + assert args.kv_cache_memory_bytes == 100000 + args = parser.parse_args(["--kv-cache-memory-bytes", "100k"]) + assert args.kv_cache_memory_bytes == 100_000 + args = parser.parse_args(["--kv-cache-memory-bytes", "1M"]) + assert args.kv_cache_memory_bytes == 2**20 + args = parser.parse_args(["--kv-cache-memory-bytes", "1m"]) + assert args.kv_cache_memory_bytes == 1_000_000 + + # Test max_num_batched_tokens (existing human-readable arg) + args = parser.parse_args(["--max-num-batched-tokens", "1024"]) + assert args.max_num_batched_tokens == 1024 + args = parser.parse_args(["--max-num-batched-tokens", "2k"]) + assert args.max_num_batched_tokens == 2_000 + args = parser.parse_args(["--max-num-batched-tokens", "4K"]) + assert args.max_num_batched_tokens == 2**10 * 4 + + def test_numa_bind_args(): parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = parser.parse_args( diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 7b480bdba67..25a9451bc2b 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json +from collections.abc import AsyncIterator from contextlib import suppress from dataclasses import dataclass, field from typing import Any @@ -19,6 +20,7 @@ from tests.entrypoints.openai.utils import ( from tests.utils import RemoteOpenAIServer from vllm._aiter_ops import is_aiter_found_and_supported from vllm.config import MultiModalConfig +from vllm.entrypoints.generate.base.serving import build_per_request_timing_metrics from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -50,9 +52,17 @@ from vllm.tokenizers import get_tokenizer from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestStateStats GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b" GPT_OSS_SPECULATOR_NAME = "RedHatAI/gpt-oss-20b-speculator.eagle3" +_PER_REQUEST_STATS = RequestStateStats( + queued_ts=1.0, + scheduled_ts=1.5, + first_token_ts=2.0, + last_token_ts=3.0, + num_generation_tokens=2, +) @pytest.fixture(scope="module") @@ -608,6 +618,169 @@ def _build_serving_chat( return serving_chat +def _build_minimal_metrics_serving_chat( + enable_per_request_metrics: bool, + enable_force_include_usage: bool = False, +) -> OpenAIServingChat: + serving = OpenAIServingChat.__new__(OpenAIServingChat) + serving.response_role = "assistant" + serving.parser_cls = None + serving.enable_auto_tools = False + serving.enable_prompt_tokens_details = False + serving.enable_log_outputs = False + serving.enable_log_deltas = False + serving.enable_force_include_usage = enable_force_include_usage + serving.request_logger = None + serving.system_fingerprint = None + serving.enable_per_request_metrics = enable_per_request_metrics + return serving + + +def _make_metrics_request_output( + metrics: RequestStateStats | None = _PER_REQUEST_STATS, + token_ids: tuple[int, ...] = (100, 101), +) -> RequestOutput: + return RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="Hello", + token_ids=list(token_ids), + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + metrics=metrics, + ) + + +async def _single_request_output( + request_output: RequestOutput, +) -> AsyncIterator[RequestOutput]: + yield request_output + + +async def _collect_metrics_stream_chunks( + serving: OpenAIServingChat, + request: ChatCompletionRequest, +) -> list[dict[str, Any]]: + chunks: list[dict[str, Any]] = [] + async for line in serving.chat_completion_stream_generator( + request, + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ): + line = line.strip() + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + if payload != "[DONE]": + chunks.append(json.loads(payload)) + return chunks + + +def test_build_per_request_timing_metrics_valid_timestamps(): + metrics = build_per_request_timing_metrics( + _PER_REQUEST_STATS, num_generation_tokens=10 + ) + + assert metrics.time_to_first_token_ms == pytest.approx(500.0) + assert metrics.generation_time_ms == pytest.approx(1000.0) + assert metrics.queue_time_ms == pytest.approx(500.0) + assert metrics.mean_itl_ms == pytest.approx(1000.0 / 9, rel=1e-4) + assert metrics.tokens_per_second == pytest.approx(10.0 / 1.5, rel=1e-4) + + +@pytest.mark.asyncio +async def test_chat_per_request_metrics_follow_server_flag(): + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=False, + ) + request_output = _make_metrics_request_output() + + disabled_serving = _build_minimal_metrics_serving_chat( + enable_per_request_metrics=False + ) + disabled_response = await disabled_serving.chat_completion_full_generator( + request, + _single_request_output(request_output), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert disabled_response.metrics is None + + enabled_serving = _build_minimal_metrics_serving_chat( + enable_per_request_metrics=True + ) + enabled_response = await enabled_serving.chat_completion_full_generator( + request, + _single_request_output(request_output), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert enabled_response.metrics is not None + assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +@pytest.mark.asyncio +async def test_chat_per_request_metrics_suppressed_for_n_greater_than_one(): + serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=True) + response = await serving.chat_completion_full_generator( + ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=False, + n=2, + ), + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ) + assert response.metrics is None + + +@pytest.mark.asyncio +async def test_chat_streaming_metrics_ride_on_usage_chunk(): + serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=True) + chunks = await _collect_metrics_stream_chunks( + serving, + ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=True, + stream_options={"include_usage": True}, + ), + ) + + usage_chunks = [chunk for chunk in chunks if chunk.get("usage")] + assert usage_chunks + assert usage_chunks[-1]["metrics"]["time_to_first_token_ms"] == pytest.approx(500.0) + + @dataclass class MockEngine: model_config: MockModelConfig = field(default_factory=MockModelConfig) @@ -1951,6 +2124,13 @@ async def test_tool_choice_validation_without_parser(): assert isinstance(response_named, ErrorResponse) assert "tool_choice" in response_named.error.message assert "--tool-call-parser" in response_named.error.message + # The function name should appear in a clean, readable form - + # guards against leaking Pydantic's internal repr of the + # ChatCompletionNamedToolChoiceParam/ChatCompletionNamedFunction + # objects directly into the client-facing error message. + assert "get_weather" in response_named.error.message + assert "ChatCompletionNamedFunction" not in response_named.error.message + assert "ChatCompletionNamedToolChoiceParam" not in response_named.error.message @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index 3e9e4850f07..4c574673817 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -183,27 +183,39 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI): """Test that thinking_token_budget limits the number of reasoning tokens. - Counts non-empty streaming ``delta.reasoning`` chunks (coarse proxy; each - chunk may represent multiple decode tokens — see - ``_count_reasoning_decode_token_ids_between_markers`` and the Qwen3.5 MTP - test for id-based checks). + Counts reasoning decode tokens by id, which is robust to how tokens are + grouped into streamed chunks (a single chunk can carry several tokens under + async scheduling / stream_interval > 1). Counting chunks under-counts. """ - reasoning_token_count = 0 + tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME) + start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False)) + end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False)) + + prompt_token_ids: list[int] = [] + decode_token_ids: list[int] = [] stream = await client.chat.completions.create( model=MODEL_NAME, messages=MESSAGES, max_tokens=100, stream=True, - extra_body={"thinking_token_budget": THINK_BUDGET}, + extra_body={"thinking_token_budget": THINK_BUDGET, "return_token_ids": True}, ) async for chunk in stream: - delta = chunk.choices[0].delta - if getattr(delta, "reasoning", None): - reasoning_token_count += 1 + if not chunk.choices: + continue + if getattr(chunk, "prompt_token_ids", None): + prompt_token_ids = list(chunk.prompt_token_ids) + delta_ids = getattr(chunk.choices[0], "token_ids", None) + if delta_ids: + decode_token_ids.extend(delta_ids) + reasoning_token_count = _count_reasoning_decode_token_ids_between_markers( + prompt_token_ids + decode_token_ids, start_ids, end_ids + ) + assert reasoning_token_count is not None, "missing reasoning start marker in ids" assert reasoning_token_count == THINK_BUDGET, ( - f"reasoning tokens ({reasoning_token_count}) exceeded " + f"reasoning tokens ({reasoning_token_count}) != " f"thinking_token_budget ({THINK_BUDGET})" ) diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index a16fa83fe32..7b628207ac1 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -9,6 +9,8 @@ import regex as re from openai import BadRequestError from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.sampling_params import SamplingParams from vllm.tokenizers import get_tokenizer # any model with a chat template should work here @@ -730,3 +732,38 @@ async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str): "structured_outputs": {"grammar": invalid_simplified_sql_grammar} }, ) + + +# Unit tests for bad_words in CompletionRequest.to_sampling_params() +def test_completion_request_bad_words_to_sampling_params(): + """bad_words should be forwarded to SamplingParams (parity with chat).""" + request = CompletionRequest( + model="test-model", + prompt="Hello", + bad_words=["foo", "bar"], + max_tokens=10, + ) + + sampling_params = request.to_sampling_params( + max_tokens=10, + default_sampling_params={}, + ) + + assert isinstance(sampling_params, SamplingParams) + assert sampling_params.bad_words == ["foo", "bar"] + + +def test_completion_request_bad_words_default_empty(): + """bad_words defaults to an empty list, matching the chat endpoint.""" + request = CompletionRequest( + model="test-model", + prompt="Hello", + max_tokens=10, + ) + + assert request.bad_words == [] + sampling_params = request.to_sampling_params( + max_tokens=10, + default_sampling_params={}, + ) + assert sampling_params.bad_words == [] diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 062c3e7583a..aa9e9c1d72e 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -11,7 +11,10 @@ from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion -from vllm.entrypoints.openai.engine.protocol import GenerationError +from vllm.entrypoints.openai.engine.protocol import ( + GenerationError, + RequestResponseMetadata, +) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.scale_out.render.serving import ServingRender @@ -20,9 +23,17 @@ from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.stats import RequestStateStats MODEL_NAME = "openai-community/gpt2" MODEL_NAME_SHORT = "gpt2" +_PER_REQUEST_STATS = RequestStateStats( + queued_ts=1.0, + scheduled_ts=1.5, + first_token_ts=2.0, + last_token_ts=3.0, + num_generation_tokens=2, +) BASE_MODEL_PATHS = [ BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME), BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT), @@ -93,6 +104,39 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion: ) +def _build_minimal_metrics_serving_completion( + enable_per_request_metrics: bool, +) -> OpenAIServingCompletion: + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.enable_prompt_tokens_details = False + serving.system_fingerprint = None + serving.enable_per_request_metrics = enable_per_request_metrics + return serving + + +def _make_metrics_request_output( + metrics: RequestStateStats | None = _PER_REQUEST_STATS, +) -> RequestOutput: + return RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="Hello", + token_ids=[100, 101], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + metrics=metrics, + ) + + def _build_renderer(model_config: MockModelConfig): return HfRenderer( MockVllmConfig(model_config, parallel_config=MockParallelConfig()), @@ -100,6 +144,58 @@ def _build_renderer(model_config: MockModelConfig): ) +def test_completion_per_request_metrics_follow_server_flag(): + request = CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10) + request_output = _make_metrics_request_output() + + disabled_serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=False + ) + disabled_response = disabled_serving.request_output_to_completion_response( + [request_output], + request, + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert disabled_response.metrics is None + + enabled_serving = _build_minimal_metrics_serving_completion( + enable_per_request_metrics=True + ) + enabled_response = enabled_serving.request_output_to_completion_response( + [request_output], + request, + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert enabled_response.metrics is not None + assert enabled_response.metrics.time_to_first_token_ms == pytest.approx(500.0) + + +def test_completion_per_request_metrics_suppressed_for_multiple_prompts(): + serving = _build_minimal_metrics_serving_completion(enable_per_request_metrics=True) + response = serving.request_output_to_completion_response( + [_make_metrics_request_output(), _make_metrics_request_output()], + CompletionRequest( + model=MODEL_NAME, + prompt=["Test prompt", "Another prompt"], + max_tokens=10, + ), + "cmpl-test-id", + 0, + MODEL_NAME, + None, + RequestResponseMetadata(request_id="cmpl-test-id"), + ) + assert response.metrics is None + + @pytest.mark.asyncio async def test_completion_error_non_stream(): """test finish_reason='error' returns 500 InternalServerError (non-streaming)""" @@ -378,3 +474,139 @@ def test_negative_prompt_token_ids_flat(): prompt=[-1], max_tokens=10, ) + + +class TestCompletionPromptListLimit: + """Regression tests for CVE: unbounded prompt list fan-out.""" + + def test_scalar_prompt_allowed(self): + request = CompletionRequest( + model=MODEL_NAME, + prompt="hello", + max_tokens=1, + ) + assert request.prompt == "hello" + + def test_single_token_list_allowed(self): + request = CompletionRequest( + model=MODEL_NAME, + prompt=[1, 2, 3], + max_tokens=1, + ) + assert request.prompt == [1, 2, 3] + + def test_bounded_text_prompt_list_allowed(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "10") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + request = CompletionRequest( + model=MODEL_NAME, + prompt=["a", "b", "c"], + max_tokens=1, + ) + assert request.prompt == ["a", "b", "c"] + + def test_bounded_token_id_prompt_list_allowed(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "10") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + request = CompletionRequest( + model=MODEL_NAME, + prompt=[[1], [2], [3]], + max_tokens=1, + ) + assert request.prompt == [[1], [2], [3]] + + def test_oversized_text_prompt_list_rejected(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + with pytest.raises( + Exception, match="prompt list length 10 exceeds the maximum" + ): + CompletionRequest( + model=MODEL_NAME, + prompt=["x"] * 10, + max_tokens=1, + ) + + def test_oversized_token_id_prompt_list_rejected(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + with pytest.raises( + Exception, match="prompt list length 10 exceeds the maximum" + ): + CompletionRequest( + model=MODEL_NAME, + prompt=[[1]] * 10, + max_tokens=1, + ) + + def test_exact_limit_allowed(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + request = CompletionRequest( + model=MODEL_NAME, + prompt=["x"] * 5, + max_tokens=1, + ) + assert len(request.prompt) == 5 + + def test_one_over_limit_rejected(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + with pytest.raises(Exception, match="prompt list length 6 exceeds the maximum"): + CompletionRequest( + model=MODEL_NAME, + prompt=["x"] * 6, + max_tokens=1, + ) + + def test_oversized_prompt_embeds_list_rejected(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + with pytest.raises(Exception, match="prompt_embeds list length 10 exceeds"): + CompletionRequest( + model=MODEL_NAME, + prompt_embeds=[b"\x00"] * 10, + max_tokens=1, + ) + + def test_bounded_prompt_embeds_list_allowed(self, monkeypatch): + monkeypatch.setenv("VLLM_MAX_COMPLETION_PROMPTS", "5") + from vllm import envs + + if hasattr(envs.__getattr__, "cache_clear"): + envs.__getattr__.cache_clear() + + request = CompletionRequest( + model=MODEL_NAME, + prompt_embeds=[b"\x00"] * 5, + max_tokens=1, + ) + assert len(request.prompt_embeds) == 5 diff --git a/tests/entrypoints/openai/completion/test_prompt_validation.py b/tests/entrypoints/openai/completion/test_prompt_validation.py index 81204b27bc0..87c6b6e1668 100644 --- a/tests/entrypoints/openai/completion/test_prompt_validation.py +++ b/tests/entrypoints/openai/completion/test_prompt_validation.py @@ -18,7 +18,7 @@ from vllm.renderers.embed_utils import safe_load_prompt_embeds @pytest.mark.asyncio async def test_empty_prompt(): - model_name = "gpt2" + model_name = "openai-community/gpt2" server_args = ["--enforce-eager"] with RemoteOpenAIServer(model_name, server_args) as remote_server: client = remote_server.get_async_client() @@ -38,7 +38,7 @@ async def test_empty_prompt(): @pytest.mark.asyncio async def test_out_of_vocab_token_ids(): - model_name = "gpt2" + model_name = "openai-community/gpt2" server_args = ["--enforce-eager"] with RemoteOpenAIServer(model_name, server_args) as remote_server: client = remote_server.get_async_client() diff --git a/tests/entrypoints/openai/correctness/test_lmeval.py b/tests/entrypoints/openai/correctness/test_lmeval.py index 5b23b423902..aad1b5e0624 100644 --- a/tests/entrypoints/openai/correctness/test_lmeval.py +++ b/tests/entrypoints/openai/correctness/test_lmeval.py @@ -71,8 +71,9 @@ def test_lm_eval_accuracy_v1_engine(): more_args = [] - # Limit compilation time for V1 - if current_platform.is_tpu(): + # Limit compilation time for V1 on TPU + # Avoid OOM on XPU + if current_platform.is_tpu() or current_platform.is_xpu(): more_args = ["--max-num-seqs", "64"] run_test(more_args) diff --git a/tests/entrypoints/openai/responses/test_namespace_tool_separator.py b/tests/entrypoints/openai/responses/test_namespace_tool_separator.py new file mode 100644 index 00000000000..c895092b45f --- /dev/null +++ b/tests/entrypoints/openai/responses/test_namespace_tool_separator.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import openai # use the official client for correctness check +import pytest + +MODEL_NAME = "Qwen/Qwen3-1.7B" +NAMESPACE = "mcp__computer_use" +TOOL_NAME = "get_app_state" +FLAT_TOOL_NAME = f"{NAMESPACE}__{TOOL_NAME}" + +tools = [ + { + "type": "namespace", + "name": NAMESPACE, + "description": "Computer control tools.", + "tools": [ + { + "type": "function", + "name": TOOL_NAME, + "description": "Get the current state of a desktop application.", + "parameters": { + "type": "object", + "properties": { + "app": { + "type": "string", + "description": "Application name, for example Chrome.", + } + }, + "required": ["app"], + "additionalProperties": False, + }, + } + ], + } +] + +prompt = [ + { + "role": "user", + "content": "Use the computer app state tool to inspect Google Chrome.", + }, +] + + +def _assert_namespace_tool_call(tool_call) -> None: + assert tool_call.type == "function_call" + assert tool_call.name == TOOL_NAME + assert tool_call.namespace == NAMESPACE + assert tool_call.name != FLAT_TOOL_NAME + + args = json.loads(tool_call.arguments) + assert args["app"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_namespace_tool_separator(client: openai.AsyncOpenAI, model_name: str): + response = await client.responses.create( + model=model_name, + input=prompt, + tools=tools, + temperature=0.0, + ) + + assert len(response.output) >= 1 + tool_call = next( + (out for out in response.output if out.type == "function_call"), None + ) + assert tool_call is not None + _assert_namespace_tool_call(tool_call) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_namespace_tool_separator_streaming( + client: openai.AsyncOpenAI, model_name: str +): + stream = await client.responses.create( + model=model_name, + input=prompt, + tools=tools, + temperature=0.0, + stream=True, + ) + events = [event async for event in stream] + + added_call = next( + ( + event.item + for event in events + if event.type == "response.output_item.added" + and getattr(event.item, "type", None) == "function_call" + ), + None, + ) + done_call = next( + ( + event.item + for event in events + if event.type == "response.output_item.done" + and getattr(event.item, "type", None) == "function_call" + ), + None, + ) + + assert added_call is not None + assert added_call.name == TOOL_NAME + assert added_call.namespace == NAMESPACE + + assert done_call is not None + _assert_namespace_tool_call(done_call) diff --git a/tests/entrypoints/openai/test_cli_args.py b/tests/entrypoints/openai/test_cli_args.py index 58dd328b325..1f764202e55 100644 --- a/tests/entrypoints/openai/test_cli_args.py +++ b/tests/entrypoints/openai/test_cli_args.py @@ -206,6 +206,14 @@ def test_chat_template_validation_for_sad_paths(serve_parser): validate_parsed_serve_args(args) +def test_per_request_metrics_requires_log_stats(serve_parser): + args = serve_parser.parse_args( + args=["--enable-per-request-metrics", "--disable-log-stats"] + ) + with pytest.raises(ValueError): + validate_parsed_serve_args(args) + + @pytest.mark.parametrize( "cli_args, expected_middleware", [ diff --git a/tests/entrypoints/openai/test_uds.py b/tests/entrypoints/openai/test_uds.py index c79a4870dea..f79e40ee413 100644 --- a/tests/entrypoints/openai/test_uds.py +++ b/tests/entrypoints/openai/test_uds.py @@ -40,4 +40,5 @@ async def test_show_version(server: RemoteOpenAIServer): response = client.get(server.url_for("version")) response.raise_for_status() - assert response.json() == {"version": VLLM_VERSION} + # Tolerate additive fields (e.g. the Rust frontend reports its own version). + assert response.json()["version"] == VLLM_VERSION diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index d5565f25d37..96555ee363a 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -369,7 +369,7 @@ async def test_chat_request( assert output.object == "list" assert len(output.data) == 1 assert output.model == MODEL_NAME - assert output.usage.prompt_tokens == 34 + assert output.usage.prompt_tokens == 33 # test continue_final_message response = requests.post( @@ -401,7 +401,7 @@ async def test_chat_request( assert output.object == "list" assert len(output.data) == 1 assert output.model == MODEL_NAME - assert output.usage.prompt_tokens == 36 + assert output.usage.prompt_tokens == 35 # test continue_final_message with add_generation_prompt response = requests.post( diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py index 56e83de3f74..df79a387afd 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import weakref +from types import SimpleNamespace import pytest import torch @@ -9,7 +10,10 @@ import torch from tests.models.utils import softmax from vllm import LLM, PoolingParams from vllm.distributed import cleanup_dist_env_and_memory +from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor +from vllm.entrypoints.pooling.scoring.typing import ScoringData from vllm.platforms import current_platform +from vllm.renderers import TokenizeParams MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls" PROMPT = "The chef prepared a delicious meal." @@ -141,6 +145,45 @@ def test_max_tokens_per_doc(llm: LLM): assert with_limit_tokens < no_limit_tokens +def test_token_type_ids_follow_post_tokenization(): + processor = object.__new__(CrossEncoderIOProcessor) + processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1) + processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt) + processor.model_config = None + processor.get_score_prompt = lambda **_: ( + "", + { + "prompt_token_ids": list(range(32)), + "token_type_ids": [0] * 16 + [1] * 16, + }, + ) + + engine_inputs, pooling_params = processor._pre_process( + ScoringData(data_1=["query"], data_2=["document"]), + TokenizeParams( + max_total_tokens=None, + truncate_prompt_tokens=16, + truncation_side="left", + ), + PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}), + ) + + assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32)) + assert pooling_params[0].extra_kwargs == { + "cache_salt": "salt", + "compressed_token_type_ids": 0, + } + + engine_inputs, pooling_params = processor._pre_process( + ScoringData(data_1=["query"], data_2=["document"]), + TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40), + PoolingParams(task="classify"), + ) + + assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8 + assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16} + + def test_pooling_params(llm: LLM): def get_outputs(use_activation): outputs = llm.score( diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index e452b7367a2..3167e00b96e 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -489,6 +489,202 @@ async def test_derender_completion_kv_transfer_params_passthrough(client): assert response.json()["kv_transfer_params"] == kv +# --------------------------------------------------------------------------- +# Resource bounds regression tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_derender_chat_bounded_payload_succeeds(client): + """Normal bounded derender payload succeeds (positive control).""" + gen_req = await _render_chat(client) + synthetic_ids = gen_req["token_ids"][:5] + + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": _make_generate_response(synthetic_ids), + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["choices"]) == 1 + assert data["choices"][0]["message"]["content"] + + +@pytest.mark.asyncio +async def test_derender_chat_oversized_token_ids_rejected(client): + """token_ids longer than max_model_len returns 400.""" + response = await client.get("/v1/models") + assert response.status_code == 200 + + # Use a token_ids list that exceeds any reasonable max_model_len. + # The tiny-random model has max_model_len of 2048. + oversized_ids = [42] * 1_000_000 + + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": _make_generate_response(oversized_ids), + }, + ) + assert response.status_code == 400 + assert "max_model_len" in response.json()["error"]["message"] + + +@pytest.mark.asyncio +async def test_derender_chat_too_many_choices_rejected(client): + """choices count exceeding VLLM_MAX_N_SEQUENCES returns 400.""" + # Default VLLM_MAX_N_SEQUENCES is 16384; use a larger count. + oversized_choices = [ + {"index": i, "token_ids": [42], "finish_reason": "stop"} for i in range(20_000) + ] + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": { + "request_id": "test-choices-bound", + "choices": oversized_choices, + }, + }, + ) + assert response.status_code == 400 + assert "choices count" in response.json()["error"]["message"] + + +@pytest.mark.asyncio +async def test_derender_completion_too_many_generate_responses_rejected(client): + """generate_responses count exceeding limit returns 400.""" + oversized_responses = [ + { + "request_id": f"gen-{i}", + "choices": [{"index": 0, "token_ids": [42], "finish_reason": "stop"}], + } + for i in range(20_000) + ] + response = await client.post( + "/v1/completions/derender", + json={ + "model": MODEL_NAME, + "generate_responses": oversized_responses, + }, + ) + assert response.status_code == 400 + assert "generate_responses count" in response.json()["error"]["message"] + + +@pytest.mark.asyncio +async def test_derender_chat_negative_token_ids_rejected(client): + """Negative token_ids are rejected at the protocol validation level.""" + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": _make_generate_response([-1, 42, 100]), + }, + ) + # vLLM's validation_exception_handler converts Pydantic errors to 400 + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_derender_chat_oversized_logprobs_rejected(client): + """logprobs.content longer than max_model_len returns 400.""" + oversized_logprobs: dict = { + "content": [ + {"token": "x", "logprob": -1.0, "bytes": None, "top_logprobs": []} + for _ in range(1_000_000) + ] + } + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": { + "request_id": "test-logprobs-bound", + "choices": [ + { + "index": 0, + "token_ids": [42], + "finish_reason": "stop", + "logprobs": oversized_logprobs, + } + ], + }, + }, + ) + assert response.status_code == 400 + assert "logprobs.content length" in response.json()["error"]["message"] + + +@pytest.mark.asyncio +async def test_derender_chat_oversized_top_logprobs_rejected(client): + """top_logprobs count exceeding max_logprobs (default 20) returns 400.""" + oversized_top_logprobs = { + "content": [ + { + "token": "x", + "logprob": -1.0, + "bytes": None, + "top_logprobs": [ + {"token": f"t{i}", "logprob": -float(i), "bytes": None} + for i in range(25) + ], + } + ] + } + response = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL_NAME, + "generate_response": { + "request_id": "test-top-logprobs-bound", + "choices": [ + { + "index": 0, + "token_ids": [42], + "finish_reason": "stop", + "logprobs": oversized_top_logprobs, + } + ], + }, + }, + ) + assert response.status_code == 400 + msg = response.json()["error"]["message"] + assert "top_logprobs count" in msg + assert "max_logprobs" in msg + + +@pytest.mark.asyncio +async def test_derender_completion_oversized_token_ids_rejected(client): + """Completion endpoint also rejects oversized token_ids.""" + oversized_ids = [42] * 1_000_000 + response = await client.post( + "/v1/completions/derender", + json={ + "model": MODEL_NAME, + "generate_responses": [ + { + "request_id": "gen-0", + "choices": [ + { + "index": 0, + "token_ids": oversized_ids, + "finish_reason": "stop", + } + ], + } + ], + }, + ) + assert response.status_code == 400 + assert "max_model_len" in response.json()["error"]["message"] + + # --------------------------------------------------------------------------- # E2E: render -> derender roundtrip with parser (reasoning + tool calls) # --------------------------------------------------------------------------- diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py index 1ab963dc180..5b00d2e578e 100644 --- a/tests/entrypoints/serve/instrumentator/test_basic.py +++ b/tests/entrypoints/serve/instrumentator/test_basic.py @@ -83,7 +83,8 @@ async def test_show_version(server: RemoteOpenAIServer): response = requests.get(server.url_for("version")) response.raise_for_status() - assert response.json() == {"version": VLLM_VERSION} + # Tolerate additive fields (e.g. the Rust frontend reports its own version). + assert response.json()["version"] == VLLM_VERSION @pytest.mark.asyncio diff --git a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py new file mode 100644 index 00000000000..0f96bf161d2 --- /dev/null +++ b/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test that http_requests_total metric records correct status codes. + +Regression test for: Prometheus http_requests_total records 4xx exceptions +(ValueError, TypeError, etc.) as 5xx because they propagate through the +PrometheusInstrumentatorMiddleware before being caught by ServerErrorMiddleware. +""" + +from argparse import Namespace +from http import HTTPStatus + +import httpx +import pytest +from fastapi import FastAPI, HTTPException, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse +from prometheus_client import CollectorRegistry +from prometheus_fastapi_instrumentator import Instrumentator + +from vllm.entrypoints.serve.utils.server_utils import exception_handler +from vllm.exceptions import VLLMNotFoundError, VLLMValidationError + + +@pytest.fixture +def registry(): + """Create a fresh Prometheus registry for each test.""" + return CollectorRegistry() + + +@pytest.fixture +def app(registry): + """Create a minimal FastAPI app that mirrors vLLM's exception handler + and Prometheus middleware setup.""" + + app = FastAPI() + + # Mock app state that exception_handler needs + app.state.args = Namespace(log_error_stack=False) + + # Register exception handlers exactly as vLLM does in build_app() + app.exception_handler(HTTPException)(_http_exception_handler) + app.exception_handler(RequestValidationError)(_validation_exception_handler) + app.exception_handler(ValueError)(exception_handler) + app.exception_handler(TypeError)(exception_handler) + app.exception_handler(OverflowError)(exception_handler) + app.exception_handler(NotImplementedError)(exception_handler) + app.exception_handler(VLLMValidationError)(exception_handler) + app.exception_handler(VLLMNotFoundError)(exception_handler) + app.exception_handler(Exception)(exception_handler) + + # Instrument with Prometheus (same as vLLM's attach_router) + Instrumentator( + excluded_handlers=["/metrics"], + registry=registry, + ).add().instrument(app) + + # Test routes that raise different exception types + @app.get("/raise_value_error") + async def raise_value_error(): + raise ValueError("invalid input value") + + @app.get("/raise_type_error") + async def raise_type_error(): + raise TypeError("wrong type") + + @app.get("/raise_overflow_error") + async def raise_overflow_error(): + raise OverflowError("number too large") + + @app.get("/raise_not_implemented_error") + async def raise_not_implemented_error(): + raise NotImplementedError("feature not supported") + + @app.get("/raise_vllm_validation_error") + async def raise_vllm_validation_error(): + raise VLLMValidationError("bad parameter", parameter="temperature") + + @app.get("/raise_vllm_not_found_error") + async def raise_vllm_not_found_error(): + raise VLLMNotFoundError("model not found") + + @app.get("/raise_http_exception_400") + async def raise_http_exception_400(): + raise HTTPException(status_code=400, detail="bad request") + + @app.get("/raise_http_exception_404") + async def raise_http_exception_404(): + raise HTTPException(status_code=404, detail="not found") + + @app.get("/raise_runtime_error") + async def raise_runtime_error(): + raise RuntimeError("unexpected server error") + + @app.get("/success") + async def success(): + return {"status": "ok"} + + return app + + +async def _http_exception_handler(req: Request, exc: HTTPException): + return JSONResponse({"error": exc.detail}, status_code=exc.status_code) + + +async def _validation_exception_handler(req: Request, exc: RequestValidationError): + return JSONResponse({"error": str(exc)}, status_code=HTTPStatus.BAD_REQUEST) + + +def _get_http_requests_total(registry, method: str, handler: str): + """Extract the http_requests_total metric values grouped by status. + + Returns a dict like {"2xx": 1.0, "5xx": 1.0} for the given handler. + """ + results: dict[str, float] = {} + for metric in registry.collect(): + if metric.name == "http_requests": + for sample in metric.samples: + if ( + sample.name == "http_requests_total" + and sample.labels.get("method") == method + and sample.labels.get("handler") == handler + ): + status = sample.labels.get("status") + results[status] = results.get(status, 0) + sample.value + return results + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "endpoint,expected_status_group,expected_http_code", + [ + # These should record as 4xx in Prometheus + ("/raise_value_error", "4xx", 400), + ("/raise_type_error", "4xx", 400), + ("/raise_overflow_error", "4xx", 400), + ("/raise_vllm_validation_error", "4xx", 400), + ("/raise_vllm_not_found_error", "4xx", 404), + ("/raise_http_exception_400", "4xx", 400), + ("/raise_http_exception_404", "4xx", 404), + # NotImplementedError returns 501 which is still 5xx group + ("/raise_not_implemented_error", "5xx", 501), + # These should record as 5xx in Prometheus (genuine server errors) + ("/raise_runtime_error", "5xx", 500), + # Successful requests should record as 2xx + ("/success", "2xx", 200), + ], + ids=[ + "ValueError->4xx", + "TypeError->4xx", + "OverflowError->4xx", + "VLLMValidationError->4xx", + "VLLMNotFoundError->4xx", + "HTTPException(400)->4xx", + "HTTPException(404)->4xx", + "NotImplementedError->5xx", + "RuntimeError->5xx", + "success->2xx", + ], +) +async def test_http_requests_total_records_correct_status( + app, + registry, + endpoint, + expected_status_group, + expected_http_code, +): + """Verify that http_requests_total records the correct status group. + + The Prometheus metric should reflect the actual HTTP status code returned + to the client, not a default 500 for all exceptions. + """ + # raise_app_exceptions=False allows the full ASGI middleware stack + # (including ServerErrorMiddleware) to handle exceptions and generate + # proper HTTP responses, just like a real server would. + transport = httpx.ASGITransport(app=app, raise_app_exceptions=False) + async with httpx.AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + response = await client.get(endpoint) + + # Verify the HTTP response code returned to the client is correct + assert response.status_code == expected_http_code, ( + f"Expected HTTP {expected_http_code} for {endpoint}, got {response.status_code}" + ) + + # Verify Prometheus recorded the correct status group + metrics = _get_http_requests_total(registry, "GET", endpoint) + assert expected_status_group in metrics, ( + f"Expected Prometheus to record '{expected_status_group}' for " + f"{endpoint}, but got: {metrics}" + ) + assert metrics[expected_status_group] == 1.0, ( + f"Expected 1 request recorded as '{expected_status_group}' for " + f"{endpoint}, but got {metrics[expected_status_group]}" + ) + + # For endpoints that should be recorded as 4xx, verify they are NOT + # incorrectly recorded as 5xx + if expected_status_group == "4xx": + assert "5xx" not in metrics, ( + f"Expected NO '5xx' recording for {endpoint} " + f"(should be '4xx'), but found: {metrics}" + ) diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index 0fa3661f8ff..1c1f6ed2359 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -74,7 +74,7 @@ class FakeHarmonyParser(HarmonyParser): self.reasoning_parser = None self.tool_parser = None self._chunk_results: list[ChunkResult] = [] - self._flush_results: list[Segment | None] = [] + self._flush_results: list[list[Segment]] = [] self.processed_chunks: list[list[int]] = [] def enqueue_chunk_result( @@ -89,7 +89,7 @@ class FakeHarmonyParser(HarmonyParser): ) ) - def enqueue_flush_result(self, segment: Segment | None) -> None: + def enqueue_flush_result(self, segment: list[Segment]) -> None: self._flush_results.append(segment) def process_chunk(self, token_ids) -> ChunkResult: @@ -98,10 +98,10 @@ class FakeHarmonyParser(HarmonyParser): return self._chunk_results.pop(0) return ChunkResult(segments=[], reasoning_token_count=0) - def flush(self) -> Segment | None: + def flush(self) -> list[Segment]: if self._flush_results: return self._flush_results.pop(0) - return None + return [] def make_harmony_context( @@ -598,13 +598,21 @@ async def test_streaming_message_synchronization(): content=[TextContent(text=response_text)], recipient=Role.USER, ) - flush_segment = Segment( - channel="commentary", - recipient=None, - delta="", - completed_message=message, - ) - parser.enqueue_flush_result(flush_segment) + flush_segments = [ + Segment( + channel="final", + recipient=None, + delta=response_text, + completed_message=None, + ), + Segment( + channel="final", + recipient=None, + delta="", + completed_message=message, + ), + ] + parser.enqueue_flush_result(flush_segments) # Create another output to trigger synchronization via flush() context.append_output( @@ -618,8 +626,9 @@ async def test_streaming_message_synchronization(): assert context.num_init_messages == 1 assert context._messages[2].content[0].text == response_text assert context.last_append_flush_status is True - assert len(context.last_append_segments) == 1 - assert context.last_append_segments[0].completed_message is message + assert len(context.last_append_segments) == 2 + assert context.last_append_segments[-2].delta == response_text + assert context.last_append_segments[-1].completed_message is message def test_turn_metrics_copy_and_reset(): diff --git a/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml new file mode 100644 index 00000000000..ba292eb9724 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "nm-testing/Qwen2-1.5B-Instruct-FP8W8" +accuracy_threshold: 0.55 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml new file mode 100644 index 00000000000..3179c3251b2 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen2-1.5B-Instruct-FP8W8-humming.yaml @@ -0,0 +1,8 @@ +model_name: "nm-testing/Qwen2-1.5B-Instruct-FP8W8" +accuracy_threshold: 0.55 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml new file mode 100644 index 00000000000..66888312ef0 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "mgoin/Qwen3-0.6B-MXFP8" +accuracy_threshold: 0.39 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml new file mode 100644 index 00000000000..b83d9e6a9e9 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-0.6B-MXFP8-humming.yaml @@ -0,0 +1,8 @@ +model_name: "mgoin/Qwen3-0.6B-MXFP8" +accuracy_threshold: 0.39 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml new file mode 100644 index 00000000000..d3de5b3792f --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-fp8.yaml @@ -0,0 +1,13 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml new file mode 100644 index 00000000000..6b76efc4e22 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml new file mode 100644 index 00000000000..310255f6bf4 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml @@ -0,0 +1,11 @@ +model_name: "QuixiAI/Qwen3-30B-A3B-AWQ" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --dtype bfloat16 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml new file mode 100644 index 00000000000..6d2702d8f4c --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml new file mode 100644 index 00000000000..8ca4777ce5e --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml new file mode 100644 index 00000000000..618e7fdcc35 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "nm-testing/Qwen3-30B-A3B-Fp8-v1" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml new file mode 100644 index 00000000000..71aabcef99d --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Fp8-v1-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nm-testing/Qwen3-30B-A3B-Fp8-v1" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml new file mode 100644 index 00000000000..251fb252cf8 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml @@ -0,0 +1,11 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml new file mode 100644 index 00000000000..d813e05d6f9 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml new file mode 100644 index 00000000000..ab89ebf5154 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-GPTQ-Int4-humming.yaml @@ -0,0 +1,9 @@ +model_name: "Qwen/Qwen3-30B-A3B-GPTQ-Int4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml new file mode 100644 index 00000000000..f77ee1173e1 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml new file mode 100644 index 00000000000..7b4f82c8458 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml @@ -0,0 +1,9 @@ +model_name: "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml index 9b77af67327..d98f91e1f99 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml @@ -5,8 +5,7 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --quantization humming - --kernel-config.enable_flashinfer_autotune=False + --linear-backend humming env: VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml index 0b1599ff94b..67725fce64a 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml @@ -5,6 +5,5 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --quantization humming - --kernel-config.enable_flashinfer_autotune=False + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml new file mode 100644 index 00000000000..a20c1433103 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml @@ -0,0 +1,9 @@ +model_name: "nvidia/Qwen3-30B-A3B-NVFP4" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml new file mode 100644 index 00000000000..b58c0d710e4 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml new file mode 100644 index 00000000000..c932091dbe3 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml new file mode 100644 index 00000000000..fa5095cc882 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3-30B-A3B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --tensor-parallel-size 1 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml new file mode 100644 index 00000000000..3af5c03a245 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-35B-A3B-FP8" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml new file mode 100644 index 00000000000..85fce244200 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-FP8-humming.yaml @@ -0,0 +1,10 @@ +model_name: "Qwen/Qwen3.5-35B-A3B-FP8" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml new file mode 100644 index 00000000000..6a85ab384f2 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --quantization experts_int8 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml new file mode 100644 index 00000000000..d282fdc7019 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-experts-int8-humming.yaml @@ -0,0 +1,10 @@ +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --quantization experts_int8 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml new file mode 100644 index 00000000000..2a118098fc0 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-fp8.yaml @@ -0,0 +1,10 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml new file mode 100644 index 00000000000..34f17a4b055 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming-act-int8.yaml @@ -0,0 +1,11 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming +env: + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml new file mode 100644 index 00000000000..b617c61eb3e --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-4B-quantized.w4a16-humming.yaml @@ -0,0 +1,9 @@ +model_name: "RedHatAI/Qwen3.5-4B-quantized.w4a16" +accuracy_threshold: 0.82 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml new file mode 100644 index 00000000000..502ab776f40 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.6-35B-A3B-NVFP4-humming.yaml @@ -0,0 +1,10 @@ +model_name: "RedHatAI/Qwen3.6-35B-A3B-NVFP4" +accuracy_threshold: 0.91 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming diff --git a/tests/evals/gsm8k/configs/humming/config-act-fp8.txt b/tests/evals/gsm8k/configs/humming/config-act-fp8.txt index 05fb6a15838..42ff6be00ef 100644 --- a/tests/evals/gsm8k/configs/humming/config-act-fp8.txt +++ b/tests/evals/gsm8k/configs/humming/config-act-fp8.txt @@ -1,2 +1,9 @@ gpt-oss-20b-humming-act-fp8.yaml Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml +Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml +Qwen3-0.6B-MXFP8-humming-act-fp8.yaml +Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml +Qwen3-30B-A3B-FP8-block-humming-act-fp8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-FP8-humming-act-fp8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-act-int8.txt b/tests/evals/gsm8k/configs/humming/config-act-int8.txt new file mode 100644 index 00000000000..b018b35cbfd --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-act-int8.txt @@ -0,0 +1,4 @@ +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt new file mode 100644 index 00000000000..2c10777a095 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt @@ -0,0 +1,3 @@ +Qwen3-30B-A3B-int5wc-hadamard-humming.yaml +Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml +Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config.txt b/tests/evals/gsm8k/configs/humming/config.txt index 821025365c7..144ed959935 100644 --- a/tests/evals/gsm8k/configs/humming/config.txt +++ b/tests/evals/gsm8k/configs/humming/config.txt @@ -1,2 +1,13 @@ gpt-oss-20b-humming.yaml Qwen3-30B-A3B-MXFP4A16-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3.5-35B-A3B-FP8-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml index 00ba9eccfda..8e0d9535030 100644 --- a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml +++ b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.yaml @@ -5,7 +5,6 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --moe-backend humming env: VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml index e2beb3739b1..7e9b6508a20 100644 --- a/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml @@ -5,5 +5,4 @@ num_fewshot: 5 server_args: >- --enforce-eager --max-model-len 8192 - --tensor-parallel-size 1 --moe-backend humming diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 89d91cb4a98..45f2a13fdd7 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -10,6 +10,7 @@ import ast import asyncio import json import os +import tempfile import time from collections.abc import Generator @@ -25,7 +26,7 @@ INVALID = -9999999 def download_and_cache_file(url: str, filename: str | None = None) -> str: """Download and cache a file from a URL.""" if filename is None: - filename = os.path.join("/tmp", url.split("/")[-1]) + filename = os.path.join(tempfile.gettempdir(), url.split("/")[-1]) if os.path.exists(filename): return filename @@ -146,6 +147,7 @@ async def call_vllm_chat_api( def _build_gsm8k_prompts( num_questions: int = 1319, num_shots: int = 5, + gen_prefix: str = "", ) -> tuple[list[str], list[int]]: """Build few-shot GSM8K completion prompts and ground-truth labels.""" if num_questions == 0: @@ -157,14 +159,15 @@ def _build_gsm8k_prompts( for i in range(num_shots): few_shot_examples += ( f"Question: {train_data[i]['question']}\n" - f"Answer: {train_data[i]['answer']}\n\n" + f"Answer:{gen_prefix} {train_data[i]['answer']}\n\n" ) prompts = [] labels = [] for i in range(num_questions): prompts.append( - few_shot_examples + f"Question: {test_data[i]['question']}\nAnswer:" + few_shot_examples + + f"Question: {test_data[i]['question']}\nAnswer:{gen_prefix}" ) labels.append(get_answer_value(test_data[i]["answer"])) @@ -213,6 +216,7 @@ def evaluate_gsm8k( temperature: float = 0.0, seed: int | None = 42, request_timeout_seconds: float = 600, + gen_prefix: str = "", ) -> dict[str, float | int]: """ Evaluate GSM8K accuracy using vLLM serve endpoint. @@ -220,7 +224,7 @@ def evaluate_gsm8k( Returns dict with accuracy, invalid_rate, latency, etc. """ base_url = f"{host}:{port}" - prompts, labels = _build_gsm8k_prompts(num_questions, num_shots) + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) num_questions = len(prompts) async def run_async_evaluation(): @@ -278,6 +282,7 @@ def evaluate_gsm8k_offline( num_shots: int = 5, max_tokens: int = 256, temperature: float = 0.0, + gen_prefix: str = "", ) -> dict[str, float | int]: """Evaluate GSM8K accuracy using an offline vllm.LLM object. @@ -286,7 +291,7 @@ def evaluate_gsm8k_offline( """ from vllm import SamplingParams - prompts, labels = _build_gsm8k_prompts(num_questions, num_shots) + prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) sampling_params = SamplingParams( temperature=temperature, diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index c9ec5ff66e5..0c48af6d3c6 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -70,6 +70,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict: host=host, port=port, request_timeout_seconds=request_timeout_seconds, + gen_prefix=eval_config.get("gen_prefix", ""), ) return results diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 0340ca9a477..01405a0e5d1 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -150,7 +150,7 @@ def _reference_index_topk( num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE pages = block_table[req_id, :num_blocks] k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1) - score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale + score = sm_scale * torch.einsum("qhd,kd->hqk", q.float(), k.float()) q_pos = prefix_len + torch.arange(q_len, device=idx_q.device) k_pos = torch.arange(k.shape[0], device=idx_q.device) @@ -621,11 +621,12 @@ def test_decode_index_topk_fp8(num_idx_heads: int): init_blocks=init_blocks, local_blocks=local_blocks, num_kv_heads=num_idx_heads, - sm_scale=head_dim**-0.5, decode_query_len=decode_query_len, + max_decode_query_len=decode_query_len, ) # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK - # in fp32, so it must match an fp32 matmul of the same e4m3 values). + # in fp32 with no scaling, so it must match an unscaled fp32 matmul of the + # same e4m3 values). expected = _reference_index_topk( idx_q.float(), index_kv_cache.float(), @@ -636,7 +637,6 @@ def test_decode_index_topk_fp8(num_idx_heads: int): topk, init_blocks, local_blocks, - head_dim**-0.5, ) _assert_topk_indices_equal_unordered(actual, expected) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py index 99b4a0e19ee..2b9a4e823c0 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -113,7 +113,7 @@ def _build_decode_metadata(): # stub with the attribute is enough for metadata construction. layer_name = "placeholder" vllm_config.compilation_config.static_forward_context[layer_name] = ( - types.SimpleNamespace(prefill_backend=None) + types.SimpleNamespace(prefill_backend=torch.empty((1,))) ) init_workspace_manager(device) diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index c02a457c98a..9e33f24ea28 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -30,8 +30,7 @@ NUM_Q_HEADS = 8 NUM_KV_HEADS = 8 HEAD_SIZES = [128, 256] BLOCK_SIZES = [16, 64] -# TODO: re-add torch.float16 once AITER reenables fp16 unified attention. -DTYPES = [torch.bfloat16] +DTYPES = [torch.bfloat16, torch.float16] FP8_DTYPE = current_platform.fp8_dtype() # (query_len, kv_len) per sequence diff --git a/tests/kernels/helion/test_silu_and_mul_per_block_quant.py b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..b8fcd9c8a67 --- /dev/null +++ b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the silu_and_mul_per_block_quant helion kernel +Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_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.silu_and_mul_per_block_quant import ( + _pick_cache, + baseline, + pick_config, + silu_and_mul_per_block_quant, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion +from vllm.utils.torch_utils import set_random_seed + +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, intermediate_size: int, group_size: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + args = ( + result, + input, + scale, + group_size, + scale_ub, + False, + ) + return args + + +class TestSiluAndMulPerBlockQuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_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( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_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( + {"intermediate_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({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_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( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 32} + ) + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestSiluAndMulPerBlockQuantCorrectness: + @pytest.mark.parametrize("num_tokens", [1, 7, 4096]) + @pytest.mark.parametrize("hidden_size", [1024, 2048, 5120]) + @pytest.mark.parametrize("group_size", [64, 128]) + @pytest.mark.parametrize("is_scale_transposed", [False, True]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8]) + @pytest.mark.parametrize("has_scale_ub", [True, False]) + @pytest.mark.parametrize("seed", [0]) + def test_silu_and_mul_per_block_quant( + self, + num_tokens: int, + hidden_size: int, + group_size: int, + is_scale_transposed: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + has_scale_ub: bool, + seed: int, + ) -> None: + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + set_random_seed(seed) + + if hidden_size % group_size != 0: + return + + if has_scale_ub and quant_dtype != FP8_DTYPE: + # skip + return + + scale = 1 / hidden_size + x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale + + if has_scale_ub: + act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:] + act_abs = act.abs().float() + scale_ub = 0.5 * (act_abs.mean() + act_abs.amax()) + else: + scale_ub = None + + ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype) + + if is_scale_transposed: + ref_scales = torch.empty( + (hidden_size // group_size, x.shape[0]), + device="cuda", + dtype=torch.float32, + ).t() + else: + ref_scales = torch.empty( + (x.shape[0], hidden_size // group_size), + device="cuda", + dtype=torch.float32, + ) + + ops_out = ref_out.clone() + ops_scales = ref_scales.clone() + + baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed) + silu_and_mul_per_block_quant( + ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed + ) + + torch.testing.assert_close(ref_scales, ops_scales) + # allow 1 ULP difference + assert ( + ref_out.view(torch.uint8).to(torch.int16) + - ops_out.view(torch.uint8).to(torch.int16) + ).abs().max() <= 1 + + +class TestSiluAndMulPerBlockQuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "silu_and_mul_per_block_quant" in registered_kernels + + kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"] + assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["out", "scales"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["silu_and_mul_per_block_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/mamba/test_cpu_short_conv.py b/tests/kernels/mamba/test_cpu_short_conv.py new file mode 100644 index 00000000000..c8e85a45511 --- /dev/null +++ b/tests/kernels/mamba/test_cpu_short_conv.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.config import CompilationConfig, VllmConfig +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.mamba.short_conv import ShortConv +from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm +from vllm.platforms import current_platform +from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + + +@pytest.fixture(autouse=True) +def mock_dist(): + with ( + patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + return_value=1, + ), + patch( + "vllm.distributed.parallel_state.model_parallel_is_initialized", + return_value=True, + ), + patch( + "vllm.distributed.parallel_state.get_tp_group", + return_value=MagicMock(rank_in_group=0), + ), + ): + yield + + +@pytest.fixture +def vllm_config(): + # ShortConv only needs compilation_config from the current vLLM config, so a + # minimal config (model_config=None) avoids mocking ModelConfig and the + # associated VllmConfig validation churn. + return VllmConfig(compilation_config=CompilationConfig()) + + +def test_short_conv_forward_native_prefill(vllm_config): + prefix = "test_layer" + config = SimpleNamespace(conv_L_cache=4, conv_bias=True) + dim = 16 + + from vllm.config import set_current_vllm_config + + with set_current_vllm_config(vllm_config): + layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) + + layer.to("cpu") + # vLLM Linear layers allocate weights with torch.empty (uninitialized). + # On ARM these come back as zero-filled pages, so in_proj output is zero and + # the prefill state stays zero. Seed + init to make the test platform-safe. + torch.manual_seed(0) + for p in layer.parameters(): + torch.nn.init.normal_(p) + dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) + dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) + + # Mock AttentionMetadata + num_prefills = 1 + num_prefill_tokens = 5 + query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32) + state_indices_tensor_p = torch.tensor([0], dtype=torch.int32) + + # ShortConvAttentionMetadata + attn_metadata = ShortConvAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=0, + num_decode_tokens=0, + num_reqs=1, + query_start_loc_p=query_start_loc_p, + has_initial_states_p=torch.tensor([False]), + state_indices_tensor_p=state_indices_tensor_p, + state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32), + num_accepted_tokens=None, + query_start_loc_d=None, + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + num_computed_tokens_p=None, + seq_lens=torch.tensor([5]), + ) + + # Mock KV cache + # conv_state shape (num_blocks, L_cache - 1, dim) + conv_state = torch.zeros((1, config.conv_L_cache - 1, dim)) + layer.kv_cache = (conv_state,) + + hidden_states = torch.randn((num_prefill_tokens, dim)) + output = torch.zeros_like(hidden_states) + + attn_metadata_dict = {prefix: attn_metadata} + with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): + layer.forward_native(hidden_states, output) + + # Check if KV cache was updated + assert not torch.allclose(conv_state, torch.zeros_like(conv_state)) + + +def test_short_conv_forward_native_decode(vllm_config): + prefix = "test_layer_decode" + config = SimpleNamespace(conv_L_cache=4, conv_bias=True) + dim = 16 + + from vllm.config import set_current_vllm_config + + with set_current_vllm_config(vllm_config): + layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) + + layer.to("cpu") + torch.manual_seed(0) + for p in layer.parameters(): + torch.nn.init.normal_(p) + dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) + dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) + + # Mock AttentionMetadata for 2 decode requests + num_decodes = 2 + state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32) + + attn_metadata = ShortConvAttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + num_reqs=num_decodes, + query_start_loc_p=None, + has_initial_states_p=None, + state_indices_tensor_p=torch.empty((0,), dtype=torch.int32), + state_indices_tensor_d=state_indices_tensor_d, + num_accepted_tokens=None, + query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32), + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + num_computed_tokens_p=None, + seq_lens=torch.tensor([1, 1]), + ) + + # Mock KV cache (2 blocks for 2 requests) + conv_state = torch.randn((2, config.conv_L_cache - 1, dim)) + layer.kv_cache = (conv_state,) + + hidden_states = torch.randn((num_decodes, dim)) + output = torch.zeros_like(hidden_states) + + old_conv_state = conv_state.clone() + + attn_metadata_dict = {prefix: attn_metadata} + with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): + layer.forward_native(hidden_states, output) + + # Check if KV cache was updated + assert not torch.allclose(conv_state, old_conv_state) + + +def test_dispatch_cpu_unquantized_gemm_conv_layer(): + # Convolution layers have >2D weights; dispatch should skip them gracefully. + # Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so + # the AMX prepack branch does not raise on AMX-capable CPUs. + class MockConvLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn(32, 1, 4, dtype=torch.bfloat16) + ) + self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16)) + + layer = MockConvLayer() + # The ndim != 2 guard returns early without raising. + dispatch_cpu_unquantized_gemm(layer, remove_weight=False) + # No cpu_linear set — conv layers are handled elsewhere. + assert not hasattr(layer, "cpu_linear") diff --git a/tests/kernels/moe/test_cpu_int4_moe.py b/tests/kernels/moe/test_cpu_int4_moe.py index 05694eb08b2..04931e386f1 100644 --- a/tests/kernels/moe/test_cpu_int4_moe.py +++ b/tests/kernels/moe/test_cpu_int4_moe.py @@ -8,19 +8,21 @@ import pytest import torch import torch.nn.functional as F -from vllm.platforms import current_platform +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import ( + CPUExpertsInt4, +) +from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import ( + convert_to_w4a8_int8_moe_format, +) +from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed -if not current_platform.is_cpu(): - pytest.skip("skipping CPU-only tests", allow_module_level=True) - -# Check if the dynamic_4bit_int_moe op is available -if not hasattr(torch.ops._C, "dynamic_4bit_int_moe"): - pytest.skip("dynamic_4bit_int_moe op not available", allow_module_level=True) - -# Check if KleidiAI ops are available -if not hasattr(torch.ops.aten, "_dyn_quant_pack_4bit_weight"): - pytest.skip("KleidiAI 4-bit ops not available", allow_module_level=True) +if ( + not current_platform.is_cpu() + or current_platform.get_cpu_architecture() != CpuArchEnum.ARM +): + pytest.skip("skipping Arm CPU-only tests", allow_module_level=True) # Tolerance for INT4 W4A8 @@ -34,49 +36,6 @@ def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: return F.silu(x[..., :d]) * x[..., d:] -def _pack_int4_weight_to_kleidi( - int4_as_int8: torch.Tensor, - scales: torch.Tensor, - bias: torch.Tensor | None, - group_size: int, - in_features: int, - out_features: int, -) -> torch.Tensor: - """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format. - - Args: - int4_as_int8: [out, in] int8 tensor with values in [-8, 7] - scales: [out, in//group_size] or [out, 1] for channel-wise - bias: [out] optional bias - group_size: Quantization group size (-1 for channel-wise) - in_features: Input dimension - out_features: Output dimension - - Returns: - Packed weight tensor in KleidiAI format - """ - # Shift to unsigned nibble [0, 15] - tmp = int4_as_int8.add(8) - # Pack pairs along input dimension - uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8) - - # Determine scale dtype based on group_size - scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16 - scales_typed = scales.to(scale_dtype) - bias_typed = None if bias is None else bias.to(torch.float32) - - # Pack using KleidiAI op - actual_group_size = in_features if group_size == -1 else group_size - return torch.ops.aten._dyn_quant_pack_4bit_weight( - uint8_nibbles, - scales_typed, - bias_typed, - actual_group_size, - in_features, - out_features, - ) - - def _make_int4_moe_weights( E: int, N: int, @@ -124,59 +83,29 @@ def _make_int4_moe_weights( w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01 w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01 - # Pack weights for each expert - w13_packed_list = [] - w2_packed_list = [] + w13_packed, w2_packed, *_ = convert_to_w4a8_int8_moe_format( + w13_weight=w13_int4, + w2_weight=w2_int4, + w13_weight_scale=w13_scales, + w2_weight_scale=w2_scales, + group_size=group_size, + w13_bias=w13_bias if has_bias else None, + w2_bias=w2_bias if has_bias else None, + ) - for e in range(E): - w13_packed_list.append( - _pack_int4_weight_to_kleidi( - w13_int4[e], - w13_scales[e], - w13_bias[e] if (has_bias and w13_bias is not None) else None, - group_size, - K, - 2 * N, - ) - ) - w2_packed_list.append( - _pack_int4_weight_to_kleidi( - w2_int4[e], - w2_scales[e], - w2_bias[e] if (has_bias and w2_bias is not None) else None, - group_size, - N, - K, - ) - ) + if group_size == -1: + w13_scale = w13_scales.float() + w2_scale = w2_scales.float() + else: + w13_scale = w13_scales.float().repeat_interleave(group_size, dim=-1) + w2_scale = w2_scales.float().repeat_interleave(group_size, dim=-1) - w13_packed = torch.stack(w13_packed_list, dim=0) - w2_packed = torch.stack(w2_packed_list, dim=0) - - # Create reference dequantized weights - w13_ref = torch.zeros(E, 2 * N, K, dtype=torch.float32) - w2_ref = torch.zeros(E, K, N, dtype=torch.float32) - - for e in range(E): - # Dequantize w13 - for i in range(2 * N): - for j in range(K): - group_idx = 0 if group_size == -1 else (j // group_size) - w13_ref[e, i, j] = ( - w13_int4[e, i, j].float() * w13_scales[e, i, group_idx].float() - ) - if has_bias and w13_bias is not None: - w13_ref[e, i, j] += w13_bias[e, i].float() - - # Dequantize w2 - for i in range(K): - for j in range(N): - group_idx = 0 if group_size == -1 else (j // group_size) - w2_ref[e, i, j] = ( - w2_int4[e, i, j].float() * w2_scales[e, i, group_idx].float() - ) - if has_bias and w2_bias is not None: - w2_ref[e, i, j] += w2_bias[e, i].float() + w13_ref = w13_int4.float() * w13_scale + w2_ref = w2_int4.float() * w2_scale + if has_bias and w13_bias is not None: + w13_ref = w13_ref + w13_bias.float().unsqueeze(-1) + if has_bias and w2_bias is not None: + w2_ref = w2_ref + w2_bias.float().unsqueeze(-1) return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias @@ -233,17 +162,20 @@ MoE_CONFIGS = [ (768, 2048, 16, 4, 64), ] SEEDS = [0, 42] +ACTIVATION_DTYPES = [torch.float32, torch.bfloat16, torch.float16] @pytest.mark.parametrize("M", NUM_TOKENS) @pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS) @pytest.mark.parametrize("seed", SEEDS) -def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): +@pytest.mark.parametrize("activation_dtype", ACTIVATION_DTYPES) +def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed, activation_dtype): """Test dynamic_4bit_int_moe kernel against dequantized torch reference.""" set_random_seed(seed) + activation = MoEActivation.SILU # Generate input activations - a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5) + a = torch.randn(M, K, dtype=activation_dtype) / (K**0.5) # Generate INT4 weights w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights( @@ -266,8 +198,6 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): ) # Test dynamic_4bit_int_moe kernel - # Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style - activation_kind = 1 apply_router_weight_on_input = False out = torch.ops._C.dynamic_4bit_int_moe( @@ -278,14 +208,14 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed): w2_packed, K, # H (hidden_size / w2_out_features) N, # I (intermediate_size / w2_in_features) - 2 * N, # I2 (2*intermediate_size / w13_out_features) group_size, apply_router_weight_on_input, - activation_kind, + CPUExpertsInt4._activation_kind(activation), ) + assert out.dtype == activation_dtype torch.testing.assert_close( - ref_out.bfloat16(), + ref_out, out, atol=INT4_W4A8_ATOL, rtol=INT4_W4A8_RTOL, diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index 5aac3784ba4..b15cbcdd812 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -8,8 +10,7 @@ from vllm.platforms import current_platform if not current_platform.is_device_capability_family(120): pytest.skip( - reason="FlashInfer CuteDSL SM12x MoE requires SM120 " - "(RTX Pro 6000 / DGX Spark).", + reason="FlashInfer B12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).", allow_module_level=True, ) @@ -18,8 +19,8 @@ from vllm.utils.flashinfer import has_flashinfer_b12x_moe if not has_flashinfer_b12x_moe(): pytest.skip( reason=( - "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout " - "not available in installed FlashInfer (needs PRs #3051 and #3066)." + "FlashInfer B12xMoEWrapper not available in installed " + "FlashInfer (needs PR #3080)." ), allow_module_level=True, ) @@ -40,7 +41,6 @@ from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( FlashInferB12xExperts, ) -from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout from vllm.utils.torch_utils import set_random_seed # Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256, @@ -59,7 +59,7 @@ def _reorder_gate_up_to_up_gate( ) -> tuple[torch.Tensor, torch.Tensor]: """Swap gate and up-projection halves along dim=1 to [up, gate] order. - The SM12x kernel expects weights in [up (w3), gate (w1)] order while the + The B12x kernel expects weights in [up (w3), gate (w1)] order while the BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``. """ @@ -70,6 +70,22 @@ def _reorder_gate_up_to_up_gate( ) +def _process_b12x_weights( + experts: FlashInferB12xExperts, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_scale_2: torch.Tensor, + w2_scale_2: torch.Tensor, +) -> None: + layer = SimpleNamespace( + w13_weight_scale=w1_scale, + w13_weight_scale_2=w1_scale_2, + w2_weight_scale=w2_scale, + w2_weight_scale_2=w2_scale_2, + ) + experts.process_weights_after_loading(layer) + + @pytest.mark.parametrize("m,n,k", MNK_FACTORS) @pytest.mark.parametrize("e", [8, 16]) @pytest.mark.parametrize("topk", [1, 2, 4]) @@ -174,22 +190,12 @@ def test_flashinfer_b12x_moe( moe_config=moe_config, quant_config=quant_config, ) - # In production, process_weights_after_loading computes these after - # normalizing block scales. In the test the scales are already in final - # form (global_scale=1.0), so we compute the MMA layouts directly. - num_experts_w1, m1, k1_sf = w1_blockscale.shape - experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( - w1_blockscale.reshape(num_experts_w1 * m1, k1_sf), - m=m1, - k=k1_sf * 16, - num_groups=num_experts_w1, - ) - num_experts_w2, m2, k2_sf = w2_blockscale.shape - experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( - w2_blockscale.reshape(num_experts_w2 * m2, k2_sf), - m=m2, - k=k2_sf * 16, - num_groups=num_experts_w2, + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, ) kernel = mk.FusedMoEKernel( @@ -224,5 +230,134 @@ def test_flashinfer_b12x_moe( torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-1) +@pytest.mark.parametrize("m,n,k", MNK_FACTORS) +@pytest.mark.parametrize("e", [8, 16]) +@pytest.mark.parametrize("topk", [1, 2, 4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.inference_mode() +def test_flashinfer_b12x_moe_relu2( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + """Test FlashInferB12xExperts with ReLU2 (non-gated) activation. + + ReLU2 is used by Nemotron-H style models. Unlike the gated SiLU + path, w1 has shape [E, N, K] (not [E, 2N, K]) and the activation + is relu(x)^2 without a gate/up split. + """ + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + + # Non-gated: w1 shape is (e, n, k), not (e, 2n, k). + w1_bf16 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + + gs = torch.ones(1, device="cuda", dtype=torch.float32) + sf_vec_size = 16 + + # W1: no gate/up reordering for non-gated. + w1_flat = w1_bf16.reshape(e * n, k) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w1_q = w1_q_flat.view(e, n, k // 2) + w1_blockscale = w1_sf_flat.view(e, n, w1_sf_flat.shape[1]) + + w2_flat = w2_bf16.reshape(e * k, n) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w2_q = w2_q_flat.view(e, k, n // 2) + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) + + ones_e = torch.ones(e, device="cuda", dtype=torch.float32) + + quant_config = nvfp4_moe_quant_config( + g1_alphas=ones_e, + g2_alphas=ones_e, + a1_gscale=ones_e, + a2_gscale=ones_e, + w1_scale=w1_blockscale, + w2_scale=w2_blockscale, + ) + + moe_config = make_dummy_moe_config( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + in_dtype=dtype, + activation=MoEActivation.RELU2_NO_MUL, + ) + + experts = FlashInferB12xExperts( + moe_config=moe_config, + quant_config=quant_config, + ) + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, + ) + + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts, + inplace=False, + ) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + b12x_output = kernel.apply( + hidden_states=a, + w1=w1_q, + w2=w2_q, + topk_weights=topk_weights, + topk_ids=topk_ids, + global_num_experts=e, + activation=MoEActivation.RELU2_NO_MUL, + apply_router_weight_on_input=False, + expert_map=None, + ) + + torch_output = torch_moe( + a, + w1_bf16, + w2_bf16, + score, + topk, + activation=MoEActivation.RELU2_NO_MUL, + ) + + torch.testing.assert_close( + b12x_output, + torch_output, + atol=2e-1, + rtol=2e-1, + ) + + if __name__ == "__main__": test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 552063988fa..cc8d9c36dc0 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -36,6 +36,9 @@ from vllm.distributed import ( get_eplb_group, tensor_model_parallel_all_gather, ) +from vllm.distributed.device_communicators.all_reduce_utils import ( + gpu_p2p_access_check, +) from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace from vllm.forward_context import set_forward_context @@ -106,6 +109,8 @@ if has_deep_ep(): if has_nixl_ep(): BACKENDS += ["nixl_ep"] +DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"} + QUANT_METHODS = [ None, "fp8", @@ -416,6 +421,22 @@ def generate_valid_test_configs( return configs +@functools.cache +def visible_devices_have_peer_access(world_size: int) -> bool: + if not current_platform.is_cuda(): + return True + + try: + return all( + gpu_p2p_access_check(src, dst) + for src in range(world_size) + for dst in range(world_size) + if src != dst + ) + except RuntimeError: + return False + + # TODO: break this up into sections def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: # routed_input_transform only makes sense with shared_experts (latent MoE) @@ -1802,6 +1823,9 @@ def test_moe_layer( if enable_eplb and not use_ep: pytest.skip("EPLB requires EP.") + if backend in DEEPEP_BACKENDS and not visible_devices_have_peer_access(world_size): + pytest.skip("DeepEP backends require peer access between visible GPUs.") + verbosity = pytestconfig.getoption("verbose") if os.environ.get("VLLM_LOGGING_LEVEL") is None: diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 8c620afbc81..d6eb488a643 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.metadata +import types from dataclasses import dataclass from importlib.util import find_spec @@ -10,7 +11,7 @@ import torch from packaging import version from tests.kernels.moe.utils import check_accuracy -from vllm._aiter_ops import is_aiter_found +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer @@ -1247,6 +1248,7 @@ def test_rocm_mxfp4_moe_oracle( num_tokens: int, hidden_size: int, intermediate_size: int, + monkeypatch: pytest.MonkeyPatch, ): """ Test ROCm MXFP4 MoE using oracle functions. @@ -1268,6 +1270,7 @@ def test_rocm_mxfp4_moe_oracle( if config["requires_gfx950"] and not ROCM_GFX950: pytest.skip(f"Backend {backend_name} requires GFX950") + import vllm.distributed.parallel_state as ps from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( @@ -1282,6 +1285,12 @@ def test_rocm_mxfp4_moe_oracle( # Initialize workspace manager (needed for modular kernels) init_workspace_manager(torch.accelerator.current_device_index()) + # Set up the TP Group to prevent failure on should_use_cdna4_mx_scale_swizzle check + monkeypatch.setattr(ps, "_TP", types.SimpleNamespace(world_size=1)) + + # AITER must be enabled or aiter_mxfp4_w4a8_moe asserts before dispatch. + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) + # Map string to enum backend = Mxfp4MoeBackend[backend_name] diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 3f3bcebd11e..5fdcb8682f5 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -55,6 +55,7 @@ def make_dummy_moe_config( intermediate_size: int = 1, in_dtype: torch.dtype = torch.bfloat16, max_num_tokens: int = 512, + activation: MoEActivation = MoEActivation.SILU, ) -> FusedMoEConfig: """ This is a dummy config for the mk constructor interface @@ -73,7 +74,7 @@ def make_dummy_moe_config( else num_experts, num_logical_experts=num_experts, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - activation=MoEActivation.SILU, + activation=activation, in_dtype=in_dtype, device="cuda", routing_method=RoutingMethodType.TopK, diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 4cb638e47af..c5eaa2f9321 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -11,6 +11,7 @@ from tests.kernels.quant_utils import ( native_per_token_group_quant_fp8, native_w8a8_block_matmul, ) +from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( @@ -93,7 +94,24 @@ def test_per_token_group_quant_fp8( tma_aligned_scales=tma_aligned_scales, ) - assert torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15) + if current_platform.is_rocm(): + # On gfx950 the Triton and PyTorch FP8 kernels can round in opposite + # directions when an element lands at the midpoint between two adjacent + # e4m3fn values (1-ULP tie-breaking). Verify: (1) no element is more + # than 1 FP8 ULP away, and (2) fewer than 0.05% of elements have any + # mismatch. Observed worst case across all parameter combos: 0.049%, + # max ULP = 1. + ulp = fp8_ulp_distance(out, ref_out) + assert (ulp <= 1).all(), ( + f"FP8 mismatch > 1 ULP: {int((ulp > 1).sum())} elements" + ) + assert float((ulp > 0).float().mean()) < 5e-4, ( + f"Too many 1-ULP mismatches: {int((ulp > 0).sum())}/{ulp.numel()}" + ) + else: + assert torch.allclose( + out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15 + ) assert torch.allclose(scale, ref_scale) if column_major_scales: diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 9a14edc4271..60fbf53d4ad 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -284,6 +284,93 @@ def test_mxfp8_native_moe(T, H, inter, E, top_k): assert _relerr(got, ref) < 5e-2 +# --------------------------------------------------------------------------- # +# Native MXFP8 grouped GEMM (dot_scaled) vs pure-PyTorch grouped matmul +# --------------------------------------------------------------------------- # +def _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, num_valid, mul_weight=None): + """Pure-PyTorch reference for ``_grouped_gemm_mxfp8``. + + For each routed (expanded) token ``tid in [0, num_valid)`` the kernel writes + ``out[tid] = a[tid // a_div] @ w[expert(tid)].T`` (fp32 accumulate), optionally + scaled by ``mul_weight[tid]``. The expert for ``tid`` is ``topk_ids.flatten() + [tid]`` (row-major expansion: ``tid = token*top_k + slot``). This is computed + here with plain ``torch.matmul`` on the dequantized operands — independent of + the Triton ``dot_scaled`` path and of the (separate) aiter backend. + """ + eids = topk_ids.reshape(-1) + n = w_deq.shape[1] + out = torch.empty(num_valid, n, dtype=torch.float32, device=a_deq.device) + for tid in range(num_valid): + e = int(eids[tid].item()) + out[tid] = a_deq[tid // a_div].float() @ w_deq[e].float().T + if mul_weight is not None: + out[tid] *= float(mul_weight[tid].item()) + return out + + +@requires_gfx950 +@pytest.mark.parametrize("T,N,K,E,top_k", [(8, 256, 128, 8, 2), (5, 512, 256, 16, 4)]) +@pytest.mark.parametrize("weighted", [False, True]) +@torch.inference_mode() +def test_mxfp8_grouped_gemm_native(T, N, K, E, top_k, weighted): + """Directly exercise ``_grouped_gemm_mxfp8`` against a non-Triton reference. + + Covers both call modes used by ``fused_moe_mxfp8_native``: + * ``weighted=False`` -> g1: ``a_div=top_k`` (a-row shared across the top_k + expansions of a token), no per-token weight. + * ``weighted=True`` -> g2: ``a_div=1`` (one a-row per expansion), output + scaled by ``topk_weights``. + """ + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( + _grouped_gemm_mxfp8, + ) + from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( + moe_align_block_size, + ) + + torch.manual_seed(0) + block_m = 64 + a_div = 1 if weighted else top_k + m_routed = T * top_k + # a-rows: g1 reads one row per token (a_div=top_k); g2 one per expansion. + a_rows = m_routed if weighted else T + a_bf16 = torch.randn(a_rows, K, device=DEVICE, dtype=torch.bfloat16) * 0.5 + w_bf16 = torch.randn(E, N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1 + a_fp8, a_scale = _mxfp8_e4m3_quantize_torch(a_bf16, is_sf_swizzled_layout=False) + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False) + + logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32) + topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1) + topk_weights = topk_weights.to(torch.float32) + topk_ids = topk_ids.to(torch.int32) + mul = topk_weights.reshape(-1) if weighted else None + + sorted_ids, expert_ids, num_post = moe_align_block_size( + topk_ids, block_m, E, None, ignore_invalid_experts=False + ) + got = _grouped_gemm_mxfp8( + a_fp8, + a_scale, + w_fp8, + w_scale, + sorted_ids, + expert_ids, + num_post, + m_routed, + top_k, + block_m, + torch.bfloat16, + a_div=a_div, + mul_weight_by=mul, + ) + # Reference: dequant the SAME bits the kernel reads, plain torch matmul. + a_deq = dequant_mxfp8_to_bf16(a_fp8, a_scale) + w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale) + ref = _ref_grouped_gemm(a_deq, w_deq, topk_ids, a_div, m_routed, mul) + assert got.shape == (m_routed, N) + assert _relerr(got, ref) < 5e-2 + + # --------------------------------------------------------------------------- # # MXFP8 linear emulation: BF16-at-load (default) vs per-step dequant + switch # --------------------------------------------------------------------------- # @@ -335,3 +422,97 @@ def test_mxfp8_linear_emulation_bf16_at_load( out = kernel.apply_weights(layer, x) assert out.dtype == act_dtype # dtype-match preserved (no tl.dot/F.linear crash) assert _relerr(out.float(), out_ref.float()) < 2e-2 + + +# ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ──────────────── +# Regression for the EP + aiter-master-switch interaction: under expert +# parallelism ``RoutedExperts.expert_map`` hands the experts either the 0/1 +# ``expert_mask`` (aiter master ON, ``rocm_aiter_fmoe_enabled``) or vLLM's -1 +# index map (master OFF). ``AiterMxfp8Experts.apply`` must forward the right 0/1 +# mask to aiter in BOTH cases. The old code always rebuilt the mask via +# ``(expert_map >= 0)``; on the already-0/1 mask that collapses to all-ones (no +# experts masked out) and EP output becomes garbage (no accuracy). +def _capture_expert_mask(expert_map, *, rocm_aiter_fmoe_enabled, global_num_experts): + """Drive the real ``AiterMxfp8Experts.apply`` mask branch and capture the + ``expert_mask`` it forwards to ``rocm_aiter_ops.fused_moe``.""" + from types import SimpleNamespace + from unittest import mock + + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( + AiterMxfp8Experts, + ) + + experts = object.__new__(AiterMxfp8Experts) # bypass heavy __init__ + experts.moe_config = SimpleNamespace( + rocm_aiter_fmoe_enabled=rocm_aiter_fmoe_enabled + ) + experts.quant_config = SimpleNamespace(gemm1_clamp_limit=None) + experts.w1_scale_val = None + experts.w2_scale_val = None + + captured = {} + + def _fake_fused_moe(hidden_states, w1, w2, tw, ti, *, expert_mask, **kw): + captured["expert_mask"] = expert_mask + return torch.zeros_like(hidden_states) + + w1 = torch.zeros(1, device=DEVICE) + w2 = torch.zeros(1, device=DEVICE) + out = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16) + hidden = torch.zeros(4, 8, device=DEVICE, dtype=torch.bfloat16) + tw = torch.ones(4, 2, device=DEVICE) + ti = torch.zeros(4, 2, dtype=torch.int32, device=DEVICE) + + with mock.patch.object(rocm_aiter_ops, "fused_moe", side_effect=_fake_fused_moe): + experts.apply( + output=out, + hidden_states=hidden, + w1=w1, + w2=w2, + topk_weights=tw, + topk_ids=ti, + activation=None, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=None, + workspace13=None, + workspace2=None, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + return captured["expert_mask"] + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_aiter_mxfp8_ep_expert_mask_both_master_modes(): + """Both aiter-master forms must yield the SAME correct 0/1 aiter mask; + guards the EP+master regression (mask must not collapse to all-ones).""" + from vllm.model_executor.layers.fused_moe.expert_map_manager import ( + determine_expert_map, + ) + + E, ep_size, ep_rank = 8, 2, 0 # rank owns global experts 0..3 + # master OFF: vLLM's -1 index map + _, idx_map, _ = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=False) + # master ON: 0/1 mask (+ trailing sentinel) that RoutedExperts forwards + _, _, ep_mask = determine_expert_map(ep_size, ep_rank, E, return_expert_mask=True) + + idx_map = idx_map.to(DEVICE) + ep_mask = ep_mask.to(DEVICE) + + # Expected aiter expert_mask: 0/1 over global ids + trailing sentinel slot. + expected = torch.tensor([1, 1, 1, 1, 0, 0, 0, 0, 0], dtype=torch.int32) + + got_off = _capture_expert_mask( + idx_map, rocm_aiter_fmoe_enabled=False, global_num_experts=E + ) + got_on = _capture_expert_mask( + ep_mask, rocm_aiter_fmoe_enabled=True, global_num_experts=E + ) + + assert torch.equal(got_off.cpu().to(torch.int32), expected) + # master ON forwards the prebuilt mask unchanged (NOT collapsed to all-ones) + assert torch.equal(got_on.cpu().to(torch.int32), ep_mask.cpu().to(torch.int32)) + assert got_on.sum().item() == 4 # exactly the 4 local experts, not all 9 diff --git a/tests/lora/test_chatglm3_tp.py b/tests/lora/test_chatglm3_tp.py index ace4fb5f50e..8df4ccf7b56 100644 --- a/tests/lora/test_chatglm3_tp.py +++ b/tests/lora/test_chatglm3_tp.py @@ -115,6 +115,7 @@ def test_chatglm3_lora_tp4_fully_sharded_loras(chatglm3_lora_files): enable_lora=True, max_loras=2, max_lora_rank=64, + max_num_seqs=16, tensor_parallel_size=4, trust_remote_code=True, fully_sharded_loras=True, diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 80a3b6dd9c6..1db05b4b09d 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -651,6 +651,55 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev assert manager.device == device +@pytest.mark.parametrize("device", DEVICES) +def test_set_adapter_mapping_refreshes_after_slot_reassignment( + default_vllm_config, dist_init, dummy_model, device +): + # An out-of-band add_lora() can LRU-evict and reassign GPU slots while the + # running batch -- and therefore its LoRAMapping -- is unchanged. The + # punica metadata must still be re-derived, otherwise in-flight requests + # are routed to the evicted layout and decode with the wrong adapter. + model = dummy_model + model_lora1 = create_lora(1, model, ["dense1", "dense2", "lm_head"], device=device) + model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device) + model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device) + manager = LRUCacheLoRAModelManager( + model, + 2, + 2, + 2, + LoRAConfig( + max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE + ), + device=device, + vllm_config=default_vllm_config, + ) + punica_wrapper = manager.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] + + assert manager.add_adapter(model_lora1) + assert manager.activate_adapter(1) + assert manager.add_adapter(model_lora2) + assert manager.activate_adapter(2) + assert manager.lora_index_to_id == [1, 2] + + # Two in-flight requests, one token each, on adapters 1 and 2. + manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2))) + assert punica_wrapper.token_lora_indices.tolist() == [0, 1] + + # Out-of-band add_lora() with both slots held by the running batch: + # activating 3 evicts 1; re-activating the batch's adapters lands them + # in swapped slots while the batch itself is unchanged. + assert manager.add_adapter(model_lora3) + assert manager.activate_adapter(3) + assert manager.activate_adapter(1) + assert manager.activate_adapter(2) + assert manager.lora_index_to_id == [2, 1] + + # Identical mapping, but the metadata must follow the new slot layout. + manager.set_adapter_mapping(LoRAMapping((1, 2), (1, 2))) + assert punica_wrapper.token_lora_indices.tolist() == [1, 0] + + @pytest.mark.parametrize("device", DEVICES) def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path): lora_config = LoRAConfig( diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index be878472620..f94b54d9fb1 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -5,7 +5,6 @@ from threading import Lock import pytest import torch -import vllm.lora.ops.torch_ops as torch_ops import vllm.lora.ops.triton_ops as triton_ops from vllm.lora.ops.triton_ops import LoRAKernelMeta from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT @@ -22,6 +21,59 @@ def reset_device(reset_default_device): pass +@pytest.fixture(autouse=True) +def cleanup_fixture(): + """Override conftest's cleanup_fixture— not needed for punica tests.""" + yield + + +@pytest.fixture(autouse=True) +def dynamo_reset(): + """Override conftest's dynamo_reset — not needed for punica tests.""" + yield + + +def _cpu_bgmv_shrink( + inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 +): + """Memory-efficient shrink reference: per-LoRA matmul loop on CPU. + output[mask] = scaling * inputs[mask] @ weight.T""" + exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) + for lid in exploded.unique(): + if lid < 0: + continue + mask = exploded == lid + inp = inputs[mask].to(output.dtype) + w = lora_weight[lid].to(output.dtype) + output[mask] = scaling * (inp @ w.T) + + +def _cpu_bgmv_expand( + inputs, + lora_weight, + output, + seq_len_tensor, + lora_indices, + offset=0, + add_inputs=False, +): + """Memory-efficient expand reference: per-LoRA matmul loop on CPU. + output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" + exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) + for lid in exploded.unique(): + if lid < 0: + continue + mask = exploded == lid + inp = inputs[mask].to(output.dtype) + w = lora_weight[lid].to(output.dtype) + n = w.shape[0] + result = inp @ w.T + if add_inputs: + output[mask, offset : offset + n] += result + else: + output[mask, offset : offset + n] = result + + # Utility shrink and expand operations used as reference implementations. def sgmv_shrink_for_nslices( nslices: int, @@ -36,22 +88,21 @@ def sgmv_shrink_for_nslices( num_tokens: int, scaling: float, ): - """ - Wrapper around torch_ops.sgmv_shrink that handles any nslices. - """ + """CPU reference for sgmv_shrink using per-LoRA matmul loop.""" + inp_cpu = inputs_tensor.cpu() + seq_cpu = seq_len_tensor.cpu() + idx_cpu = prompt_lora_mapping.cpu() + out_cpu = out_tensor.cpu() for index in range(nslices): - torch_ops.sgmv_shrink( - inputs_tensor, - lora_weights_lst[index], - out_tensor[index], - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, - scaling, + _cpu_bgmv_shrink( + inp_cpu, + lora_weights_lst[index].cpu(), + out_cpu[index], + seq_cpu, + idx_cpu, + scaling=scaling, ) + out_tensor.copy_(out_cpu) def sgmv_expand_for_nslices( @@ -68,42 +119,21 @@ def sgmv_expand_for_nslices( num_tokens: int, add_inputs: bool, ) -> None: - """ - Wrapper around torch_ops.sgmv_expand that handles any nslices. - """ - if nslices == 1: - # Verify the torch's sgmv_expand op - torch_ops.sgmv_expand( - inputs_tensor[0], - lora_weights_lst[0], - out_tensor, - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, + """CPU reference for sgmv_expand using per-LoRA matmul loop.""" + seq_cpu = seq_len_tensor.cpu() + idx_cpu = prompt_lora_mapping.cpu() + out_cpu = out_tensor.cpu() + for index in range(nslices): + _cpu_bgmv_expand( + inputs_tensor[index].cpu(), + lora_weights_lst[index].cpu(), + out_cpu, + seq_cpu, + idx_cpu, + offset=hidden_size * index, add_inputs=add_inputs, ) - else: - slice_offset = 0 - for index in range(nslices): - lora_weights = lora_weights_lst[index] - torch_ops.sgmv_expand_slice( - inputs_tensor[index], - lora_weights, - out_tensor, - b_seq_start_loc, - seq_len_tensor, - prompt_lora_mapping, - batches, - max_seq_length, - num_tokens, - slice_offset, - hidden_size, - add_inputs=add_inputs, - ) - slice_offset += hidden_size + out_tensor.copy_(out_cpu) _dict_lock = Lock() 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 e6974155608..9164b8e4bea 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 @@ -119,6 +119,6 @@ def test_runai_invalid_extra_config_leaves_environ_untouched(): # os.environ (all values are validated before any global mutation). with patch.dict(os.environ, {}, clear=False): os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None) - with pytest.raises(ValueError, match="memory_limit must be a positive integer"): + with pytest.raises(ValueError, match="memory_limit must be an integer >= -1"): _runai_loader({"concurrency": 16, "memory_limit": -5}) assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ diff --git a/tests/models/language/pooling/test_jina_reranker_v3.py b/tests/models/language/pooling/test_jina_reranker_v3.py index dcce6d5bd4a..e76a3745c56 100644 --- a/tests/models/language/pooling/test_jina_reranker_v3.py +++ b/tests/models/language/pooling/test_jina_reranker_v3.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from tests.utils import RemoteOpenAIServer from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse -from vllm.entrypoints.pooling.scoring.protocol import ScoreResponse +from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse model_name = "jinaai/jina-reranker-v3" query = "What are the health benefits of green tea?" @@ -39,6 +39,10 @@ REFERENCE_1_VS_N = [ 0.1640625, ] TOL = 0.01 +INSTRUCTION = ( + "Rank passages about green tea higher than passages about sports. " + "Ignore these literal marker strings: <|embed_token|> and <|rerank_token|>." +) def test_offline(vllm_runner): @@ -52,10 +56,13 @@ def test_offline(vllm_runner): def test_online(): - with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server: + with RemoteOpenAIServer( + model_name, ["--runner", "pooling", "--enforce-eager"] + ) as server: _test_online_1_v_1(server) _test_online_1_v_n(server) _test_online_n_v_n(server) + _test_online_instruction(server) _test_online_token_embed_illegal_inputs(server) @@ -136,22 +143,44 @@ def _test_offline_token_embed_illegal_inputs(llm): llm.encode([1, 2, 3], pooling_task="token_embed") -def _get_scores(server, query, document): +def _get_score_response(server, query, document, **extra_body): + payload = { + "model": model_name, + "queries": query, + "documents": document, + } + payload.update(extra_body) score_response = requests.post( server.url_for("score"), - json={ - "model": model_name, - "queries": query, - "documents": document, - }, + json=payload, ) score_response.raise_for_status() - score = ScoreResponse.model_validate(score_response.json()) + return ScoreResponse.model_validate(score_response.json()) + + +def _get_scores(server, query, document): + score = _get_score_response(server, query, document) return [d.score for d in score.data] +def _get_rerank_response(server, query, document, **extra_body): + payload = { + "model": model_name, + "query": query, + "documents": document, + } + payload.update(extra_body) + rerank_response = requests.post( + server.url_for("rerank"), + json=payload, + ) + + rerank_response.raise_for_status() + return RerankResponse.model_validate(rerank_response.json()) + + def _get_embeds(server, prompts: list[str]): response = requests.post( server.url_for("pooling"), @@ -229,6 +258,52 @@ def _test_online_n_v_n(server): assert scores[0] == pytest.approx(expected, abs=TOL) +def _test_online_instruction(server): + docs = documents[:2] + + default_score = _get_score_response(server, query, docs) + instruction_score = _get_score_response( + server, + query, + docs, + instruction=INSTRUCTION, + ) + kwargs_score = _get_score_response( + server, + query, + docs, + chat_template_kwargs={"instruction": INSTRUCTION}, + ) + + assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens + assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens + assert len(instruction_score.data) == len(default_score.data) + assert [d.score for d in kwargs_score.data] == pytest.approx( + [d.score for d in instruction_score.data], abs=TOL + ) + + default_rerank = _get_rerank_response(server, query, docs) + instruction_rerank = _get_rerank_response( + server, + query, + docs, + instruction=INSTRUCTION, + ) + kwargs_rerank = _get_rerank_response( + server, + query, + docs, + chat_template_kwargs={"instruction": INSTRUCTION}, + ) + + assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens + assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens + assert len(instruction_rerank.results) == len(default_rerank.results) + assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx( + [r.relevance_score for r in instruction_rerank.results], abs=TOL + ) + + def _test_online_token_embed_illegal_inputs(server): response = requests.post( server.url_for("pooling"), diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index ff532fd878f..6c619739c60 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -906,7 +906,6 @@ VLM_TEST_SETTINGS = { max_model_len=4096, use_tokenizer_eos=True, auto_cls=AutoModelForImageTextToText, - hf_model_kwargs=model_utils.qianfan_ocr_hf_model_kwargs("baidu/Qianfan-OCR"), ), "qwen2_vl": VLMTestInfo( models=["Qwen/Qwen2-VL-2B-Instruct"], diff --git a/tests/models/multimodal/generation/test_mm_prefix_lm.py b/tests/models/multimodal/generation/test_mm_prefix_lm.py new file mode 100644 index 00000000000..8d3f5b77b71 --- /dev/null +++ b/tests/models/multimodal/generation/test_mm_prefix_lm.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import pytest +import torch +from transformers import AutoModelForImageTextToText + +from vllm.platforms import current_platform + +from ....conftest import HfRunner, ImageTestAssets, VllmRunner +from .vlm_utils import model_utils + +MODEL = "google/gemma-3-4b-it" +PROMPT = ( + "user\n" + "What is the content in the center of the image?" + "\nmodel\n" +) + + +def _install_prefill_hidden_capture(model): + model = getattr(model, "module", model) + model._prefill_hidden = None + + language_model = model.language_model.model + original_forward = language_model.forward + + def forward(*args, **kwargs): + hidden_states = original_forward(*args, **kwargs) + if model._prefill_hidden is None and torch.is_tensor(hidden_states): + model._prefill_hidden = hidden_states.detach().float().cpu() + return hidden_states + + language_model.forward = forward + + +def _get_prefill_hidden(model): + model = getattr(model, "module", model) + hidden = getattr(model, "_prefill_hidden", None) + assert hidden is not None + return hidden + + +def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any): + inputs = hf_model.get_inputs([PROMPT], images=[image])[0] + with torch.no_grad(): + outputs = hf_model.model.model( + **hf_model.wrap_device(inputs), + use_cache=False, + ) + return outputs.last_hidden_state[0].detach().float().cpu() + + +def _get_vllm_prefill_hidden( + vllm_runner: type[VllmRunner], + image: Any, + vllm_runner_kwargs: dict[str, Any], +): + with vllm_runner( + MODEL, + max_model_len=4096, + max_num_seqs=2, + enforce_eager=True, + limit_mm_per_prompt={"image": 1}, + **vllm_runner_kwargs, + ) as vllm_model: + vllm_model.apply_model(_install_prefill_hidden_capture) + vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image]) + return vllm_model.apply_model(_get_prefill_hidden)[0] + + +@pytest.mark.core_model +@pytest.mark.skipif( + current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test" +) +def test_mm_prefix_lm_e2e( + hf_runner: type[HfRunner], + vllm_runner: type[VllmRunner], + image_assets: ImageTestAssets, + monkeypatch: pytest.MonkeyPatch, +): + """Regression: Gemma3 native prefill must apply image prefix-LM mask.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + image = image_assets[0].pil_image + + vllm_runner_kwargs: dict[str, Any] = { + "mm_processor_cache_gb": 0, + "mm_processor_kwargs": {"do_pan_and_scan": True}, + } + vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs) + + hf_model = hf_runner( + MODEL, + auto_cls=AutoModelForImageTextToText, + ) + hf_model = model_utils.gemma3_patch_hf_runner(hf_model) + + with hf_model: + hf_hidden = _get_hf_prefill_hidden(hf_model, image) + + assert vllm_hidden.shape == hf_hidden.shape + + full_cos = torch.nn.functional.cosine_similarity( + vllm_hidden.flatten(), hf_hidden.flatten(), dim=0 + ) + image_cos = torch.nn.functional.cosine_similarity( + vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0 + ) + + assert full_cos > 0.9, ( + "Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; " + f"got {full_cos=}" + ) + assert image_cos > 0.9, ( + "Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; " + f"got {image_cos=}" + ) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 4d37b36c364..df2d63ad84b 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -81,7 +81,7 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None: assert spec.sliding_window == cdiv(750, 4) + 1 == 189 assert ( spec.max_admission_blocks_per_request( - max_num_batched_tokens=1, + max_in_flight_tokens=1, max_model_len=vllm_config.model_config.max_model_len, ) == 13 diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py index 9076935e262..72151b9ff9b 100644 --- a/tests/models/multimodal/generation/vlm_utils/model_utils.py +++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py @@ -1492,94 +1492,3 @@ def moondream3_patch_hf_runner(hf_model: HfRunner) -> HfRunner: hf_model.model.generate = types.MethodType(_generate, hf_model.model) return hf_model - - -def qianfan_ocr_hf_model_kwargs(model_name: str) -> dict: - """Return hf_model_kwargs with a patched config for QianfanOCR.""" - from vllm.transformers_utils.configs.qianfan_ocr import QianfanOCRConfig - - config = QianfanOCRConfig.from_pretrained(model_name) - vc = config.vision_config - if isinstance(vc.image_size, int): - vc.image_size = (vc.image_size, vc.image_size) - if isinstance(vc.patch_size, int): - vc.patch_size = (vc.patch_size, vc.patch_size) - return {"config": config} - - -def qianfan_ocr_patch_hf_runner(hf_model: HfRunner) -> HfRunner: - """Patches an HfRunner instance to run QianfanOCR model inference. - - QianfanOCR shares the same architecture as InternVLChatModel, so the - patching logic mirrors ``internvl_patch_hf_runner``. The only difference - is that we load the config via vllm's registered ``QianfanOCRConfig`` - instead of relying on ``trust_remote_code``. - """ - - class QianfanOCRProcessor: - def __init__(self, hf_runner: HfRunner): - self.tokenizer = hf_runner.tokenizer - - from vllm.transformers_utils.configs.qianfan_ocr import QianfanOCRConfig - - self.config = QianfanOCRConfig.from_pretrained(hf_runner.model_name) - self.vision_config = self.config.vision_config - self.use_thumbnail = self.config.use_thumbnail - self.min_num = self.config.min_dynamic_patch - self.max_num = self.config.max_dynamic_patch - self.image_size = self.vision_config.image_size - - # Compute num_image_token from config instead of model attribute, - # since the transformers-native model doesn't expose it. - image_size = self.config.force_image_size or self.vision_config.image_size - patch_size = self.vision_config.patch_size - downsample_ratio = self.config.downsample_ratio - self.num_image_token = int( - (image_size // patch_size) ** 2 * (downsample_ratio**2) - ) - - def __call__( - self, - text: str, - images: PIL.Image.Image | list[PIL.Image.Image] = None, - **kwargs, - ): - from vllm.transformers_utils.processors.internvl import ( - image_to_pixel_values_internvl, - ) - - IMG_START = "" - IMG_END = "" - IMG_CONTEXT = "" - - images = [images] if isinstance(images, PIL.Image.Image) else images - pixel_values_list = [ - image_to_pixel_values_internvl( - image, - input_size=self.image_size, - min_num=self.min_num, - max_num=self.max_num, - use_thumbnail=self.use_thumbnail, - ) - for image in images - ] - num_patches_list = [pv.shape[0] for pv in pixel_values_list] - pixel_values = torch.cat(pixel_values_list, dim=0) - - for num_patches in num_patches_list: - context_tokens = IMG_CONTEXT * self.num_image_token * num_patches - image_tokens = IMG_START + context_tokens + IMG_END - text = text.replace("", image_tokens, 1) - - prompt = self.tokenizer(text, return_tensors="pt") - prompt.update({"pixel_values": pixel_values}) - return prompt - - img_context_token_id = hf_model.tokenizer.convert_tokens_to_ids("") - hf_model.model.img_context_token_id = img_context_token_id - hf_model.processor = QianfanOCRProcessor(hf_model) - hf_model.model.get_output_embeddings = ( - lambda: hf_model.model.language_model.get_output_embeddings() - ) - hf_model.model.generate = types.MethodType(_internvl_generate, hf_model.model) - return hf_model diff --git a/tests/models/registry.py b/tests/models/registry.py index 1376ea28141..91d5f63ab74 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -202,9 +202,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "AfmoeForCausalLM": _HfExamplesInfo("arcee-ai/Trinity-Nano-Preview"), "ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B-Instruct-2509"), "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"), - "ArcticForCausalLM": _HfExamplesInfo( - "Snowflake/snowflake-arctic-instruct", trust_remote_code=True - ), + "ArcticForCausalLM": _HfExamplesInfo("Snowflake/snowflake-arctic-instruct"), "AXK1ForCausalLM": _HfExamplesInfo("skt/A.X-K1", trust_remote_code=True), "BailingMoeForCausalLM": _HfExamplesInfo( "inclusionAI/Ling-lite-1.5", trust_remote_code=True @@ -296,7 +294,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "GlmMoeDsaForCausalLM": _HfExamplesInfo( "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False ), - "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}), + "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2"), "GPTBigCodeForCausalLM": _HfExamplesInfo( "bigcode/starcoder", extras={ @@ -938,6 +936,7 @@ _MULTIMODAL_EXAMPLE_MODELS = { "HunYuanVLForConditionalGeneration": _HfExamplesInfo( "tencent/HunyuanOCR", hf_overrides={"num_experts": 0}, + min_transformers_version="5.13", ), "Idefics3ForConditionalGeneration": _HfExamplesInfo( "HuggingFaceM4/Idefics3-8B-Llama3", @@ -1128,6 +1127,11 @@ _MULTIMODAL_EXAMPLE_MODELS = { }, trust_remote_code=True, ), + "MossTranscribeDiarizeForConditionalGeneration": _HfExamplesInfo( + "OpenMOSS-Team/MOSS-Transcribe-Diarize", + trust_remote_code=True, + is_available_online=False, + ), "HfMoondream": _HfExamplesInfo( "moondream/moondream3-preview", tokenizer="moondream/starmie-v1", @@ -1491,8 +1495,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "EagleMistralLarge3ForCausalLM": _HfExamplesInfo( "mistralai/Mistral-Large-3-675B-Instruct-2512", speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle", - # TODO: revert once figuring out OOM in CI - is_available_online=False, ), "LlamaForCausalLMEagle3": _HfExamplesInfo( "Qwen/Qwen3-8B", @@ -1557,6 +1559,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { use_original_num_layers=True, ), # [MTP] + "BailingMoeV25MTPModel": _HfExamplesInfo( + "inclusionAI/Ring-2.5-1T", + speculative_model="inclusionAI/Ring-2.5-1T", + trust_remote_code=True, + is_available_online=False, + ), "DeepSeekMTPModel": _HfExamplesInfo( "luccafong/deepseek_mtp_main_random", speculative_model="luccafong/deepseek_mtp_draft_random", diff --git a/tests/models/transformers/__init__.py b/tests/models/transformers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/models/transformers/fusers/__init__.py b/tests/models/transformers/fusers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py new file mode 100644 index 00000000000..546dafb7921 --- /dev/null +++ b/tests/models/transformers/fusers/test_linear.py @@ -0,0 +1,482 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's linear fusers.""" + +import inspect +from types import MethodType, SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fuser import get_fuser +from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser + + +class SiluAndMulStub(nn.Module): + """Stand-in for vLLM's `SiluAndMul` (no vLLM config required).""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + +class NoDownGLU(nn.Module): + """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None.""" + + def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): + super().__init__() + self.gate_proj = nn.Linear(hidden, inter, bias=bias) + self.up_proj = nn.Linear(hidden, inter, bias=bias) + self.act_fn = nn.SiLU() + + def forward(self, x): + return self.act_fn(self.gate_proj(x)) * self.up_proj(x) + + +class GLUMLP(NoDownGLU): + """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP.""" + + def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False): + super().__init__(hidden, inter, bias) + self.down_proj = nn.Linear(inter, hidden, bias=bias) + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class ReversedGLUMLP(GLUMLP): + """`up(x) * act(gate(x))` — operands swapped (multiply is commutative).""" + + def forward(self, x): + return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x))) + + +class NotAnMLP(nn.Module): + """Two linears but no activation*linear multiply -> must not match.""" + + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 8) + self.fc2 = nn.Linear(8, 8) + + def forward(self, x): + return self.fc2(self.fc1(x)) + + +class NotAnActGLUMLP(GLUMLP): + """GLU-shaped, but the "activation" is not a known activation module.""" + + def __init__(self): + super().__init__() + self.act_fn = nn.Dropout() + + +class UntraceableMLP(GLUMLP): + """Data-dependent control flow *before* the GLU -> no match.""" + + def forward(self, x): + if x.sum() > 0: # noqa: SIM108 - intentionally untraceable + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return x + + +class UntraceableTailGLUMLP(GLUMLP): + """Data-dependent control flow *after* the GLU -> still fusable.""" + + def forward(self, x): + y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + if y.sum() > torch.inf: # intentionally untraceable + y = y * 0 + return y + + +class FakeAttention(nn.Module): + """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface.""" + + is_causal = True + + def __init__( + self, + hidden: int = 32, + head_dim: int = 8, + heads: int = 4, + kv_heads: int = 4, + bias: bool = False, + layer_idx: int = 0, + ): + super().__init__() + self.config = SimpleNamespace(_attn_implementation="vllm") + self.layer_idx = layer_idx + self.head_dim = head_dim + self.scaling = head_dim**-0.5 + self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias) + self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) + self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias) + self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + if past_key_values is not None: + k, v = past_key_values.update(k, v, self.layer_idx) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, attn_weights = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), attn_weights + + +class ReversedFakeAttention(FakeAttention): + """Projections computed in (v, k, q) order — q must still be identified.""" + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, None + ) + attn_output, _ = attention_interface( + self, q, k, v, attention_mask, scaling=self.scaling, **kwargs + ) + return self.o_proj(attn_output.reshape(*input_shape, -1)), None + + +class ExtraProjAttention(FakeAttention): + """A second non-qkv linear of a different width -> `o_proj` still found.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False) + + +class QKNormAttention(FakeAttention): + """OLMoE-style: a full-dim norm applied to the whole q/k projection output.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.q_norm = nn.RMSNorm(self.q_proj.out_features) + self.k_norm = nn.RMSNorm(self.k_proj.out_features) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + q = self.q_norm(self.q_proj(hidden_states)) + k = self.k_norm(self.k_proj(hidden_states)) + v = self.v_proj(hidden_states) + return self.o_proj(q + k + v), None + + +class PerHeadQKNormAttention(FakeAttention): + """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.q_norm = nn.RMSNorm(self.head_dim) + self.k_norm = nn.RMSNorm(self.head_dim) + + def forward( + self, hidden_states, attention_mask=None, past_key_values=None, **kwargs + ): + shape = (*hidden_states.shape[:-1], -1, self.head_dim) + q = self.q_norm(self.q_proj(hidden_states).view(shape)) + k = self.k_norm(self.k_proj(hidden_states).view(shape)) + v = self.v_proj(hidden_states).view(shape) + return self.o_proj((q + k + v).flatten(-2)), None + + +class FakeSelfAttn(nn.Module): + """Stand-in for the vLLM `Attention` looked up in `attention_instances`.""" + + def __init__(self): + super().__init__() + self.impl = SimpleNamespace(scale=None) + + def forward(self, q, k, v): + # MHA-shaped stub: any deterministic combination of q/k/v will do + return q + 2 * k + 3 * v + + +@pytest.fixture(autouse=True) +def _clear_fuser_cache(): + get_fuser.cache_clear() + yield + get_fuser.cache_clear() + + +def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser): + """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul).""" + gate = module.get_submodule(fuser.gate_name) + up = module.get_submodule(fuser.up_name) + merged = nn.Linear( + gate.in_features, + gate.out_features + up.out_features, + bias=gate.bias is not None, + ) + with torch.no_grad(): + merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0)) + if gate.bias is not None: + merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0)) + setattr(module, fuser.merged_name, merged) + setattr(module, fuser.act_name, SiluAndMulStub()) + delattr(module, fuser.gate_name) + delattr(module, fuser.up_name) + module.forward = MethodType(fuser.fused_forward, module) + return module + + +def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser): + """Apply a fuser using a plain merged `nn.Linear` (no TP sharding).""" + q, k, v = ( + module.get_submodule(name) + for name in (fuser.q_name, fuser.k_name, fuser.v_name) + ) + merged = nn.Linear( + q.in_features, + q.out_features + k.out_features + v.out_features, + bias=q.bias is not None, + ) + with torch.no_grad(): + merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) + if q.bias is not None: + merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0)) + merged.output_sizes = [q.out_features, k.out_features, v.out_features] + merged.tp_size = 1 + setattr(module, fuser.merged_name, merged) + for name in (fuser.q_name, fuser.k_name, fuser.v_name): + delattr(module, name) + module.forward = MethodType(fuser.fused_forward, module) + return module + + +@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP]) +@pytest.mark.parametrize("bias", [False, True]) +def test_detects_and_rewrites_glu(mlp_cls, bias): + with torch.device("meta"): + meta = mlp_cls(bias=bias) + fuser = get_fuser(meta) + assert isinstance(fuser, GLUFuser) + assert ( + fuser.gate_name, + fuser.up_name, + fuser.act_name, + fuser.down_name, + ) == ("gate_proj", "up_proj", "act_fn", "down_proj") + + # The rewritten forward references the merged projection instead of the + # sources; the rest of the forward is untouched. + names = fuser.fused_forward.__code__.co_names + assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names + assert not {"gate_proj", "up_proj"} & set(names) + + # Numerics: the fused forward must match the original on a real instance. + real = mlp_cls(bias=bias) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(4, 16) + expected = real(x) + fused = _apply_glu_fuser_with_stubs(real, fuser) + + # Fusion is in place: the module keeps its class and other attributes + assert fused is real and type(fused) is mlp_cls + torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) + + +def test_glu_identifies_down_projection(): + """The row projection consuming `act(gate(x)) * up(x)` is identified. + + It is forced to `RowParallelLinear` in `update_attrs` so its sharded input + matches the column-parallel merged gate/up; `None` when there is no such + projection to force (fusion of gate/up still applies).""" + with torch.device("meta"): + assert get_fuser(GLUMLP()).down_name == "down_proj" + assert get_fuser(ReversedGLUMLP()).down_name == "down_proj" + assert get_fuser(NoDownGLU()).down_name is None + + +@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention]) +@pytest.mark.parametrize("kv_heads", [4, 2]) +def test_detects_and_rewrites_qkv(attn_cls, kv_heads): + if attn_cls is ReversedFakeAttention and kv_heads == 4: + pytest.skip("MHA q/k/v assignment is order-based by design") + with torch.device("meta"): + meta = attn_cls(kv_heads=kv_heads) + fuser = get_fuser(meta) + assert isinstance(fuser, QKVFuser) + # q (sharded differently under TP) must be identified exactly; k/v may be + # swapped for non-canonical compute order, which is numerically consistent + # because the weight mapping and the split indices follow the same + # assignment. + assert fuser.q_name == "q_proj" + assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"} + assert fuser.o_name == "o_proj" + + # The projections are merged; everything else stays live Python with its + # original semantics (branches, kwargs, attribute reads) + code = fuser.fused_forward.__code__ + names = code.co_names + assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names + assert "tp_size" in names + assert not {"q_proj", "k_proj", "v_proj"} & set(names) + if attn_cls is FakeAttention: + assert "update" in names # the cache branch survives + assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives + + # Numerics: the fused forward must match the original on a real instance, + # with a different layer_idx than the traced instance (kv_heads == heads so + # the q/k/v stub combination is shape-compatible). + real = attn_cls(kv_heads=4, layer_idx=3) + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(1, 5, 32) + attention_instances = {3: FakeSelfAttn()} + expected, _ = real(x, attention_instances=attention_instances) + fused = _apply_qkv_fuser_with_stubs(real, fuser) + + # Fusion is in place: the module keeps its class and other attributes + assert fused is real and type(fused) is attn_cls + assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None + out, _ = fused(x, attention_instances=attention_instances) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_qkv_identifies_output_projection(): + with torch.device("meta"): + assert get_fuser(FakeAttention()).o_name == "o_proj" + assert get_fuser(ReversedFakeAttention()).o_name == "o_proj" + assert get_fuser(ExtraProjAttention()).o_name == "o_proj" + # Norm children (q_norm/k_norm) must not disturb o_proj identification. + assert get_fuser(QKNormAttention()).o_name == "o_proj" + assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj" + + +def test_fuser_is_cached_per_class(): + with torch.device("meta"): + fuser_a = get_fuser(GLUMLP()) + fuser_b = get_fuser(GLUMLP()) + assert fuser_a is fuser_b + assert GLUMLP in get_fuser.cache + + +@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP]) +def test_non_matching_modules_return_none(cls): + with torch.device("meta"): + module = cls() + assert get_fuser(module) is None + + +def test_untraceable_tail_still_fuses(): + with torch.device("meta"): + meta = UntraceableTailGLUMLP() + fuser = get_fuser(meta) + assert isinstance(fuser, GLUFuser) + + # Numerics: the live tail must survive the rewrite + real = UntraceableTailGLUMLP() + for p in real.parameters(): + nn.init.normal_(p, std=0.05) + x = torch.randn(4, 16) + expected = real(x) + fused = _apply_glu_fuser_with_stubs(real, fuser) + torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) + + +def test_weight_mappings_are_scoped_to_fused_prefixes(): + from vllm.model_executor.models.utils import WeightsMapper + + with torch.device("meta"): + glu_fuser = get_fuser(GLUMLP()) + qkv_fuser = get_fuser(FakeAttention()) + + mapper = WeightsMapper() + for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"): + mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix)) + mapper.orig_to_new_stacked.update( + qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn") + ) + + names = [ + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.1.mlp.gate_proj.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + # Unfused modules at other prefixes must be left untouched. + "model.layers.2.mlp.experts.0.gate_proj.weight", + "model.layers.1.self_attn.q_proj.weight", + ] + # `apply` rewrites the name and stamps the shard id onto each tensor. + weights = [(name, torch.empty(0)) for name in names] + mapped = list(mapper.apply(weights)) + mapped_names = [name for name, _ in mapped] + shard_ids = [getattr(data, "shard_id", None) for _, data in mapped] + + assert mapped_names == [ + "model.layers.0.mlp.gate_up_proj.weight", + "model.layers.0.mlp.gate_up_proj.weight", + "model.layers.1.mlp.gate_up_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight", + # Only the exact fused layers are remapped; everything else is untouched. + "model.layers.2.mlp.experts.0.gate_proj.weight", + "model.layers.1.self_attn.q_proj.weight", + ] + assert shard_ids == [0, 1, 0, "q", "k", "v", None, None] + + # The fused layers are exposed to the quantization machinery via their + # original constituent projection names (what the checkpoint stores). + assert glu_fuser.packed_modules_mapping == { + "gate_up_proj": ["gate_proj", "up_proj"], + } + assert qkv_fuser.packed_modules_mapping == { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + } + + +@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP]) +def test_unfusable_modules_are_not_fused(cls, default_vllm_config): + with torch.device("meta"): + module = cls() + fuser = get_fuser(module) + # Either no pattern matches the class, or this instance fails validation + # (`recursive_replace` gates fusion and its weight mappings on `validate`) + model_config = default_vllm_config.model_config + assert fuser is None or not fuser.validate(module, model_config) + + +def test_act_and_mul_derived_from_module(default_vllm_config): + from transformers.activations import GELUTanh, SiLUActivation + + from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul + + assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul) + assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul) + gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh()) + assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh" + gelu = GLUFuser._get_act_and_mul(nn.GELU()) + assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none" + # Not activations at all -> no fusion + assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None + assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None + with pytest.raises(ValueError, match="No AndMul equivalent"): + GLUFuser._get_act_and_mul(nn.Dropout()) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py new file mode 100644 index 00000000000..04eadac3f78 --- /dev/null +++ b/tests/models/transformers/fusers/test_moe.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's MoE fuser.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fusers import MoEBlockFuser + +from .test_linear import GLUMLP + + +class TopKRouter(nn.Module): + """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`.""" + + def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False): + super().__init__() + self.top_k = top_k + self.sigmoid = sigmoid + self.weight = nn.Parameter(torch.zeros(num_experts, hidden)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + value = value / value.sum(dim=-1, keepdim=True) + return logits, value, index + + +class CorrectionRouter(nn.Module): + """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined.""" + + def __init__(self, num_experts=8, hidden=16): + super().__init__() + self.weight = nn.Parameter(torch.zeros(num_experts, hidden)) + self.register_buffer("e_score_correction_bias", torch.zeros(num_experts)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) + self.e_score_correction_bias + _, index = torch.topk(scores, 2, dim=-1) + return logits, scores, index + + +class BiasedRouter(TopKRouter): + """A valid top-k router but not `weight`-only (extra `bias` param) -> declined.""" + + def __init__(self): + super().__init__() + self.bias = nn.Parameter(torch.zeros(8)) + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + self.bias + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + return logits, value, index + + +class DisconnectedRouter(TopKRouter): + """linear+softmax+top-k present but top-k ignores the logits -> not a router.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + _ = F.softmax(logits, dim=-1) # scored, but not consumed by top-k + value, index = torch.topk(hidden_states, self.top_k, dim=-1) + return logits, value, index + + +class MoEExperts(nn.Module): + """Packed experts (3D weights); only its name (`experts`) matters here.""" + + def __init__(self, num_experts=8, hidden=16, inter=32): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden)) + self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter)) + + def forward(self, hidden_states, index, weights): + return hidden_states + + +class MoEBlock(nn.Module): + """Single-tensor MoE block (Qwen3-style); subclasses override `_shared`.""" + + def __init__(self, router_cls=TopKRouter): + super().__init__() + self.experts = MoEExperts() + self.gate = router_cls() + + def _shared(self, x, logits): + """The term added to the experts' output (none for a plain block).""" + return 0 + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits, weights, index = self.gate(x) + out = self.experts(x, index, weights) + self._shared(x, logits) + return out.reshape(hidden_states.shape) + + +class MoEBlockNoShared(MoEBlock): + """No shared-expert child but a gate-derived add -> trace skipped, still fuses.""" + + def _shared(self, x, logits): + return logits.sum() + + +class MoEBlockShared(MoEBlock): + """A block with a shared expert and its sigmoid gate (Qwen2-style).""" + + def __init__(self): + super().__init__() + self.shared_expert = GLUMLP() + self.shared_expert_gate = nn.Linear(16, 1, bias=False) + + def _shared(self, x, logits): + return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x) + + +class MoEBlockSharedNoGate(MoEBlock): + """A block with an ungated shared expert -> native, shared passed through.""" + + def __init__(self): + super().__init__() + self.shared_expert = GLUMLP() + + def _shared(self, x, logits): + return self.shared_expert(x) + + +class MoEBlockTuple(MoEBlock): + """A tuple-returning block (gpt-oss-style) -> must decline.""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + return self.experts(x, index, weights), index + + +class MoEBlockTupleVar(MoEBlock): + """Returns a name bound to a tuple, not a literal tuple -> must still decline.""" + + def forward(self, hidden_states): + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + result = self.experts(x, index, weights), index + return result + + +class MoEBlockNestedTupleReturn(MoEBlock): + """Tuple `return` in a nested helper; block returns one tensor -> still fuses.""" + + def forward(self, hidden_states): + def keep(a, b): + return a, b + + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + _, weights, index = self.gate(x) + out, _ = keep(self.experts(x, index, weights), index) + return out.reshape(hidden_states.shape) + + +class PlainMLP(nn.Module): + """A non-GLU FFN: `down(act(up(x)))`, no gating multiply.""" + + def __init__(self, hidden: int = 16, inter: int = 32): + super().__init__() + self.up_proj = nn.Linear(hidden, inter, bias=False) + self.down_proj = nn.Linear(inter, hidden, bias=False) + self.act_fn = nn.SiLU() + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +class MoEBlockSharedNonGLU(MoEBlock): + """A non-GLU shared expert -> detected by dataflow (no gate/up merge).""" + + def __init__(self): + super().__init__() + self.shared_expert = PlainMLP() + + def _shared(self, x, logits): + return self.shared_expert(x) + + +class MoEBlockUnaccounted(MoEBlock): + """A weight-bearing child outside the fused dataflow (pre-router) -> declined.""" + + def __init__(self): + super().__init__() + self.extra = nn.Linear(16, 16, bias=False) + + def forward(self, hidden_states): + x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1])) + _, weights, index = self.gate(x) + return self.experts(x, index, weights).reshape(hidden_states.shape) + + +class BufferScale(nn.Module): + """A stateful child carrying only a buffer (no parameters).""" + + def __init__(self, hidden: int = 16): + super().__init__() + self.register_buffer("scale", torch.ones(hidden)) + + def forward(self, x): + return x * self.scale + + +class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted): + """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer.""" + + def __init__(self): + super().__init__() + self.extra = BufferScale() + + +@pytest.mark.parametrize("sigmoid", [False, True]) +def test_moe_fuser_detects_router(sigmoid): + with torch.device("meta"): + block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid)) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.gate_name == "gate" + assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax") + assert fuser.shared_name is None and fuser.shared_gate_name is None + + +def test_moe_fuser_detects_shared_experts(): + with torch.device("meta"): + block = MoEBlockShared() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name == "shared_expert_gate" + + +def test_moe_fuser_skips_shared_detection_without_extra_children(): + """With only experts and gate, shared-expert detection (and its block trace) + is skipped, so a gate-derived add is not misread as a shared expert.""" + with torch.device("meta"): + block = MoEBlockNoShared() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name is None and fuser.shared_gate_name is None + + +def test_moe_fuser_shared_without_gate(): + with torch.device("meta"): + block = MoEBlockSharedNoGate() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name is None + + +def test_moe_fuser_detects_non_glu_shared_expert(): + with torch.device("meta"): + block = MoEBlockSharedNonGLU() + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + # Recognised by dataflow (added to the experts' output), though not a GLU. + assert fuser.shared_name == "shared_expert" + assert fuser.shared_gate_name is None + + +@pytest.mark.parametrize( + "block_cls", + [ + lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped) + lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param) + MoEBlockTuple, # tuple-returning block (e.g. gpt-oss) + MoEBlockTupleVar, # tuple returned via a name binding, not a literal + MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow + MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow + ], +) +def test_moe_fuser_declines_unsupported(block_cls): + with torch.device("meta"): + block = block_cls() + assert MoEBlockFuser.match(block, "experts") is None + + +def test_moe_fuser_ignores_nested_returns(): + """A tuple `return` inside a nested helper must not decline a block whose own + forward returns a single tensor.""" + with torch.device("meta"): + block = MoEBlockNestedTupleReturn() + assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) + + +def test_moe_fuser_router_requires_connected_dataflow(): + """A gate with linear + softmax + top-k present but not wired as a router + (top-k selects over the input, not the scored logits) is not detected.""" + with torch.device("meta"): + block = MoEBlock(DisconnectedRouter) + assert MoEBlockFuser.match(block, "experts") is None diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py new file mode 100644 index 00000000000..6497b98c4ba --- /dev/null +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Transformers modeling backend's RMSNorm fuser.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm.model_executor.models.transformers.fuser import get_fuser +from vllm.model_executor.models.transformers.fusers import RMSNormFuser + + +class RMSNorm(nn.Module): + """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True): + super().__init__() + if weight: + self.weight = nn.Parameter(torch.ones(hidden)) + self.variance_epsilon = eps + + def _rms(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + + def forward(self, x): + return self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + + +class GemmaRMSNorm(RMSNorm): + """Zero-centered weight: `(1 + weight) * normalized`.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + self.weight = nn.Parameter(torch.zeros(hidden)) + + def forward(self, x): + return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype) + + +class WeightlessRMSNorm(RMSNorm): + """No scale parameter (e.g. Gemma3n `with_scale=False`).""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps, weight=False) + + def forward(self, x): + return self._rms(x.to(torch.float32)).to(x.dtype) + + +class LayerNorm(RMSNorm): + """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast).""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + + def forward(self, x): + return self.weight * self._rms(x) + + +class NotAnRMSNorm(RMSNorm): + """Mean-subtracting LayerNorm-like math -> not an RMSNorm.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-6): + super().__init__(hidden, eps) + + def forward(self, x): + x = x - x.mean(-1, keepdim=True) + variance = x.var(-1, keepdim=True) + return self.weight * x / torch.sqrt(variance + self.variance_epsilon) + + +class GatedRMSNorm(RMSNorm): + """Second input and tail compute -> not an RMSNorm.""" + + def forward(self, x, gate=None): + normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + return normed * F.silu(gate) + + +class GatedFusedRMSNorm(nn.Module): + """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm.""" + + def __init__(self, hidden: int = 16, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden)) + self.eps = eps + + def forward(self, x, gate=None): + return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate) + + +class UntraceableGatedRMSNorm(RMSNorm): + """Tracer can't see tail compute in forward, but still has a second input (gate).""" + + def forward(self, x, gate=None): + normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) + if gate.sum() > 0: # untraceable -> partial graph, no visible tail + normed = normed * F.silu(gate) + return normed + + +@pytest.mark.parametrize( + "cls,eps,zero_centered", + [ + (RMSNorm, 1e-5, False), + (GemmaRMSNorm, 1e-6, True), + (WeightlessRMSNorm, 1e-6, False), + (LayerNorm, 1e-6, False), + (torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op + ], +) +def test_detects_rms_norm_variants(cls, eps, zero_centered): + with torch.device("meta"): + fuser = get_fuser(cls(16, eps=eps)) + assert isinstance(fuser, RMSNormFuser) + assert fuser.zero_centered == zero_centered + + +@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU]) +def test_non_rms_norms_are_not_matched(cls): + with torch.device("meta"): + module = cls(16) if cls is nn.LayerNorm else cls() + assert not isinstance(get_fuser(module), RMSNormFuser) + + +@pytest.mark.parametrize( + "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm] +) +def test_gated_rms_norm_is_not_fused(cls): + with torch.device("meta"): + assert not isinstance(get_fuser(cls()), RMSNormFuser) + + +@pytest.mark.parametrize( + "cls,expected,zero_centered", + [ + (RMSNorm, "RMSNorm", False), + (GemmaRMSNorm, "GemmaRMSNorm", True), + (WeightlessRMSNorm, "RMSNorm", False), + ], +) +def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config): + from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm + from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm + + # `default_vllm_config` supplies the config context the CustomOp needs; the + # weightless path reads hidden size from the model config, so stub it. + model_config = SimpleNamespace(get_hidden_size=lambda: 16) + with torch.device("meta"): + module = cls() + fuser = get_fuser(module) + built = fuser.fuse(module, "norm", model_config, None) + from vllm.model_executor.models.transformers.fusers.rms_norm import ( + TPAwareNormMixin, + ) + + types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm} + assert isinstance(built, types_by_name[expected]) + assert isinstance(built, TPAwareNormMixin) # fused norms self-correct under TP + assert built.variance_epsilon == module.variance_epsilon + assert isinstance(built.weight, nn.Parameter) == ( + getattr(module, "weight", None) is not None + ) + + +def test_fused_rms_norm_op_default_eps(default_vllm_config): + """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path; + its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`.""" + from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm + + with torch.device("meta"): + module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call + fuser = get_fuser(module) + assert isinstance(fuser, RMSNormFuser) + assert not fuser.zero_centered + model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32) + built = fuser.fuse(module, "norm", model_config, None) + assert isinstance(built, VLLMRMSNorm) + assert built.variance_epsilon == torch.finfo(torch.float32).eps + + +def test_eps_is_derived_per_instance(default_vllm_config): + """Two instances of the same norm class with different eps must fuse to their + own eps: the type-cached fuser holds only structure, not this value.""" + model_config = SimpleNamespace(get_hidden_size=lambda: 16) + with torch.device("meta"): + for eps in (1e-5, 1e-6): + module = RMSNorm(16, eps=eps) + built = get_fuser(module).fuse(module, "norm", model_config, None) + assert built.variance_epsilon == eps + + +def test_fused_norm_is_gather_capable(default_vllm_config): + """Every fused norm is emitted gather-capable, so a norm on a head-sharded + projection (OLMoE-style) self-corrects at runtime with no QKV-specific + plumbing. A full-width input skips the gather and equals a plain norm.""" + from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm + from vllm.model_executor.models.transformers.fusers import rms_norm + + torch.manual_seed(0) + x = torch.randn(4, 16) + for gathered_cls, plain_cls in [ + (rms_norm.TPAwareRMSNorm, RMSNorm), + (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm), + ]: + gathered = gathered_cls(hidden_size=16, eps=1e-6) + assert isinstance(gathered, rms_norm.TPAwareNormMixin) + plain = plain_cls(hidden_size=16, eps=1e-6) + with torch.no_grad(): + weight = torch.randn(16) + gathered.weight.copy_(weight) + plain.weight.copy_(weight) + torch.testing.assert_close(gathered(x), plain(x)) + + +def test_gathered_norm_rejects_uneven_sharding(default_vllm_config): + """A sharded input (narrower than the full-width weight) that does not tile + the weight evenly across ranks is rejected before any collective.""" + from vllm.model_executor.models.transformers.fusers import rms_norm + + norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6) + norm.tp_size = 2 # emulate TP=2 without a real process group + with pytest.raises(ValueError, match="does not tile it evenly"): + norm(torch.randn(2, 3)) # 3 * 2 != 8 diff --git a/tests/models/test_transformers.py b/tests/models/transformers/test_backend.py similarity index 82% rename from tests/models/test_transformers.py rename to tests/models/transformers/test_backend.py index eadc3534c37..a3eea1783f5 100644 --- a/tests/models/test_transformers.py +++ b/tests/models/transformers/test_backend.py @@ -6,10 +6,16 @@ from typing import Any import pytest -from ..conftest import HfRunner, VllmRunner -from ..utils import multi_gpu_test, prep_prompts -from .registry import HF_EXAMPLE_MODELS -from .utils import check_embeddings_close, check_logprobs_close +from ...conftest import HfRunner, VllmRunner +from ...utils import multi_gpu_test, prep_prompts +from ..registry import HF_EXAMPLE_MODELS +from ..utils import check_embeddings_close, check_logprobs_close + + +@pytest.fixture(scope="function", autouse=True) +def enable_pickle(monkeypatch): + """`LLM.apply_model` requires pickling a function.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") def get_model(arch: str) -> str: @@ -18,6 +24,17 @@ def get_model(arch: str) -> str: return model_info.default +def get_num_fused(model) -> tuple[int, int]: + from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ) + + glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules()) + qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules()) + return glu, qkv + + def check_implementation( runner_ref: type[HfRunner | VllmRunner], runner_test: type[VllmRunner], @@ -25,6 +42,7 @@ def check_implementation( model: str, kwargs_ref: dict[str, Any] | None = None, kwargs_test: dict[str, Any] | None = None, + num_fused: tuple[int, int] = (1, 1), **kwargs, ): if kwargs_ref is None: @@ -41,6 +59,12 @@ def check_implementation( model_config = model_test.llm.llm_engine.model_config assert model_config.using_transformers_backend() + num_layers = model_config.hf_config.get_text_config().num_hidden_layers + expected_glu, expected_qkv = num_fused + for num_glu, num_qkv in model_test.apply_model(get_num_fused): + assert num_glu == expected_glu * num_layers + assert num_qkv == expected_qkv * num_layers + outputs_test = model_test.generate_greedy_logprobs(*args) with runner_ref(model, **kwargs_ref) as model_ref: @@ -58,11 +82,11 @@ def check_implementation( @pytest.mark.parametrize( - "model,model_impl", + "model,model_impl,num_fused", [ - ("meta-llama/Llama-3.2-1B-Instruct", "transformers"), - ("hmellor/Ilama-3.2-1B", "auto"), # CUSTOM CODE - ("allenai/OLMoE-1B-7B-0924", "transformers"), # MoE + ("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)), + ("hmellor/Ilama-3.2-1B", "auto", (1, 1)), # CUSTOM CODE + ("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)), # MoE ], ) # trust_remote_code=True by default def test_models( @@ -71,6 +95,7 @@ def test_models( example_prompts: list[str], model: str, model_impl: str, + num_fused: tuple[int, int], ) -> None: import transformers from packaging.version import Version @@ -84,7 +109,12 @@ def test_models( ) check_implementation( - hf_runner, vllm_runner, example_prompts, model, model_impl=model_impl + hf_runner, + vllm_runner, + example_prompts, + model, + num_fused=num_fused, + model_impl=model_impl, ) diff --git a/tests/models/utils.py b/tests/models/utils.py index 1c4b9c93b09..91d76d5f243 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -486,6 +486,7 @@ def dummy_hf_overrides( "Gemma3nForConditionalGeneration", "Gemma4ForCausalLM", "Gemma4ForConditionalGeneration", + "Gemma4MTPModel", "DiffusionGemmaForBlockDiffusion", ) else 1 diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py new file mode 100644 index 00000000000..8be70383b8a --- /dev/null +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for VLLMUnprocessableEntityError and media fetch error handling. + +Verifies that unprocessable image URLs (404, 403, DNS failures, etc.) return +HTTP 422 instead of 500. +""" + +from http import HTTPStatus +from unittest.mock import AsyncMock, MagicMock, patch + +import aiohttp +import pytest + +from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.exceptions import VLLMUnprocessableEntityError +from vllm.multimodal.media import MediaConnector + + +class TestVLLMUnprocessableEntityError: + """Tests for VLLMUnprocessableEntityError exception.""" + + def test_creation(self): + exc = VLLMUnprocessableEntityError("Test error") + assert str(exc) == "Test error" + assert exc.parameter is None + + def test_creation_with_parameter_and_value(self): + exc = VLLMUnprocessableEntityError( + "Test error", + parameter="image_url", + value="https://example.com/image.jpg", + ) + assert "parameter=image_url" in str(exc) + assert "value=https://example.com/image.jpg" in str(exc) + + def test_is_value_error_subclass(self): + exc = VLLMUnprocessableEntityError("Test") + assert isinstance(exc, ValueError) + + +class TestMediaConnectorErrorHandling: + """Tests for MediaConnector error handling.""" + + @pytest.mark.asyncio + async def test_fetch_image_async_404(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=404, + message="Not Found", + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + await connector.fetch_image_async("https://example.com/missing.jpg") + + assert exc_info.value.parameter == "image_url" + + @pytest.mark.asyncio + async def test_fetch_image_async_dns_error(self): + """DNS errors are transient and should remain as-is for retry.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientConnectorDNSError( + connection_key=MagicMock(), + os_error=MagicMock(), + ) + + with pytest.raises(aiohttp.ClientConnectorDNSError) as exc_info: + await connector.fetch_image_async( + "https://nonexistent.example/image.jpg" + ) + + assert isinstance(exc_info.value, aiohttp.ClientConnectorDNSError) + + @pytest.mark.asyncio + async def test_fetch_image_async_500_preserved(self): + """5xx errors should remain as server errors.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "async_get_bytes", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=500, + message="Internal Server Error", + ) + + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + await connector.fetch_image_async("https://example.com/image.jpg") + + assert exc_info.value.status == 500 + + def test_fetch_image_404(self): + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientResponseError( + request_info=MagicMock(), + history=(), + status=404, + message="Not Found", + ) + + with pytest.raises(VLLMUnprocessableEntityError) as exc_info: + connector.fetch_image("https://example.com/missing.jpg") + + assert exc_info.value.parameter == "image_url" + + def test_fetch_image_connection_error(self): + """Connection errors are transient and should remain as-is for retry.""" + connector = MediaConnector() + + with patch.object( + connector.connection, "get_bytes", new_callable=MagicMock + ) as mock_get: + mock_get.side_effect = aiohttp.ClientConnectionError("Connection refused") + + with pytest.raises(aiohttp.ClientConnectionError) as exc_info: + connector.fetch_image("https://example.com/image.jpg") + + assert isinstance(exc_info.value, aiohttp.ClientConnectionError) + + +class TestErrorResponse: + """Tests for error response creation.""" + + def test_unprocessable_entity_returns_422(self): + exc = VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Cannot connect", + parameter="image_url", + value="https://example.com/image.jpg", + ) + + response = create_error_response(exc) + + assert response.error.code == HTTPStatus.UNPROCESSABLE_ENTITY.value + assert response.error.type == "UnprocessableEntityError" + assert response.error.param == "image_url" + + def test_unprocessable_entity_message(self): + exc = VLLMUnprocessableEntityError("Test error message") + response = create_error_response(exc) + + assert response.error.message == "Test error message" + assert response.error.code == 422 diff --git a/tests/multimodal/media/test_video.py b/tests/multimodal/media/test_video.py index e4b3afff084..671abd7b077 100644 --- a/tests/multimodal/media/test_video.py +++ b/tests/multimodal/media/test_video.py @@ -16,7 +16,11 @@ from vllm.assets.video import ( video_to_pil_images_list, ) from vllm.multimodal.media import ImageMediaIO, VideoMediaIO -from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader +from vllm.multimodal.video import ( + PYNVVIDEOCODEC_VIDEO_BACKEND, + VIDEO_LOADER_REGISTRY, + VideoLoader, +) from ..utils import cosine_similarity, create_video_from_image, normalize_image @@ -357,3 +361,95 @@ def test_load_base64_jpeg_raises_on_zero_num_frames(): with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"): videoio.load_base64("video/jpeg", data) + + +# --------------------------------------------------------------------------- +# GPU video backend policy tests +# --------------------------------------------------------------------------- + + +class TestMergeKwargsGpuBackendPolicy: + """Verify that merge_kwargs blocks request-level GPU backend selection + when the static (engine-level) config did not configure that backend.""" + + def test_pynvvideocodec_requires_gpu(self): + assert VIDEO_LOADER_REGISTRY.backend_requires_gpu(PYNVVIDEOCODEC_VIDEO_BACKEND) + + def test_strips_video_backend_pynv_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"video_backend": "pynvvideocodec"}, + ) + assert "video_backend" not in result + + def test_strips_backend_pynv_when_not_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"num_frames": 16}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result.get("backend") != "pynvvideocodec" + + def test_preserves_video_backend_pynv_when_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"video_backend": "pynvvideocodec", "num_frames": 8}, + ) + assert result["video_backend"] == "pynvvideocodec" + assert result["num_frames"] == 8 + + def test_preserves_backend_pynv_when_static(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"backend": "pynvvideocodec"}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result["backend"] == "pynvvideocodec" + + @pytest.mark.parametrize("backend", ["opencv", "pyav", "torchcodec"]) + def test_software_video_backend_passes_through(self, backend: str): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"video_backend": backend}, + ) + assert result["video_backend"] == backend + + @pytest.mark.parametrize("backend", ["opencv", "pyav"]) + def test_software_codec_backend_passes_through(self, backend: str): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={"backend": backend}, + ) + assert result["backend"] == backend + + def test_strips_both_keys_independently(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs=None, + runtime_kwargs={ + "video_backend": "pynvvideocodec", + "backend": "pynvvideocodec", + "num_frames": 4, + }, + ) + assert "video_backend" not in result + assert result.get("backend") != "pynvvideocodec" + assert result["num_frames"] == 4 + + def test_other_kwargs_preserved_when_gpu_backend_stripped(self): + result = VideoMediaIO.merge_kwargs( + default_kwargs={"fps": 2}, + runtime_kwargs={ + "video_backend": "pynvvideocodec", + "num_frames": 16, + }, + ) + assert "video_backend" not in result + assert result["num_frames"] == 16 + + def test_static_pynv_with_different_runtime_gpu_backend(self): + """If static sets pynv via video_backend but runtime tries to set it + via the codec-level 'backend' key (without a static match), strip it.""" + result = VideoMediaIO.merge_kwargs( + default_kwargs={"video_backend": "pynvvideocodec"}, + runtime_kwargs={"backend": "pynvvideocodec"}, + ) + assert result.get("backend") != "pynvvideocodec" + assert result["video_backend"] == "pynvvideocodec" diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 6fccc926a21..6aeb7dc486e 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -768,6 +768,118 @@ def test_pyav_backend_returns_target_frames_not_keyframes(): ) +# ============================================================================ +# TorchCodec Backend Tests +# ============================================================================ + + +def test_torchcodec_backend_loads_frames( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """Test that the torchcodec codec backend can load frames.""" + pytest.importorskip("torchcodec") + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_data, num_frames=8, backend="torchcodec" + ) + + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] == 8 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec" + assert "total_num_frames" in metadata + assert "fps" in metadata + assert "duration" in metadata + + +def test_torchcodec_dynamic_backend_loads_frames( + dummy_video_path, monkeypatch: pytest.MonkeyPatch +): + """Test that the torchcodec codec with dynamic sampling can load frames.""" + pytest.importorskip("torchcodec") + with monkeypatch.context() as m: + m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic") + + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv_dynamic") + frames, metadata = loader.load_bytes( + video_data, fps=2, max_duration=10, backend="torchcodec" + ) + + assert frames.ndim == 4 + assert frames.shape[3] == 3 # RGB + assert frames.shape[0] > 0 + assert frames.shape[0] == len(metadata["frames_indices"]) + assert metadata["video_backend"] == "torchcodec_dynamic" + + +def test_torchcodec_backend_rejects_frame_recovery(dummy_video_path): + """frame_recovery is OpenCV-only; torchcodec must reject it.""" + pytest.importorskip("torchcodec") + with open(dummy_video_path, "rb") as f: + video_data = f.read() + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + with pytest.raises(AssertionError): + loader.load_bytes( + video_data, num_frames=8, backend="torchcodec", frame_recovery=True + ) + + +def test_torchcodec_backend_returns_target_frames_not_keyframes(): + """Regression test: torchcodec must return the requested frames, not the + GOP keyframe they seek back to. + + Mirrors ``test_pyav_backend_returns_target_frames_not_keyframes``: a long + GOP (single keyframe at frame 0) with a per-frame green-channel marker. + With ``seek_mode="exact"`` torchcodec resolves each index to the exact + frame, so the returned markers must be distinct, ordered, and match the + requested indices. + """ + pytest.importorskip("torchcodec") + num_frames = 50 + num_sampled = 4 + height, width = 64, 64 + + video_bytes = create_long_gop_video( + num_frames=num_frames, width=width, height=height + ) + + loader = VIDEO_LOADER_REGISTRY.load("opencv") + frames, metadata = loader.load_bytes( + video_bytes, num_frames=num_sampled, backend="torchcodec" + ) + assert frames.shape == (num_sampled, height, width, 3) + + requested = list(metadata["frames_indices"]) + assert len(requested) == num_sampled + + actual = [int(f[height // 2, width // 2, 1]) for f in frames] + + assert len(set(actual)) == num_sampled, ( + f"torchcodec returned only {len(set(actual))} distinct frames for " + f"{num_sampled} requested indices: markers={actual}, " + f"requested={requested}. Keyframe-snap regression." + ) + + assert actual == sorted(actual), f"Returned frames out of order: markers={actual}" + + for marker, want_idx in zip(actual, requested): + assert abs(marker - want_idx) <= 10, ( + f"Frame mismatch: requested index {want_idx}, " + f"got marker {marker} (tolerance ±10)" + ) + + @pytest.mark.parametrize( "loader_key, kwargs, expected_num_frames", [ @@ -854,6 +966,42 @@ def test_pyav_backend_returns_target_frames_not_keyframes(): 120, id="glm46v-pyav-60s", ), + # uniform sampling + torchcodec codec (same frame counts as opencv) + pytest.param( + "opencv", + {"num_frames": 32, "backend": "torchcodec"}, + 32, + id="torchcodec-num_frames", + ), + pytest.param( + "opencv", {"fps": 2, "backend": "torchcodec"}, 120, id="torchcodec-fps" + ), + pytest.param( + "opencv", + {"num_frames": 500, "fps": 2, "backend": "torchcodec"}, + 120, + id="torchcodec-num_frames_wins_fps", + ), + # dynamic sampling + torchcodec codec + pytest.param( + "opencv_dynamic", + {"fps": 1, "max_duration": 60, "backend": "torchcodec"}, + 60, + id="torchcodec_dynamic-within_max_duration", + ), + pytest.param( + "opencv_dynamic", + {"fps": 2, "max_duration": 30, "backend": "torchcodec"}, + 60, + id="torchcodec_dynamic-exceeds_max_duration", + ), + # glm46v dynamic FPS + torchcodec codec + pytest.param( + "glm46v", + {"backend": "torchcodec"}, + 120, + id="glm46v-torchcodec-60s", + ), ], ) def test_video_loader_frames_sampling( @@ -864,6 +1012,8 @@ def test_video_loader_frames_sampling( expected_num_frames: int, ): """Test video loader frames sampling functionality.""" + if kwargs.get("backend") == "torchcodec": + pytest.importorskip("torchcodec") monkeypatch.setenv("VLLM_VIDEO_LOADER_BACKEND", loader_key) loader = VIDEO_LOADER_REGISTRY.load(loader_key) diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index ba90252fe1c..3e9cff64aa1 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -7,7 +7,6 @@ from collections.abc import Sequence import pytest from openai_harmony import ( Conversation, - HarmonyError, Message, RenderConversationConfig, Role, @@ -51,6 +50,15 @@ def chat_request(): ) +@pytest.fixture +def malformed_msgs_str() -> list[str]: + return [ + "<|channel|>analysis<|message|>thinking<|end|>", + "<|start|>assistant<|channel|>commentary<|message|>thinking<|end|>", + '<|start|>assistant<|channel|>final {"answer": "hi"}<|return|>', + ] + + def encode_output(harmony_str: str) -> list[int]: return get_encoding().encode(harmony_str, allowed_special="all") @@ -131,13 +139,22 @@ def tool_call_entries(delta_message) -> list[tuple[int, str | None, str | None]] ] +def assert_parser_is_reset(harmony_parser: HarmonyParser): + assert harmony_parser._parser is None + assert harmony_parser._num_processed_messages == 0 + assert harmony_parser._current_message_tokens == [] + + class TestFlush: def test_flush(self, harmony_parser): harmony_parser.process_chunk( encode_output("<|channel|>analysis<|message|>Think") ) - flushed = harmony_parser.flush() + flushed_segments = harmony_parser.flush() + assert flushed_segments is not None + assert len(flushed_segments) == 1 + flushed = flushed_segments[0] assert flushed is not None assert flushed.channel == "analysis" @@ -145,15 +162,27 @@ class TestFlush: assert flushed.delta == "" assert flushed.completed_message is not None assert get_text(flushed.completed_message) == "Think" - assert harmony_parser._parser is None + assert_parser_is_reset(harmony_parser) - def test_flush_raises_and_resets_on_non_terminal_eos(self, harmony_parser): - harmony_parser.process_chunk(encode_output("<|channel|>analysis")) + def test_flush_recovers_invalid_output(self, harmony_parser, malformed_msgs_str): + for msg_str in malformed_msgs_str[:-1]: + chunk = harmony_parser.process_chunk(encode_output(msg_str)) + assert "".join(segment.delta for segment in chunk.segments) == "thinking" - with pytest.raises(HarmonyError): - harmony_parser.flush() + last_msg_str = malformed_msgs_str[-1] + harmony_parser.process_chunk(encode_output(last_msg_str)) + flushed_segments = harmony_parser.flush() + assert len(flushed_segments) == 2 + delta_segment = flushed_segments[0] + message_segment = flushed_segments[1] - assert harmony_parser._parser is None + assert delta_segment.channel == "final" + assert delta_segment.recipient is None + assert delta_segment.delta == last_msg_str + assert message_segment.channel == "final" + assert message_segment.recipient is None + assert get_text(message_segment.completed_message) == last_msg_str + assert_parser_is_reset(harmony_parser) class TestParse: @@ -364,7 +393,7 @@ class TestParse: assert reasoning is None assert content == "I'm in the middle of answering" assert tool_calls is None - assert harmony_parser._parser is None + assert_parser_is_reset(harmony_parser) def test_interrupted_reasoning_first_message(self, harmony_parser, chat_request): reasoning, content, tool_calls = harmony_parser.parse( @@ -378,7 +407,7 @@ class TestParse: assert reasoning == "I'm in the middle of thinking" assert content is None assert tool_calls is None - assert harmony_parser._parser is None + assert_parser_is_reset(harmony_parser) def test_truncated_output(self, harmony_parser, chat_request): reasoning, content, tool_calls = harmony_parser.parse( @@ -394,24 +423,23 @@ class TestParse: assert reasoning == "I'm thinking." assert content == "I'm in the middle of answering" assert tool_calls is None - assert harmony_parser._parser is None + assert_parser_is_reset(harmony_parser) - def test_malformed_final_recovers_raw_content(self, harmony_parser, chat_request): - raw_output = ( - "<|channel|>analysis<|message|>thinking<|end|>" - '<|start|>assistant<|channel|>final {"answer": "hi"}<|return|>' - ) + def test_malformed_msgs_recovers_raw_content( + self, harmony_parser, chat_request, malformed_msgs_str + ): + combined_output = "".join(malformed_msgs_str) reasoning, content, tool_calls = harmony_parser.parse( - raw_output, + "", chat_request, - model_output_token_ids=encode_output(raw_output), + model_output_token_ids=encode_output(combined_output), ) - assert content == raw_output - assert reasoning is None + assert reasoning == "thinking" + assert content == "thinking\n" + malformed_msgs_str[-1] assert tool_calls is None - assert harmony_parser._parser is None + assert_parser_is_reset(harmony_parser) @pytest.mark.parametrize( ("harmony_str", "expected_content"), @@ -489,7 +517,7 @@ class TestParseDelta: assert second_delta is not None assert second_delta.content == "Answer" assert second_delta.reasoning is None - assert parser._parser is None + assert_parser_is_reset(parser) def test_multi_token(self, gpt_oss_tokenizer, chat_request): parser = HarmonyParser(gpt_oss_tokenizer) @@ -506,25 +534,33 @@ class TestParseDelta: assert delta.reasoning is None assert not delta.tool_calls - def test_malformed_final_recovers_raw_content( - self, gpt_oss_tokenizer, chat_request + def test_malformed_msgs_recovers_raw_content( + self, gpt_oss_tokenizer, chat_request, malformed_msgs_str ): parser = HarmonyParser(gpt_oss_tokenizer) - delta = parser.parse_delta( - delta_text='final {"answer": "hi"}', - delta_token_ids=encode_output( - '<|channel|>final {"answer": "hi"}<|return|>' - ), + for msg_str in malformed_msgs_str[:-1]: + delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output(msg_str), + request=chat_request, + finished=False, + ) + assert delta.reasoning or delta.content == "thinking" + assert not delta.tool_calls + + last_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output(malformed_msgs_str[-1]), request=chat_request, finished=True, ) - assert delta is not None - assert delta.content == 'final {"answer": "hi"}' - assert delta.reasoning is None - assert not delta.tool_calls - assert parser._parser is None + assert last_delta is not None + assert last_delta.content == malformed_msgs_str[-1] + assert last_delta.reasoning is None + assert not last_delta.tool_calls + assert_parser_is_reset(parser) @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"]) def test_tool_call_split_across_deltas( diff --git a/tests/plugins/vllm_add_dummy_endpoint_plugin/setup.py b/tests/plugins/vllm_add_dummy_endpoint_plugin/setup.py new file mode 100644 index 00000000000..04a4936327d --- /dev/null +++ b/tests/plugins/vllm_add_dummy_endpoint_plugin/setup.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from setuptools import setup + +setup( + name="vllm_add_dummy_endpoint_plugin", + version="0.1", + packages=["vllm_add_dummy_endpoint_plugin"], + entry_points={ + "vllm.endpoint_plugins": [ + "dummy_admin_endpoint_plugin = vllm_add_dummy_endpoint_plugin:DummyAdminEndpointPlugin" # noqa + ] + }, +) diff --git a/tests/plugins/vllm_add_dummy_endpoint_plugin/vllm_add_dummy_endpoint_plugin/__init__.py b/tests/plugins/vllm_add_dummy_endpoint_plugin/vllm_add_dummy_endpoint_plugin/__init__.py new file mode 100644 index 00000000000..9b0aca0c046 --- /dev/null +++ b/tests/plugins/vllm_add_dummy_endpoint_plugin/vllm_add_dummy_endpoint_plugin/__init__.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worked example `vllm.endpoint_plugins` entry point. + +Reports scheduler config via `collective_rpc`. Demonstrates the full +contract: `attach_router` registers the route at Phase A (`build_app`) and +`init_state` stashes the `EngineClient` the route handler needs at Phase B +(`init_app_state`). + +`required_tasks` is `None`, so this plugin is also eligible on the CPU only +render server which has no `EngineClient`. `init_state` is called with +`engine_client=None` in that case and the route handler returns 503 rather +than reaching for a client that doesn't exist. +""" + +from fastapi import FastAPI, HTTPException, Request + + +class DummyAdminEndpointPlugin: + name = "dummy_admin_endpoint_plugin" + required_tasks: tuple[str, ...] | None = None + + def attach_router(self, app: FastAPI) -> None: + @app.get("/v1/admin/scheduler_config") + async def scheduler_config(raw_request: Request): + engine_client = raw_request.app.state.dummy_engine_client + if engine_client is None: + raise HTTPException( + status_code=503, + detail="scheduler_config requires an engine, which this " + "server does not have", + ) + results = await engine_client.collective_rpc("get_scheduler_config") + return {"scheduler_config": results} + + async def init_state(self, engine_client, state, args) -> None: + state.dummy_engine_client = engine_client diff --git a/tests/plugins_tests/test_endpoint_plugins.py b/tests/plugins_tests/test_endpoint_plugins.py new file mode 100644 index 00000000000..9385d98f2e3 --- /dev/null +++ b/tests/plugins_tests/test_endpoint_plugins.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the `vllm.endpoint_plugins` framework (RFC #46565). + +Uses the worked in repo example plugin (`vllm_add_dummy_endpoint_plugin`, +installed via `tests/plugins/vllm_add_dummy_endpoint_plugin`) exercising both +`EndpointPlugin` hooks against a fake `EngineClient`, unit tests for the +`load_endpoint_plugins` gating matrix and an e2e test that drives a real HTTP +request through the plugin's route. +""" + +from argparse import Namespace +from typing import Any + +import httpx +import pytest +from fastapi import FastAPI +from vllm_add_dummy_endpoint_plugin import DummyAdminEndpointPlugin + +from vllm.entrypoints.openai.api_server import ( + _attach_endpoint_plugins, + _init_endpoint_plugins_state, + build_app, +) +from vllm.entrypoints.openai.cli_args import make_arg_parser +from vllm.plugins import load_endpoint_plugins +from vllm.plugins.endpoint_plugins.interface import EndpointPlugin +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +class _RaisingEndpointPlugin: + """Factory that raises to exercise the "instantiation fails" path.""" + + name = "raising_endpoint_plugin" + required_tasks = None + + def __init__(self): + raise RuntimeError("boom") + + +class _FakeEngineClient: + """Minimal stand in exercising `collective_rpc`. Not a real engine.""" + + def __init__(self, rpc_result: Any = None): + self.rpc_result = rpc_result + self.rpc_calls: list[tuple[str, tuple, dict]] = [] + + async def collective_rpc(self, method, timeout=None, args=(), kwargs=None): + self.rpc_calls.append((method, args, kwargs or {})) + return self.rpc_result + + +def _build_args() -> Namespace: + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers() + serve_parser = subparsers.add_parser("serve") + make_arg_parser(serve_parser) + return serve_parser.parse_args([]) + + +def _fake_loader(factories: dict[str, Any]): + def _load_plugins_by_group(group: str) -> dict[str, Any]: + assert group == "vllm.endpoint_plugins" + return factories + + return _load_plugins_by_group + + +def test_dummy_plugin_satisfies_protocol(): + assert isinstance(DummyAdminEndpointPlugin(), EndpointPlugin) + + +def test_no_plugins_loaded_when_allowlist_unset(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("VLLM_PLUGINS", raising=False) + + assert load_endpoint_plugins(("generate",)) == [] + + +def test_no_plugins_loaded_when_allowlist_is_empty_string( + monkeypatch: pytest.MonkeyPatch, +): + """`VLLM_PLUGINS=""` parses to `[""]`, not `None` (see `vllm.envs`), so it + must be treated as a (non strict) allowlist matching no plugin name, not + as "unset".""" + monkeypatch.setenv("VLLM_PLUGINS", "") + + assert load_endpoint_plugins(("generate",)) == [] + + +def test_plugin_loaded_when_allowlisted_and_task_matches( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + + plugins = load_endpoint_plugins(("generate",)) + + assert len(plugins) == 1 + assert isinstance(plugins[0], DummyAdminEndpointPlugin) + + +def test_plugin_skipped_when_required_tasks_miss(monkeypatch: pytest.MonkeyPatch): + class _GenerateOnlyPlugin(DummyAdminEndpointPlugin): + required_tasks = ("generate",) + + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + monkeypatch.setattr( + "vllm.plugins.load_plugins_by_group", + _fake_loader({"dummy_admin_endpoint_plugin": _GenerateOnlyPlugin}), + ) + + assert load_endpoint_plugins(("embed",)) == [] + assert len(load_endpoint_plugins(("generate",))) == 1 + + +def test_plugin_loaded_when_required_tasks_is_none(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + + assert len(load_endpoint_plugins(supported_tasks=None)) == 1 + + +def test_plugin_skipped_when_required_tasks_set_but_supported_tasks_none( + monkeypatch: pytest.MonkeyPatch, +): + class _GenerateOnlyPlugin(DummyAdminEndpointPlugin): + required_tasks = ("generate",) + + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + monkeypatch.setattr( + "vllm.plugins.load_plugins_by_group", + _fake_loader({"dummy_admin_endpoint_plugin": _GenerateOnlyPlugin}), + ) + + assert load_endpoint_plugins(supported_tasks=None) == [] + + +def test_factory_raising_is_logged_and_skipped(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv( + "VLLM_PLUGINS", "raising_endpoint_plugin,dummy_admin_endpoint_plugin" + ) + monkeypatch.setattr( + "vllm.plugins.load_plugins_by_group", + _fake_loader( + { + "raising_endpoint_plugin": _RaisingEndpointPlugin, + "dummy_admin_endpoint_plugin": DummyAdminEndpointPlugin, + } + ), + ) + + plugins = load_endpoint_plugins(("generate",)) + + assert len(plugins) == 1 + assert isinstance(plugins[0], DummyAdminEndpointPlugin) + + +def test_attach_is_noop_when_nothing_discovered(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("VLLM_PLUGINS", raising=False) + + app = FastAPI() + _attach_endpoint_plugins(app, ("generate",)) + + assert app.state.endpoint_plugins == [] + + +@pytest.mark.asyncio +async def test_init_state_is_noop_without_phase_a(monkeypatch: pytest.MonkeyPatch): + """`init_app_state` callers that never ran `build_app` (e.g. + `run_batch.py`, which builds a bare `State()`) must not crash just + because `state.endpoint_plugins` was never set.""" + from starlette.datastructures import State + + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + + state = State() + await _init_endpoint_plugins_state(_FakeEngineClient(), state, _build_args()) + + assert not hasattr(state, "dummy_engine_client") + + +@pytest.mark.asyncio +async def test_render_server_attaches_endpoint_plugins_with_no_engine_client( + monkeypatch: pytest.MonkeyPatch, +): + """The CPU only render server has no `EngineClient` but a plugin eligible + for the `render` task (`required_tasks` is `None` or includes `"render"`) + still gets its routes attached at Phase A. Phase B passes `None` for + `engine_client` and it's up to the plugin to handle that.""" + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + + args = _build_args() + app = build_app(args, ("render",)) + + assert len(app.state.endpoint_plugins) == 1 + assert any( + getattr(route, "path", None) == "/v1/admin/scheduler_config" + for route in app.routes + ) + + await _init_endpoint_plugins_state(None, app.state, args) + + assert app.state.dummy_engine_client is None + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/v1/admin/scheduler_config") + + assert response.status_code == 503 + + +@pytest.mark.asyncio +async def test_endpoint_plugin_end_to_end(monkeypatch: pytest.MonkeyPatch): + """Phase A (attach) + Phase B (init) wired through `build_app` then + exercised with a real HTTP request against the worked example plugin.""" + monkeypatch.setenv("VLLM_PLUGINS", "dummy_admin_endpoint_plugin") + + args = _build_args() + app = build_app(args, supported_tasks=()) + + assert len(app.state.endpoint_plugins) == 1 + assert any( + getattr(route, "path", None) == "/v1/admin/scheduler_config" + for route in app.routes + ) + + fake_engine_client = _FakeEngineClient(rpc_result=["cfg-a", "cfg-b"]) + await _init_endpoint_plugins_state(fake_engine_client, app.state, args) + + assert app.state.dummy_engine_client is fake_engine_client + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + response = await client.get("/v1/admin/scheduler_config") + + assert response.status_code == 200 + assert response.json() == {"scheduler_config": ["cfg-a", "cfg-b"]} + assert fake_engine_client.rpc_calls == [("get_scheduler_config", (), {})] diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index f5a38ddb51d..2cd70e4cf0f 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -61,7 +61,7 @@ MODELS = [ ) @pytest.mark.parametrize("model", MODELS) def test_auto_round_model(vllm_runner, model): - with vllm_runner(model, enforce_eager=True) as llm: + with vllm_runner(model) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=8) assert output @@ -336,7 +336,7 @@ def test_wna16_xpu_prefers_ark_when_available(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: True) monkeypatch.setattr(current_platform, "is_cpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (True, None, object(), DummyQuantLinear), ) @@ -355,7 +355,7 @@ def test_wna16_xpu_falls_back_when_ark_unavailable(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: True) monkeypatch.setattr(current_platform, "is_cpu", lambda: False) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (False, "missing", None, None), ) @@ -377,7 +377,7 @@ def test_wna16_cpu_gptq_prefers_ark_when_available(monkeypatch) -> None: monkeypatch.setattr(current_platform, "is_xpu", lambda: False) monkeypatch.setattr(current_platform, "is_cpu", lambda: True) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (True, None, object(), DummyQuantLinear), ) @@ -398,7 +398,7 @@ def test_wna16_cpu_gptq_raises_when_ark_and_marlin_unavailable( monkeypatch.setattr(current_platform, "is_xpu", lambda: False) monkeypatch.setattr(current_platform, "is_cpu", lambda: True) monkeypatch.setattr( - "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state", + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state", lambda: (False, "missing", None, None), ) monkeypatch.setattr( diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 32450231487..a08e14c53c8 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -133,6 +133,81 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head(): assert isinstance(method, ModelOptNvFp4LinearMethod) +def test_modelopt_mixed_precision_resolves_declared_packed_projection(): + config = _mixed_precision_config( + { + "model.layers.0.self_attn.q_proj": {"quant_algo": "MXFP8"}, + "model.layers.0.self_attn.k_proj": {"quant_algo": "MXFP8"}, + "model.layers.0.self_attn.v_proj": {"quant_algo": "MXFP8"}, + } + ) + config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + + assert config._resolve_quant_algo("model.layers.0.self_attn.qkv_proj") == "MXFP8" + + +def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling(): + config = _mixed_precision_config( + { + "model.layers.0.linear_attn.in_proj_qkv": {"quant_algo": "FP8"}, + "model.layers.0.linear_attn.in_proj_z": {"quant_algo": "FP8"}, + "model.layers.0.linear_attn.out_proj": {"quant_algo": "FP8"}, + } + ) + config.packed_modules_mapping = { + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + } + + assert ( + config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_qkvz") == "FP8" + ) + assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None + + +def test_modelopt_mixed_precision_infers_fused_gate_up_projection(): + from vllm.model_executor.layers.linear import LinearBase + + config = _mixed_precision_config( + { + "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"}, + "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"}, + } + ) + + fake_layer = MagicMock(spec=LinearBase) + with patch( + "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel" + ): + method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj") + + assert isinstance(method, ModelOptNvFp4LinearMethod) + + +@pytest.mark.parametrize( + ("quantized_prefix", "missing_prefix"), + [ + ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"), + ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"), + ], +) +def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear( + quantized_prefix, missing_prefix +): + from vllm.model_executor.layers.linear import LinearBase + + config = _mixed_precision_config( + { + quantized_prefix: {"quant_algo": "NVFP4"}, + } + ) + + fake_layer = MagicMock(spec=LinearBase) + method = config.get_quant_method(fake_layer, missing_prefix) + + assert isinstance(method, UnquantizedLinearMethod) + + def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale(): holder = Mock() scale = torch.nn.Parameter(torch.empty(1)) @@ -249,10 +324,11 @@ def test_modelopt_fp8_pc_pt_checkpoint_setup(default_vllm_config, vllm_runner): assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod) assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod) - assert qkv_proj.weight.dtype == torch.float8_e4m3fn - assert o_proj.weight.dtype == torch.float8_e4m3fn - assert gate_up_proj.weight.dtype == torch.float8_e4m3fn - assert down_proj.weight.dtype == torch.float8_e4m3fn + fp8_dtype = current_platform.fp8_dtype() + assert qkv_proj.weight.dtype == fp8_dtype + assert o_proj.weight.dtype == fp8_dtype + assert gate_up_proj.weight.dtype == fp8_dtype + assert down_proj.weight.dtype == fp8_dtype # Per-channel scales; activations are dynamically scaled per token. assert hasattr(qkv_proj, "weight_scale") diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py index e33cc304710..2238c41f498 100644 --- a/tests/renderers/test_chat_utils_prompt_embeds.py +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -40,7 +40,7 @@ from vllm.renderers.hf import ( # Qwen2TokenizerFast (SentencePiece BPE variant) # BertTokenizerFast (WordPiece) TOKENIZER_IDS: Final[list[str]] = [ - "gpt2", + "openai-community/gpt2", "Qwen/Qwen2.5-1.5B-Instruct", "bert-base-uncased", ] diff --git a/tests/test_config.py b/tests/test_config.py index 3837057658b..1e93b610da5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -649,6 +649,30 @@ def test_nested_hf_overrides(): assert model_config.hf_config.vision_config.hidden_size == 512 +def test_model_class_overrides_registers_target(): + """`model_class_overrides` redirects an architecture to a custom class.""" + from vllm.model_executor.models import ModelRegistry + + arch = "_TestModelClassOverrideArch" + target = "vllm.model_executor.models.llama:LlamaForCausalLM" + assert arch not in ModelRegistry.models + + model_config = ModelConfig( + "facebook/opt-125m", + model_class_overrides={arch: target}, + ) + try: + # Accessing `.registry` is the chokepoint that applies the overrides; + # it has already run during construction. + registered = model_config.registry.models[arch] + assert registered.module_name == "vllm.model_executor.models.llama" + assert registered.class_name == "LlamaForCausalLM" + # Idempotent: a second access does not re-register or error out. + assert model_config.registry.models[arch] is registered + finally: + ModelRegistry.models.pop(arch, None) + + @pytest.mark.skipif( current_platform.is_rocm(), reason="Encoder Decoder models not supported on ROCm." ) diff --git a/tests/test_sampling_params.py b/tests/test_sampling_params.py new file mode 100644 index 00000000000..e5d811fbb13 --- /dev/null +++ b/tests/test_sampling_params.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass + +import pytest + +from vllm import SamplingParams + + +@dataclass +class MockModelConfig: + is_diffusion: bool = False + max_logprobs: int = 20 + logits_processors: list | None = None + + def get_vocab_size(self) -> int: + return 1024 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"temperature": 0.7}, + {"temperature": 0.0}, + {"min_p": 0.1}, + {"seed": 42}, + {"min_tokens": 5}, + {"logit_bias": {0: 1.0}}, + {"bad_words": ["foo"]}, + {"allowed_token_ids": [0, 1]}, + ], +) +def test_diffusion_rejects_unsupported_params(kwargs: dict): + params = SamplingParams(**kwargs) + with pytest.raises(ValueError, match="not yet supported with diffusion"): + params.verify(MockModelConfig(is_diffusion=True), None, None, None) + + +def test_diffusion_accepts_default_params(): + SamplingParams().verify(MockModelConfig(is_diffusion=True), None, None, None) + + +def test_diffusion_accepts_top_k_top_p(): + params = SamplingParams(top_p=0.9, top_k=10) + params.verify(MockModelConfig(is_diffusion=True), None, None, None) + + +def test_non_diffusion_models_unaffected(): + params = SamplingParams(temperature=0.7, top_k=10, seed=42) + params.verify(MockModelConfig(), None, None, None) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index fc4da3f8fec..4da6381a0bc 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -5,7 +5,7 @@ from typing import _get_protocol_attrs # type: ignore import pytest from transformers import ( PreTrainedTokenizerBase, - PreTrainedTokenizerFast, + TokenizersBackend, ) from vllm.tokenizers import TokenizerLike, get_tokenizer @@ -23,8 +23,8 @@ def _assert_tokenizer_like(tokenizer: object): def test_tokenizer_like_protocol(): - tokenizer = get_tokenizer("gpt2", use_fast=True) - assert isinstance(tokenizer, PreTrainedTokenizerFast) + tokenizer = get_tokenizer("openai-community/gpt2", use_fast=True) + assert isinstance(tokenizer, TokenizersBackend) _assert_tokenizer_like(tokenizer) tokenizer = get_tokenizer( @@ -38,12 +38,14 @@ def test_tokenizer_like_protocol(): assert isinstance(tokenizer, HfTokenizer) # Verify it's a fast tokenizer (required for FastIncrementalDetokenizer) - assert isinstance(tokenizer, PreTrainedTokenizerFast) + assert isinstance(tokenizer, TokenizersBackend) assert "DSV32" in tokenizer.__class__.__name__ _assert_tokenizer_like(tokenizer) -@pytest.mark.parametrize("tokenizer_name", ["facebook/opt-125m", "gpt2"]) +@pytest.mark.parametrize( + "tokenizer_name", ["facebook/opt-125m", "openai-community/gpt2"] +) def test_tokenizer_revision(tokenizer_name: str): # Assume that "main" branch always exists tokenizer = get_tokenizer(tokenizer_name, revision="main") diff --git a/tests/tokenizers_/test_detokenize.py b/tests/tokenizers_/test_detokenize.py index 2f173bec80c..23eaca9fc36 100644 --- a/tests/tokenizers_/test_detokenize.py +++ b/tests/tokenizers_/test_detokenize.py @@ -5,7 +5,7 @@ from collections.abc import Generator from typing import Any import pytest -from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast +from transformers import AutoTokenizer, PythonBackend, TokenizersBackend from vllm.sampling_params import SamplingParams from vllm.tokenizers.mistral import MistralTokenizer @@ -33,7 +33,7 @@ TRUTH = [ TOKENIZERS = [ "facebook/opt-125m", - "gpt2", + "openai-community/gpt2", "bigcode/tiny_starcoder_py", "EleutherAI/gpt-j-6b", "EleutherAI/pythia-70m", @@ -153,13 +153,13 @@ def test_decode_streaming( spaces_between_special_tokens, fast, ): - if fast and not isinstance(tokenizer, PreTrainedTokenizerFast): + if fast and not isinstance(tokenizer, TokenizersBackend): pytest.skip() if skip_special_tokens and not spaces_between_special_tokens: pytest.skip() - if not fast and isinstance(tokenizer, PreTrainedTokenizerFast): + if not fast and isinstance(tokenizer, TokenizersBackend): # Fix up inconsistency in fast/slow tokenizer behaviour. tokenizer.add_special_tokens( { @@ -173,7 +173,7 @@ def test_decode_streaming( extra_decode_args = ( {} - if not isinstance(tokenizer, PreTrainedTokenizer) + if not isinstance(tokenizer, PythonBackend) else {"spaces_between_special_tokens": spaces_between_special_tokens} ) @@ -225,7 +225,7 @@ def test_decode_streaming( @pytest.mark.parametrize("tokenizer_name", TOKENIZERS) @pytest.mark.parametrize("fast", (True, False)) def test_oov_decode(tokenizer, fast): - if fast and not isinstance(tokenizer, PreTrainedTokenizerFast): + if fast and not isinstance(tokenizer, TokenizersBackend): pytest.skip() decoded_text, out_ids = _run_incremental_decode( diff --git a/tests/tokenizers_/test_hf.py b/tests/tokenizers_/test_hf.py index 3ccbbd73e7a..61c81302f07 100644 --- a/tests/tokenizers_/test_hf.py +++ b/tests/tokenizers_/test_hf.py @@ -14,7 +14,7 @@ from vllm.tokenizers.hf import ( ) -@pytest.mark.parametrize("model_id", ["gpt2", "zai-org/chatglm3-6b"]) +@pytest.mark.parametrize("model_id", ["openai-community/gpt2", "zai-org/chatglm3-6b"]) def test_cached_tokenizer(model_id: str): reference_tokenizer = AutoTokenizer.from_pretrained( model_id, trust_remote_code=True @@ -47,7 +47,7 @@ def _check_consistency(target: TokenizerLike, expected: TokenizerLike): assert target.encode("prompt") == expected.encode("prompt") -@pytest.mark.parametrize("model_id", ["gpt2"]) +@pytest.mark.parametrize("model_id", ["openai-community/gpt2"]) def test_thread_pool_tokenizer_pickle(model_id: str): """Regression test for issue #45433: the thread-pool tokenizer wrapper reconstructs through maybe_make_thread_pool on unpickling, which used to diff --git a/tests/tool_parsers/conftest.py b/tests/tool_parsers/conftest.py index 89609b257c3..23e0eff98a2 100644 --- a/tests/tool_parsers/conftest.py +++ b/tests/tool_parsers/conftest.py @@ -9,4 +9,4 @@ from vllm.tokenizers import TokenizerLike @pytest.fixture(scope="module") def default_tokenizer() -> TokenizerLike: - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") diff --git a/tests/tool_parsers/test_gigachat3_tool_parser.py b/tests/tool_parsers/test_gigachat3_tool_parser.py index b00b410b2fa..00a97095134 100644 --- a/tests/tool_parsers/test_gigachat3_tool_parser.py +++ b/tests/tool_parsers/test_gigachat3_tool_parser.py @@ -19,7 +19,7 @@ from vllm.tool_parsers import ToolParser, ToolParserManager def default_tokenizer() -> TokenizerLike: """Override module-scoped default_tokenizer because gigachat tests mutate the tokenizer via ``add_tokens``.""" - return AutoTokenizer.from_pretrained("gpt2") + return AutoTokenizer.from_pretrained("openai-community/gpt2") MSG_SEP_TOKEN = "<|message_sep|>\n\n" diff --git a/tests/tool_parsers/test_glm47_moe_tool_parser.py b/tests/tool_parsers/test_glm47_moe_tool_parser.py index c9767f6f62f..224fea08d74 100644 --- a/tests/tool_parsers/test_glm47_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm47_moe_tool_parser.py @@ -7,12 +7,16 @@ import json from unittest.mock import Mock import pytest +from openai.types.responses import ResponseFunctionToolCall from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionToolsParam, FunctionDefinition, ) +from vllm.entrypoints.openai.engine.protocol import FunctionCall +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.utils import build_response_output_items from vllm.tokenizers import get_tokenizer from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser @@ -58,7 +62,69 @@ def mock_request(sample_tools) -> ChatCompletionRequest: return request +@pytest.fixture +def namespace_tool_request() -> ResponsesRequest: + return ResponsesRequest.model_validate( + { + "input": "hi", + "tools": [ + { + "type": "namespace", + "name": "mcp__computer_use", + "description": "Computer use tools.", + "tools": [ + { + "type": "function", + "name": "get_app_state", + "description": "Get app state.", + "parameters": { + "type": "object", + "properties": { + "app": {"type": "string"}, + }, + }, + } + ], + } + ], + } + ) + + class TestGlm47ExtractToolCalls: + def test_namespace_tool_call_round_trip_to_responses_output( + self, glm47_tokenizer, namespace_tool_request + ): + parser = Glm47MoeModelToolParser( + glm47_tokenizer, tools=namespace_tool_request.tools + ) + out = ( + "mcp__computer_use__get_app_state" + "app" + "Google Chrome" + "" + ) + + result = parser.extract_tool_calls(out, request=namespace_tool_request) + + assert result.tools_called + tool_call = result.tool_calls[0].function + assert tool_call == FunctionCall( + name="mcp__computer_use__get_app_state", + arguments='{"app": "Google Chrome"}', + ) + + output_items = build_response_output_items( + reasoning=None, + content=None, + tool_calls=[tool_call], + tools=namespace_tool_request.tools, + ) + output_tool_call = output_items[0] + assert isinstance(output_tool_call, ResponseFunctionToolCall) + assert output_tool_call.name == "get_app_state" + assert output_tool_call.namespace == "mcp__computer_use" + def test_no_tool_call(self, glm47_tool_parser, mock_request): out = "This is a plain response." r = glm47_tool_parser.extract_tool_calls(out, request=mock_request) diff --git a/tests/tool_parsers/test_granite_tool_parser.py b/tests/tool_parsers/test_granite_tool_parser.py index 2046c11c5d2..af3386112f5 100644 --- a/tests/tool_parsers/test_granite_tool_parser.py +++ b/tests/tool_parsers/test_granite_tool_parser.py @@ -2,13 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + import pytest from tests.tool_parsers.common_tests import ( ToolParserTestConfig, ToolParserTests, ) -from tests.tool_parsers.utils import run_tool_extraction +from tests.tool_parsers.utils import ( + run_tool_extraction, + run_tool_extraction_streaming, + split_string_into_token_deltas, +) +from vllm.tokenizers import get_tokenizer +from vllm.tool_parsers.granite_tool_parser import GraniteToolParser class TestGraniteToolParser(ToolParserTests): @@ -116,3 +124,38 @@ I'll get that information.""", f"Expected 1 tool call from string format, got {len(tool_calls)}" ) assert tool_calls[0].function.name == "get_weather" + + +# granite emits arguments before name and its own tokenizer (not gpt2) is used +# here so the token boundaries match production; get_tokenizer only fetches the +# small tokenizer files, not the model weights. +@pytest.fixture(scope="module") +def granite_tokenizer(): + return get_tokenizer(tokenizer_name="ibm-granite/granite-3.1-8b-instruct") + + +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5]) +def test_streaming_parallel_calls_batched_deltas(granite_tokenizer, chunk_size): + """A batched delta (multiple tokens) spanning the boundary between two + parallel calls must not drop the first call's name. granite streams + arguments before name, so the name only completes as the next call appears. + """ + parser = GraniteToolParser(granite_tokenizer) + model_output = ( + '<|tool_call|> [{"arguments": {"city": "Tokyo"}, "name": "get_weather"}, ' + '{"arguments": {"timezone": "Asia/Tokyo"}, "name": "get_time"}]' + ) + token_deltas = split_string_into_token_deltas(granite_tokenizer, model_output) + batched = [ + "".join(token_deltas[i : i + chunk_size]) + for i in range(0, len(token_deltas), chunk_size) + ] + reconstructor = run_tool_extraction_streaming( + parser, batched, assert_one_tool_per_delta=False + ) + names = [tc.function.name for tc in reconstructor.tool_calls] + assert names == ["get_weather", "get_time"] + # trailing args of the final call are flushed by the serving layer + assert json.loads(reconstructor.tool_calls[0].function.arguments) == { + "city": "Tokyo" + } diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/tool_parsers/test_mistral_tool_parser.py index 03a10ef0991..4f75e2f576f 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/tool_parsers/test_mistral_tool_parser.py @@ -144,6 +144,7 @@ def stream_delta_message_generator( mistral_tokenizer: TokenizerLike, model_output: str | None, tools: list[tuple[str, str]] | None, + chunk_size: int = 1, ) -> Generator[DeltaMessage, None, None]: if ( isinstance(mistral_tokenizer, MistralTokenizer) @@ -182,15 +183,13 @@ def stream_delta_message_generator( previous_tokens = None prefix_offset = 0 read_offset = 0 + pending_text = "" + pending_token_ids: list[int] = [] for i, delta_token in enumerate(all_token_ids): - delta_token_ids = [delta_token] - previous_token_ids = all_token_ids[:i] - current_token_ids = all_token_ids[: i + 1] - (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( detokenize_incrementally( tokenizer=mistral_tokenizer, - all_input_ids=current_token_ids, + all_input_ids=all_token_ids[: i + 1], prev_tokens=previous_tokens, prefix_offset=prefix_offset, read_offset=read_offset, @@ -198,27 +197,39 @@ def stream_delta_message_generator( spaces_between_special_tokens=True, ) ) + previous_tokens = ( + previous_tokens + new_tokens if previous_tokens else new_tokens + ) + prefix_offset = new_prefix_offset + read_offset = new_read_offset - current_text = previous_text + delta_text + # Buffer tokens so each streamed delta can carry ``chunk_size`` tokens, + # reproducing the multi-token deltas produced by async scheduling / + # stream_interval > 1. + pending_text += delta_text + pending_token_ids.append(delta_token) + if len(pending_token_ids) < chunk_size and i != len(all_token_ids) - 1: + continue + + previous_token_ids = all_token_ids[: i + 1 - len(pending_token_ids)] + current_token_ids = all_token_ids[: i + 1] + current_text = previous_text + pending_text delta_message = mistral_tool_parser.extract_tool_calls_streaming( previous_text, current_text, - delta_text, + pending_text, previous_token_ids, current_token_ids, - delta_token_ids, + pending_token_ids, request=_DUMMY_REQUEST, ) if delta_message: yield delta_message previous_text = current_text - previous_tokens = ( - previous_tokens + new_tokens if previous_tokens else new_tokens - ) - prefix_offset = new_prefix_offset - read_offset = new_read_offset + pending_text = "" + pending_token_ids = [] @pytest.mark.parametrize( @@ -1572,3 +1583,39 @@ def test_grammar_from_tool_parser_set_by_adjust_request( request = _make_request() result = mistral_tool_parser.adjust_request(request) assert result._grammar_from_tool_parser is True + + +@pytest.mark.parametrize("chunk_size", [2, 3, 4, 5]) +def test_streaming_pre_v11_parallel_calls_batched_deltas( + mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size +): + """A batched delta spanning the boundary between two parallel calls must + keep them on distinct indices (the bug collapsed both onto index 0).""" + model_output = ( + '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, ' + '{"name": "get_current_weather", "arguments": ' + '{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]' + ) + names: list[str] = [] + args: list[str] = [] + idx = -1 + for delta_message in stream_delta_message_generator( + mistral_pre_v11_tool_parser, + mistral_pre_v11_tokenizer, + model_output, + tools=None, + chunk_size=chunk_size, + ): + for tool_call in delta_message.tool_calls or []: + if tool_call.index != idx: + idx = tool_call.index + args.append("") + if tool_call.function and tool_call.function.name: + names.append(tool_call.function.name) + if tool_call.function and tool_call.function.arguments: + args[tool_call.index] += tool_call.function.arguments + + assert names == ["add", "get_current_weather"] + assert len(args) == 2 + # trailing args of the final call are flushed by the serving layer + assert json.loads(args[0]) == {"a": 3.5, "b": 4} diff --git a/tests/tool_use/test_parallel_tool_calls.py b/tests/tool_use/test_parallel_tool_calls.py index 0f7f6893162..4cfd165f1a8 100644 --- a/tests/tool_use/test_parallel_tool_calls.py +++ b/tests/tool_use/test_parallel_tool_calls.py @@ -115,14 +115,12 @@ async def test_parallel_tool_calls( assert not role_name or role_name == "assistant" role_name = "assistant" - # if a tool call is streamed make sure there's exactly one - # (based on the request parameters + # a chunk may carry >1 tool-call delta at a parallel-call boundary streamed_tool_calls = chunk.choices[0].delta.tool_calls - if streamed_tool_calls and len(streamed_tool_calls) > 0: - # make sure only one diff is present - correct even for parallel - assert len(streamed_tool_calls) == 1 - tool_call = streamed_tool_calls[0] + for tool_call in streamed_tool_calls or []: + # deltas arrive in non-decreasing index order + assert tool_call.index >= tool_call_idx # if a new tool is being called, set up empty arguments if tool_call.index != tool_call_idx: diff --git a/tests/utils.py b/tests/utils.py index 08579f99e4d..2a3bdb91fe0 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1459,6 +1459,46 @@ def multi_process_parallel( ray.shutdown() +def assert_rocm_custom_allreduce_backend_state( + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> None: + from vllm.distributed.parallel_state import get_tp_group + + device_communicator = get_tp_group().device_communicator + aiter_ar_comm = device_communicator.aiter_ar_comm + if use_aiter_custom_ar: + assert aiter_ar_comm is not None, "AITER CustomAllreduce was not initialized." + assert not aiter_ar_comm.disabled, "AITER CustomAllreduce is disabled." + assert device_communicator.ca_comm is None, ( + "vLLM CustomAllreduce should not be initialized when AITER CA is used." + ) + else: + assert aiter_ar_comm is None, ( + "AITER CustomAllreduce should not be initialized when disabled." + ) + assert device_communicator.ca_comm is not None, ( + "vLLM CustomAllreduce should be initialized when AITER CA is disabled." + ) + + qr_comm = device_communicator.qr_comm + assert qr_comm is not None, "QuickReduce communicator was not initialized." + if quick_reduce_quantization == "NONE": + assert qr_comm.disabled, "QuickReduce should be disabled." + else: + assert not qr_comm.disabled, "QuickReduce should be enabled." + + +def assert_rocm_custom_allreduce_backend_state_on_worker( + _worker, + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> None: + assert_rocm_custom_allreduce_backend_state( + use_aiter_custom_ar, quick_reduce_quantization + ) + + @contextmanager def error_on_warning(category: type[Warning] = Warning): """ diff --git a/tests/utils_/test_async_utils.py b/tests/utils_/test_async_utils.py index 03d116bdfd8..cd41cdaf264 100644 --- a/tests/utils_/test_async_utils.py +++ b/tests/utils_/test_async_utils.py @@ -40,3 +40,25 @@ async def test_merge_async_iterators(): print("Iterator was cancelled normally") except (Exception, asyncio.CancelledError) as e: raise AssertionError() from e + + +@pytest.mark.asyncio +async def test_merge_async_iterators_single_closes_underlying(): + # The single-iterator fast path must close the underlying generator when + # the merged generator is closed, matching the multi-iterator path. On the + # buggy fast path the underlying generator is left running. + closed = False + + async def gen(): + nonlocal closed + try: + while True: + yield "x" + await asyncio.sleep(0.01) + finally: + closed = True + + merged = merge_async_iterators(gen()) + assert await anext(merged) == (0, "x") + await merged.aclose() + assert closed diff --git a/tests/v1/attention/test_linear_attention_metadata_builder.py b/tests/v1/attention/test_linear_attention_metadata_builder.py new file mode 100644 index 00000000000..3ef811b3a66 --- /dev/null +++ b/tests/v1/attention/test_linear_attention_metadata_builder.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import CUDAGraphMode, SpeculativeConfig +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionMetadataBuilder, + LinearAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") + + +def _create_mamba_spec(num_speculative_blocks: int = 1) -> MambaSpec: + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=num_speculative_blocks, + ) + + +def test_bailing_linear_attention_reports_uniform_batch_cudagraph_support(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + + support = BailingLinearAttentionMetadataBuilder.get_cudagraph_support( + vllm_config, _create_mamba_spec() + ) + + assert support == AttentionCGSupport.UNIFORM_BATCH + + +def test_non_bailing_linear_attention_keeps_single_token_cudagraph_support(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["MiniMaxText01ForCausalLM"], + "model_type": "minimax_text_01", + } + ) + + support = LinearAttentionMetadataBuilder.get_cudagraph_support( + vllm_config, _create_mamba_spec() + ) + + assert support == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + + +def test_linear_attention_spec_decode_full_graph_metadata_pads_cache_slots(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + + builder = BailingLinearAttentionMetadataBuilder( + kv_cache_spec=_create_mamba_spec(), + layer_names=["model.layers.0.self_attn"], + vllm_config=vllm_config, + device=DEVICE, + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]), + BLOCK_SIZE, + DEVICE, + ) + common.block_table_tensor[2].fill_(-1) + + metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32), + ) + + assert metadata.num_decodes == 3 + assert metadata.num_prefills == 0 + assert metadata.num_decode_tokens == 4 + assert metadata.state_indices_tensor_d is not None + assert metadata.state_indices_tensor_d.shape == (3, 2) + assert torch.equal( + metadata.state_indices_tensor_d[2], + torch.full((2,), PAD_SLOT_ID, dtype=torch.int32), + ) + assert torch.equal( + metadata.state_indices_tensor[2], + torch.tensor(PAD_SLOT_ID, dtype=torch.int32), + ) + assert metadata.query_start_loc_d is not None + assert metadata.query_start_loc_d.tolist() == [0, 2, 4, 4] + assert metadata.num_accepted_tokens is not None + assert metadata.num_accepted_tokens.tolist() == [1, 2, 1] + + +def test_linear_attention_full_graph_metadata_uses_stable_decode_buffers(): + vllm_config = create_vllm_config( + hf_config_override={ + "architectures": ["BailingMoeV2_5ForCausalLM"], + "model_type": "bailing_hybrid", + } + ) + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + + builder = BailingLinearAttentionMetadataBuilder( + kv_cache_spec=_create_mamba_spec(), + layer_names=["model.layers.0.self_attn"], + vllm_config=vllm_config, + device=DEVICE, + ) + + common = create_common_attn_metadata( + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ) + common.block_table_tensor = torch.tensor( + [[10, 11], [12, 13], [-1, -1]], + dtype=torch.int32, + device=DEVICE, + ) + + first = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32), + ) + assert first.state_indices_tensor_d is not None + assert first.query_start_loc_d is not None + assert first.num_accepted_tokens is not None + state_ptr = first.state_indices_tensor_d.data_ptr() + query_ptr = first.query_start_loc_d.data_ptr() + accepted_ptr = first.num_accepted_tokens.data_ptr() + + common2 = create_common_attn_metadata( + BatchSpec(seq_lens=[36, 0, 0], query_lens=[2, 0, 0]), + BLOCK_SIZE, + DEVICE, + arange_block_indices=True, + ) + common2.block_table_tensor = torch.tensor( + [[20, 21], [-1, -1], [-1, -1]], + dtype=torch.int32, + device=DEVICE, + ) + second = builder.build( + common_prefix_len=0, + common_attn_metadata=common2, + num_accepted_tokens=torch.tensor([2, 1, 1], dtype=torch.int32), + ) + + assert second.state_indices_tensor_d is not None + assert second.query_start_loc_d is not None + assert second.num_accepted_tokens is not None + assert second.state_indices_tensor_d.data_ptr() == state_ptr + assert second.query_start_loc_d.data_ptr() == query_ptr + assert second.num_accepted_tokens.data_ptr() == accepted_ptr + assert second.state_indices_tensor_d.tolist() == [ + [20, 21], + [PAD_SLOT_ID, PAD_SLOT_ID], + [PAD_SLOT_ID, PAD_SLOT_ID], + ] + assert second.query_start_loc_d.tolist() == [0, 2, 2, 2] + assert second.num_accepted_tokens.tolist() == [2, 1, 1] diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 48c6de8f8bd..8f9e8acac60 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -136,17 +136,9 @@ def test_standard_attention_backend_selection( # Get the backend class path from vllm.platforms.rocm import RocmPlatform - # The AITER unified attention kernel only supports BF16/FP8 KV caches - # (its 3D kernel asserts on fp16), so it must be selected with bf16. - dtype = ( - torch.bfloat16 - if selected_backend == "ROCM_AITER_UNIFIED_ATTN" - else torch.float16 - ) - attn_selector_config = AttentionSelectorConfig( head_size=128, - dtype=dtype, + dtype=torch.float16, kv_cache_dtype="auto", block_size=16, use_mla=False, diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 8253a8422e2..9b6f6458961 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -324,3 +324,48 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): assert request.status == RequestStatus.FINISHED_ERROR assert request.request_id not in scheduler.requests assert not scheduler.running + + +def test_no_placeholder_underflow_on_discarded_spec_frame(): + num_spec = 5 + scheduler = create_scheduler( + async_scheduling=True, + num_speculative_tokens=num_spec, + speculative_method="ngram_gpu", + ) + req = create_requests(num_requests=1, max_tokens=20)[0] + req.num_computed_tokens = req.num_tokens + scheduler.requests[req.request_id] = req + scheduler.running.append(req) + req.status = RequestStatus.RUNNING + + req.num_output_placeholders = 1 + req.async_tokens_to_discard = num_spec + computed_before = req.num_computed_tokens + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={req.request_id: num_spec + 1}, + total_num_scheduled_tokens=num_spec + 1, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={req.request_id: [10] * num_spec}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_runner_output = ModelRunnerOutput( + req_ids=[req.request_id], + req_id_to_index={req.request_id: 0}, + sampled_token_ids=[[999]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(scheduler_output, model_runner_output) + + assert req.num_output_placeholders == 1 + assert req.num_computed_tokens == computed_before + assert req.async_tokens_to_discard == num_spec - 1 + assert req.status == RequestStatus.RUNNING diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 95237fa2723..947672c48af 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1404,12 +1404,15 @@ def test_get_max_concurrency_for_kv_cache_config(): enable_chunked_prefill=True, max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, + # Pin to sync: SWA per-request bounds grow with overlapping batches. + async_scheduling=False, ) vllm_config = VllmConfig( model_config=model_config, scheduler_config=scheduler_config, ) + assert vllm_config.max_concurrent_batches == 1 full_attention_spec = FullAttentionSpec( block_size=16, @@ -2484,6 +2487,91 @@ def test_unify_hybrid_kv_cache_specs(): kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) +def test_unify_kv_cache_spec_page_size_mamba(): + """Regression test for https://github.com/vllm-project/vllm/issues/43626. + + MambaSpec's page_size_bytes is determined by its state shapes and does not + change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba + page instead of scaling its block_size. This situation arises when a layer + with a page larger than the (already platform-aligned) Mamba page joins the + specs, e.g. a dense draft model with more KV heads than the hybrid main + model. + """ + # 1. Hybrid main model (Mamba + full attention, pages already aligned at + # 16KB) plus a dense draft model layer with a 2x larger page (32KB). + # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so + # the old code scaled the Mamba block_size, which left page_size_bytes + # unchanged at 16384. + mamba_spec = new_mamba_spec() # page_size_bytes = 16384 + main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 + draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 + assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 + assert draft_attn_spec.page_size_bytes == 32768 + + unified = kv_cache_utils.unify_kv_cache_spec_page_size( + { + "mamba_layer": mamba_spec, + "main_attn_layer": main_attn_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + # Mamba page is padded; block_size (caching granularity) is unchanged. + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + assert unified["mamba_layer"].block_size == mamba_spec.block_size + # Attention layer with smaller page still unifies by scaling block_size. + assert unified["main_attn_layer"].page_size_bytes == 32768 + assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size + assert unified["main_attn_layer"].page_size_padded is None + # Layer already at max page size is unchanged. + assert unified["draft_attn_layer"] == draft_attn_spec + + # 2. Mamba page already padded by the platform (state smaller than the + # padded page); the padding is re-applied at the new maximum. + padded_mamba_spec = new_mamba_spec( + shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 + ) + assert padded_mamba_spec.page_size_bytes == 16384 + unified = kv_cache_utils.unify_kv_cache_spec_page_size( + { + "mamba_layer": padded_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + + # 3. Mamba page that does not evenly divide the maximum page size is + # padded as well (the divisibility constraint only applies to block_size + # scaling). + odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) + assert odd_mamba_spec.page_size_bytes == 24576 + assert 32768 % odd_mamba_spec.page_size_bytes != 0 + unified = kv_cache_utils.unify_kv_cache_spec_page_size( + { + "mamba_layer": odd_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + + # 4. Attention layers with non-divisible page sizes still raise. + with pytest.raises(NotImplementedError): + kv_cache_utils.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + + # 5. Uniform page sizes are returned unchanged. + specs = { + "mamba_layer": new_mamba_spec(), + "attn_layer": new_kv_cache_spec(), + } + assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs + + def test_hma_not_disabled_when_kv_events_enabled(): """ Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 6246a233290..59260a499ef 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -3453,7 +3453,8 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): manager = make_kv_cache_manager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) @@ -3509,7 +3510,8 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): manager = make_kv_cache_manager( config, max_model_len=max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + # Single (sync) batch in flight, so in-flight tokens == batched tokens. + max_in_flight_tokens=max_num_batched_tokens, enable_caching=True, hash_block_size=block_size, ) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3e2b7dc5832..900f8a9b06a 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1356,6 +1356,41 @@ def test_spec_decode_padding_first_decode_step(): assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec +def test_spec_decode_padding_skipped_for_diffusion(): + """Diffusion spec tokens are the fixed-size denoising canvas, not + rejectable drafts: a first-decode-step request must keep its 1-token span + instead of being padded to 1 + num_spec_tokens, which would overflow the + canvas. + """ + num_spec = 3 + scheduler = create_scheduler( + num_speculative_tokens=num_spec, + enable_prefix_caching=True, + block_size=16, + ) + # Diffusion schedulers initialize this to 0 (model_config.is_diffusion). + scheduler.num_sampled_tokens_per_step = 0 + r1, r2 = create_requests( + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16 + ) + + scheduler.add_request(r1) + out = scheduler.schedule() + assert out.num_scheduled_tokens[r1.request_id] == 33 + _model_output(scheduler, out, [[100]]) + scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]])) + + # r2 arrives; its whole prompt is a prefix-cache hit -> needs exactly + # 1 token while r1 is a running speculative decode. + scheduler.add_request(r2) + out = scheduler.schedule() + + assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3] + # r2 keeps its true 1-token span; no placeholder drafts are attached. + assert out.num_scheduled_tokens[r2.request_id] == 1 + assert r2.request_id not in out.scheduled_spec_decode_tokens + + def test_spec_decode_padding_skipped_with_prefill_in_batch(): """Padding is skipped when the batch contains a prefill chunk: the batch is already mixed/non-uniform, so padding a new decode request buys nothing. diff --git a/tests/v1/core/test_swa_inflight_window_free.py b/tests/v1/core/test_swa_inflight_window_free.py new file mode 100644 index 00000000000..737796a4bf7 --- /dev/null +++ b/tests/v1/core/test_swa_inflight_window_free.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Out-of-window block frees vs in-flight GPU steps. + +With async scheduling / PP, `num_computed_tokens` optimistically includes +tokens of unprocessed steps whose attention windows still read the blocks +just below the optimistic boundary (and rejected spec tokens can roll it +back), so `allocate_slots` frees on the processed-token basis: +`num_computed_tokens - num_in_flight_tokens`. +""" + +import torch + +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, SlidingWindowSpec +from vllm.v1.outputs import ModelRunnerOutput + +from .utils import create_requests, create_scheduler, mock_kv + +NUM_PROMPT_TOKENS = 100 +BLOCK_SIZE = 16 +SLIDING_WINDOW = 16 +# Tokens 0..84 are outside the window of the next token to compute +# (100 - 16 + 1 = 85), i.e. 5 full blocks. +NUM_OUT_OF_WINDOW_BLOCKS = 85 // BLOCK_SIZE +# Chunked-local skips whole chunks left of the current one: +# (100 // 32) * 32 = 96 settled tokens -> 6 full blocks. +CHUNK_SIZE = 32 +NUM_OUT_OF_CHUNK_BLOCKS = (NUM_PROMPT_TOKENS // CHUNK_SIZE) * CHUNK_SIZE // BLOCK_SIZE + + +def _make_model_runner_output( + scheduler_output: SchedulerOutput, + token_id: int = 0, +) -> ModelRunnerOutput: + req_ids = list(scheduler_output.num_scheduled_tokens.keys()) + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[token_id] for _ in req_ids], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + +def _create_swa_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + + +def _create_chunked_scheduler(async_scheduling: bool): + return create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=async_scheduling, + kv_cache_spec=ChunkedLocalAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=CHUNK_SIZE, + ), + ) + + +def _num_null_blocks(scheduler, request_id: str) -> int: + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + null_block = manager._null_block + return sum(1 for b in manager.req_to_blocks[request_id] if b is null_block) + + +def test_num_in_flight_tokens_accounting(): + scheduler = create_scheduler(async_scheduling=True) + request = create_requests(num_requests=1, num_tokens=NUM_PROMPT_TOKENS)[0] + scheduler.add_request(request) + + out0 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + # Async: decode scheduled before the prefill output is processed. + out1 = scheduler.schedule() + assert request.num_in_flight_tokens == NUM_PROMPT_TOKENS + 1 + + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 1 + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + assert request.num_in_flight_tokens == 0 + + +def test_swa_free_waits_for_in_flight_step(): + """Async: out-of-window blocks stay allocated until the step that still + reads them has been processed.""" + scheduler = _create_swa_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-window blocks: + # they must not be freed yet. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-window + # blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + assert ( + block_pool.get_num_free_blocks() + == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_swa_free_immediate_when_sync(): + """Sync: no in-flight step at schedule time, frees happen at the first + decode allocation as before.""" + scheduler = _create_swa_scheduler(async_scheduling=False) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert request.num_in_flight_tokens == 0 + + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS + + +def test_swa_admission_cap_accounts_for_overlapping_batches(): + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment. + assert base == 129 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens. + assert overlapped == 193 + + +def test_chunked_local_free_waits_for_in_flight_step(): + """Chunked-local attention frees whole chunks left of the current one, and + is exposed to the same load-WAR: with async scheduling those chunks must + stay allocated until the in-flight step that still reads them settles.""" + scheduler = _create_chunked_scheduler(async_scheduling=True) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + block_pool = scheduler.kv_cache_manager.block_pool + + out0 = scheduler.schedule() # prefill, in flight from here on + free_after_prefill = block_pool.get_num_free_blocks() + + # Decode scheduled while the prefill still reads the out-of-chunk blocks. + out1 = scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == 0 + assert block_pool.get_num_free_blocks() == free_after_prefill + + # Prefill output processed; the next allocate frees the out-of-chunk blocks. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + assert ( + block_pool.get_num_free_blocks() == free_after_prefill + NUM_OUT_OF_CHUNK_BLOCKS + ) + # Not double-freed on the following steps. + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + scheduler.schedule() + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_CHUNK_BLOCKS + + +def test_chunked_local_admission_cap_accounts_for_overlapping_batches(): + spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=1024, + ) + base = spec.max_admission_blocks_per_request( + max_in_flight_tokens=1024, max_model_len=16384 + ) + # (1024 + 1024) tokens -> 128 blocks. + assert base == 128 + overlapped = spec.max_admission_blocks_per_request( + max_in_flight_tokens=2 * 1024, max_model_len=16384 + ) + # One extra in-flight chunk is held back: (1024 + 2 * 1024) tokens. + assert overlapped == 192 + + +def test_connector_finish_frees_on_settled_basis(): + """The out-of-window prune done at request finish, before the block table + is handed to a KV connector (simple CPU offload / NIXL store), must use the + same processed-token basis. Otherwise the connector reads/hands off a block + the still-in-flight step is optimistically counted as done with, which is + the load-path WAR this fix closes.""" + scheduler = create_scheduler( + block_size=BLOCK_SIZE, + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + kv_cache_spec=SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=SLIDING_WINDOW, + ), + ) + request = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE + )[0] + scheduler.add_request(request) + req_id = request.request_id + + out0 = scheduler.schedule() # prefill, in flight from here on + scheduler.schedule() # decode over-scheduled: num_computed_tokens optimistic + + # Finishing now (connector store) must NOT prune out-of-window blocks the + # in-flight prefill still reads. + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == 0 + + # Once the in-flight step settles, the same prune releases them. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler._connector_finished(request) + assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2450b23669a..19beba1a53d 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -30,6 +30,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -54,12 +55,16 @@ def create_scheduler( block_size: int = 16, max_model_len: int | None = None, num_speculative_tokens: int | None = None, + speculative_method: str | None = None, skip_tokenizer_init: bool = False, async_scheduling: bool = False, pipeline_parallel_size: int = 1, + data_parallel_size: int = 1, + num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None, use_ec_connector: bool = False, ec_role: str | None = None, use_v2_model_runner: bool | None = None, + kv_cache_spec: KVCacheSpec | None = None, ) -> Scheduler | AsyncScheduler: """Create scheduler under test. @@ -126,9 +131,18 @@ def create_scheduler( speculative_config: SpeculativeConfig | None = None if num_speculative_tokens is not None: - speculative_config = SpeculativeConfig( + spec_kwargs: dict = dict( model="ngram", num_speculative_tokens=num_speculative_tokens ) + if num_speculative_tokens_per_batch_size is not None: + spec_kwargs["num_speculative_tokens_per_batch_size"] = ( + num_speculative_tokens_per_batch_size + ) + if speculative_method is not None: + spec_kwargs["method"] = speculative_method + spec_kwargs["prompt_lookup_max"] = num_speculative_tokens + spec_kwargs["prompt_lookup_min"] = 1 + speculative_config = SpeculativeConfig(**spec_kwargs) ec_transfer_config = ( ECTransferConfig( @@ -144,25 +158,25 @@ def create_scheduler( scheduler_config=scheduler_config, model_config=model_config, cache_config=cache_config, - parallel_config=ParallelConfig(pipeline_parallel_size=pipeline_parallel_size), + parallel_config=ParallelConfig( + pipeline_parallel_size=pipeline_parallel_size, + data_parallel_size=data_parallel_size, + ), kv_transfer_config=kv_transfer_config, speculative_config=speculative_config, ec_transfer_config=ec_transfer_config, ) + if kv_cache_spec is None: + kv_cache_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) kv_cache_config = KVCacheConfig( num_blocks=num_blocks, # A large number of blocks to hold all requests kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=1, - dtype=torch.float32, - ), - ) - ], + kv_cache_groups=[KVCacheGroupSpec(["layer"], kv_cache_spec)], ) cache_config.num_gpu_blocks = num_blocks register_all_kvcache_specs(vllm_config) diff --git a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py index 4bb8d63a8a2..11f77492d2e 100644 --- a/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py +++ b/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py @@ -45,7 +45,9 @@ def test_prompts(): use_fork_for_test = ( - fork_new_process_for_each_test if not current_platform.is_rocm() else lambda x: x + fork_new_process_for_each_test + if not (current_platform.is_rocm() or current_platform.is_xpu()) + else lambda x: x ) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index dd9efc66960..5a7af6f22c5 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -39,13 +39,18 @@ from vllm.v1.worker.mamba_utils import get_mamba_groups class StepAction: num_computed_tokens_start: int num_scheduled_tokens: int - kv_cache_block_ids: list[int] # [] to follow last step + kv_cache_block_ids: list[int] # per-block mask: 1=held, 0=freed/nulled preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy num_speculative_tokens = 3 +# Whether the run under test uses async scheduling. Set by each test entrypoint +# before generation; consulted where the scheduler's optimistic token count must +# be corrected for in-flight (possibly-rejected) speculative tokens. +async_scheduling_mode = False + num_accepted_tokens = 1 prompt_token_ids: list[int] = [] MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" @@ -249,7 +254,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable): scheduler_output.scheduled_cached_reqs.num_computed_tokens[0] ) if ( - self.num_spec_tokens + async_scheduling_mode + and self.num_spec_tokens and num_prompt_tokens is not None and num_computed_tokens > num_prompt_tokens ): @@ -496,7 +502,10 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn) -def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: +def get_mamba_prefix_cache_step_configs( + async_scheduling: bool = False, +) -> dict[str, TestConfig]: + a = async_scheduling tests = { "accept_1": TestConfig( num_prompt_tokens=554, @@ -504,14 +513,24 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=1, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), - StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(556, 4, [], (-1, -1), (-1, -1)), - StepAction(557, 4, [], (0, 1), (-1, -1)), - StepAction(558, 4, [], (-1, -1), (-1, -1)), - StepAction(559, 4, [], (-1, -1), (1, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 555, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), + StepAction(558, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), # test case 2.1: no hit, accept 2 tokens @@ -521,11 +540,19 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=2, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(558, 4, [], (1, 1), (2, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), # test case 2.2: no hit, accept 2 tokens @@ -535,10 +562,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=2, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)), - StepAction(559, 4, [], (-1, -1), (1, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction(559, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -548,10 +581,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [], (-1, -1), (-1, -1)), - StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(559, 4, [], (2, 1), (1, 0)), - StepAction(562, 4, [], (-1, -1), (-1, -1)), + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 556, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1) + ), + StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -561,10 +602,16 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)), - StepAction(560, 4, [], (-1, -1), (-1, -1)), - StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction(560, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 563, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "accept_3_3": TestConfig( @@ -573,9 +620,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=3, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -585,9 +638,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(553, 4, [], (-1, -1), (-1, -1)), + StepAction(553, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)), - StepAction(561, 4, [], (-1, -1), (-1, -1)), + StepAction( + 561, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -597,9 +656,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(554, 4, [], (-1, -1), (-1, -1)), + StepAction(554, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)), - StepAction(562, 4, [], (-1, -1), (-1, -1)), + StepAction( + 562, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -609,9 +674,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(555, 4, [], (-1, -1), (-1, -1)), + StepAction(555, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)), - StepAction(563, 4, [], (-1, -1), (-1, -1)), + StepAction( + 563, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), ], ), @@ -621,9 +692,15 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: num_accepted_tokens=4, step_actions=[ StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)), - StepAction(556, 4, [], (-1, -1), (3, 0)), + StepAction(556, 4, [1, 1, 1, 1], (-1, -1), (3, 0)), StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), - StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 564, + 4, + [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_block_size": TestConfig( @@ -642,7 +719,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)), - StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), + StepAction( + 560 * 2, + 4, + [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1], + (1, 2), + (-1, -1), + ), ], ), "prompt_2_block_size_10": TestConfig( @@ -652,7 +735,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)), - StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 560 * 2 + 10, + 4, + [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_3_block_size": TestConfig( @@ -662,7 +751,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)), - StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)), + StepAction( + 560 * 3, + 4, + [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1], + (2, 3), + (-1, -1), + ), ], ), "prompt_3_block_size_10": TestConfig( @@ -672,7 +767,13 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: step_actions=[ StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)), StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)), - StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)), + StepAction( + 560 * 3 + 10, + 4, + [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1], + (-1, -1), + (-1, -1), + ), ], ), "prompt_10_block_size": TestConfig( @@ -691,14 +792,18 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: StepAction( 560 * 9, 560, - [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], (8, 9), (-1, -1), ), StepAction( 560 * 10, 4, - [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], (9, 10), (-1, -1), ), @@ -720,34 +825,28 @@ def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]: StepAction( 560 * 9, 560 + 10, - [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1] + if a + else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1], (8, 10), (-1, -1), ), ], ), } - return tests -def fill_following_kv_cache_block_ids(test_config: TestConfig) -> None: - for step_action_prev, step_action_next in zip( - test_config.step_actions[:-1], test_config.step_actions[1:] - ): - if len(step_action_next.kv_cache_block_ids) == 0: - step_action_next.kv_cache_block_ids = ( - step_action_prev.kv_cache_block_ids.copy() - ) - - -@create_new_process_for_each_test() -def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): +def _run_mamba_prefix_cache_mrv1( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + global async_scheduling_mode + async_scheduling_mode = async_scheduling run_ref_mamba_state_in_subprocess() apply_patch(monkeypatch) prompt_dataset = datasets.load_dataset("heheda/a_long_article") full_prompt = prompt_dataset["train"][0]["text"] - tests = get_mamba_prefix_cache_step_configs() + tests = get_mamba_prefix_cache_step_configs(async_scheduling) engine = LLM( model=MODEL, @@ -761,6 +860,7 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): }, max_num_batched_tokens=3072, hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, + async_scheduling=async_scheduling, seed=42, ) global prompt_token_ids @@ -777,7 +877,6 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): ) global cur_step_action_idx cur_step_action_idx = 0 - fill_following_kv_cache_block_ids(test_config) global step_actions step_actions = test_config.step_actions _ = engine.generate( @@ -800,7 +899,20 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): @create_new_process_for_each_test() -def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): +def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False) + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True) + + +def _run_mamba_prefix_cache_mrv2( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + global async_scheduling_mode + async_scheduling_mode = async_scheduling monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") envs.disable_envs_cache() @@ -998,11 +1110,12 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): max_num_batched_tokens=3072, max_model_len=BLOCK_SIZE * 12, hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, + async_scheduling=async_scheduling, seed=42, ) try: - tests = get_mamba_prefix_cache_step_configs() + tests = get_mamba_prefix_cache_step_configs(async_scheduling) global step_actions global cur_step_action_idx @@ -1010,7 +1123,6 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): for test_name, test_config in tests.items(): num_accepted_tokens = test_config.num_accepted_tokens cur_step_action_idx = 0 - fill_following_kv_cache_block_ids(test_config) step_actions = test_config.step_actions sampling_params = SamplingParams( temperature=0.0, @@ -1053,3 +1165,13 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): del engine torch.accelerator.empty_cache() cleanup_dist_env_and_memory() + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=False) + + +@create_new_process_for_each_test() +def test_mamba_prefix_cache_mrv2_async(monkeypatch: pytest.MonkeyPatch): + _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=True) diff --git a/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py b/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py new file mode 100644 index 00000000000..18e51584fa8 --- /dev/null +++ b/tests/v1/e2e/general/test_rocm_aiter_custom_ar.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode +from vllm.envs import disable_envs_cache +from vllm.platforms import current_platform + +from ....conftest import VllmRunner +from ....utils import ( + assert_rocm_custom_allreduce_backend_state_on_worker, + multi_gpu_test, +) + +PROMPTS = ["Hello, my name is", "The capital of France is"] + + +def _run_generation( + vllm_runner: type[VllmRunner], + monkeypatch: pytest.MonkeyPatch, + compilation_config: CompilationConfig, + *, + model: str, + max_tokens: int, + use_aiter_custom_ar: bool, + quick_reduce_quantization: str, +) -> list[tuple[list[int], str]]: + with monkeypatch.context() as m: + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv( + "VLLM_ROCM_USE_AITER_CUSTOM_AR", + "1" if use_aiter_custom_ar else "0", + ) + m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quick_reduce_quantization) + disable_envs_cache() + rocm_aiter_ops.refresh_env_variables() + + with vllm_runner( + model, + dtype="half", + tensor_parallel_size=2, + compilation_config=compilation_config, + max_model_len=256, + max_num_seqs=len(PROMPTS), + gpu_memory_utilization=0.7, + ) as llm: + llm.get_llm().collective_rpc( + assert_rocm_custom_allreduce_backend_state_on_worker, + args=(use_aiter_custom_ar, quick_reduce_quantization), + ) + + return llm.generate_greedy(PROMPTS, max_tokens) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only") +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed") +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "quick_reduce_quantization", + [ + pytest.param("FP", id="quick-reduce-on"), + pytest.param("NONE", id="quick-reduce-off"), + ], +) +@pytest.mark.parametrize( + "cudagraph_mode", + [ + pytest.param(CUDAGraphMode.NONE, id="cudagraph-none"), + pytest.param(CUDAGraphMode.FULL, id="cudagraph-full"), + ], +) +@pytest.mark.parametrize( + "model,max_tokens", + [ + pytest.param("facebook/opt-125m", 8, id="opt-125m"), + ], +) +def test_rocm_aiter_custom_ar_e2e( + vllm_runner: type[VllmRunner], + monkeypatch: pytest.MonkeyPatch, + cudagraph_mode: CUDAGraphMode, + quick_reduce_quantization: str, + model: str, + max_tokens: int, +): + compilation_mode = ( + CompilationMode.NONE + if cudagraph_mode == CUDAGraphMode.NONE + else CompilationMode.VLLM_COMPILE + ) + compilation_config = CompilationConfig( + mode=compilation_mode, + cudagraph_mode=cudagraph_mode, + ) + + baseline_generations = _run_generation( + vllm_runner, + monkeypatch, + compilation_config, + model=model, + max_tokens=max_tokens, + use_aiter_custom_ar=False, + quick_reduce_quantization=quick_reduce_quantization, + ) + aiter_custom_ar_generations = _run_generation( + vllm_runner, + monkeypatch, + compilation_config, + model=model, + max_tokens=max_tokens, + use_aiter_custom_ar=True, + quick_reduce_quantization=quick_reduce_quantization, + ) + + assert aiter_custom_ar_generations == baseline_generations diff --git a/tests/v1/engine/utils.py b/tests/v1/engine/utils.py index 013e73bd8e4..324c9c9ad56 100644 --- a/tests/v1/engine/utils.py +++ b/tests/v1/engine/utils.py @@ -7,14 +7,14 @@ from typing import TypeAlias import numpy as np import torch -from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast +from transformers import PythonBackend, TokenizersBackend from vllm.engine.arg_utils import EngineArgs from vllm.v1.engine import EngineCoreOutput, FinishReason from vllm.v1.metrics.stats import PrefillStats from vllm.v1.outputs import LogprobsLists, LogprobsTensors -GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast +GeneralTokenizerType: TypeAlias = PythonBackend | TokenizersBackend # Number of sample logprobs to request when testing sample logprobs NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5 @@ -193,7 +193,7 @@ def _create_random_top_token_test_matrix( def decode_token( tok_id: int, - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, ) -> str: """Reproduce the process of detokenizing a token for testing purposes. @@ -210,7 +210,7 @@ def decode_token( def generate_dummy_sample_logprobs( sampled_tokens_list: list, num_logprobs: int, - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, ) -> list[tuple[list[int], list[float], int]]: """Generate dummy sample logprobs @@ -259,7 +259,7 @@ def generate_dummy_sample_logprobs( def generate_dummy_prompt_logprobs_tensors( prompt_tokens_list: list, num_logprobs: int, - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, ) -> LogprobsTensors: """Generate dummy prompt logprobs tensors diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 9ce225c3a49..9732f467377 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -28,6 +28,7 @@ hybrid_ssm_configs=( # GDN (Qwen3.5) "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENFORCE_EAGER=0 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--spec-method,mtp,--spec-tokens,1" # Mamba1 (Jamba) "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ai21labs/AI21-Jamba2-3B VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" ) diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index 0e7f6af7e38..9ae15ad0e8e 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -76,6 +76,7 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2} PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128} DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128} +ENFORCE_EAGER=${ENFORCE_EAGER:-1} # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048) VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} @@ -157,11 +158,13 @@ run_tests_for_model() { VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ vllm serve $model_name \ --port $PORT \ - --enforce-eager \ --block-size ${PREFILL_BLOCK_SIZE} \ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ --tensor-parallel-size $PREFILLER_TP_SIZE \ --kv-transfer-config '$KV_CONFIG_P'" + if [[ "$ENFORCE_EAGER" == "1" ]]; then + BASE_CMD="${BASE_CMD} --enforce-eager" + fi if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" for arg in "${extra_args[@]}"; do @@ -206,10 +209,12 @@ run_tests_for_model() { VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \ vllm serve $model_name \ --port $PORT \ - --enforce-eager \ --block-size ${DECODE_BLOCK_SIZE} \ --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ --kv-transfer-config '$KV_CONFIG_D'" + if [[ "$ENFORCE_EAGER" == "1" ]]; then + BASE_CMD="${BASE_CMD} --enforce-eager" + fi if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" for arg in "${extra_args[@]}"; do diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index a34f07edc97..d3fee8d6ba5 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -9,8 +9,10 @@ PREFILL_GPU_ID=${PREFILL_GPU_ID:-0} DECODE_GPU_ID=${DECODE_GPU_ID:-1} MODEL=${MODEL:-"ibm-granite/granite-4.0-h-tiny"} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8} +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-FLASHINFER} -echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL)" +echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL, backend=$ATTENTION_BACKEND)" KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}' @@ -36,6 +38,14 @@ cleanup_instances() { cleanup_instances +EXTRA_ARGS=() +if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a EXTRA_ARGS <<< "$VLLM_SERVE_EXTRA_ARGS" +fi +if [[ -n "$ATTENTION_BACKEND" ]]; then + EXTRA_ARGS+=(--attention-backend "$ATTENTION_BACKEND") +fi + # Start prefill instance PREFILL_PORT=8001 CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \ @@ -51,8 +61,8 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ - --attention-backend FLASHINFER \ - --kv-transfer-config "$KV_CONFIG" & + --kv-transfer-config "$KV_CONFIG" \ + "${EXTRA_ARGS[@]}" & # Start decode instance DECODE_PORT=8002 @@ -69,8 +79,8 @@ vllm serve $MODEL \ --trust-remote-code \ --enable-prefix-caching \ --mamba-cache-mode all \ - --attention-backend FLASHINFER \ - --kv-transfer-config "$KV_CONFIG" & + --kv-transfer-config "$KV_CONFIG" \ + "${EXTRA_ARGS[@]}" & echo "Waiting for prefill instance on port $PREFILL_PORT..." wait_for_server "$PREFILL_PORT" diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh index 2e71858983e..dae632dfce8 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh @@ -18,6 +18,7 @@ # Environment variables: # MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) # GPU_MEMORY_UTILIZATION - GPU memory fraction (default: 0.6) +# ATTENTION_BACKEND - optional attention backend for vllm serve # VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve # SKIP_CROSS_LAYERS - set to 1 to skip the cross-layer layout test # SKIP_NORMAL_LAYOUT - set to 1 to skip the normal layout test @@ -34,9 +35,11 @@ fi GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.6} BLOCK_SIZE=${BLOCK_SIZE:-128} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-} VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") # ── KV transfer configs ───────────────────────────────────────────────── @@ -139,6 +142,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Start decode instance ── @@ -161,6 +167,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Wait for servers ── diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh index a80950b3413..de6c9abcc6b 100755 --- a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh @@ -19,6 +19,7 @@ # MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) # KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB) # BLOCK_SIZE - KV cache block size (default: 128) +# ATTENTION_BACKEND - optional attention backend for vllm serve # VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve set -xe @@ -34,9 +35,11 @@ fi KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048} BLOCK_SIZE=${BLOCK_SIZE:-128} +ATTENTION_BACKEND=${ATTENTION_BACKEND:-} VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # ── KV transfer config ────────────────────────────────────────────────── @@ -110,6 +113,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Start decode instance ── @@ -133,6 +139,9 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} $arg" done fi + if [[ -n "$ATTENTION_BACKEND" ]]; then + BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND" + fi eval "$BASE_CMD &" # ── Wait for servers ── diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index 9e5f564bba5..beb639724e8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -7,7 +7,7 @@ import torch from tests.v1.kv_connector.unit.utils import create_vllm_config from vllm.config import KVEventsConfig, KVTransferConfig -from vllm.distributed.kv_events import BlockRemoved, BlockStored +from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( OffloadingEventGroupSpec, OffloadingEventsTracker, @@ -28,10 +28,9 @@ from vllm.v1.kv_offload.base import ( OffloadKey, make_offload_key, ) -from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec -_CPU_MEDIUM = CPULoadStoreSpec.medium() +_CPU_MEDIUM = MEDIUM_CPU _FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec( kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value, kv_cache_spec_sliding_window=None, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index 6f36a6c8149..fb6b785c886 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -10,9 +10,11 @@ from prometheus_client import Counter, Gauge, Histogram from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, OffloadPromMetrics, + _ConnectorMetricName, _MetricType, _StatsKey, _TransferMetricName, + get_connector_metric_definitions, ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import ( OffloadingConnector, @@ -37,6 +39,42 @@ MY_COUNTER = "my_counter" MY_LABEL = "my_label" +def test_connector_metric_histogram_buckets(): + metadata = get_connector_metric_definitions() + + sync_delay = metadata[_ConnectorMetricName.LOOKUP_SYNC_DELAY] + assert isinstance(sync_delay, OffloadingHistogramMetadata) + assert sync_delay.buckets == ( + 0.00001, + 0.00005, + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + ) + + async_delay = metadata[_ConnectorMetricName.LOOKUP_ASYNC_DELAY] + assert isinstance(async_delay, OffloadingHistogramMetadata) + assert async_delay.buckets == ( + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + 5, + 10, + ) + + class _FakeMetric: def __init__(self, **kwargs): self.kwargs = kwargs 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 ad5792e6c3c..6c55b91d8da 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -11,6 +11,10 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import ( to_keys, ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, + _ConnectorMetricName, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, @@ -32,6 +36,72 @@ from vllm.v1.kv_offload.base import ( from vllm.v1.request import RequestStatus +def _reduce_kv_connector_stats(runner): + reduced: dict[str, int | float] = {} + for payload in runner.kv_connector_stats: + stats = ( + payload + if hasattr(payload, "reduce") + else OffloadingConnectorStats(data=payload) + ) + for key, value in stats.reduce().items(): + reduced[key] = reduced.get(key, 0) + value + return reduced + + +def test_scheduler_reports_allocation_failure(request_runner): + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + runner.new_request(token_ids=[0] * 4) + runner.manager.prepare_store.side_effect = lambda keys, req_context: None + + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + reduced = _reduce_kv_connector_stats(runner) + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 + + +def test_scheduler_reports_lookup_sync_delay(request_runner): + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + runner.new_request(token_ids=[1] * 4) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + reduced = _reduce_kv_connector_stats(runner) + assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_count"] == 1 + assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0 + + +def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): + """A deferred lookup reports its async delay once it resolves.""" + runner = request_runner( + block_size=4, + num_gpu_blocks=10, + async_scheduling=False, + ) + runner.manager.lookup.side_effect = [LookupResult.RETRY, LookupResult.MISS] + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + + runner.new_request(token_ids=[1] * 4) + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + + reduced = _reduce_kv_connector_stats(runner) + assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_count"] == 1 + assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0 + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_offloading_connector(request_runner, async_scheduling: bool): block_size = 4 @@ -2146,15 +2216,17 @@ class TestEagle: # ------------------------------------------------------------------- @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_full_attn_store_excludes_trailing_block( + def test_full_attn_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): - """Eagle full-attention group stores all blocks except the trailing - one. + """Eagle full-attention group excludes the trailing block only while + decoding. Setup: 2 groups — group 0 is normal full-attention, group 1 is - eagle full-attention. With a 3-block prompt, group 1 should store - only blocks 0 and 1, skipping block 2 (the volatile tail). + eagle full-attention. With a 3-block prompt, group 1 stores all 3 + prompt blocks at the end of prefill (the trailing prompt block is + stable), but skips block 3 once it fills with decoded tokens (its + draft-layer KV is volatile until the next block starts). """ block_size = 4 block_size_factor = 1 @@ -2200,23 +2272,27 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) + # 4 decoded tokens fill block 3 entirely with decode tokens (one + # extra token so the block is stored under async scheduling too). runner.run( - decoded_tokens=[EOS_TOKEN_ID], + decoded_tokens=[1, 1, 1, 1, 1, EOS_TOKEN_ID], expected_stored=( (0, 0), (0, 1), (0, 2), + (0, 3), (1, 0), (1, 1), + (1, 2), ), ) @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_sw_store_excludes_trailing_block( + def test_sw_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): - """Eagle sliding-window group stores all blocks except the trailing - one.""" + """Eagle sliding-window group stores all prompt blocks but excludes + the trailing block while decoding.""" block_size = 4 sliding_window = 8 num_gpu_blocks = 100 @@ -2251,15 +2327,18 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) + # 4 decoded tokens fill block 3 entirely with decode tokens. runner.run( - decoded_tokens=[EOS_TOKEN_ID], - expected_stored=((0, 0), (0, 1)), + decoded_tokens=[1, 1, 1, 1, EOS_TOKEN_ID], + expected_stored=((0, 0), (0, 1), (0, 2)), ) @pytest.mark.parametrize("async_scheduling", [True, False]) - def test_single_block_nothing_stored(self, request_runner, async_scheduling: bool): - """An eagle group with only one block stores nothing: that block is - the tail.""" + def test_single_block_stored_at_end_of_prefill( + self, request_runner, async_scheduling: bool + ): + """An eagle group with a single-block prompt stores it at the end of + prefill: prompt blocks are stable, so no tail is held back.""" block_size = 4 block_size_factor = 1 offloaded_block_size = block_size * block_size_factor @@ -2290,17 +2369,80 @@ class TestEagle: runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) - runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=()) - runner.manager.prepare_store.assert_not_called() + runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=((0, 0),)) + + @pytest.mark.parametrize("async_scheduling", [True, False]) + def test_multichunk_store_no_interior_holes( + self, request_runner, async_scheduling: bool + ): + """Eagle store must not drop interior blocks across prefill chunks. + + Regression: the trailing-block exclusion (num_blocks - 1) was applied + when collecting keys, but next_stored_block_idx advanced by the + non-decremented count, so the trailing block of every chunked-prefill + chunk was skipped and never re-considered. With the harness chunk budget + (1000 tokens) and block_size 4, a prompt longer than one chunk lost the + block at the chunk boundary, leaving a permanent gap that caps prefix + reuse at the first hole. Only the trailing decode block may be held + back; all other blocks must be stored exactly once (no duplicates from + next_stored_block_idx regressing at the prefill->decode transition). + """ + block_size = 4 + block_size_factor = 1 + offloaded_block_size = block_size * block_size_factor + num_gpu_blocks = 1000 + + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=True, + ), + ] + runner = request_runner( + block_size=block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + block_size_factor=block_size_factor, + ) + assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group + + # Prompt spans more than one prefill chunk (chunk budget 1000 tokens). + num_blocks = 256 + runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + # Decode a few non-EOS tokens so prefill completes across both chunks + # before the request finishes. + runner._run([1, 1, 1, 1, EOS_TOKEN_ID], complete_transfers=True) + + offsets = sorted( + b.request_block_offset + for t in runner.completed_stores + for b in t.gpu_blocks + ) + # The stored blocks must be contiguous from 0: no interior block is + # dropped at a chunk boundary. (The bug left a gap at offloaded block + # 249, the tail of the first 1000-token chunk.) + assert offsets == list(range(len(offsets))), ( + f"interior hole in stored blocks: {offsets}" + ) @pytest.mark.parametrize("async_scheduling", [True, False]) def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool): """Eagle group constrains load: convergence tightens both groups. - Store 3 offloaded blocks per group (eagle group skips tail → stores - 2). Then a new request loads from CPU. The eagle group's post-pop hit - (2) does not tighten below group 0's hit (3), so both groups load - normally. + Store 3 offloaded blocks per group (all prompt blocks, so the eagle + group stores all 3 as well). Then a new request loads from CPU. The + eagle group pops its trailing hit block on load, tightening the hit + to 2 blocks for both groups. """ block_size = 4 block_size_factor = 1 @@ -2349,6 +2491,7 @@ class TestEagle: (0, 2), (1, 0), (1, 1), + (1, 2), ), ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index c2884649bdd..73ea5e2be1d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -73,10 +73,6 @@ class MockLoadStoreSpec(LoadStoreSpec): def __init__(self, offload_keys: Iterable[OffloadKey]): self.offload_keys: list[OffloadKey] = list(offload_keys) - @staticmethod - def medium() -> str: - return "Mock" - def __repr__(self) -> str: return repr(self.offload_keys) @@ -133,6 +129,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager = MagicMock(spec=OffloadingManager) self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.manager.lookup.return_value = LookupResult.MISS + self.manager.get_stats.return_value = None self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() @@ -335,6 +332,7 @@ class RequestRunner: self.completed_loads: list[TransferSummary] = [] self.completed_stores: list[TransferSummary] = [] self.flushed_gpu_blocks: set[GPUBlock] = set() + self.kv_connector_stats: list[Any] = [] # block_id -> GPUBlock self.gpu_blocks: dict[int, GPUBlock] = {} @@ -348,6 +346,12 @@ class RequestRunner: slot_mapping={}, ) + def _record_kv_connector_stats(self, engine_outputs: dict[int, Any]) -> None: + for output in engine_outputs.values(): + scheduler_stats = output.scheduler_stats + if scheduler_stats is not None and scheduler_stats.kv_connector_stats: + self.kv_connector_stats.append(scheduler_stats.kv_connector_stats) + def new_request( self, token_ids: list[int], @@ -524,13 +528,17 @@ class RequestRunner: if self.async_scheduling: # in async scheduling we update the output of the previous step if prev_model_runner_output is not None: - self.scheduler.update_from_output( + engine_outputs = self.scheduler.update_from_output( prev_scheduler_output, prev_model_runner_output ) + self._record_kv_connector_stats(engine_outputs) prev_scheduler_output = scheduler_output prev_model_runner_output = model_runner_output else: - self.scheduler.update_from_output(scheduler_output, model_runner_output) + engine_outputs = self.scheduler.update_from_output( + scheduler_output, model_runner_output + ) + self._record_kv_connector_stats(engine_outputs) if post_step_fn is not None: post_step_fn() @@ -546,9 +554,10 @@ class RequestRunner: if token_id is None: if self.async_scheduling: # sample last token - self.scheduler.update_from_output( + engine_outputs = self.scheduler.update_from_output( prev_scheduler_output, prev_model_runner_output ) + self._record_kv_connector_stats(engine_outputs) break self._parse_transfers() diff --git a/tests/v1/kv_connector/unit/test_hf3fs_connector.py b/tests/v1/kv_connector/unit/test_hf3fs_connector.py index cd525e23b14..94bb94c6fbd 100644 --- a/tests/v1/kv_connector/unit/test_hf3fs_connector.py +++ b/tests/v1/kv_connector/unit/test_hf3fs_connector.py @@ -33,7 +33,7 @@ def hf3fs_stats(): def _make_cuda_event(): """Return a real CUDA event when available, otherwise a MagicMock.""" if torch.cuda.is_available(): - return torch.Event() + return torch.cuda.Event() return MagicMock() diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 89a0374b462..e16d6d378ea 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -6,6 +6,7 @@ from dataclasses import dataclass import numpy as np import pytest +from vllm.distributed.kv_events import MEDIUM_CPU from vllm.v1.kv_offload.base import ( LoadStoreSpec, LookupResult, @@ -99,7 +100,7 @@ def verify_events( stores: list[set[OffloadKey]] = [] evictions: list[set[OffloadKey]] = [] for event in events: - assert event.medium == CPULoadStoreSpec.medium() + assert event.medium == MEDIUM_CPU if event.removed: evictions.append(set(event.keys)) else: @@ -114,6 +115,25 @@ def verify_events( assert tuple(stores) == to_key_sets(expected_stores) +def test_cpu_eviction_removed_precedes_stored(): + """An eviction is announced before the store that reuses its capacity.""" + manager = make_cpu_manager(num_blocks=2, enable_events=True) + + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + list(manager.take_events()) + + manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) + manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX) + + events = list(manager.take_events()) + removed_idx = [i for i, event in enumerate(events) if event.removed] + stored_idx = [i for i, event in enumerate(events) if not event.removed] + assert removed_idx and stored_idx, events + assert max(removed_idx) < min(stored_idx) + assert all(event.medium == manager.medium for event in events) + + @pytest.mark.parametrize("eviction_policy", ["lru", "arc"]) def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): """ @@ -224,6 +244,64 @@ def test_cpu_manager_reports_cache_usage_gauge(): check_usage_stats(manager, 0.0) +def test_cpu_manager_reports_allocation_size_histogram(): + manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) + manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) + + stats = manager.get_stats() + + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 2 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3 + + # The cache-usage gauge is always reported, so get_stats() never returns + # None, but the histogram has nothing new once its samples are consumed. + second_stats = manager.get_stats() + assert second_stats is not None + assert f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count" not in ( + second_stats.reduce() + ) + + +def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch): + manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + + def fail_allocate_blocks(keys): + raise RuntimeError("allocation failed") + + monkeypatch.setattr(manager, "_allocate_blocks", fail_allocate_blocks) + + with pytest.raises(RuntimeError, match="allocation failed"): + manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) + + stats = manager.get_stats() + + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3 + + +def test_cpu_manager_reports_allocation_size_on_eviction_failure(): + manager = make_cpu_manager(num_blocks=1, cache_policy="lru") + + manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) + manager.get_stats() + + assert manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) is None + + stats = manager.get_stats() + + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1 + assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1 + + def test_cpu_manager(): """ Tests CPUOffloadingManager with lru policy. @@ -819,3 +897,26 @@ def test_evictable_cache_block_count(): manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX) # cache state [10, 11, 14, 15] <- all blocks idle assert manager._num_evictable_cache_blocks == 4 + + +def test_touch_forwards_req_context_to_policy(monkeypatch): + """Regression: CPUOffloadingManager.touch forwards ReqContext to policy.""" + manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + received = [] + + def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + received.append((list(keys), req_context)) + + monkeypatch.setattr(manager._policy, "touch", spy_touch) + + keys = to_keys([1, 2]) + ctx = make_req_context( + req_id="test-req", + kv_transfer_params={"test_param": "test_value"}, + ) + + manager.touch(keys, ctx) + + assert len(received) == 1 + assert received[0][0] == keys + assert received[0][1] is ctx diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 543ee44330d..b051f44ecfd 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -21,7 +21,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheTensor, ) -from vllm.v1.kv_offload.base import OffloadingSpec +from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec from vllm.v1.kv_offload.factory import OffloadingSpecFactory from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec @@ -248,8 +248,8 @@ def test_duplicate_registration_raises(): # --------------------------------------------------------------------------- -def test_build_metric_definitions_empty_below_threshold(): - """store_threshold < 2 → only base metric (no stores_skipped).""" +def test_build_metric_definitions_below_threshold(): + """store_threshold < 2 keeps stores_skipped disabled.""" from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics config = _make_vllm_config(store_threshold=1) @@ -258,6 +258,32 @@ def test_build_metric_definitions_empty_below_threshold(): config.kv_transfer_config.kv_connector_extra_config ) assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics + assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics + + +def test_build_metric_definitions_allocation_size_histogram(): + """CPU allocation size is always reported as a histogram.""" + from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics + + config = _make_vllm_config(store_threshold=0) + spec_cls = OffloadingSpecFactory.get_spec_cls(config) + metrics = spec_cls.build_metric_definitions( + config.kv_transfer_config.kv_connector_extra_config + ) + metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE] + assert isinstance(metadata, OffloadingHistogramMetadata) + assert metadata.buckets == ( + 1, + 4, + 16, + 64, + 256, + 1024, + 4096, + 16384, + 65536, + 262144, + ) def test_build_metric_definitions_returns_counter_at_threshold(): diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 3df9d30691d..f429d6cf62d 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -25,6 +25,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult +from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager # --------------------------------------------------------------------------- @@ -418,3 +419,49 @@ class TestMockObjTierShutdown: tier, _ = _make_tier(num_blocks=4) tier.shutdown() tier.shutdown() # must not raise + + +class TestObjStoreConfig: + def test_explicit_credentials_included(self): + cfg = ObjStoreConfig( + bucket="b", + endpoint_override="ep", + access_key="ak", + secret_key="sk", + ) + params = cfg.to_nixl_params() + assert params["access_key"] == "ak" + assert params["secret_key"] == "sk" + + def test_credentials_omitted_when_empty(self): + cfg = ObjStoreConfig(bucket="b", endpoint_override="ep") + params = cfg.to_nixl_params() + assert "access_key" not in params + assert "secret_key" not in params + assert "session_token" not in params + assert "region" not in params + assert params["bucket"] == "b" + assert params["endpoint_override"] == "ep" + + def test_session_token_and_region_included(self): + cfg = ObjStoreConfig( + bucket="b", + endpoint_override="ep", + access_key="ak", + secret_key="sk", + session_token="tok", + region="us-east-1", + ) + params = cfg.to_nixl_params() + assert params["session_token"] == "tok" + assert params["region"] == "us-east-1" + + def test_ca_bundle_included_when_set(self): + cfg = ObjStoreConfig( + bucket="b", + endpoint_override="ep", + ca_bundle="/path/to/ca.pem", + ) + params = cfg.to_nixl_params() + assert params["ca_bundle"] == "/path/to/ca.pem" + assert "access_key" not in params diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index fe9283a8119..13460f68d7c 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( from vllm.v1.kv_offload.base import ( LookupResult, OffloadingCounterMetadata, + OffloadingEvent, OffloadKey, OffloadPolicy, ReqContext, @@ -242,6 +243,24 @@ class TestTieringOffloadingManager: if req_context.req_id not in self.manager._req_state: self.manager.on_new_request(req_context) + def test_take_events_aggregates_tier_owned_events(self, manager_setup): + primary_event = OffloadingEvent(to_keys([1]), "CPU", removed=False) + secondary_event1 = OffloadingEvent(to_keys([2]), "tier-1", removed=False) + secondary_event2 = OffloadingEvent(to_keys([3]), "tier-2", removed=True) + + self.primary_tier.take_events = MagicMock(return_value=[primary_event]) + self.secondary_tier1.take_events = MagicMock(return_value=[secondary_event1]) + self.secondary_tier2.take_events = MagicMock(return_value=[secondary_event2]) + + assert list(self.manager.take_events()) == [ + primary_event, + secondary_event1, + secondary_event2, + ] + self.primary_tier.take_events.assert_called_once_with() + self.secondary_tier1.take_events.assert_called_once_with() + self.secondary_tier2.take_events.assert_called_once_with() + def test_basic_store_to_primary(self, manager_setup): """Test basic store operation to primary tier.""" blocks = to_keys(range(3)) diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index ec150272792..fba240fea6a 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -85,6 +85,7 @@ def _model_config(vocab_size: int = 10): return SimpleNamespace( max_logprobs=20, logits_processors=None, + is_diffusion=False, get_vocab_size=lambda: vocab_size, ) diff --git a/tests/v1/spec_decode/test_dynamic_sd.py b/tests/v1/spec_decode/test_dynamic_sd.py index fe9f30ba25f..8d46b241c62 100644 --- a/tests/v1/spec_decode/test_dynamic_sd.py +++ b/tests/v1/spec_decode/test_dynamic_sd.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Regression tests for the Dynamic SD batch-size schedule helpers.""" +import logging + import pytest from tests.v1.core.utils import create_requests, create_scheduler @@ -187,6 +189,33 @@ def test_scheduler_falls_back_to_static_k_when_dsd_not_configured(): assert output.num_spec_tokens_to_schedule == 3 +def test_dynamic_sd_is_disabled_with_data_parallel(caplog_vllm): + with caplog_vllm.at_level(logging.WARNING, logger="vllm"): + scheduler = create_scheduler( + max_num_seqs=256, + max_num_batched_tokens=2560, + num_speculative_tokens=3, + num_speculative_tokens_per_batch_size=[ + (1, 16, 3), + (64, 128, 2), + (256, 4096, 0), + ], + data_parallel_size=2, + ) + + speculative_config = scheduler.vllm_config.speculative_config + assert speculative_config is not None + assert speculative_config.num_speculative_tokens_per_batch_size is None + assert scheduler.dynamic_sd_lookup is None + assert "Dynamic speculative decoding is not supported with data parallelism" in ( + caplog_vllm.text + ) + + output = _add_requests_and_schedule(scheduler, 256) + assert len(output.num_scheduled_tokens) == 256 + assert output.num_spec_tokens_to_schedule == 3 + + def test_scheduler_uses_static_k_when_no_requests_are_scheduled(): scheduler = _make_scheduler_with_dynamic_sd( [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], diff --git a/tests/v1/structured_output/test_backend_guidance.py b/tests/v1/structured_output/test_backend_guidance.py index ca8c9b0d785..edcdd983c0e 100644 --- a/tests/v1/structured_output/test_backend_guidance.py +++ b/tests/v1/structured_output/test_backend_guidance.py @@ -17,7 +17,7 @@ from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.structured_output.backend_guidance import GuidanceBackend from vllm.v1.structured_output.backend_types import StructuredOutputOptions -TOKENIZER = "gpt2" +TOKENIZER = "openai-community/gpt2" @pytest.fixture(scope="module") diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 00722ccc244..7490cbb27be 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -18,6 +18,7 @@ from vllm.config import ( VllmConfig, set_current_vllm_config, ) +from vllm.config.reasoning import ReasoningConfig from vllm.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, @@ -255,6 +256,31 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids( + dist_init, +): + vllm_config = get_vllm_config() + assert vllm_config.model_config.logits_processors is None + reasoning_config = ReasoningConfig( + reasoning_start_str="", reasoning_end_str="" + ) + reasoning_config._reasoning_start_token_ids = [1] + reasoning_config._reasoning_end_token_ids = [2] + vllm_config.reasoning_config = reasoning_config + + with set_current_vllm_config(vllm_config): + model_config = vllm_config.model_config + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config) + head_size = model_config.get_head_size() + vllm_config.compilation_config.static_forward_context["layer.0"] = Attention( + num_heads, head_size, 0.1 + ) + runner = GPUModelRunner(vllm_config, torch.device("cpu")) + + assert runner.input_batch.thinking_budget_state_holder is not None + assert runner.input_batch.logitsprocs_need_output_token_ids is False + + @pytest.mark.skip_global_cleanup @pytest.mark.parametrize( ("world_size", "is_last_rank", "expected_calls"), diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 31be4a8402f..cdaa644b62e 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import patch import pytest @@ -13,7 +14,12 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_VIDEO_BACKEND, ) from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.worker import startup_plan from vllm.v1.worker.gpu_worker import Worker +from vllm.v1.worker.startup_plan import ( + maybe_apply_startup_plan, + maybe_save_startup_plan, +) def _worker_with_mm_config( @@ -114,3 +120,69 @@ def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers( assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == ( available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3) ) + + +# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and +# saved by Worker.determine_available_memory / compile_or_warm_up_model. + + +def _plan_worker(config_hash="abc123", free_memory=78 * GiB_bytes, kv_bytes=None): + """The minimal Worker surface the startup-plan entry points touch.""" + return SimpleNamespace( + vllm_config=SimpleNamespace(compute_hash=lambda: config_hash), + rank=0, + parallel_config=SimpleNamespace(world_size=1), + init_snapshot=SimpleNamespace(free_memory=free_memory), + cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_bytes), + ) + + +def _plan_platform(name="NVIDIA H100 PCIe"): + return SimpleNamespace( + get_device_name=lambda device_id=0: name, + get_device_total_memory=lambda device_id=0: 80 * GiB_bytes, + get_device_capability=lambda device_id=0: (9, 0), + ) + + +@pytest.fixture +def plan_env(monkeypatch: pytest.MonkeyPatch, tmp_path): + """Enable the startup plan, isolated under a tmp cache root.""" + monkeypatch.setenv("VLLM_ENABLE_STARTUP_PLAN", "1") + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path)) + with patch.object(startup_plan, "current_platform", _plan_platform()): + yield + + +def test_startup_plan_fingerprint_sensitivity(plan_env): + """The fingerprint is the OOM-safety key: stable for identical inputs, + different for anything the profiled value depends on.""" + fp = startup_plan.compute_plan_fingerprint + base = fp(_plan_worker().vllm_config, 0, 1) + assert base == fp(_plan_worker().vllm_config, 0, 1) + assert base != fp(_plan_worker("other").vllm_config, 0, 1) + assert base != fp(_plan_worker().vllm_config, 1, 2) + with patch.object(startup_plan, "current_platform", _plan_platform("NVIDIA A100")): + assert base != fp(_plan_worker().vllm_config, 0, 1) + with patch("vllm.__version__", "0.0.0+plan-test"): + assert base != fp(_plan_worker().vllm_config, 0, 1) + + +def test_startup_plan_apply_gate(plan_env): + """Only a fingerprint-matching, memory-safe plan is ever applied.""" + maybe_save_startup_plan(_plan_worker(), 50 * GiB_bytes) + + applied = _plan_worker() + maybe_apply_startup_plan(applied) + assert applied.cache_config.kv_cache_memory_bytes == 50 * GiB_bytes + + less_memory = _plan_worker(free_memory=60 * GiB_bytes) + other_config = _plan_worker(config_hash="zzz999") + for refused in (less_memory, other_config): + maybe_apply_startup_plan(refused) + assert refused.cache_config.kv_cache_memory_bytes is None + + # An explicit --kv-cache-memory is never overridden. + explicit = _plan_worker(kv_bytes=7 * GiB_bytes) + maybe_apply_startup_plan(explicit) + assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes diff --git a/tests/v1/worker/test_xpu_model_runner.py b/tests/v1/worker/test_xpu_model_runner.py new file mode 100644 index 00000000000..5ddf490c9b4 --- /dev/null +++ b/tests/v1/worker/test_xpu_model_runner.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ``vllm.v1.worker.xpu_model_runner`` (XPU worker / CUDA shims).""" + +import pytest +import torch +from torch._dynamo.variables.torch import TorchInGraphFunctionVariable + +from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + +# XPU-only: needs distinct torch.cuda vs torch.xpu current_stream symbols. +pytestmark = pytest.mark.skipif( + not hasattr(torch, "xpu") or not hasattr(torch.xpu, "current_stream"), + reason="torch.xpu.current_stream is required", +) + + +# Child process: patched torch.cuda must not leak to other tests in the session. +@pytest.mark.forked +def test_torch_cuda_wrapper_allows_dynamo_handler_registration() -> None: + """Guard against XPU CUDA shim breaking Torch Dynamo during AOT compile. + + Before the fix, ``_torch_cuda_wrapper`` assigned + ``torch.cuda.current_stream = torch.xpu.current_stream`` (same function object). + On the first AOT/profile run, Dynamo builds its in-graph handler table and + registers ``torch.cuda.current_stream`` and ``torch.xpu.current_stream`` + separately; duplicate identity triggers:: + + AssertionError: Handler already registered for + + That surfaced as EngineCore failing in ``profile_run`` / ``_get_handlers()``. + The fix uses distinct shim callables so both can be registered. + + This test replays the post-init state (wrapper applied, patches left on + ``torch.cuda``) and checks that Dynamo's real ``_get_handlers()`` succeeds. + """ + # Same entry point as XPUModelRunner.__init__ (patches persist after exit). + with _torch_cuda_wrapper(): + pass + + # Fresh handler table build, as on first torch.compile / AOT in the worker. + # Registers torch.cuda.current_stream and torch.xpu.current_stream separately; + # if they are the same object (pre-fix alias), raises Handler already registered. + TorchInGraphFunctionVariable._get_handlers.cache_clear() + TorchInGraphFunctionVariable._get_handlers() diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index 9a67a013f1b..aec7b85d59c 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/tools/pre_commit/check_torch_cuda.py @@ -9,7 +9,7 @@ import regex as re # --------------------------------------------------------------------------- # _TORCH_CUDA_PATTERNS = [ r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device|device\()\b", - r"\btorch\.cuda\.(manual_seed|manual_seed_all|Event)\b", + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b", r"\bwith\storch\.cuda\.device\b", # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally r"\bcuda_device_count_stateless\(\)\b", @@ -21,7 +21,6 @@ ALLOWED_FILES = { "vllm/device_allocator/", "vllm/distributed/weight_transfer/ipc_engine.py", "tests/distributed/test_packed_tensor.py", - "tools/pre_commit/check_torch_cuda.py", } @@ -40,13 +39,6 @@ def scan_file(path: str) -> int: f"Found {matched_text} API call. Use set_random_seed instead." ) return 1 - if matched_text == "torch.cuda.Event": - print( - f"{path}:{line_num}: " - "\033[91merror:\033[0m " - "Found torch.cuda.Event API call. Use torch.Event instead." - ) - return 1 print( f"{path}:{line_num}: " "\033[91merror:\033[0m " # red color diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 40ccbcf9a28..bed4d8254a0 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2,12 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools from collections.abc import Callable -from contextlib import contextmanager -from typing import Protocol import torch from torch._ops import OpOverload -from torch.distributed import ProcessGroup import vllm.envs as envs from vllm.platforms import current_platform @@ -52,42 +49,6 @@ def is_aiter_found() -> bool: IS_AITER_FOUND = is_aiter_found() -class AiterCustomAllreduceProto(Protocol): - max_size: int - world_size: int - fully_connected: bool - - @contextmanager - def capture(self): ... - def close(self) -> None: ... - def fused_ar_rms( - self, - inp: torch.Tensor, - res_inp: torch.Tensor, - *, - w: torch.Tensor, - eps: float, - registered: bool = False, - use_1stage: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor]: ... - def fused_ar_rms_per_group_quant( - self, - inp: torch.Tensor, - res_inp: torch.Tensor, - *, - w: torch.Tensor, - eps: float, - group_size: int = 128, - registered: bool = False, - use_1stage: bool = False, - emit_bf16: bool = False, - ) -> ( - tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] - ): ... - def should_custom_ar(self, inp: torch.Tensor) -> bool: ... - - def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. @@ -830,6 +791,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( ) -> tuple[torch.Tensor, torch.Tensor]: aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -840,8 +802,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -854,12 +816,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms( + result = ca.custom_fused_ar_rms( input_, residual, - w=weight, - eps=epsilon, - registered=torch.cuda.is_current_stream_capturing(), + weight, + epsilon, use_1stage=use_1stage, ) assert result is not None @@ -890,6 +851,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( """ aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -900,8 +862,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -914,7 +876,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms_per_group_quant( + result = ca.fused_ar_rms_per_group_quant( input_, residual, w=weight, @@ -962,6 +924,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( """ aiter_ar = rocm_aiter_ops.get_aiter_allreduce() assert aiter_ar is not None, "aiter allreduce must be initialized" + ca = aiter_ar.aiter_ca total_bytes = input_.numel() * input_.element_size() hidden_dim = input_.shape[-1] @@ -972,8 +935,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( else: hidden_ok = False token_ok = token_num <= 80 - world_size = aiter_ar.world_size - full_nvlink = aiter_ar.fully_connected + world_size = ca.world_size + full_nvlink = ca.fully_connected if world_size == 2: size_ok = True @@ -986,7 +949,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( use_1stage = hidden_ok and token_ok and size_ok - result = aiter_ar.fused_ar_rms_per_group_quant( + result = ca.fused_ar_rms_per_group_quant( input_, residual, w=weight, @@ -1577,6 +1540,7 @@ class rocm_aiter_ops: # Check if the env variable is set _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + _CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA @@ -1598,9 +1562,6 @@ class rocm_aiter_ops: # num_shared_experts / shared_expert_scoring_func args (7-arg form). _TOPK_SOFTMAX_FUSED_SIGMOID: bool | None = None - _ALL_REDUCE_MAX_SIZE: int = 8192 * 1024 * 8 * 2 - _CUSTOM_ALL_REDUCE: AiterCustomAllreduceProto | None = None - @classmethod def refresh_env_variables(cls): """ @@ -1611,6 +1572,7 @@ class rocm_aiter_ops: you can call this function to reload the env variables. """ cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER + cls._CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA @@ -1770,6 +1732,11 @@ class rocm_aiter_ops: def is_mha_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._MHA_ENABLED + @classmethod + @if_aiter_supported + def is_custom_all_reduce_enabled(cls) -> bool: + return cls._AITER_ENABLED and cls._CUSTOM_ALL_REDUCE_ENABLED + @classmethod @if_aiter_supported def is_shuffle_kv_cache_enabled(cls) -> bool: @@ -1824,33 +1791,20 @@ class rocm_aiter_ops: return cls.is_linear_enabled() and on_gfx950() @classmethod - def initialize_aiter_allreduce( - cls, group: ProcessGroup, device: torch.device - ) -> None: - try: - from aiter.dist.device_communicators.custom_all_reduce import ( - CustomAllreduce as AiterCustomAllreduce, - ) + def get_aiter_allreduce(cls): + """Return the TP device communicator's AITER custom-allreduce if it has + one, return None otherwise + """ + from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( + AiterCustomAllreduce, + ) + from vllm.distributed.parallel_state import get_tp_group - cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device) - except Exception: - cls._CUSTOM_ALL_REDUCE = None - - @classmethod - def get_aiter_allreduce(cls) -> AiterCustomAllreduceProto | None: - return cls._CUSTOM_ALL_REDUCE - - @classmethod - def destroy_aiter_allreduce(cls) -> None: - if cls._CUSTOM_ALL_REDUCE is not None: - cls._CUSTOM_ALL_REDUCE.close() - cls._CUSTOM_ALL_REDUCE = None - - @classmethod - def get_aiter_allreduce_max_size(cls) -> int | None: - # effective max input size (based on upstream aiter version: v0.1.10.post3) - # https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/aiter/dist/device_communicators/custom_all_reduce.py#L272-L273 - return int(cls._ALL_REDUCE_MAX_SIZE / 2) + device_comm = get_tp_group().device_communicator + aiter_ar_comm = getattr(device_comm, "aiter_ar_comm", None) + return ( + aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None + ) @classmethod @if_aiter_supported @@ -2165,21 +2119,6 @@ class rocm_aiter_ops: def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501 return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501 - # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823). - @classmethod - def has_fused_allreduce_rmsnorm_quant_per_group(cls) -> bool: - """True if the running AITER build exposes the per-group AR+RMS+quant - kernel (added in ROCm/aiter PR #2823). - - The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off - this so vLLM degrades to the AR+RMS-only fusion when run against an - older aiter that lacks the per-group launcher. - """ - aiter_ar = cls.get_aiter_allreduce() - return aiter_ar is not None and hasattr( - aiter_ar, "fused_ar_rms_per_group_quant" - ) - @staticmethod def get_fused_mla_dual_rms_norm_op() -> OpOverload: return torch.ops.vllm.fused_mla_dual_rms_norm.default diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2a32a38a695..4de43e96a88 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2244,7 +2244,20 @@ def topk_sigmoid( gating_output: torch.Tensor, renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float = 1.0, ) -> None: + if current_platform.is_xpu(): + # xpu doesn't support routed_scaling_factor currently, will revert + # in next vllm-xpu-kernels bumpup + torch.ops._moe_C.topk_sigmoid( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize, + e_score_correction_bias, + ) + return torch.ops._moe_C.topk_sigmoid( topk_weights, topk_ids, @@ -2252,6 +2265,7 @@ def topk_sigmoid( gating_output, renormalize, e_score_correction_bias, + routed_scaling_factor, ) @@ -2974,6 +2988,24 @@ if hasattr(torch.ops._C, "fused_experts_cpu"): return torch.empty_like(hidden_states) +if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): + + @register_fake("_C::dynamic_4bit_int_moe") + def dynamic_4bit_int_moe_fake( + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w13_packed: torch.Tensor, + w2_packed: torch.Tensor, + hidden_size: int, + intermediate_size: int, + group_size: int, + apply_router_weight_on_input: bool, + activation_kind: int, + ) -> torch.Tensor: + return x.new_empty((x.size(0), hidden_size)) + + def fused_experts_cpu( hidden_states: torch.Tensor, w1: torch.Tensor, diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index ee706037abb..ab400028925 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -19,7 +19,6 @@ from vllm.compilation.passes.fusion.rms_quant_fusion import ( from vllm.config import VllmConfig from vllm.config.utils import Range from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce -from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -129,6 +128,29 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = { }, } +MiB = 1024 * 1024 + + +def _select_flashinfer_allreduce_use_oneshot( + workspace_backend: str, + device_capability: int | None, + world_size: int, + current_tensor_size: int, +) -> bool | None: + if workspace_backend == "mnnvl": + # FlashInfer sizes MNNVL workspaces around its own AUTO strategy. + # Forcing vLLM's per-rank threshold can request one-shot for tensors + # larger than the MNNVL one-shot workspace. + return None + + if device_capability is None: + max_one_shot_size = None + else: + max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get( + device_capability, {} + ).get(world_size) + return max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB + if flashinfer_comm is not None: from vllm.distributed.device_communicators.flashinfer_all_reduce import ( @@ -139,8 +161,6 @@ if flashinfer_comm is not None: ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern - MiB = 1024 * 1024 - def call_trtllm_fused_allreduce_norm( allreduce_in: torch.Tensor, residual: torch.Tensor, @@ -175,16 +195,6 @@ if flashinfer_comm is not None: ) curr_device = current_platform.get_device_capability() device_capability = curr_device.to_int() if curr_device is not None else None - # Get one shot input size limit for the current world size - # for the current device capability - max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get( - device_capability, # type: ignore[arg-type, unused-ignore] - {}, - ).get(world_size, None) - # Use one shot if no max size is specified - use_oneshot = ( - max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB - ) # Select workspace based on pattern: quant patterns use the # trtllm quant workspace, non-quant patterns use the primary workspace. @@ -206,6 +216,12 @@ if flashinfer_comm is not None: assert workspace is not None, ( "Flashinfer allreduce workspace must be initialized when using flashinfer" ) + use_oneshot = _select_flashinfer_allreduce_use_oneshot( + workspace.backend, + device_capability, + world_size, + current_tensor_size, + ) assert flashinfer_comm is not None if norm_out is None: norm_out = allreduce_in @@ -249,7 +265,7 @@ if flashinfer_comm is not None: # the end for the one-shot path; the two-shot path is synchronized # and keeps the early completion. Related one-shot instability in # the same kernel: flashinfer-ai/flashinfer#1223. - trigger_completion_at_end=use_oneshot + trigger_completion_at_end=(use_oneshot is True) or num_tokens > PDL_ADVANCE_LAUNCH_TOKENS, ) @@ -1473,39 +1489,23 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): ) return - device_comm = get_tp_group().device_communicator - if device_comm is None: - logger.warning_once("Device communicator is required.") - return - - ca_comm = getattr(device_comm, "ca_comm", None) + ca_comm = rocm_aiter_ops.get_aiter_allreduce() if ca_comm is None: - logger.warning_once("Custom Allreduce is required.") + logger.warning_once( + "AITER allreduce fusions are disabled " + "because AITER Custom All Reduce is not enabled. " + "Set VLLM_ROCM_USE_AITER_CUSTOM_AR=1 " + "to enable it." + ) return self.ca_comm = ca_comm - assert isinstance(ca_comm, CustomAllreduce) - - group = get_tp_group().cpu_group - rocm_aiter_ops.initialize_aiter_allreduce(group, self.device) hidden_dim = config.model_config.get_hidden_size() element_size = torch.tensor([], dtype=self.model_dtype).element_size() - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() - if max_size is None: - logger.warning("AITER allreduce fusion must be initialized") - return - - # Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim. - # Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM - # and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12 - # hidden_dim is a runtime argument. Detect the older API via the missing - # `_pool` attribute and skip fusion for unsupported sizes. - # Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590 - aiter_ar = rocm_aiter_ops.get_aiter_allreduce() + max_size = ca_comm.effective_max_size() _AITER_OLD_FUSED_AR_RMS_HIDDEN = (512, 1024, 2048, 4096) if ( - aiter_ar is not None - and not hasattr(aiter_ar, "_pool") + not ca_comm.supports_dynamic_hidden_dim and hidden_dim not in _AITER_OLD_FUSED_AR_RMS_HIDDEN ): logger.warning_once( @@ -1515,10 +1515,6 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): _AITER_OLD_FUSED_AR_RMS_HIDDEN, hidden_dim, ) - # Tear down aiter's custom-allreduce so its IPC handles don't - # race with vllm's ca_comm on the unfused fallback path. - with contextlib.suppress(Exception): - rocm_aiter_ops.destroy_aiter_allreduce() return max_token_num = max_size // (hidden_dim * element_size) @@ -1532,9 +1528,7 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): # fall back to the AR+RMS-only fusion paired with PR #41825's # standalone RMS+quant fusion -- still correct, just leaves the # post-AR quant as a standalone kernel. - supports_per_group_quant = ( - rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group() - ) + supports_per_group_quant = ca_comm.supports_per_group_quant if not supports_per_group_quant: logger.warning_once( "AITER AR+RMS+per-group-FP8-quant fusion disabled: aiter " @@ -1609,9 +1603,3 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): logger.debug( "%s Replaced %s patterns", self.__class__.__name__, self.matched_count ) - - def __del__(self) -> None: - if getattr(self, "disabled", True): - return - with contextlib.suppress(Exception): - rocm_aiter_ops.destroy_aiter_allreduce() diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 70a58004fc5..1091ec5f505 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -161,13 +161,11 @@ class CacheConfig: """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. - - In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254), + """In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254), some layers can skip tokens corresponding to prefill. This flag enables attention metadata for eligible layers to be overridden with metadata necessary for implementing this optimization in some models (e.g. Gemma3n) + NOTE: KV cache sharing is not supported for MRv2 (v2 model runner). """ kv_cache_memory_bytes: int | None = None @@ -205,6 +203,7 @@ class CacheConfig: ignored_factors = { # Runtime/derived knobs that don't affect compiled graph shape "gpu_memory_utilization", + "kv_cache_memory_bytes", "is_attention_free", "num_gpu_blocks_override", "enable_prefix_caching", diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index e98ab1b3e08..aff94fe03d2 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -147,6 +147,7 @@ LinearBackend = Literal[ "flashinfer_cudnn", "flashinfer_b12x", "marlin", + "humming", "triton", "deep_gemm", "torch", diff --git a/vllm/config/model.py b/vllm/config/model.py index e19c9408914..6e3fef0dcca 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -74,6 +74,12 @@ else: logger = init_logger(__name__) +# Process-local record of which (arch, target) model-class overrides have been +# registered in *this* process. Must not live on ModelConfig: that instance is +# pickled to each worker, so an instance flag would arrive already "registered" +# while the worker's own global ModelRegistry is still untouched. +_REGISTERED_MODEL_CLASS_OVERRIDES: set[tuple[str, str]] = set() + RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] @@ -274,6 +280,13 @@ class ModelConfig: hf_overrides: HfOverrides = field(default_factory=dict) """If a dictionary, contains arguments to be forwarded to the Hugging Face config. If a callable, it is called to update the HuggingFace config.""" + model_class_overrides: dict[str, str] = field(default_factory=dict) + """Override the model class used for one or more architectures, mapping the + architecture name to a `"module:class"` target (the same format accepted by + `ModelRegistry.register_model`). This registers the target class at runtime, + e.g. `{"GlmMoeDsaForCausalLM": + "vllm.models.deepseek_v32.nvidia.model:DeepseekV32ForCausalLM"}`. This + argument is for development and debugging purposes only.""" generation_config: str = "auto" """The folder path to the generation config. Defaults to `"auto"`, the generation config will be loaded from model path. If set to `"vllm"`, no @@ -812,8 +825,34 @@ class ModelConfig: @property def registry(self): + self._maybe_register_model_class_overrides() return me_models.ModelRegistry + def _maybe_register_model_class_overrides(self) -> None: + # Apply ``model_class_overrides`` here because this property is the + # single chokepoint through which every model-class inspect/resolve + # passes, in both the engine front-end and every worker process. The + # guard is process-local (see ``_REGISTERED_MODEL_CLASS_OVERRIDES``), so + # each worker re-registers into its own ModelRegistry exactly once + # rather than trusting a pickled-in instance flag. + if not self.model_class_overrides: + return + pending = [ + (arch, target) + for arch, target in self.model_class_overrides.items() + if (arch, target) not in _REGISTERED_MODEL_CLASS_OVERRIDES + ] + if not pending: + return + logger.warning_once( + "Applying model_class_overrides %s. This is intended for " + "development/debugging.", + str(self.model_class_overrides), + ) + for arch, target in pending: + me_models.ModelRegistry.register_model(arch, target) + _REGISTERED_MODEL_CLASS_OVERRIDES.add((arch, target)) + @property def architectures(self) -> list[str]: return self.model_arch_config.architectures @@ -999,6 +1038,7 @@ class ModelConfig: "modelopt", "modelopt_fp4", "modelopt_mxfp8", + "mxfp8", "modelopt_mixed", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) @@ -1252,7 +1292,7 @@ class ModelConfig: def is_deepseek_mla(self) -> bool: return self.model_arch_config.is_deepseek_mla - @property + @cached_property def is_mm_prefix_lm(self) -> bool: return self.model_arch_config.is_mm_prefix_lm diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 5bd528f4c9c..7c270b0c0eb 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -632,8 +632,8 @@ class ParallelConfig: # The all_reduce at the end of attention (during o_proj) means that # inputs are replicated across each rank of the tensor parallel group. - # If using expert-parallelism with DeepEP All2All ops, replicated - # tokens results in useless duplicate computation and communication. + # If using expert-parallelism, replicated tokens results in useless + # duplicate computation and communication. # # In this case, ensure the input to the experts is sequence parallel # to avoid the excess work. diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 983759bc9a5..3dc8d33cf9f 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +import functools +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Literal, get_args from pydantic import Field, SkipValidation, field_validator, model_validator @@ -46,6 +48,7 @@ MTPModelTypes = Literal[ "qwen3_5_mtp", "longcat_flash_mtp", "minimax_m3_mtp", + "bailing_hybrid_mtp", "mtp", "pangu_ultra_moe_mtp", "step3p5_mtp", @@ -303,8 +306,10 @@ class SpeculativeConfig: ) factors.append(uses_aux_hidden_states) - # The specific layers used also affect the computation graph if uses_aux_hidden_states and self.draft_model_config is not None: + factors.append(self.draft_model_config.compute_hash()) + + # The specific layers used also affect the computation graph. layer_ids = getattr( self.draft_model_config.hf_config, "eagle_aux_hidden_state_layer_ids", @@ -463,6 +468,21 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]} ) + architectures = getattr(hf_config, "architectures", []) or [] + if ( + hf_config.model_type == "bailing_hybrid" + or "BailingMoeV2_5ForCausalLM" in architectures + ): + hf_config.model_type = "bailing_hybrid_mtp" + if hf_config.model_type == "bailing_hybrid_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + { + "n_predict": n_predict, + "architectures": ["BailingMoeV25MTPModel"], + } + ) + if hf_config.model_type == "exaone_moe": hf_config.model_type = "exaone_moe_mtp" if hf_config.model_type == "exaone_moe_mtp": @@ -572,6 +592,40 @@ class SpeculativeConfig: return hf_config + @staticmethod + def _apply_composed_hf_override( + target_hf_overrides: Callable[[PretrainedConfig], PretrainedConfig], + hf_config: PretrainedConfig, + ) -> PretrainedConfig: + hf_config = SpeculativeConfig.hf_config_override(hf_config) + return target_hf_overrides(hf_config) + + @staticmethod + def compose_draft_hf_overrides( + target_hf_overrides: HfOverrides | None, + ) -> Callable[[PretrainedConfig], PretrainedConfig]: + """Build the ``hf_overrides`` for the draft ``ModelConfig``. + + Callable overrides on the target are config-to-config transforms + (e.g. test harnesses shrinking ``num_hidden_layers``) and must also + reach the draft config — otherwise a draft belonging to a large + target is instantiated at full size even when the target is shrunk. + Dict overrides are target-specific key patches and are not applied + to the draft. + + The composed override must stay picklable: the draft ``ModelConfig`` + is sent to spawned engine-core processes, so a local closure would + fail with ``Can't get local object`` during pickling. Bind the + target via ``functools.partial`` over a module-referenceable static + method instead. + """ + if not callable(target_hf_overrides): + return SpeculativeConfig.hf_config_override + + return functools.partial( + SpeculativeConfig._apply_composed_hf_override, target_hf_overrides + ) + def __post_init__(self): # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter @@ -736,7 +790,12 @@ class SpeculativeConfig: if self.method == "medusa": draft_hf_overrides = {"model_type": "medusa"} else: - draft_hf_overrides = SpeculativeConfig.hf_config_override + # Compose any callable hf_overrides set on the target so the + # draft config receives the same transform (e.g. the test + # shrink). Dict overrides stay target-only. + draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides( + self.target_model_config.hf_overrides + ) self.draft_model_config = ModelConfig( model=self.model, runner="draft", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c3f689c0a9c..062877b1694 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -501,6 +501,17 @@ class VllmConfig: return 2 return pp_size + @property + def max_in_flight_tokens(self) -> int: + # Upper bound on tokens that are scheduled but not yet settled (freed): + # every concurrent batch may hold up to a full `max_num_batched_tokens`. + # Recycling-aware KV cache specs (sliding-window, chunked-local) reserve + # for this because out-of-window blocks are freed on the processed-token + # basis, so in-flight steps transiently keep their blocks. + return ( + self.max_concurrent_batches * self.scheduler_config.max_num_batched_tokens + ) + @property def num_speculative_tokens(self) -> int: if ( @@ -531,6 +542,11 @@ class VllmConfig: ): return True + # Mixed sliding/full DFlash drafts need multiple KV groups (V2 only); + # force V2 as for dspark, since a hybrid target otherwise defaults to V1. + if self._dflash_needs_multi_kv_group(): + return True + if self.model_config is not None and self.model_config.is_diffusion: return True @@ -554,6 +570,18 @@ class VllmConfig: return True + def _dflash_needs_multi_kv_group(self) -> bool: + """Whether a DFlash draft mixes sliding-window and full attention.""" + spec = self.speculative_config + if spec is None or spec.method != "dflash": + return False + draft_config = getattr(spec, "draft_model_config", None) + if draft_config is None: + return False + layer_types = getattr(draft_config.hf_config, "layer_types", None) or [] + num_sliding = sum(lt == "sliding_attention" for lt in layer_types) + return 0 < num_sliding < len(layer_types) + def _is_default_v2_model_runner_model(self) -> bool: model_config = self.model_config if model_config is None: @@ -787,6 +815,25 @@ class VllmConfig: ) self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + def _maybe_disable_dynamic_sd_for_data_parallel(self) -> None: + speculative_config = self.speculative_config + if ( + speculative_config is None + or not speculative_config.uses_dynamic_speculative_decoding() + or self.parallel_config.data_parallel_size <= 1 + ): + return + + logger.warning_once( + "Dynamic speculative decoding is not supported with data " + "parallelism because data-parallel ranks can select different " + "speculative-token counts, causing DP divergence and deadlocks. " + "Disabling num_speculative_tokens_per_batch_size and falling back " + "to static num_speculative_tokens=%d.", + speculative_config.num_speculative_tokens, + ) + speculative_config.num_speculative_tokens_per_batch_size = None + def _post_init_kv_transfer_config(self) -> None: """Update KVTransferConfig based on top-level configs in VllmConfig. @@ -1201,6 +1248,7 @@ class VllmConfig: "optimization level defaults." ) + self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() if ( @@ -1850,8 +1898,13 @@ class VllmConfig: tp_size = self.parallel_config.tensor_parallel_size from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_enabled(): - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size() + max_size: int | None = None + if rocm_aiter_ops.is_custom_all_reduce_enabled(): + from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( # noqa: E501 + AiterCustomAllreduce, + ) + + max_size = AiterCustomAllreduce.effective_max_size() else: max_size = compilation_config.pass_config.flashinfer_max_size(tp_size) if max_size is not None and self.model_config is not None: @@ -1936,12 +1989,21 @@ class VllmConfig: if architecture is None: return + from vllm.model_executor.models import ModelRegistry from vllm.model_executor.models.config import ( MODELS_CONFIG_MAP, HybridAttentionMambaModelConfig, ) cls = MODELS_CONFIG_MAP.get(architecture, None) + if cls is None: + # `architecture` may be an HF base-model name (e.g. "Mamba2Model" + # when `architectures` is omitted); normalize to the resolved arch + # so per-arch config hooks are not skipped. + architecture = ModelRegistry._normalize_arch( + architecture, self.model_config + ) + cls = MODELS_CONFIG_MAP.get(architecture, None) if cls is not None: cls.verify_and_update_config(self) diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py new file mode 100644 index 00000000000..63e06b77a77 --- /dev/null +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-owned wrapper over AITER's ``CustomAllreduce``. + +vLLM's ``CudaCommunicator`` stores one of these as ``aiter_ar_comm`` (when +``VLLM_ROCM_USE_AITER_CUSTOM_AR`` is set) so the plain allreduce and +the fused allreduce+RMSNorm path share a single AITER instance with its IPC buffers. + +""" + +import torch +from torch.distributed import ProcessGroup + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class AiterCustomAllreduce: + # Default IPC buffer size for AITER's CustomAllreduce. + MAX_SIZE: int = 8192 * 1024 * 8 * 2 + + @classmethod + def effective_max_size(cls) -> int: + """ + Max input byte size eligible for AITER custom allreduce. + """ + return cls.MAX_SIZE // 2 + + def __init__( + self, + group: ProcessGroup, + device: int | str | torch.device, + max_size: int | None = None, + ): + from aiter.dist.device_communicators.custom_all_reduce import ( + CustomAllreduce as _AiterCustomAllreduce, + ) + + if max_size is None: + max_size = self.MAX_SIZE + + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + + @property + def aiter_ca(self): + return self._impl + + @property + def disabled(self) -> bool: + return self._impl.disabled + + def should_custom_ar(self, inp: torch.Tensor) -> bool: + return self._impl.should_custom_ar(inp) + + def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None: + return self._impl.custom_all_reduce(inp) + + def capture(self): + return self._impl.capture() + + def close(self) -> None: + self._impl.close() + + @property + def supports_dynamic_hidden_dim(self) -> bool: + """Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim. + Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM + and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12 + hidden_dim is a runtime argument. Older builds are detected via + AiterCustomAllreduce.supports_dynamic_hidden_dim; This function is used to + skip fusion for unsupported sizes on them. + Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590 + """ + return hasattr(self._impl, "_pool") + + @staticmethod + def build_supports_per_group_quant() -> bool: + """True if the running AITER build exposes the per-group AR+RMS+quant + kernel (added in ROCm/aiter PR #2823). + + The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off + this so vLLM degrades to the AR+RMS-only fusion when run against an + older aiter that lacks the per-group launcher. + """ + from aiter.dist.device_communicators.custom_all_reduce import ( + CustomAllreduce as _AiterCustomAllreduce, + ) + + return hasattr(_AiterCustomAllreduce, "fused_ar_rms_per_group_quant") + + # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823). + @property + def supports_per_group_quant(self) -> bool: + return self.build_supports_per_group_quant() diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 6fd889daeb0..9a443b7fc16 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -175,10 +175,11 @@ class DeviceCommunicatorBase: config = get_current_vllm_config_or_none() if config is not None: - # as long as we use data parallel (coupled data parallel - # where all data parallel ranks execute forward together), - # we initialize the all2all manager used in expert parallel. - use_ep = config.parallel_config.data_parallel_size > 1 + # initialize the all2all manager for DP or sequence-parallel EP. + use_ep = ( + config.parallel_config.data_parallel_size > 1 + or config.parallel_config.use_sequence_parallel_moe + ) all2all_backend = config.parallel_config.all2all_backend self.is_ep_communicator = unique_name.split(":")[0] == "ep" diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index b92015b1880..555fd0ec948 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -6,6 +6,7 @@ import torch from torch.distributed import ProcessGroup import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.device_communicators.all_reduce_utils import ( NCCL_SYMM_MEM_ALL_REDUCE_CONFIG, should_nccl_symm_mem_ag_rs, @@ -19,6 +20,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from ..utils import StatelessProcessGroup +from .aiter_custom_all_reduce import AiterCustomAllreduce from .base_device_communicator import DeviceCommunicatorBase logger = init_logger(__name__) @@ -48,16 +50,21 @@ class CudaCommunicator(DeviceCommunicatorBase): use_custom_allreduce = False use_torch_symm_mem = False use_flashinfer_allreduce = False + use_aiter_allreduce = False else: from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER + use_aiter_allreduce = use_custom_allreduce and bool( + rocm_aiter_ops.is_custom_all_reduce_enabled() + ) self.use_custom_allreduce = use_custom_allreduce self.use_torch_symm_mem = use_torch_symm_mem self.use_flashinfer_allreduce = use_flashinfer_allreduce + self.use_aiter_allreduce = use_aiter_allreduce # lazy import to avoid documentation build error from vllm.distributed.device_communicators.custom_all_reduce import ( @@ -85,6 +92,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None + self.aiter_ar_comm: AiterCustomAllreduce | None = None if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( @@ -98,7 +106,13 @@ class CudaCommunicator(DeviceCommunicatorBase): device=self.device, ) - if use_custom_allreduce and self.world_size > 1: + if self.use_aiter_allreduce and self.world_size > 1: + self.aiter_ar_comm = AiterCustomAllreduce( + group=self.cpu_group, + device=self.device, + ) + + if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: # Initialize a custom fast all-reduce implementation. self.ca_comm = CustomAllreduce( group=self.cpu_group, @@ -108,13 +122,14 @@ class CudaCommunicator(DeviceCommunicatorBase): ), ) - if current_platform.is_rocm(): - # Initialize a custom quick all-reduce implementation for AMD. - # Quick reduce is designed as a complement to custom allreduce. - # Based on quickreduce (https://github.com/mk1-project/quickreduce). - # If it's a rocm, 'use_custom_allreduce==True' means it must - # currently be an MI300 series. - self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) + if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm(): + # Initialize a custom quick all-reduce implementation for AMD. + # Quick reduce is designed as a complement to custom allreduce + # (vLLM's or AITER's), so it is initialized for either backend. + # Based on quickreduce (https://github.com/mk1-project/quickreduce). + # On ROCm, 'use_custom_allreduce==True' means it must currently be + # an MI300 series. + self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device) if self.world_size > 1: self._log_all_reduce_backend_selection() @@ -203,6 +218,7 @@ class CudaCommunicator(DeviceCommunicatorBase): "NCCL_SYMM_MEM", "QUICK_REDUCE", "FLASHINFER", + "AITER_CUSTOM", "CUSTOM", "SYMM_MEM", "PYNCCL", @@ -236,6 +252,8 @@ class CudaCommunicator(DeviceCommunicatorBase): enabled_ar_backends.append("QUICK_REDUCE") if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled: enabled_ar_backends.append("FLASHINFER") + if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled: + enabled_ar_backends.append("AITER_CUSTOM") if self.ca_comm is not None and not self.ca_comm.disabled: enabled_ar_backends.append("CUSTOM") if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled: @@ -261,8 +279,8 @@ class CudaCommunicator(DeviceCommunicatorBase): out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) if out is not None: return out - # always try quick reduce first, then flashinfer, then custom allreduce, - # and then pynccl. (quick reduce just for ROCM MI3*) + # always try quick reduce first, then flashinfer, then the AITER or vLLM + # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*) qr_comm = self.qr_comm if ( qr_comm is not None @@ -281,6 +299,15 @@ class CudaCommunicator(DeviceCommunicatorBase): out = fi_ar_comm.all_reduce(input_) assert out is not None return out + aiter_ar_comm = self.aiter_ar_comm + if ( + aiter_ar_comm is not None + and not aiter_ar_comm.disabled + and aiter_ar_comm.should_custom_ar(input_) + ): + out = aiter_ar_comm.custom_all_reduce(input_) + assert out is not None + return out ca_comm = self.ca_comm if ( ca_comm is not None @@ -509,6 +536,9 @@ class CudaCommunicator(DeviceCommunicatorBase): self.pynccl_comm = None if self.ca_comm is not None: self.ca_comm = None + if self.aiter_ar_comm is not None: + self.aiter_ar_comm.close() + self.aiter_ar_comm = None if self.fi_ar_comm is not None: self.fi_ar_comm.destroy() self.fi_ar_comm = None diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index e98c765f537..feacb03d28b 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -742,8 +742,8 @@ class EplbState: is_main_rank = ep_rank == 0 if is_main_rank: if not self.is_async or is_profile: - start_event = torch.Event(enable_timing=True) - end_event = torch.Event(enable_timing=True) + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) start_event.record() logger.info( "Rearranging experts %s %s...", diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 21a7ee68fa9..dee19749745 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -31,7 +31,7 @@ class CpuGpuEvent: """ def __init__(self): - self._event = torch.Event() + self._event = torch.cuda.Event() self._recorded = threading.Event() def wait(self, stream: torch.cuda.Stream | None = None): @@ -56,7 +56,7 @@ class CpuGpuEvent: "CpuGpuEvent.record() called before the previous event was " "consumed by wait()" ) - self._event = torch.Event() + self._event = torch.cuda.Event() self._event.record(stream) self._recorded.set() diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index adc8b082699..a7d83bb378f 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -43,6 +43,7 @@ class KVCacheEvent( MEDIUM_GPU = "GPU" +MEDIUM_CPU = "CPU" class BlockStored(KVCacheEvent): @@ -458,15 +459,12 @@ class ZmqEventPublisher(EventPublisher): for seq, buf in self._buffer: if seq >= start_seq: - # [identity, empty_delim, seq_bytes, payload] - # (identity, empty_delim) are stripped off by the router - # receiving payload is (seq_bytes, payload) + # Subscriber receives (topic, seq_bytes, payload) self._replay.send_multipart( - (client_id, b"", seq.to_bytes(8, "big"), buf) + (client_id, b"", self._topic_bytes, seq.to_bytes(8, "big"), buf) ) # Send end of sequence marker - # receiving payload is (-1, b""") - self._replay.send_multipart((client_id, b"", self.END_SEQ, b"")) + self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b"")) @staticmethod def offset_endpoint_port( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index a604bd5528f..299ff037ad2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -239,7 +239,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # this event is complete the request is considered "done sending" # by get_finished; clients block on the per-file flock to wait for # the disk write itself. - self._req_copy_events: dict[str, torch.Event] = {} + self._req_copy_events: dict[str, torch.cuda.Event] = {} # req_ids reported as finished-generating by the scheduler, # accumulated across get_finished calls. self._accumulated_finished_req_ids: set[str] = set() @@ -320,7 +320,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): @staticmethod def _write_tensors( tensors: dict[str, torch.Tensor], - event: torch.Event, + event: torch.cuda.Event, filename: str, lock_fd: int | None, ) -> None: @@ -375,7 +375,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): copy_stream = self._get_copy_stream() # Ensure the copy stream sees all prior writes on the default stream. - ready_event = torch.Event() + ready_event = torch.cuda.Event() ready_event.record() copy_stream.wait_event(ready_event) @@ -396,7 +396,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): pinned_hs.copy_(hidden_states_gpu, non_blocking=True) # Record completion of this copy on the copy stream. - copy_done = torch.Event() + copy_done = torch.cuda.Event() copy_done.record(copy_stream) # token_ids is already on CPU (created in request_finished). diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py index edcf83b1925..a54233453bb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py @@ -221,7 +221,7 @@ class Hf3fsClient: @wsynchronized() def batch_write( - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event ) -> list[int]: """Write data from tensors to the file at specified offsets. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py index 55a8b5a161e..526375952fe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py @@ -133,7 +133,7 @@ class AsyncOperationManager: # CUDA streams for async operations self._save_stream = torch.cuda.Stream() self._load_stream = torch.cuda.Stream() - self._save_event = torch.Event() + self._save_event = torch.cuda.Event() # Buffer allocators for data copying self._save_buffer_allocator = CopyBufferAllocator( @@ -171,7 +171,7 @@ class AsyncOperationManager: def submit_save_operation(self, request_id: str, block_ids, block_hashes) -> Future: """Submit a save operation for async execution.""" future: Future[Any] = Future() - main_stream_event = torch.Event() + main_stream_event = torch.cuda.Event() main_stream_event.record() task = (request_id, block_ids, block_hashes, future, main_stream_event) self._save_queue.put(task) @@ -304,7 +304,7 @@ class AsyncOperationManager: block_ids, buffers, "gather" ) - save_stream_event = torch.Event() + save_stream_event = torch.cuda.Event() save_stream_event.record(self._save_stream) # Record gather completion # Step3: Write data in batches diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py index e2718d1faa1..3914663a62d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py @@ -75,7 +75,7 @@ class Hf3fsClient: return torch.frombuffer(buffer_data, dtype=dtype) def batch_write( - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event ) -> list[int]: """Write data from tensors to file at specified offsets.""" results = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py index 6d83380cad3..2e75519df12 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py @@ -430,7 +430,7 @@ class LMCacheMPWorkerAdapter: @_lmcache_nvtx_annotate def submit_store_request( - self, request_id: str, op: LoadStoreOp, event: torch.Event + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): """ Submit a KV cache store request to LMCache @@ -464,7 +464,7 @@ class LMCacheMPWorkerAdapter: @_lmcache_nvtx_annotate def submit_retrieve_request( - self, request_id: str, op: LoadStoreOp, event: torch.Event + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event ): """ Submit a KV cache retrieve request to LMCache @@ -501,7 +501,7 @@ class LMCacheMPWorkerAdapter: self, request_ids: list[str], ops: list[LoadStoreOp], - event: torch.Event, + event: torch.cuda.Event, ): """ Submit a batched store request to LMCache @@ -550,7 +550,7 @@ class LMCacheMPWorkerAdapter: self, request_ids: list[str], ops: list[LoadStoreOp], - event: torch.Event, + event: torch.cuda.Event, ): """ Submit a batched retrieve request to LMCache diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index 2ca35be2b51..8786e91a5a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -589,7 +589,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1): return with torch.cuda.stream(torch.cuda.current_stream()): - event = torch.Event(interprocess=True) + event = torch.cuda.Event(interprocess=True) event.record() self.worker_adapter.batched_submit_retrieve_requests( @@ -663,7 +663,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1): return with torch.cuda.stream(torch.cuda.current_stream()): - event = torch.Event(interprocess=True) + event = torch.cuda.Event(interprocess=True) event.record() self.worker_adapter.batched_submit_store_requests( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 3b69ce9a177..ef98ec0d4e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -323,7 +323,7 @@ class ReqMeta: can_save: bool | None = None load_spec: LoadSpec | None = None is_last_chunk: bool | None = None - current_event: torch.Event | None = None + current_event: torch.cuda.Event | None = None token_ids: list[int] | None = None num_prompt_tokens: int | None = None 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 818c2479a14..e60d2f47a4e 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 @@ -1357,7 +1357,7 @@ class MooncakeStoreWorker: current_event = None for request in meta.requests: if request.can_save: - current_event = torch.Event() + current_event = torch.cuda.Event() current_event.record() break diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 7b8ab566058..15585123e5c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -56,7 +56,7 @@ class WriteTask: local_block_ids: list[int] remote_block_ids_hint: list[int] | None layer_name: str - event: torch.Event + event: torch.cuda.Event remote_notify_port: int remote_ip: str enqueue_time: float = field(default_factory=time.perf_counter) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 172836f5cc3..de21a1398e0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1061,7 +1061,7 @@ class MoRIIOConnectorWorker: # when mori-io supports ibgda functionality stream = torch.cuda.current_stream() - event = torch.Event() + event = torch.cuda.Event() event.record(stream) task = WriteTask( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index a90250d285e..baa168e1708 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -30,6 +30,14 @@ class _TransferMetricName: STORE_SIZE = "vllm:kv_offload_store_size" +class _ConnectorMetricName: + """Connector-side metrics emitted by scheduler-side offloading code.""" + + LOOKUP_SYNC_DELAY = "vllm:kv_offload_lookup_sync_delay_seconds" + LOOKUP_ASYNC_DELAY = "vllm:kv_offload_lookup_async_delay_seconds" + ALLOCATION_FAILURE = "vllm:kv_offload_allocation_failure" + + class _TransferType: """Transfer direction labels for deprecated CPU offload metrics.""" @@ -74,6 +82,50 @@ def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]: documentation="Histogram of KV offload store operation size, in bytes.", buckets=TRANSFER_SIZE_BUCKETS, ), + _ConnectorMetricName.LOOKUP_SYNC_DELAY: OffloadingHistogramMetadata( + documentation=( + "Histogram of the time spent in a single offload lookup call, " + "in seconds." + ), + buckets=( + 0.00001, + 0.00005, + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + ), + ), + _ConnectorMetricName.LOOKUP_ASYNC_DELAY: OffloadingHistogramMetadata( + documentation=( + "Histogram of time between a request's offload lookup first " + "deferring and the following lookup resolving, or request " + "finish, in seconds." + ), + buckets=( + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + 5, + 10, + ), + ), + _ConnectorMetricName.ALLOCATION_FAILURE: OffloadingCounterMetadata( + documentation=( + "Number of KV offload store allocation attempts that failed." + ), + ), } @@ -223,7 +275,7 @@ class OffloadingConnectorStats(KVConnectorStats): def increase_counter( self, counter_name: str, - counter_increase_value: int | float, + counter_increase_value: int | float = 1, labelvalues: tuple[str, ...] = (), ) -> None: """Increase a counter on the stats payload.""" 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 284098ce945..75a9696b72c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from itertools import islice @@ -21,6 +22,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, + _ConnectorMetricName, _TransferMetricName, ) from vllm.logger import init_logger @@ -246,6 +248,9 @@ class RequestOffloadState: # In-flight job IDs. Per the connector's invariant, at any given time # this contains either a single load job, or one or more store jobs. transfer_jobs: set[int] = field(default_factory=set) + # time.monotonic() of this request's first deferred offload lookup; + # None once consumed (observed) or while no lookup is pending. + deferred_lookup_start_time: float | None = None def __post_init__(self) -> None: self.group_states = tuple( @@ -293,12 +298,39 @@ class RequestOffloadState: for group_state, new_blocks in zip(self.group_states, new_block_id_groups): group_state.block_ids.extend(new_blocks) + def storable_blocks( + self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int + ) -> int: + """Number of leading offloaded blocks eligible for store. + + For eagle/MTP groups the volatile trailing block of the offloadable + range is excluded while decoding: the draft-layer KV of the last + accepted position may be rewritten after spec-token rejection. During + prefill the trailing block is stable (the draft input for a chunk's + last position is the next prompt token), so it is stored immediately. + The exclusion must be applied consistently everywhere + ``next_stored_block_idx`` is derived: otherwise the trailing block of + each step is skipped on collection but jumped over by + ``next_stored_block_idx``, so it is never re-considered and a + permanent hole breaks prefix-reuse lookup. + """ + num_blocks = num_offloadable_tokens // group_config.offloaded_block_size + is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens + if group_config.is_eagle_group and is_decoding: + num_blocks = max(0, num_blocks - 1) + return num_blocks + def advance_stored_idx(self, num_offloadable_tokens: int) -> None: + # max(): at the prefill->decode transition of a block-aligned prompt, + # storable_blocks drops by one (the eagle exclusion kicks in), and the + # index must not move backwards past already-stored blocks. for group_config, group_state in zip( self.config.kv_group_configs, self.group_states ): - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size - group_state.next_stored_block_idx = num_blocks + group_state.next_stored_block_idx = max( + group_state.next_stored_block_idx, + self.storable_blocks(group_config, num_offloadable_tokens), + ) def update_num_hit_blocks(self, num_cached_tokens: int) -> None: for group_config, group_state in zip( @@ -325,7 +357,7 @@ class OffloadingConnectorScheduler: ): self.config = SchedulerOffloadConfig.from_spec(spec) self.manager: OffloadingManager = spec.get_manager() - self._connector_stats: OffloadingConnectorStats | None = None + self._connector_stats = OffloadingConnectorStats() full_attention_groups: list[int] = [] sliding_window_groups: list[int] = [] @@ -375,6 +407,18 @@ class OffloadingConnectorScheduler: self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + def _maybe_observe_lookup_async_delay( + self, req_status: RequestOffloadState + ) -> None: + start_time = req_status.deferred_lookup_start_time + if start_time is None: + return + req_status.deferred_lookup_start_time = None + self._connector_stats.observe_histogram( + _ConnectorMetricName.LOOKUP_ASYNC_DELAY, + time.monotonic() - start_time, + ) + def _generate_job_id(self) -> int: job_id = self._job_counter self._job_counter += 1 @@ -685,7 +729,17 @@ class OffloadingConnectorScheduler: if request.skip_reading_prefix_cache: num_hit_tokens = 0 else: + lookup_start = time.monotonic() num_hit_tokens = self._lookup(req_status) + self._connector_stats.observe_histogram( + _ConnectorMetricName.LOOKUP_SYNC_DELAY, + time.monotonic() - lookup_start, + ) + if num_hit_tokens is None: + if req_status.deferred_lookup_start_time is None: + req_status.deferred_lookup_start_time = lookup_start + else: + self._maybe_observe_lookup_async_delay(req_status) req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) @@ -876,9 +930,9 @@ class OffloadingConnectorScheduler: for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size - if group_config.is_eagle_group: - num_blocks = max(0, num_blocks - 1) + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) start_block_idx = group_state.next_stored_block_idx if num_blocks <= start_block_idx: @@ -925,6 +979,9 @@ class OffloadingConnectorScheduler: new_offload_keys, req_status.req_context ) if store_output is None: + self._connector_stats.increase_counter( + _ConnectorMetricName.ALLOCATION_FAILURE + ) logger.warning("Request %s: cannot store blocks", req_id) continue @@ -947,7 +1004,9 @@ class OffloadingConnectorScheduler: is_sliding_window = ( group_config.sliding_window_size_in_blocks is not None ) - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) start_block_idx = group_state.next_stored_block_idx block_ids = group_state.block_ids num_group_blocks = 0 @@ -980,7 +1039,9 @@ class OffloadingConnectorScheduler: group_sizes.append(num_group_blocks) block_indices.append(start_gpu_block_idx or 0) - group_state.next_stored_block_idx = num_blocks + group_state.next_stored_block_idx = max( + group_state.next_stored_block_idx, num_blocks + ) src_spec = GPULoadStoreSpec( src_block_ids, group_sizes=group_sizes, block_indices=block_indices @@ -1115,10 +1176,7 @@ class OffloadingConnectorScheduler: transfer_stats.observe_histogram( _TransferMetricName.STORE_SIZE, size ) - if self._connector_stats is None: - self._connector_stats = transfer_stats - else: - self._connector_stats.aggregate(transfer_stats) + self._connector_stats.aggregate(transfer_stats) for job_id, count in meta.completed_jobs.items(): assert count > 0 @@ -1159,8 +1217,10 @@ class OffloadingConnectorScheduler: del self._req_status[job_status.req_id] def get_stats(self) -> OffloadingConnectorStats | None: - stats = self._connector_stats - self._connector_stats = None + stats: OffloadingConnectorStats | None = None + if not self._connector_stats.is_empty(): + stats = self._connector_stats + self._connector_stats = OffloadingConnectorStats() manager_stats = self.manager.get_stats() if manager_stats is not None: @@ -1198,7 +1258,7 @@ class OffloadingConnectorScheduler: return False, None self.manager.on_request_finished(req_status.req_context) - + self._maybe_observe_lookup_async_delay(req_status) if not req_status.transfer_jobs: # No in-flight jobs: no later complete_store()/complete_load() calls # need this request's state. diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 11b9e24e864..162ed03d23b 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -269,11 +269,17 @@ def _create_subgroups_split_group( must enter with the same ``split_ranks`` definition. Each rank receives the subgroup it belongs to. """ + from vllm.distributed.utils import ( + get_cpu_distributed_timeout_or_none, + get_distributed_timeout_or_none, + ) + device_backend_str = _device_backend_str(torch_distributed_backend) self_device_group = torch.distributed.split_group( split_ranks=group_ranks, group_desc=f"{group_name}:device", backend=device_backend_str, + timeout=get_distributed_timeout_or_none(), ) # CPU subgroup: split_group requires the requested backend filter to # include the parent's default device type (= the device the parent PG @@ -284,6 +290,7 @@ def _create_subgroups_split_group( split_ranks=group_ranks, group_desc=f"{group_name}:cpu", backend=f"cpu:gloo,{device_backend_str}", + timeout=get_cpu_distributed_timeout_or_none(), ) return self_device_group, self_cpu_group @@ -417,13 +424,19 @@ class GroupCoordinator: self.rank_in_group = ranks.index(self.rank) break else: - from vllm.distributed.utils import get_cpu_distributed_timeout_or_none + from vllm.distributed.utils import ( + get_cpu_distributed_timeout_or_none, + get_distributed_timeout_or_none, + ) timeout = get_cpu_distributed_timeout_or_none() + device_timeout = get_distributed_timeout_or_none() for ranks in group_ranks: device_group = torch.distributed.new_group( - ranks, backend=torch_distributed_backend + ranks, + backend=torch_distributed_backend, + timeout=device_timeout, ) # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. @@ -504,10 +517,16 @@ class GroupCoordinator: This is a collective call: every world rank must invoke it. Used where we want to issue ops that can run concurrently with ops on `device_group`. """ + from vllm.distributed.utils import get_distributed_timeout_or_none + + device_timeout = get_distributed_timeout_or_none() sibling: ProcessGroup | None = None for ranks in self.group_ranks: pg = torch.distributed.new_group( - ranks, backend=self.torch_distributed_backend, group_desc=group_desc + ranks, + backend=self.torch_distributed_backend, + group_desc=group_desc, + timeout=device_timeout, ) if self.rank in ranks: sibling = pg diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index ef3c11ff64e..eec9890f028 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -533,6 +533,16 @@ def get_cpu_distributed_timeout_or_none() -> timedelta | None: return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None +def get_distributed_timeout_or_none() -> timedelta | None: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + timeout_seconds = vllm_config.parallel_config.distributed_timeout_seconds + return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None + + def init_gloo_process_group( prefix_store: PrefixStore, group_rank: int, @@ -616,6 +626,10 @@ def stateless_init_torch_distributed_process_group( gloo_timeout = get_cpu_distributed_timeout_or_none() if gloo_timeout is not None: timeout = gloo_timeout + else: + device_timeout = get_distributed_timeout_or_none() + if device_timeout is not None: + timeout = device_timeout if listen_socket is not None: store = create_tcp_store( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index efdd7696fdc..d80ecb9c855 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -354,14 +354,17 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: elif contains_type(type_hints, set): kwargs[name].update(collection_to_kwargs(type_hints, set)) elif contains_type(type_hints, int): + # Arguments that accept human-readable integer strings (e.g., 1K, 2M, 1G) + human_readable_int_args = { + "max_num_batched_tokens", + "max_num_scheduled_tokens", + "kv_cache_memory_bytes", + "safetensors_prefetch_block_size", + } if name == "max_model_len": kwargs[name]["type"] = human_readable_int_or_auto kwargs[name]["help"] += f"\n\n{human_readable_int_or_auto.__doc__}" - elif name in ( - "max_num_batched_tokens", - "kv_cache_memory_bytes", - "safetensors_prefetch_block_size", - ): + elif name in human_readable_int_args: kwargs[name]["type"] = human_readable_int kwargs[name]["help"] += f"\n\n{human_readable_int.__doc__}" else: @@ -523,6 +526,7 @@ class EngineArgs: gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None + max_num_scheduled_tokens: int | None = None max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold @@ -536,6 +540,9 @@ class EngineArgs: code_revision: str | None = ModelConfig.code_revision hf_token: bool | str | None = ModelConfig.hf_token hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides") + model_class_overrides: dict[str, str] = get_field( + ModelConfig, "model_class_overrides" + ) tokenizer_revision: str | None = ModelConfig.tokenizer_revision quantization: QuantizationMethods | str | None = ModelConfig.quantization quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None @@ -851,6 +858,9 @@ class EngineArgs: model_group.add_argument("--config-format", **model_kwargs["config_format"]) model_group.add_argument("--hf-token", **model_kwargs["hf_token"]) model_group.add_argument("--hf-overrides", **model_kwargs["hf_overrides"]) + model_group.add_argument( + "--model-class-overrides", **model_kwargs["model_class_overrides"] + ) model_group.add_argument("--pooler-config", **model_kwargs["pooler_config"]) model_group.add_argument( "--generation-config", **model_kwargs["generation_config"] @@ -1406,6 +1416,13 @@ class EngineArgs: "default": None, }, ) + scheduler_group.add_argument( + "--max-num-scheduled-tokens", + **{ + **scheduler_kwargs["max_num_scheduled_tokens"], + "default": None, + }, + ) scheduler_group.add_argument( "--max-num-seqs", **{ @@ -1622,6 +1639,7 @@ class EngineArgs: code_revision=self.code_revision, hf_token=self.hf_token, hf_overrides=self.hf_overrides, + model_class_overrides=self.model_class_overrides, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization, @@ -2144,6 +2162,7 @@ class EngineArgs: scheduler_config = SchedulerConfig( runner_type=model_config.runner_type, max_num_batched_tokens=self.max_num_batched_tokens, + max_num_scheduled_tokens=self.max_num_scheduled_tokens, max_num_seqs=self.max_num_seqs, max_model_len=model_config.max_model_len, enable_chunked_prefill=self.enable_chunked_prefill, diff --git a/vllm/entrypoints/generate/api_router.py b/vllm/entrypoints/generate/api_router.py index 5a26e475b06..062a7947e79 100644 --- a/vllm/entrypoints/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -130,6 +130,7 @@ async def init_generate_state( enable_force_include_usage=args.enable_force_include_usage, enable_log_outputs=args.enable_log_outputs, enable_log_deltas=args.enable_log_deltas, + enable_per_request_metrics=args.enable_per_request_metrics, ) state.openai_serving_chat = ( OpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None @@ -150,6 +151,7 @@ async def init_generate_state( return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, enable_force_include_usage=args.enable_force_include_usage, + enable_per_request_metrics=args.enable_per_request_metrics, ) if "generate" in supported_tasks else None diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 196b6e3f87f..78f84ca872b 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -18,6 +18,7 @@ from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, GenerationError, + PerRequestTimingMetrics, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -34,6 +35,7 @@ from vllm.tracing import ( extract_trace_headers, log_tracing_disabled_warning, ) +from vllm.v1.metrics.stats import RequestStateStats logger = init_logger(__name__) @@ -41,6 +43,61 @@ RequestT = TypeVar("RequestT", bound=AnyRequest) _T = TypeVar("_T") +def build_per_request_timing_metrics( + metrics: RequestStateStats | None, + num_generation_tokens: int, +) -> PerRequestTimingMetrics: + """Build per-request timing metrics from ``RequestStateStats``. + + ``generation_time_ms`` is the decode interval only (first output token to + last output token); it excludes both queue wait and prefill/TTFT. + ``tokens_per_second`` is overall output throughput: all generated tokens + over the inference interval (scheduling to last output token), so it counts + the prefill/TTFT phase and is not simply the reciprocal of ``mean_itl_ms``. + Each field is left ``None`` when the timestamps it depends on are + unavailable. + """ + if metrics is None: + return PerRequestTimingMetrics() + + queued_ts = metrics.queued_ts + scheduled_ts = metrics.scheduled_ts + first_token_ts = metrics.first_token_ts + last_token_ts = metrics.last_token_ts + + time_to_first_token_ms: float | None = None + generation_time_ms: float | None = None + queue_time_ms: float | None = None + mean_itl_ms: float | None = None + tokens_per_second: float | None = None + + if scheduled_ts > 0 and first_token_ts > 0: + time_to_first_token_ms = (first_token_ts - scheduled_ts) * 1000 + + if first_token_ts > 0 and last_token_ts > 0: + generation_time_ms = (last_token_ts - first_token_ts) * 1000 + + if queued_ts > 0 and scheduled_ts > 0: + queue_time_ms = (scheduled_ts - queued_ts) * 1000 + + if first_token_ts > 0 and last_token_ts > 0 and num_generation_tokens > 1: + decode_time = last_token_ts - first_token_ts + mean_itl_ms = decode_time / (num_generation_tokens - 1) * 1000 + + if scheduled_ts > 0 and last_token_ts > 0: + inference_time_ms = (last_token_ts - scheduled_ts) * 1000 + if inference_time_ms > 0: + tokens_per_second = num_generation_tokens / inference_time_ms * 1000 + + return PerRequestTimingMetrics( + time_to_first_token_ms=time_to_first_token_ms, + generation_time_ms=generation_time_ms, + queue_time_ms=queue_time_ms, + mean_itl_ms=mean_itl_ms, + tokens_per_second=tokens_per_second, + ) + + @dataclass(kw_only=True) class ServeContext(Generic[RequestT]): request: RequestT diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 6fb27c365d9..59c7ee84cae 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -51,7 +51,11 @@ from vllm.entrypoints.serve.utils.server_utils import ( log_response, validation_exception_handler, ) -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import ( + VLLMNotFoundError, + VLLMUnprocessableEntityError, + VLLMValidationError, +) from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_derenderer import OnlineDerenderer @@ -74,6 +78,41 @@ logger = init_logger("vllm.entrypoints.openai.api_server") _FALLBACK_SUPPORTED_TASKS: tuple[SupportedTask, ...] = ("generate",) +def _attach_endpoint_plugins( + app: FastAPI, supported_tasks: tuple["SupportedTask", ...] +) -> None: + """Phase A of endpoint plugin wiring: discover, gate and attach routes. + + Attached last after all core routers. This is so endpoint plugin routes can + shadow core routes with the same path (see `EndpointPlugin.attach_router` + docstring). No-ops when no plugins are discovered/allowlisted. + """ + from vllm.plugins import load_endpoint_plugins + + endpoint_plugins = load_endpoint_plugins(supported_tasks) + for plugin in endpoint_plugins: + plugin.attach_router(app) + app.state.endpoint_plugins = endpoint_plugins + + +async def _init_endpoint_plugins_state( + engine_client: EngineClient | None, state: State, args: Namespace +) -> None: + """Phase B of endpoint plugin wiring: initialize per app plugin state. + + `state.endpoint_plugins` is set by `_attach_endpoint_plugins` (Phase A) + in `build_app`. Some `init_app_state` callers (e.g. `run_batch.py`) + build their own bare `State` without going through `build_app`. As a result + `endpoint_plugins` may be absent and are treated that the same as "none attached". + + `engine_client` is `None` for the CPU only render server which has no + engine (see `init_render_app_state`). Plugins must handle a `None` + `engine_client` themselves (see `EndpointPlugin.init_state`). + """ + for plugin in getattr(state, "endpoint_plugins", []): + await plugin.init_state(engine_client, state, args) + + @asynccontextmanager async def build_async_engine_client( args: Namespace, @@ -230,6 +269,12 @@ def build_app( register_pooling_api_routers(app, supported_tasks, model_config) + # Endpoint plugins are attached last so their routes are registered after all core + # routers. This runs even for the CPU only render server. A plugin eligible for + # the `render` task still gets its routes registered. It receives + # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). + _attach_endpoint_plugins(app, supported_tasks) + app.root_path = args.root_path app.add_middleware( CORSMiddleware, @@ -244,7 +289,18 @@ def build_app( app.exception_handler(EngineGenerateError)(engine_error_handler) app.exception_handler(EngineDeadError)(engine_error_handler) app.exception_handler(GenerationError)(generation_error_handler) + # Register specific exception types so they are handled by + # ExceptionMiddleware (inside the Prometheus middleware) rather than + # ServerErrorMiddleware (outside it). Without this, these exceptions + # propagate through Prometheus as unhandled and get recorded as 5xx + # even though they result in 4xx responses to the client. app.exception_handler(VLLMValidationError)(exception_handler) + app.exception_handler(VLLMUnprocessableEntityError)(exception_handler) + app.exception_handler(VLLMNotFoundError)(exception_handler) + app.exception_handler(ValueError)(exception_handler) + app.exception_handler(TypeError)(exception_handler) + app.exception_handler(OverflowError)(exception_handler) + app.exception_handler(NotImplementedError)(exception_handler) app.exception_handler(Exception)(exception_handler) # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY @@ -416,6 +472,8 @@ async def init_app_state( init_pooling_state(engine_client, state, args, request_logger, supported_tasks) + await _init_endpoint_plugins_state(engine_client, state, args) + state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 @@ -507,6 +565,10 @@ async def init_render_app_state( state.enable_server_load_tracking = False state.server_load_metrics = 0 + # No `EngineClient` exists for the render server, so plugins get `None` and + # must handle it themselves (see `EndpointPlugin.init_state`). + await _init_endpoint_plugins_state(None, state, args) + def create_server_socket( addr: tuple[str, int], diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index ab905677ab5..cce51157f84 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -26,6 +26,7 @@ from vllm.entrypoints.openai.engine.protocol import ( FunctionDefinition, LegacyStructuralTagResponseFormat, OpenAIBaseModel, + PerRequestTimingMetrics, StreamOptions, StructuralTagResponseFormat, ToolCall, @@ -133,6 +134,7 @@ class ChatCompletionResponse(OpenAIBaseModel): kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters." ) + metrics: PerRequestTimingMetrics | None = None class ChatCompletionResponseStreamChoice(OpenAIBaseModel): @@ -160,6 +162,7 @@ class ChatCompletionStreamResponse(OpenAIBaseModel): # Rendered prompt text from chat templating (only set when # ``return_prompt_text=True`` on the request); only sent on the first chunk. prompt_text: str | None = None + metrics: PerRequestTimingMetrics | None = None class ChatCompletionToolsParam(OpenAIBaseModel): @@ -1000,8 +1003,8 @@ class BatchChatCompletionRequest(OpenAIBaseModel): response_format: Any | None = None seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX) stop: str | list[str] | None = Field(default_factory=list) - temperature: float | None = 0.7 - top_p: float | None = 1.0 + temperature: float | None = None + top_p: float | None = None user: str | None = None tool_choice: Literal["none"] | None = "none" include_reasoning: bool = True @@ -1010,8 +1013,8 @@ class BatchChatCompletionRequest(OpenAIBaseModel): best_of: int | None = None use_beam_search: bool = False top_k: int | None = None - min_p: float | None = 0.0 - repetition_penalty: float | None = 1.0 + min_p: float | None = None + repetition_penalty: float | None = None length_penalty: float | None = 1.0 early_stopping: bool = False structured_outputs: StructuredOutputsParams | None = None diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9ef144e6a14..caa3f724da4 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -22,6 +22,7 @@ from vllm.entrypoints.chat_utils import ( from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, GenerationError, + build_per_request_timing_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -41,6 +42,7 @@ from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ErrorResponse, FunctionCall, + PerRequestTimingMetrics, PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, @@ -123,6 +125,7 @@ class OpenAIServingChat(GenerateBaseServing): enable_log_outputs: bool = False, enable_log_deltas: bool = True, default_chat_template_kwargs: dict[str, Any] | None = None, + enable_per_request_metrics: bool = False, ) -> None: super().__init__( engine_client=engine_client, @@ -161,6 +164,7 @@ class OpenAIServingChat(GenerateBaseServing): self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage + self.enable_per_request_metrics = enable_per_request_metrics self.default_sampling_params = self.model_config.get_diff_sampling_param() mc = self.model_config self.override_max_tokens = ( @@ -461,8 +465,10 @@ class OpenAIServingChat(GenerateBaseServing): stream_options, self.enable_force_include_usage ) + last_res: RequestOutput | None = None try: async for res in result_generator: + last_res = res if res.prompt_token_ids is not None: num_prompt_tokens = len(res.prompt_token_ids) if res.encoder_prompt_token_ids is not None: @@ -752,6 +758,21 @@ class OpenAIServingChat(GenerateBaseServing): mm_token_counts, ) + # In streaming, metrics ride on this final usage chunk, which is + # only emitted when usage reporting is enabled (i.e. + # ``stream_options.include_usage=true`` or + # ``--enable-force-include-usage``). + stream_per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # See note in chat_completion_full_generator: suppress for n>1. + and (request.n or 1) == 1 + ): + last_metrics = last_res.metrics if last_res is not None else None + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, completion_tokens + ) + final_usage_chunk = ChatCompletionStreamResponse( id=request_id, object=chunk_object_type, @@ -760,6 +781,7 @@ class OpenAIServingChat(GenerateBaseServing): model=model_name, usage=final_usage, system_fingerprint=self.system_fingerprint, + metrics=stream_per_request_metrics, ) final_usage_data = final_usage_chunk.model_dump_json( exclude_unset=True, exclude_none=True @@ -1003,6 +1025,18 @@ class OpenAIServingChat(GenerateBaseServing): request_metadata.final_usage_info = usage + per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # Timing metrics describe a single generation stream. For n>1 the + # returned stats belong to only one of the n sequences, so they + # cannot be accurately attributed to the request; suppress instead. + and (request.n or 1) == 1 + ): + per_request_metrics = build_per_request_timing_metrics( + final_res.metrics, num_generated_tokens + ) + # ``final_res.prompt`` is the rendered chat-templated prompt text prompt_text = final_res.prompt if request.return_prompt_text else None @@ -1019,6 +1053,7 @@ class OpenAIServingChat(GenerateBaseServing): ), prompt_text=prompt_text, kv_transfer_params=final_res.kv_transfer_params, + metrics=per_request_metrics, ) # Log complete response if output logging is enabled diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 1533895edcd..8dbb6994390 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -25,12 +25,9 @@ from vllm.entrypoints.serve.utils.constants import ( H11_MAX_HEADER_COUNT_DEFAULT, H11_MAX_INCOMPLETE_EVENT_SIZE_DEFAULT, ) -from vllm.logger import init_logger from vllm.tool_parsers import ToolParserManager from vllm.utils.argparse_utils import FlexibleArgumentParser -logger = init_logger(__name__) - class LoRAParserAction(argparse.Action): def __call__( @@ -134,6 +131,8 @@ class BaseFrontendArgs: log. The default of None means unlimited.""" enable_prompt_tokens_details: bool = False """If set to True, enable prompt_tokens_details in usage.""" + enable_per_request_metrics: bool = False + """If set to True, include per-request timing metrics in API responses.""" enable_server_load_tracking: bool = False """If set to True, enable tracking server_load_metrics in the app state.""" enable_force_include_usage: bool = False @@ -398,6 +397,14 @@ def validate_parsed_serve_args(args: argparse.Namespace): if args.enable_log_outputs and not args.enable_log_requests: raise TypeError("Error: --enable-log-outputs requires --enable-log-requests") + if getattr(args, "enable_per_request_metrics", False) and getattr( + args, "disable_log_stats", False + ): + raise ValueError( + "Error: --enable-per-request-metrics requires engine statistics " + "logging; remove --disable-log-stats to enable per-request metrics." + ) + if args.data_parallel_multi_port_external_lb: from vllm.entrypoints.openai.dp_supervisor import ( validate_multi_port_external_lb_args, diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index b96d4f3c0c7..a7b7996fed7 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -9,12 +9,14 @@ from typing import Annotated, Any, Literal from pydantic import Field, model_validator +import vllm.envs as envs from vllm.config import ModelConfig from vllm.config.utils import replace from vllm.entrypoints.openai.engine.protocol import ( AnyResponseFormat, LegacyStructuralTagResponseFormat, OpenAIBaseModel, + PerRequestTimingMetrics, StreamOptions, StructuralTagResponseFormat, UsageInfo, @@ -34,6 +36,7 @@ from vllm.sampling_params import ( ThinkingTokenBudget, ) from vllm.utils import random_uuid +from vllm.utils.collection_utils import is_list_of logger = init_logger(__name__) @@ -92,6 +95,7 @@ class CompletionRequest(OpenAIBaseModel): ) allowed_token_ids: list[int] | None = None prompt_logprobs: int | None = None + bad_words: list[str] = Field(default_factory=list) # --8<-- [end:completion-sampling-params] # --8<-- [start:completion-extra-params] @@ -368,6 +372,7 @@ class CompletionRequest(OpenAIBaseModel): structured_outputs=self.structured_outputs, logit_bias=self.logit_bias, allowed_token_ids=self.allowed_token_ids, + bad_words=self.bad_words, extra_args=extra_args or None, skip_clone=True, # Created fresh per request, safe to skip clone repetition_detection=self.repetition_detection, @@ -496,6 +501,38 @@ class CompletionRequest(OpenAIBaseModel): return data + @model_validator(mode="before") + @classmethod + def validate_prompt_list_length(cls, data): + max_prompts = envs.VLLM_MAX_COMPLETION_PROMPTS + + prompt = data.get("prompt") + if ( + isinstance(prompt, list) + and len(prompt) > 0 + and not is_list_of(prompt, int) + and len(prompt) > max_prompts + ): + raise VLLMValidationError( + f"prompt list length {len(prompt)} exceeds the maximum " + f"allowed count of {max_prompts}. To increase this " + "limit, set the VLLM_MAX_COMPLETION_PROMPTS " + "environment variable.", + parameter="prompt", + ) + + prompt_embeds = data.get("prompt_embeds") + if isinstance(prompt_embeds, list) and len(prompt_embeds) > max_prompts: + raise VLLMValidationError( + f"prompt_embeds list length {len(prompt_embeds)} exceeds " + f"the maximum allowed count of {max_prompts}. To increase " + "this limit, set the VLLM_MAX_COMPLETION_PROMPTS " + "environment variable.", + parameter="prompt_embeds", + ) + + return data + @model_validator(mode="before") @classmethod def check_cache_salt_support(cls, data): @@ -558,6 +595,7 @@ class CompletionResponse(OpenAIBaseModel): kv_transfer_params: dict[str, Any] | None = Field( default=None, description="KVTransfer parameters." ) + metrics: PerRequestTimingMetrics | None = None class CompletionResponseStreamChoice(OpenAIBaseModel): @@ -589,3 +627,4 @@ class CompletionStreamResponse(OpenAIBaseModel): # Set only on the final chunk of a stream to mirror non-streaming responses # without the per-chunk serialization overhead. system_fingerprint: str | None = None + metrics: PerRequestTimingMetrics | None = None diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index aeade306465..d26a455cc8e 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -16,6 +16,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, GenerationError, + build_per_request_timing_metrics, clamp_prompt_logprobs, format_token_id_placeholder, ) @@ -29,6 +30,7 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, + PerRequestTimingMetrics, PromptTokenUsageInfo, RequestResponseMetadata, UsageInfo, @@ -61,6 +63,7 @@ class OpenAIServingCompletion(GenerateBaseServing): return_tokens_as_token_ids: bool = False, enable_prompt_tokens_details: bool = False, enable_force_include_usage: bool = False, + enable_per_request_metrics: bool = False, ): super().__init__( engine_client=engine_client, @@ -72,6 +75,7 @@ class OpenAIServingCompletion(GenerateBaseServing): self.online_renderer = online_renderer self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage + self.enable_per_request_metrics = enable_per_request_metrics self.default_sampling_params = self.model_config.get_diff_sampling_param() mc = self.model_config @@ -300,8 +304,10 @@ class OpenAIServingCompletion(GenerateBaseServing): stream_options, self.enable_force_include_usage ) + last_res: RequestOutput | None = None try: async for prompt_idx, res in result_generator: + last_res = res prompt_token_ids = res.prompt_token_ids prompt_logprobs = res.prompt_logprobs @@ -448,6 +454,23 @@ class OpenAIServingCompletion(GenerateBaseServing): ) if include_usage: + # In streaming, metrics ride on this final usage chunk, which is + # only emitted when usage reporting is enabled (i.e. + # ``stream_options.include_usage=true`` or + # ``--enable-force-include-usage``). + stream_per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # See note in request_output_to_completion_response: suppress + # when not attributable to one stream (multi-prompt or n>1). + and num_prompts == 1 + and (request.n or 1) == 1 + ): + last_metrics = last_res.metrics if last_res is not None else None + stream_per_request_metrics = build_per_request_timing_metrics( + last_metrics, total_completion_tokens + ) + final_usage_chunk = CompletionStreamResponse( id=request_id, created=created_time, @@ -455,6 +478,7 @@ class OpenAIServingCompletion(GenerateBaseServing): choices=[], usage=final_usage_info, system_fingerprint=self.system_fingerprint, + metrics=stream_per_request_metrics, ) final_usage_data = final_usage_chunk.model_dump_json( exclude_unset=False, exclude_none=True @@ -589,6 +613,23 @@ class OpenAIServingCompletion(GenerateBaseServing): ) request_metadata.final_usage_info = usage + + per_request_metrics: PerRequestTimingMetrics | None = None + if ( + self.enable_per_request_metrics + # Metrics describe a single generation stream, so suppress them when + # they cannot be attributed to one: multiple prompts (timestamps + # span prompts) or n>1 (stats belong to one of the n sequences). + and len(final_res_batch) == 1 + and (request.n or 1) == 1 + ): + last_metrics = ( + last_final_res.metrics if last_final_res is not None else None + ) + per_request_metrics = build_per_request_timing_metrics( + last_metrics, num_generated_tokens + ) + if final_res_batch: kv_transfer_params = final_res_batch[0].kv_transfer_params return CompletionResponse( @@ -599,6 +640,7 @@ class OpenAIServingCompletion(GenerateBaseServing): usage=usage, system_fingerprint=self.system_fingerprint, kv_transfer_params=kv_transfer_params, + metrics=per_request_metrics, ) def _create_completion_logprobs( diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 084d8d429a6..2c32fcf20c6 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -115,6 +115,14 @@ class UsageInfo(OpenAIBaseModel): prompt_tokens_details: PromptTokenUsageInfo | None = None +class PerRequestTimingMetrics(OpenAIBaseModel): + time_to_first_token_ms: float | None = None + generation_time_ms: float | None = None + queue_time_ms: float | None = None + mean_itl_ms: float | None = None + tokens_per_second: float | None = None + + class RequestResponseMetadata(BaseModel): request_id: str final_usage_info: UsageInfo | None = None diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index d75e5d5a548..ed989e2ba9f 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -18,7 +18,7 @@ from openai.types.responses.response_output_item import McpCall from openai.types.responses.response_output_message import ResponseOutputMessage from openai.types.responses.response_output_text import ResponseOutputText from openai.types.responses.tool import Mcp -from openai_harmony import Author, HarmonyError, Message, Role, TextContent +from openai_harmony import Author, Message, Role, TextContent from vllm import envs from vllm.entrypoints.chat_utils import ( @@ -349,6 +349,7 @@ class ParsableContext(ConversationContext): reasoning=reasoning, content=content, tool_calls=tool_calls, + tools=self.request.tools, ) ) elif completion.text: @@ -616,7 +617,7 @@ class HarmonyContext(ConversationContext): self.num_tool_output_tokens = 0 self.last_append_segments: list[Segment] = [] - self.last_append_flush_status: bool | HarmonyError = False + self.last_append_flush_status: bool = False # Turn tracking - replaces multiple individual tracking variables self.current_turn_metrics = TurnMetrics() @@ -643,10 +644,10 @@ class HarmonyContext(ConversationContext): if output.finished: self.finish_reason = output.outputs[0].finish_reason - flushed = self.response_parser.flush() - if flushed is not None: - segments.append(flushed) - self.last_append_flush_status = flushed is not None + flushed_segments = self.response_parser.flush() + if flushed_segments: + segments.extend(flushed_segments) + self.last_append_flush_status = len(flushed_segments) > 0 self.all_turn_metrics.append(self.current_turn_metrics.copy()) self.current_turn_metrics.reset() diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index eb2d66bdd8f..ba8bc5a40f1 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -592,10 +592,19 @@ class ResponsesRequest(OpenAIBaseModel): ) elif is_named_tool_choice and tools is not None: tool_name = tool_choice.get("name") - tool_names = { - t.get("name") if isinstance(t, dict) else getattr(t, "name", None) - for t in tools - } + tool_names = set() + for tool in tools: + if isinstance(tool, dict): + if tool.get("type") == "namespace": + namespace = tool.get("name") + for namespaced_tool in tool.get("tools", []): + namespaced_name = namespaced_tool.get("name") + tool_names.add(namespaced_name) + tool_names.add(f"{namespace}__{namespaced_name}") + else: + tool_names.add(tool.get("name")) + else: + tool_names.add(getattr(tool, "name", None)) if not tool_name or tool_name not in tool_names: raise VLLMValidationError( "Tool choice 'function' not found in 'tools' parameter.", diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index a62c4623992..40a52012792 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -1073,6 +1073,7 @@ class OpenAIServingResponses(GenerateBaseServing): content=content, tool_calls=tool_calls, logprobs=logprobs, + tools=request.tools, ) # Fallback when no parser is configured @@ -1339,7 +1340,7 @@ class OpenAIServingResponses(GenerateBaseServing): [StreamingResponsesResponse], StreamingResponsesResponse ], ) -> AsyncGenerator[StreamingResponsesResponse, None]: - processor = SimpleStreamingEventProcessor() + processor = SimpleStreamingEventProcessor(tools=request.tools) def _get_logprobs( output: CompletionOutput, diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 35021caf2ab..531a35c5722 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -58,6 +58,7 @@ from openai.types.responses.response_output_item import McpCall from openai.types.responses.response_reasoning_item import ( Content as ResponseReasoningTextContent, ) +from openai.types.responses.tool import Tool from openai_harmony import Message as HarmonyMessage from vllm.entrypoints.mcp.tool_server import ToolServer @@ -71,6 +72,10 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponseReasoningPartDoneEvent, StreamingResponsesResponse, ) +from vllm.entrypoints.openai.responses.utils import ( + build_responses_tool_call_name_map, + resolve_responses_tool_call_name, +) from vllm.outputs import CompletionOutput from vllm.parser.harmony import Segment from vllm.utils import random_uuid @@ -815,6 +820,7 @@ class SimpleStreamingState: accumulated_text: str = "" tool_call_id: str = "" tool_call_name: str = "" + tool_call_namespace: str | None = None tool_call_index: int | None = None has_emitted_tool_call_delta: bool = False current_state: _StateType = field(default_factory=lambda: _StateType.NONE) @@ -1016,11 +1022,13 @@ def emit_simple_tool_call_open( state: SimpleStreamingState, name: str, index: int | None, + namespace: str | None = None, ) -> list[StreamingResponsesResponse]: state.current_state = _StateType.TOOL_CALL state.current_item_id = random_uuid() state.tool_call_id = f"call_{random_uuid()}" state.tool_call_name = name + state.tool_call_namespace = namespace state.tool_call_index = index state.accumulated_text = "" state.has_emitted_tool_call_delta = False @@ -1034,6 +1042,7 @@ def emit_simple_tool_call_open( id=state.current_item_id, call_id=state.tool_call_id, name=name, + namespace=namespace, arguments="", status="in_progress", ), @@ -1081,6 +1090,7 @@ def emit_simple_tool_call_done( item=ResponseFunctionToolCall( type="function_call", name=state.tool_call_name, + namespace=state.tool_call_namespace, arguments=state.accumulated_text, status="completed", id=state.current_item_id, @@ -1089,6 +1099,7 @@ def emit_simple_tool_call_done( ), ) state.output_index += 1 + state.tool_call_namespace = None state.current_state = _StateType.NONE return events @@ -1166,8 +1177,13 @@ class SimpleStreamingEventProcessor: ), } - def __init__(self, state: SimpleStreamingState | None = None) -> None: + def __init__( + self, + state: SimpleStreamingState | None = None, + tools: list[Tool] | None = None, + ) -> None: self.state = state or SimpleStreamingState() + self.tool_call_name_map = build_responses_tool_call_name_map(tools) def resolve_target_state( self, delta_message: DeltaMessage @@ -1224,8 +1240,15 @@ class SimpleStreamingEventProcessor: handlers = self._STATE_HANDLERS[target_state] if target_state == _StateType.TOOL_CALL: assert tool_call is not None + call_name = resolve_responses_tool_call_name( + tool_call.function.name, + tool_call_name_map=self.tool_call_name_map, + ) return handlers.open_fn( - self.state, tool_call.function.name, tool_call.index + self.state, + call_name.name, + tool_call.index, + call_name.namespace, ) return handlers.open_fn(self.state) diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index a2f35dca235..07a9704f9b5 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -35,6 +35,12 @@ from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionMessa from vllm.entrypoints.openai.engine.protocol import FunctionCall from vllm.entrypoints.openai.responses.protocol import ResponseInputOutputItem from vllm.logger import init_logger +from vllm.tool_parsers.utils import ( + build_responses_tool_call_name_map, + flat_namespace_tool_name, + iter_response_function_tool_dicts, + resolve_responses_tool_call_name, +) from vllm.utils import random_uuid logger = init_logger(__name__) @@ -45,8 +51,10 @@ def build_response_output_items( content: str | None, tool_calls: list[FunctionCall] | None, logprobs: list[Logprob] | None = None, + tools: list[Tool] | None = None, ) -> list[ResponseOutputItem]: outputs: list[ResponseOutputItem] = [] + tool_call_name_map = build_responses_tool_call_name_map(tools) if reasoning: outputs.append( @@ -81,6 +89,9 @@ def build_response_output_items( if tool_calls: for idx, tool_call in enumerate(tool_calls): + call_name = resolve_responses_tool_call_name( + tool_call.name, tool_call_name_map=tool_call_name_map + ) outputs.append( ResponseFunctionToolCall( id=f"fc_{random_uuid()}", @@ -88,7 +99,8 @@ def build_response_output_items( or make_tool_call_id(func_name=tool_call.name, idx=idx), type="function_call", status="completed", - name=tool_call.name, + name=call_name.name, + namespace=call_name.namespace, arguments=tool_call.arguments, ) ) @@ -219,10 +231,13 @@ def _construct_message_from_response_item( ) if isinstance(item, ResponseFunctionToolCall): + tool_name = item.name + if item.namespace: + tool_name = flat_namespace_tool_name(item.namespace, item.name) tool_call = ChatCompletionMessageToolCallParam( id=item.call_id, function=FunctionCallTool( - name=item.name, + name=tool_name, arguments=item.arguments, ), type="function", @@ -318,7 +333,17 @@ def _construct_message_from_response_item( def extract_function_tool_names(tools: list[Tool]) -> frozenset[str]: - return frozenset(tool.name for tool in tools if tool.type == "function") + names = [] + for tool in tools: + if tool.type == "function": + names.append(tool.name) + elif tool.type == "namespace": + names.extend( + flat_namespace_tool_name(tool.name, namespaced_tool.name) + for namespaced_tool in tool.tools + if namespaced_tool.type == "function" + ) + return frozenset(names) def extract_tool_types(tools: list[Tool]) -> set[str]: @@ -358,7 +383,7 @@ def construct_tool_dicts( tool_dicts = None else: tool_dicts = [ - convert_tool_responses_to_completions_format(tool.model_dump()) - for tool in tools + convert_tool_responses_to_completions_format(tool) + for tool in iter_response_function_tool_dicts(tools) ] return tool_dicts diff --git a/vllm/entrypoints/openai/run_batch.py b/vllm/entrypoints/openai/run_batch.py index 58975b4f86b..6ae608da0ad 100644 --- a/vllm/entrypoints/openai/run_batch.py +++ b/vllm/entrypoints/openai/run_batch.py @@ -28,6 +28,7 @@ from urllib3.util import parse_url import vllm.envs as envs from vllm.config import config +from vllm.connections import global_http_connection from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient from vllm.entrypoints.openai.api_server import init_app_state @@ -493,18 +494,9 @@ async def download_bytes_from_url( # between urllib3 and aiohttp (e.g. backslash-@ attacks). url = url_spec.url - async with ( - aiohttp.ClientSession() as session, - session.get( - url, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) as resp, - ): - if resp.status != 200: - raise Exception( - f"Failed to download data from URL: {url}. Status: {resp.status}" - ) - return await resp.read() + return await global_http_connection.async_get_bytes( + url, allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS + ) else: raise ValueError( diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index b7b9ed5bc3e..db872b7f304 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -36,6 +36,39 @@ from .utils import ( ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest] +def _apply_post_tokenization_to_token_type_ids( + tokenizer: Any, + tok_params: TokenizeParams, + token_type_ids: list[int], +) -> list[int]: + pad_length = tok_params.pad_prompt_tokens + if pad_length is not None and pad_length < 0: + pad_length = tok_params.max_input_tokens + + if pad_length is not None and pad_length > len(token_type_ids): + pad_token_type_id = token_type_ids[-1] if token_type_ids else 0 + token_type_ids = token_type_ids + [pad_token_type_id] * ( + pad_length - len(token_type_ids) + ) + + max_length = tok_params.truncate_prompt_tokens + if max_length is not None and max_length < 0: + max_length = tok_params.max_input_tokens + + if max_length is None or max_length >= len(token_type_ids): + return token_type_ids + if max_length == 0: + return token_type_ids[:0] + + side = tok_params.truncation_side or ( + tokenizer.truncation_side if tokenizer is not None else None + ) + if side == "left": + return token_type_ids[-max_length:] + + return token_type_ids[:max_length] + + class ScoringIOProcessor(PoolingIOProcessor): name: str pooling_task: PoolingTask @@ -469,9 +502,16 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): else None, ) - if token_type_ids := engine_prompt.pop("token_type_ids", None): + token_type_ids = engine_prompt.pop("token_type_ids", None) + tok_params.apply_post_tokenization(self.tokenizer, engine_prompt) + + if token_type_ids is not None: params = pooling_params.clone() - compressed = compress_token_type_ids(token_type_ids) + compressed = compress_token_type_ids( + _apply_post_tokenization_to_token_type_ids( + self.tokenizer, tok_params, token_type_ids + ) + ) params.extra_kwargs = { **(params.extra_kwargs or {}), "compressed_token_type_ids": compressed, @@ -480,7 +520,6 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): else: pooling_params_list.append(pooling_params) - tok_params.apply_post_tokenization(self.tokenizer, engine_prompt) if engine_prompt_extras: target_prompt = extract_target_prompt(self.model_config, engine_prompt) target_prompt.update(engine_prompt_extras) @@ -618,12 +657,6 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): class JinaRankingIOProcessorMixin: - @staticmethod - def sanitize_input(text: str, special_tokens: dict[str, str]) -> str: - for token in special_tokens.values(): - text = text.replace(token, "") - return text - @staticmethod def format_docs_prompts_func( query: str, @@ -641,11 +674,13 @@ class JinaRankingIOProcessorMixin: if special_tokens is None: special_tokens = default_special_tokens - query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens) - docs = [ - JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens) - for doc in docs - ] + def sanitize_input(text: str) -> str: + for token in special_tokens.values(): + text = text.replace(token, "") + return text + + query = sanitize_input(query) + docs = [sanitize_input(doc) for doc in docs] prefix = ( "<|im_start|>system\n" @@ -668,6 +703,7 @@ class JinaRankingIOProcessorMixin: ) if instruction: + instruction = sanitize_input(instruction) prompt += f"\n{instruction}\n\n" doc_prompts = [ @@ -703,12 +739,24 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM ) -> Sequence[EngineInput]: queries = self.ensure_str(scoring_data.data_1) docs = self.ensure_str(scoring_data.data_2) + chat_template_kwargs = ( + prompt_extras.get("chat_template_kwargs") if prompt_extras else None + ) + instruction = ( + chat_template_kwargs.get("instruction") if chat_template_kwargs else None + ) if len(queries) == 1: - prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)] + prompts = [ + self.format_docs_prompts_func( + query=queries[0], docs=docs, instruction=instruction + ) + ] else: prompts = [ - self.format_docs_prompts_func(query=q, docs=[d]) + self.format_docs_prompts_func( + query=q, docs=[d], instruction=instruction + ) for q, d in zip(queries, docs) ] diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index e125007a549..613ff65ad0e 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -3,6 +3,7 @@ import time from typing import cast +import vllm.envs as envs from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse from vllm.entrypoints.openai.completion.protocol import CompletionResponse from vllm.entrypoints.openai.engine.protocol import ( @@ -28,6 +29,7 @@ from ..token_in_token_out.mm_serde import encode_mm_kwargs_item from ..token_in_token_out.protocol import ( DerenderChatRequest, DerenderCompletionRequest, + GenerateResponse, MultiModalFeatures, PlaceholderRangeInfo, ) @@ -51,6 +53,71 @@ class ServingDerender(BaseServing): self.online_derenderer = online_derenderer + def _validate_derender_bounds( + self, + generate_responses: list[GenerateResponse], + ) -> ErrorResponse | None: + """Reject derender payloads that exceed resource bounds. + + Runs before any tokenizer.decode() or parser invocation to prevent + CPU/memory exhaustion from oversized caller-supplied token structures. + """ + max_n = envs.VLLM_MAX_N_SEQUENCES + max_model_len = self.model_config.max_model_len + # See ModelConfig.max_logprobs for semantics and default value. + max_logprobs = self.model_config.max_logprobs + + if len(generate_responses) > max_n: + return self.create_error_response( + f"generate_responses count ({len(generate_responses)}) " + f"exceeds server maximum ({max_n}). " + f"Set VLLM_MAX_N_SEQUENCES to increase this limit." + ) + + for gen in generate_responses: + if len(gen.choices) > max_n: + return self.create_error_response( + f"choices count ({len(gen.choices)}) in response " + f"'{gen.request_id}' exceeds server maximum ({max_n})." + ) + + for choice in gen.choices: + if choice.token_ids and len(choice.token_ids) > max_model_len: + return self.create_error_response( + f"token_ids length ({len(choice.token_ids)}) in " + f"choice {choice.index} exceeds " + f"max_model_len ({max_model_len})." + ) + if choice.logprobs and choice.logprobs.content: + if len(choice.logprobs.content) > max_model_len: + return self.create_error_response( + f"logprobs.content length " + f"({len(choice.logprobs.content)}) in " + f"choice {choice.index} exceeds " + f"max_model_len ({max_model_len})." + ) + for entry in choice.logprobs.content: + if ( + max_logprobs >= 0 + and entry.top_logprobs + and len(entry.top_logprobs) > max_logprobs + ): + return self.create_error_response( + f"top_logprobs count " + f"({len(entry.top_logprobs)}) in " + f"choice {choice.index} exceeds " + f"max_logprobs ({max_logprobs})." + ) + + if gen.prompt_logprobs and len(gen.prompt_logprobs) > max_model_len: + return self.create_error_response( + f"prompt_logprobs length ({len(gen.prompt_logprobs)}) " + f"in response '{gen.request_id}' exceeds " + f"max_model_len ({max_model_len})." + ) + + return None + async def derender_chat_response( self, request: DerenderChatRequest, @@ -68,6 +135,10 @@ class ServingDerender(BaseServing): if error_check_ret is not None: return error_check_ret + bounds_error = self._validate_derender_bounds([request.generate_response]) + if bounds_error is not None: + return bounds_error + try: choices = await self.online_derenderer.derender_chat( request.generate_response, request.chat_request @@ -117,6 +188,13 @@ class ServingDerender(BaseServing): if error_check_ret is not None: return error_check_ret + if not request.generate_responses: + return self.create_error_response("generate_responses must not be empty") + + bounds_error = self._validate_derender_bounds(request.generate_responses) + if bounds_error is not None: + return bounds_error + ( choices, total_prompt_tokens, @@ -125,9 +203,6 @@ class ServingDerender(BaseServing): request.generate_responses, request.prompt_tokens ) - if not request.generate_responses: - return self.create_error_response("generate_responses must not be empty") - first = request.generate_responses[0] kv_params = first.kv_transfer_params if any( diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 48a1c4722dd..233ebf070c5 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -190,6 +190,13 @@ class GenerateResponseChoice(BaseModel): # or (b) ``enable_return_routed_experts`` is off server-side. routed_experts: str | None = None + @field_validator("token_ids") + @classmethod + def validate_token_ids(cls, v: list[int] | None) -> list[int] | None: + if v is not None and any(t < 0 for t in v): + raise ValueError("token_ids must not contain negative values") + return v + class GenerateResponseStreamChoice(BaseModel): index: int diff --git a/vllm/entrypoints/serve/utils/error_response.py b/vllm/entrypoints/serve/utils/error_response.py index 4dea1513a42..fc17a75c75a 100644 --- a/vllm/entrypoints/serve/utils/error_response.py +++ b/vllm/entrypoints/serve/utils/error_response.py @@ -27,12 +27,20 @@ def create_error_response( "create_error_response called with %s: %s", type(exc).__name__, exc ) - from vllm.exceptions import VLLMNotFoundError, VLLMValidationError + from vllm.exceptions import ( + VLLMNotFoundError, + VLLMUnprocessableEntityError, + VLLMValidationError, + ) if isinstance(exc, VLLMValidationError): err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = exc.parameter + elif isinstance(exc, VLLMUnprocessableEntityError): + err_type = "UnprocessableEntityError" + status_code = HTTPStatus.UNPROCESSABLE_ENTITY + param = exc.parameter elif isinstance(exc, VLLMNotFoundError): err_type = "NotFoundError" status_code = HTTPStatus.NOT_FOUND diff --git a/vllm/envs.py b/vllm/envs.py index 1f94be8ac5f..9305bb9cbe7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -95,12 +95,17 @@ if TYPE_CHECKING: VLLM_USE_PRECOMPILED_RUST: bool = False VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: bool = False VLLM_DOCKER_BUILD_CONTEXT: bool = False + VLLM_BUILD_COMMIT: str = "unknown" + VLLM_BUILD_PIPELINE: str = "local" + VLLM_BUILD_URL: str = "" + VLLM_IMAGE_TAG: str = "" VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False CMAKE_BUILD_TYPE: Literal["Debug", "Release", "RelWithDebInfo"] | None = None VERBOSE: bool = False VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds VLLM_MAX_N_SEQUENCES: int = 16384 + VLLM_MAX_COMPLETION_PROMPTS: int = 1024 VLLM_PLUGINS: list[str] | None = None VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None VLLM_LORA_RESOLVER_HF_REPO_LIST: str | None = None @@ -119,6 +124,7 @@ if TYPE_CHECKING: VLLM_USE_OINK_OPS: bool = False VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True VLLM_ROCM_USE_AITER: bool = False + VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False VLLM_ROCM_USE_AITER_LINEAR: bool = True VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False @@ -232,6 +238,7 @@ if TYPE_CHECKING: VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True VLLM_ALLREDUCE_USE_FLASHINFER: bool = False VLLM_TUNED_CONFIG_FOLDER: str | None = None + VLLM_ENABLE_STARTUP_PLAN: bool = False VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False @@ -618,6 +625,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_DOCKER_BUILD_CONTEXT": lambda: ( os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in ("1", "true") ), + # Build provenance metadata embedded in official vllm-openai images. + # Set via Docker ENV at image build time; informational only. + "VLLM_BUILD_COMMIT": lambda: os.environ.get("VLLM_BUILD_COMMIT", "unknown"), + "VLLM_BUILD_PIPELINE": lambda: os.environ.get("VLLM_BUILD_PIPELINE", "local"), + "VLLM_BUILD_URL": lambda: os.environ.get("VLLM_BUILD_URL", ""), + "VLLM_IMAGE_TAG": lambda: os.environ.get("VLLM_IMAGE_TAG", ""), # CMake build type # If not set, defaults to "Debug" or "RelWithDebInfo" # Available options: "Debug", "Release", "RelWithDebInfo" @@ -1057,6 +1070,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MAX_N_SEQUENCES": lambda: int( os.environ.get("VLLM_MAX_N_SEQUENCES", "16384") ), + # Maximum number of prompts allowed in a single /v1/completions request + # when the prompt field is a list. Prevents unbounded fan-out of engine + # requests from a single API call. Default: 1024. + "VLLM_MAX_COMPLETION_PROMPTS": lambda: int( + os.environ.get("VLLM_MAX_COMPLETION_PROMPTS", "1024") + ), # a list of plugin names to load, separated by commas. # if this is not set, it means all plugins will be loaded # if this is set to an empty string, no plugins will be loaded @@ -1146,6 +1165,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER": lambda: ( os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1") ), + # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's + # CudaCommunicator on ROCm. + "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1") + ), # Whether to use aiter paged attention. # By default is disabled. "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: ( @@ -1715,6 +1739,16 @@ environment_variables: dict[str, Callable[[], Any]] = { # Each component first checks this folder, then the configs shipped with # vLLM (if any). If no JSON matches, it uses a hard-coded heuristic. "VLLM_TUNED_CONFIG_FOLDER": lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None), + # Opt-in persistence of the startup plan. When enabled, each worker + # saves the memory-profiling result (the suggested --kv-cache-memory value + # and the free-memory baseline) under VLLM_CACHE_ROOT/startup_plan/, + # keyed by a hardware+config fingerprint, and later boots auto-apply it + # -- skipping memory profiling -- when the fingerprint matches and + # current free memory >= the recorded baseline. + # See vllm/v1/worker/startup_plan.py. + "VLLM_ENABLE_STARTUP_PLAN": lambda: bool( + int(os.getenv("VLLM_ENABLE_STARTUP_PLAN", "0")) + ), # Valid values are container,code_interpreter,web_search_preview # ex VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter # If the server_label of your mcp tool is not in this list it will @@ -2049,6 +2083,8 @@ def compile_factors() -> dict[str, object]: "VLLM_DEBUG_DUMP_PATH", "VLLM_PORT", "VLLM_CACHE_ROOT", + # Runtime memory-plan persistence; does not affect compiled graphs. + "VLLM_ENABLE_STARTUP_PLAN", "LD_LIBRARY_PATH", "VLLM_SERVER_DEV_MODE", "VLLM_DP_MASTER_IP", diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 931040b8ceb..4112c3de24b 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -64,3 +64,37 @@ class LoRAAdapterNotFoundError(VLLMNotFoundError): def __str__(self): return self.message + + +class VLLMUnprocessableEntityError(ValueError): + """vLLM-specific error for unprocessable entity requests. + + This exception is raised when the request content is invalid or cannot be + processed, such as when an image URL points to a non-existent or inaccessible + resource (404, 403, DNS failure, etc.). + + Args: + message: The error message describing the unprocessable entity. + parameter: Optional parameter name that failed validation. + value: Optional value that was rejected during validation. + """ + + def __init__( + self, + message: str, + *, + parameter: str | None = None, + value: Any = None, + ) -> None: + super().__init__(message) + self.parameter = parameter + self.value = value + + def __str__(self): + base = super().__str__() + extras = [] + if self.parameter is not None: + extras.append(f"parameter={self.parameter}") + if self.value is not None: + extras.append(f"value={self.value}") + return f"{base} ({', '.join(extras)})" if extras else base diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 10f400364ee..f57fc2c950b 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -83,7 +83,10 @@ class DPMetadata: num_tokens_across_dp_cpu: torch.Tensor, ) -> "DPMetadata": assert num_tokens_across_dp_cpu is not None - assert parallel_config.data_parallel_size > 1 + assert ( + parallel_config.data_parallel_size > 1 + or parallel_config.use_sequence_parallel_moe + ) assert parallel_config.is_moe_model is not False dp_rank = parallel_config.data_parallel_rank batchsize = num_tokens @@ -277,14 +280,20 @@ def set_forward_context( dp_metadata: DPMetadata | None = None if ( - vllm_config.parallel_config.data_parallel_size > 1 + ( + vllm_config.parallel_config.data_parallel_size > 1 + or vllm_config.parallel_config.use_sequence_parallel_moe + ) and vllm_config.parallel_config.is_moe_model is not False and (attn_metadata is not None or num_tokens is not None) ): # If num_tokens_across_dp hasn't already been initialized, then # initialize it here. Both DP padding and Microbatching will be # disabled. - if num_tokens_across_dp is None: + if ( + num_tokens_across_dp is None + and vllm_config.parallel_config.data_parallel_size > 1 + ): assert ubatch_slices is None assert num_tokens is not None _, num_tokens_across_dp, _ = coordinate_batch_across_dp( @@ -293,6 +302,9 @@ def set_forward_context( allow_microbatching=False, ) assert num_tokens_across_dp is not None + elif num_tokens_across_dp is None: + assert num_tokens is not None + num_tokens_across_dp = torch.tensor([num_tokens], dtype=torch.int32) dp_metadata = DPMetadata.make( vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp ) diff --git a/vllm/kernels/helion/config_manager.py b/vllm/kernels/helion/config_manager.py index ca37a68e810..052b515c37f 100644 --- a/vllm/kernels/helion/config_manager.py +++ b/vllm/kernels/helion/config_manager.py @@ -162,13 +162,6 @@ class ConfigSet: config_key, ) - def has_config(self, platform: str, config_key: CaseKey) -> bool: - platform = platform.lower() - platform_dict = self._configs.get(platform) - if platform_dict is None: - return False - return config_key in platform_dict - class ConfigManager: """File-level configuration management for Helion kernels (global singleton).""" @@ -327,15 +320,3 @@ class ConfigManager: logger.info("Saved config to: %s", platform_path) return platform_path - - def config_exists( - self, - kernel_name: str, - platform: str, - config_key: CaseKey, - ) -> bool: - platform_data = self._load_platform_file(kernel_name, platform) - if not platform_data: - return False - target = dict(config_key) - return any(entry["key"] == target for entry in platform_data) diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json new file mode 100644 index 00000000000..b50cabb0da9 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json @@ -0,0 +1,2099 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [ + null + ], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "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", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "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", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 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": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "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": [ + "last", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "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", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "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", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 256 + }, + "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": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "", + "first", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 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": [ + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "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", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "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", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "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", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "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": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "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": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "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": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "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": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "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": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4096 + }, + "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": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4096 + }, + "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": [ + "", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "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", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "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": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "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": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json new file mode 100644 index 00000000000..ca411c69470 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json @@ -0,0 +1,2207 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "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": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "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": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "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": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "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": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "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": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "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": [ + "first", + "first", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "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", + "", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "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": [ + "", + "first", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "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": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "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": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "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": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "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", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "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": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "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", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "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": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "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", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "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", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "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": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "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": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "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": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "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": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 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", + "first", + "" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 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", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 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", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "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": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "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": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 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", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "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", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "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": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 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", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "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": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "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": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "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", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "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": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 16 + ], + "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": [ + "last", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..f3aaf226b04 --- /dev/null +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -0,0 +1,252 @@ +# 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.kernels.helion.utils import ( + get_fp8_dtype, + get_int8_min_max, + get_int8_min_scaling_factor, +) +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] + intermediate_size_list = [6144, 12288, 25600] + + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + group_size_list = [128] + inputs = {} + for intermediate_size, group_size, num_tokens in product( + intermediate_size_list, group_size_list, num_tokens_list + ): + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + + config_key = CaseKey( + { + "intermediate_size": intermediate_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = (result, input, scale, group_size, scale_ub, 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 intermediate_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 intermediate_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 + + result, _, _, group_size, *_ = args + num_tokens, intermediate_size = result.shape + + cache_key = (num_tokens, group_size, intermediate_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["intermediate_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size)) + best_group_size = min( + configs[best_intermediate_size], key=lambda s: abs(s - group_size) + ) + available_num_tokens = sorted(configs[best_intermediate_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( + { + "intermediate_size": best_intermediate_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + return + + +def baseline( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + torch.ops._C.silu_and_mul_per_block_quant( + out, input, scales, group_size, scale_ub, is_scale_transposed + ) + + +@register_kernel( + mutates_args=["out", "scales"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + ignore_warnings=[helion.exc.TensorOperationInWrapper], + ), +) # type: ignore[misc] +def silu_and_mul_per_block_quant( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, # dummy +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, two_intermediate_size = input.shape + hl.specialize(two_intermediate_size) + + assert two_intermediate_size % 2 == 0 + intermediate_size = two_intermediate_size // 2 + + assert out.shape[0] == num_tokens + assert out.shape[1] == intermediate_size + fp8_dtype = get_fp8_dtype() + assert out.dtype in [fp8_dtype, torch.int8] + + if scale_ub is not None: + assert out.dtype == fp8_dtype + assert scale_ub.dtype == torch.float32 + + assert scales.ndim == 2 and scales.dtype == torch.float32 + + assert scales.shape[0] == num_tokens + groups_per_row = scales.shape[1] + hl.specialize(groups_per_row) + assert ( + intermediate_size % group_size == 0 + and intermediate_size // group_size == groups_per_row + ) + + assert group_size in [64, 128] + hl.specialize(group_size) + + assert input.stride()[-1] == 1 + assert out.stride()[-1] == 1 + + quant_dtype = out.dtype + qtype_traits_min: int | float + qtype_traits_max: int | float + if quant_dtype == torch.int8: + qtype_traits_min, qtype_traits_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_traits_min, qtype_traits_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_traits_max * 512.0) + + qtype_max = float(qtype_traits_max) + + input = input.view(num_tokens, -1, group_size) + out = out.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_a_blk = input[tile_m, tile_gn, tile_n].to(torch.float32) + x_b_blk = hl.load( + input, + [tile_m, tile_gn.index + groups_per_row, tile_n], + extra_mask=(tile_gn.index + groups_per_row < 2 * groups_per_row)[ + None, :, None + ], + ).to(torch.float32) + x_blk = x_a_blk * torch.sigmoid(x_a_blk) * x_b_blk + s_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32) + + if scale_ub is not None: + scale_ub_s = hl.load(scale_ub, []) + s_blk = s_blk.clamp(max=scale_ub_s) + s_blk = s_blk * (1.0 / qtype_max) + s_blk = s_blk.clamp(min=min_scaling_factor) + + scales[tile_m, tile_gn] = s_blk + if quant_dtype == torch.int8: + y_blk = (x_blk * (1.0 / s_blk[:, :, None])).round() + else: + y_blk = x_blk / s_blk[:, :, None] + + out[tile_m, tile_gn, tile_n] = y_blk.clamp( + qtype_traits_min, qtype_traits_max + ).to(out.dtype) diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index bff3c0cf454..5c8e829b299 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -89,7 +89,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): vllm_config = get_current_vllm_config() self._lora_stream = _get_lora_aux_cuda_stream() assert current_platform.is_cuda_alike() - self._events = [torch.Event(), torch.Event()] + self._events = [torch.cuda.Event(), torch.cuda.Event()] # lora_linear avoids prefix conflicts with the base layer self.layer_name = self.base_layer.prefix + ".lora_linear_async" compilation_config = vllm_config.compilation_config diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 39f60aad5db..63a4ea9a829 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -118,7 +118,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): def _init_lora_stream_context(self) -> None: self._lora_stream: torch.cuda.Stream | None = None - self._events: tuple[torch.Event, ...] | None = None + self._events: tuple[torch.cuda.Event, ...] | None = None if not self._enable_aux_cuda_stream: return if not current_platform.is_cuda_alike(): @@ -127,7 +127,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse # the same event objects; reuse-within-a-pair is fine because the # second pair starts only after intermediate_cache1.add_() has joined. - self._events = tuple(torch.Event() for _ in range(4)) + self._events = tuple(torch.cuda.Event() for _ in range(4)) def _build_lora_context(self): use_dual_stream = ( diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index bc3d278af4e..8b5d97c1e2e 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -117,6 +117,7 @@ class LoRAModelManager: self.packed_modules: dict[str, list[str]] = {} self.modules: dict[str, BaseLayerWithLoRA] = {} self._last_mapping: LoRAMapping | None = None + self._last_slot_layout: tuple[int | None, ...] | None = None is_moe = is_moe_model(self.model) self._is_moe = is_moe @@ -1147,9 +1148,14 @@ class LoRAModelManager: return True def set_adapter_mapping(self, mapping: LoRAMapping) -> None: - if self._last_mapping != mapping: + # The punica metadata derives from the slot layout as well as the + # mapping: an out-of-band add_lora() can LRU-evict and reassign slots + # while the running batch, and thus the mapping, is unchanged. + slot_layout = tuple(self.lora_index_to_id) + if self._last_mapping != mapping or self._last_slot_layout != slot_layout: self._set_adapter_mapping(mapping) self._last_mapping = mapping + self._last_slot_layout = slot_layout def remove_adapter(self, adapter_id: int) -> bool: self.deactivate_adapter(adapter_id) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index f5e3a16d71b..5d1d3707308 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -74,6 +74,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import ( from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( FlashInferMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.humming import ( + HummingMxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.marlin import ( MarlinMxFp4LinearKernel, ) @@ -91,6 +94,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp8.humming import ( + HummingMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8.marlin import ( MarlinMxfp8LinearKernel, ) @@ -120,6 +126,9 @@ from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.nvfp4.humming import ( + HummingNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.marlin import ( MarlinNvFp4LinearKernel, ) @@ -154,6 +163,10 @@ from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( FlashInferFp8DeepGEMMDynamicBlockScaledKernel, FlashInferFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.humming import ( + HummingFP8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( MarlinFP8ScaledMMLinearKernel, ) @@ -225,6 +238,14 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "flashinfer_b12x": { FlashInferB12xNvFp4LinearKernel, }, + "humming": { + HummingFP8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, + HummingLinearKernel, + HummingMxfp8LinearKernel, + HummingMxFp4LinearKernel, + HummingNvFp4LinearKernel, + }, "marlin": { MarlinFP8ScaledMMLinearKernel, MarlinLinearKernel, @@ -292,6 +313,7 @@ _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] PlatformEnum.CUDA: [ CutlassInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, + HummingInt8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [AiterInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel], } @@ -304,6 +326,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = CutlassFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, + HummingFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ AiterHipbMMPerTokenFp8ScaledMMLinearKernel, @@ -335,6 +358,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ CutlassFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, + HummingFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ AiterFp8BlockScaledMMKernel, @@ -351,6 +375,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ _POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { PlatformEnum.CUDA: [ + HummingFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ], PlatformEnum.ROCM: [ @@ -371,10 +396,10 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { MacheteLinearKernel, AllSparkLinearKernel, MarlinLinearKernel, - HummingLinearKernel, ConchLinearKernel, ExllamaLinearKernel, TritonW4A16LinearKernel, + HummingLinearKernel, ], PlatformEnum.ROCM: [ RDNA3W4A16LinearKernel, @@ -400,6 +425,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { FlashInferCutlassMxfp8LinearKernel, MarlinMxfp8LinearKernel, EmulationMxfp8LinearKernel, + HummingMxfp8LinearKernel, ], PlatformEnum.ROCM: [ # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and @@ -426,6 +452,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { FlashInferCudnnNvFp4LinearKernel, FbgemmNvFp4LinearKernel, EmulationNvFp4LinearKernel, + HummingNvFp4LinearKernel, ], PlatformEnum.ROCM: [ EmulationNvFp4LinearKernel, @@ -436,6 +463,7 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { PlatformEnum.CUDA: [ FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, + HummingMxFp4LinearKernel, ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py index c0b8c35bbd5..f9b4bd6e435 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py @@ -19,6 +19,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import ( from vllm.model_executor.kernels.linear.mixed_precision.exllama import ( ExllamaLinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.humming import ( + HummingLinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.machete import ( MacheteLinearKernel, ) @@ -52,6 +55,7 @@ __all__ = [ "CutlassW4A8LinearKernel", "Dynamic4bitLinearKernel", "ExllamaLinearKernel", + "HummingLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", "RDNA3W4A16LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 764c0f4227f..7f1a9024ee1 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -23,8 +23,6 @@ class HummingLinearKernel(MPLinearKernel): return False, "Humming is not installed" if c.has_g_idx: return False, "Humming does not support act-order (g_idx)" - if c.zero_points: - return False, "Humming linear kernel only supports symmetric weights" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -41,6 +39,11 @@ class HummingLinearKernel(MPLinearKernel): "group_size": 0 if group_size == -1 else group_size, } + if self.config.zero_points: + assert self.w_zp_name is not None + name_map["zero_point"] = self.w_zp_name + quant_config["has_zero_point"] = True + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) prepare_humming_layer(layer, quant_config) diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py new file mode 100644 index 00000000000..d93f5d48158 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + + +class HummingMxFp4LinearKernel(MxFp4LinearKernel): + """Humming GEMM Kernel for MXFP4.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + + quant_config = { + "quant_method": "humming", + "dtype": "float4e2m1", + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/mxfp8/humming.py b/vllm/model_executor/kernels/linear/mxfp8/humming.py new file mode 100644 index 00000000000..ed1cd39cbd6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/humming.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class HummingMxfp8LinearKernel(Mxfp8LinearKernel): + """Humming GEMM Kernel for MXFP8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + + quant_config = { + "quant_method": "humming", + "dtype": "float8e4m3", + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 364608a806a..9714f92dffd 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -89,13 +89,7 @@ def _mxfp8_dot_scaled_linear( N = w.shape[0] x_q, x_scale = mxfp8_e4m3_quantize(x) out = torch.empty((M, N), dtype=x.dtype, device=x.device) - # Regime-gated launch tiles for gfx950, tuned at MiniMax-M3 shapes: - # for example, 8k/1k, 1k/1k - if M >= 1024: - BLOCK_M, BLOCK_N, num_warps, num_stages = 128, 256, 8, 2 - else: - BLOCK_M, BLOCK_N, num_warps, num_stages = 64, 64, 4, 2 - BLOCK_K = 128 + BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K) grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _mxfp8_linear_kernel[grid]( x_q, @@ -125,6 +119,81 @@ def _mxfp8_dot_scaled_linear( return out +def _select_cfg(M, N, K): + """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950. + + The M-bucketed, shape-adaptive tile selection here is the speedup over the + upstream 2-bucket launcher. Tiles are pipelined (num_stages>=2, larger BLOCK_K) + and occupancy- and shape-aware: keyed on the LOCAL (M, N, K), so it adapts to the + TP-sharded shapes (e.g. MiniMax-M3 TP=4 vs TP=8, where local N and K differ) — + large-K prefill uses BLOCK_K=256; short-K (K=768) widens N. BLOCK_K must divide K + (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible + (served K: 384/768/1024/2048/6144). + """ + if M <= 64: + # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The + # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in + # parallel) + LARGE BLOCK_K (fewer K-iters, bigger coalesced weight loads). + # Tuned by CUDA-graph replay latency. Optimal at both TP=4 and TP=8. + if K % 1024 == 0: # K=2048, 6144 -> graph-best 16x16x1024 (all M) + return 16, 16, 1024, 2, 2 + if K % 512 == 0: + return 16, 16, 512, 2, 3 + if K % 256 == 0: # K=768 (shared_down) -> graph-best 16x32x256 + return 16, 32, 256, 4, 3 + return 16, 32, 128, 4, 3 + # mid-M (65..256) on SMALL local-N: still occupancy-bound (a 64x64 tile makes too + # few N-tiles), so the narrow-BLOCK_N decode-style tile fills the CUs better. + # N<=1536 covers the real fused-qkv local N at TP=8: q heads shard but the GQA KV + # (4) + sparse-indexer (4) heads are < TP=8, so vLLM replicates them to 1/rank -> + # N = 1024 + 4*128 = 1536 (not 2560/2=1280). For the wider 1280= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256: + return 128, 128, 256, 8, 3 + return 128, 128, 128, 8, 3 + # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked): + # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since + # the entry requires K%128==0) for any other K to stay correct. + if M <= 256: # conc~128 decode + small prefill chunk: occupancy tile + if K % 512 == 0: + return 64, 64, 512, 8, 2 + return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2) + if M <= 1024: # medium prefill chunk: BN=64 keeps small-N occupied + return (128, 64, 256, 8, 3) if K % 256 == 0 else (128, 64, 128, 8, 3) + # large prefill (M > 1024). The previously graph-tuned tall 256x128x256 and deep + # 128x128x512 tiles won only on triton 3.6; on triton 3.7 their larger BLOCK_M / + # BLOCK_K register+LDS footprint spills (or hits "out of resources" on stricter + # ROCm/triton builds). The 128x128x256 tile is equal-or-faster on triton 3.7 (the + # M=4096,N=2048,K=6144 shape: ~104us vs the 256x128x256 tile's ~184us), within ~5% + # on 3.6, and inside the footprint of the tiles used elsewhere in this selector. + # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K + # / very-large-M shapes. + if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))): + return 128, 128, 256, 8, 3 + # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the + # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there). + if N <= 1024 and K % 256 == 0: + return 128, 64, 256, 8, 3 + return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3) + + class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py new file mode 100644 index 00000000000..7e390343854 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + +logger = init_logger(__name__) + + +class HummingNvFp4LinearKernel(NvFp4LinearKernel): + """Humming GEMM Kernel for NVFP4.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Route through humming's compressed-tensors nvfp4 loader (same path as + # the MoE oracle); the native group_tensor schema mishandles a scalar + # global scale. + quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "type": "float", + "num_bits": 4, + "strategy": "group", + "group_size": 16, + } + # CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`. + if not hasattr(layer, "weight_packed"): + layer.weight_packed = layer.weight + del layer.weight + # The CT linear scheme inverts the global scale (1/scale) for + # marlin/cutlass; humming wants the original. + layer.weight_global_scale = torch.nn.Parameter( + 1.0 / layer.weight_global_scale, requires_grad=False + ) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/humming.py b/vllm/model_executor/kernels/linear/scaled_mm/humming.py new file mode 100644 index 00000000000..7b8ed21fbd6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/humming.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_linear_layer_to_humming_standard, + prepare_humming_layer, +) +from vllm.platforms import current_platform + +from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, + Int8ScaledMMLinearKernel, + Int8ScaledMMLinearLayerConfig, +) + +logger = init_logger(__name__) + + +class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Humming GEMM Kernel for FP8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from vllm.utils.humming import dtypes + + name_map = {"weight": "weight", "weight_scale": "weight_scale"} + scale_torch_dtype = self.config.weight_quant_key.scale.dtype + scale_dtype = dtypes.DataType.from_torch_dtype(scale_torch_dtype) + + quant_config = { + "quant_method": "humming", + "dtype": "float8e4m3", + "scale_dtype": scale_dtype, + } + + assert self.config.weight_quant_key.scale2 is None + scale_group_shape = self.config.weight_quant_key.scale.group_shape + if scale_group_shape.is_per_tensor(): + quant_config["weight_scale_type"] = "tensor" + if not hasattr(layer, "global_scale") and hasattr(layer, "weight_scale"): + del name_map["weight_scale"] + name_map["global_scale"] = "weight_scale" + elif scale_group_shape.is_per_channel(): + quant_config["weight_scale_type"] = "channel" + elif scale_group_shape.is_per_group(): + quant_config["weight_scale_type"] = "group" + quant_config["group_size"] = scale_group_shape.col + else: + assert scale_group_shape.row > 0 and scale_group_shape.col > 0 + quant_config["weight_scale_type"] = "block" + quant_config["weight_scale_group_size_n"] = scale_group_shape.row + quant_config["weight_scale_group_size"] = scale_group_shape.col + + if hasattr(layer, "weight_scale_inv"): + name_map["weight_scale"] = "weight_scale_inv" + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + pass + + +class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): + """Humming GEMM Kernel for INT8.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement( + cls, config: Int8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_name, weight_scale_name, *_ = self.layer_param_names + name_map = {"weight": weight_name, "weight_scale": weight_scale_name} + quant_config = {"quant_method": "humming", "dtype": "int8"} + weight = getattr(layer, weight_name) + weight.data = weight.data + 128 + + convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) + prepare_humming_layer(layer, quant_config) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.utils.humming import HummingMethod + + flatten_inputs = x.view(-1, x.size(-1)) + output = HummingMethod.forward_layer( + layer=layer, + inputs=flatten_inputs, + compute_config=layer.compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index 3d23b02bea8..f30d6ced1d3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -197,29 +197,6 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): return False, "XPUFp8BlockScaledMM only support on XPU" return True, None - def process_weights_after_loading(self, layer: torch.nn.Module): - super().process_weights_after_loading(layer) - scale_attr = ( - "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale" - ) - scale = getattr(layer, scale_attr) - # Models with scale_fmt=ue8m0 (e.g. DeepSeek-V4) store weight scales - # as float8_e8m0fnu. The oneDNN fp8_gemm kernel dispatches to its - # "block quant" path only when NEITHER scale is e8m0: - # - # is_block_quant = (m1_sc != e8m0) && (m2_sc != e8m0) && ... - # - # Since activation scales are always float32 (use_ue8m0=False on XPU, - # DeepGEMM requires Hopper/Blackwell), an e8m0 weight scale causes - # is_block_quant=false and falls into the wrong per-channel path, - # producing NaN. Converting e8m0→float32 here at load time (one-time, - # negligible overhead for small scale tensors) ensures the kernel sees - # matching dtypes and correctly enters the block-quant path with the - # actual group_size derived from scale tensor shapes. - if scale.dtype == torch.float8_e8m0fnu: - scale = scale.to(torch.float32) - replace_parameter(layer, scale_attr, scale.data.t().contiguous()) - def apply_block_scaled_mm( self, A: torch.Tensor, @@ -228,11 +205,12 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): Bs: torch.Tensor, ) -> torch.Tensor: # Weight is [N, K]. Use .t() to create a [K, N] view without copying. + # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN. return torch.ops._xpu_C.fp8_gemm( A, B.t(), self.config.out_dtype, As, - Bs, + Bs.t().contiguous(), torch.Tensor(), ) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 0115912ce4c..7c3fe4407d6 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -129,10 +129,12 @@ class SiluAndMul(CustomOp): def __init__(self, *, compile_native: bool = True): super().__init__(compile_native=compile_native) - if current_platform.is_cuda_alike() or current_platform.is_xpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ): self.op = torch.ops._C.silu_and_mul - elif current_platform.is_cpu(): - self._forward_method = self.forward_native @staticmethod def forward_native(x: torch.Tensor) -> torch.Tensor: @@ -150,6 +152,11 @@ class SiluAndMul(CustomOp): def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return self.forward_cuda(x) + return self.forward_native(x) + @CustomOp.register("silu_and_mul_with_clamp") class SiluAndMulWithClamp(CustomOp): @@ -417,14 +424,14 @@ class GeluAndMul(CustomOp): self.op(out, x) return out - def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: - if self.op: - return self.forward_cuda(x) - return self.native(x) - def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return self.forward_cuda(x) + return self.forward_native(x) + def extra_repr(self) -> str: return f"approximate={repr(self.approximate)}" @@ -526,6 +533,11 @@ class NewGELU(CustomOp): def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return self.forward_cuda(x) + return self.forward_native(x) + # --8<-- [start:gelu_fast] @CustomOp.register("gelu_fast") @@ -553,6 +565,11 @@ class FastGELU(CustomOp): def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return self.forward_cuda(x) + return self.forward_native(x) + # --8<-- [start:quick_gelu] @CustomOp.register("quick_gelu") @@ -581,6 +598,11 @@ class QuickGELU(CustomOp): def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC: + return self.forward_cuda(x) + return self.forward_native(x) + # --8<-- [start:relu2] @CustomOp.register("relu2") diff --git a/vllm/model_executor/layers/fla/ops/layernorm_guard.py b/vllm/model_executor/layers/fla/ops/layernorm_guard.py index 8b9e275737e..279b06b5ebc 100644 --- a/vllm/model_executor/layers/fla/ops/layernorm_guard.py +++ b/vllm/model_executor/layers/fla/ops/layernorm_guard.py @@ -331,46 +331,6 @@ def rmsnorm_fn( ) -class LayerNormGated(nn.Module): - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = True, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - torch.nn.init.zeros_(self.bias) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" - return layernorm_fn( - x, - self.weight, - self.bias, - z=z, - group_size=self.group_size, - eps=self.eps, - norm_before_gate=self.norm_before_gate, - ) - - class RMSNormGated(nn.Module): def __init__( self, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 55c7238f642..ad2ed510b33 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1038,7 +1038,7 @@ class FusedMoEParallelConfig: @property def use_all2all_kernels(self): - return self.dp_size > 1 and self.use_ep + return self.use_ep and (self.dp_size > 1 or self.is_sequence_parallel) @property def use_deepep_ht_kernels(self): diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index 3cbab0a0d54..c5330f3b438 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -78,9 +78,6 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): @staticmethod def _supports_parallel_config(moe_parallel_config) -> bool: - # Both TP (expert_map=None) and EP are supported: apply() forwards the - # expert_map as aiter's ``expert_mask`` (the per-rank local-expert - # selection), mirroring the native rocm_aiter_moe path. return True @staticmethod @@ -129,17 +126,21 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): limit = self.quant_config.gemm1_clamp_limit swiglu_limit = 0.0 if limit is None else float(limit) - # Under EP, aiter expects ``expert_mask`` as a 0/1 *local-expert* mask - # over global ids with a trailing fake-expert sentinel slot - # (shape ``[global_num_experts + 1]``), NOT vLLM's expert_map (a - # global->local index map with -1 for non-local). Convert it; aiter - # derives the global->local compaction from the mask itself. ``None`` - # under pure TP. - if expert_map is not None: + # Under EP, aiter expects ``expert_mask``: a 0/1 *local-expert* mask over + # global ids with a trailing fake-expert sentinel slot (shape + # ``[global_num_experts + 1]``), from which it derives the global->local + # compaction. What ``RoutedExperts.expert_map`` hands us depends on the + # aiter master switch (``rocm_aiter_fmoe_enabled``). + # Branching on the (static) master flag — not the tensor contents — + # keeps this HIP-graph/torch.compile safe (no data-dependent sync). + # ``None`` under pure TP. + if expert_map is None: + expert_mask = None + elif self.moe_config.rocm_aiter_fmoe_enabled: + expert_mask = expert_map + else: local_mask = (expert_map >= 0).to(torch.int32) expert_mask = torch.cat([local_mask, local_mask.new_zeros(1)]) - else: - expert_mask = None # Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the # call is captured under HIP graphs / torch.compile (a direct diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index 275f80c68fe..90972fc7f0c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -210,7 +210,7 @@ def persistent_masked_m_silu_mul_quant( DeepGemmQuantScaleFMT.UE8M0, ] - device_capability = current_platform.get_device_capability(device_id=y.device.index) + device_capability = current_platform.get_device_capability() assert device_capability is not None cuda_arch = device_capability.to_int() diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py index c21072676ec..f4933e1bd76 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py @@ -14,12 +14,13 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + kInt4W4A8StaticChannelSym, kInt4W4A8StaticGroup32Sym, kInt4W4A8StaticGroup64Sym, kInt4W4A8StaticGroup128Sym, kInt4W4A8StaticGroupSym, ) -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): @@ -48,6 +49,48 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + if ( + not current_platform.is_cpu() + or current_platform.get_cpu_architecture() != CpuArchEnum.ARM + ): + return False, "kernel only supports Arm CPU" + + if moe_config.in_dtype not in ( + torch.float32, + torch.bfloat16, + torch.float16, + ): + return ( + False, + f"kernel does not support {moe_config.in_dtype} input/output dtype", + ) + + try: + _ = torch.ops.aten._dyn_quant_matmul_4bit + _ = torch.ops.aten._dyn_quant_pack_4bit_weight + except AttributeError: + return ( + False, + f"PyTorch {torch.__version__} does not support " + "_dyn_quant_* 4bit ops. Install a newer version", + ) + + return mk.FusedMoEExperts.is_supported_config( + cls, + moe_config, + weight_key, + activation_key, + activation_format, + ) + @staticmethod def _supports_current_device() -> bool: return current_platform.is_cpu() @@ -86,8 +129,9 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): Can be channel-wise or group-wise quantization - Activations: dynamic per-token 8-bit integer quantization """ - # group size must be multiple of 32 + # channelwise or groupwise with group size being a multiple of 32 SUPPORTED_W_A = [ + (kInt4W4A8StaticChannelSym, None), (kInt4W4A8StaticGroup128Sym, None), (kInt4W4A8StaticGroup64Sym, None), (kInt4W4A8StaticGroup32Sym, None), @@ -120,7 +164,8 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): """Expert parallelism not yet supported.""" return False - def _activation_kind(self, activation: MoEActivation) -> int: + @staticmethod + def _activation_kind(activation: MoEActivation) -> int: """Convert MoEActivation to kernel activation kind integer. Returns: @@ -200,29 +245,22 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): e_score_correction_bias=e_score_correction_bias, ) - # Extract dimensions from weight tensors - # w1 is w13_packed: [num_experts, packed_data...] - # w2 is w2_packed: [num_experts, packed_data...] - # These dimensions should be available from the layer - # For now, we'll extract from moe_config - K = self.moe_config.hidden_dim - N = self.moe_config.intermediate_size_per_partition + hidden_size = self.moe_config.hidden_dim + intermediate_size = self.moe_config.intermediate_size_per_partition assert self.quant_config.block_shape is not None - if self.quant_config.is_per_act_token: + # C++ kernel expects an int: -1 for channelwise, and group size for groupwise + if self.quant_config.block_shape == [-1, 1]: group_size = -1 else: group_size = self.quant_config.block_shape[1] - - # Call the dynamic 4-bit int MoE kernel return torch.ops._C.dynamic_4bit_int_moe( hidden_states, topk_ids.to(torch.long), topk_weights, w1, # w13_weight_packed w2, # w2_weight_packed - K, # hidden_size (w2_out_features) - N, # intermediate_size (w2_in_features) - N * 2, # 2*intermediate_size (w13_out_features) + hidden_size, + intermediate_size, group_size, apply_router_weight_on_input, self._activation_kind(activation), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py index 38200d9d090..ca23e127249 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -20,7 +22,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( - flashinfer_b12x_fused_moe, flashinfer_convert_sf_to_mma_layout, has_flashinfer_b12x_moe, ) @@ -42,6 +43,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. """ + _ACTIVATION_MAP: dict[MoEActivation, str] = { + MoEActivation.SILU: "silu", + MoEActivation.RELU2_NO_MUL: "relu2", + } + def __init__( self, moe_config: FusedMoEConfig, @@ -60,6 +66,30 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): # one. Holding it on the instance keeps apply() alloc-free. self._fc2_input_scale: torch.Tensor | None = None + # Shape params for B12xMoEWrapper construction. + self.global_num_experts = moe_config.num_experts + self.topk = moe_config.experts_per_token + self.hidden_dim = moe_config.hidden_dim + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.max_num_tokens = moe_config.max_num_tokens + self.local_expert_offset = self.ep_rank * self.num_local_experts + + activation = moe_config.activation + if activation not in self._ACTIVATION_MAP: + raise ValueError( + f"FlashInferB12xExperts does not support " + f"activation {activation!r}. " + f"Supported: {list(self._ACTIVATION_MAP.keys())}" + ) + self._activation_str = self._ACTIVATION_MAP[activation] + + # Lazily created on first apply() call. + self._wrapper: Any | None = None + self.w1_sf_mma: torch.Tensor | None = None + self.w2_sf_mma: torch.Tensor | None = None + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Normalise block scales to absorb the per-expert weight global scale # (w_gs). vLLM's NVFP4 convention stores: @@ -141,7 +171,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( @@ -158,11 +188,13 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU + return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return True + # B12xMoEWrapper does not yet support expert parallelism: its local + # expert count must equal the global expert count. + return not moe_parallel_config.use_ep def supports_expert_map(self) -> bool: return False @@ -190,13 +222,29 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: - # b12x_fused_moe expects BF16 hidden states and performs its own FP4 + # B12xMoEWrapper expects BF16 hidden states and performs its own FP4 # quantization internally. Returning True prevents the modular kernel - # from pre-quantizing activations, which would produce an FP4-packed - # tensor with size(-1)=k//2 and break the scale-factor conversion that - # expects size(-1)=k. + # from pre-quantizing activations. return True + def _ensure_wrapper(self) -> None: + """Lazily create B12xMoEWrapper on first use.""" + if self._wrapper is not None: + return + + from flashinfer.fused_moe import B12xMoEWrapper + + self._wrapper = B12xMoEWrapper( + num_experts=self.global_num_experts, + top_k=self.topk, + hidden_size=self.hidden_dim, + intermediate_size=self.intermediate_size_per_partition, + use_cuda_graph=True, + max_num_tokens=self.max_num_tokens, + num_local_experts=self.num_local_experts, + activation=self._activation_str, + ) + def apply( self, output: torch.Tensor, @@ -224,13 +272,16 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): assert self._fc2_input_scale is not None, ( "_fc2_input_scale must be set by process_weights_after_loading" ) + assert self.w1_sf_mma is not None and self.w2_sf_mma is not None, ( + "process_weights_after_loading must run before FlashInferB12xExperts.apply" + ) - top_k = topk_ids.shape[1] + self._ensure_wrapper() + wrapper = self._wrapper + assert wrapper is not None - flashinfer_b12x_fused_moe( + wrapper_output = wrapper.run( x=hidden_states, - token_selected_experts=topk_ids.to(torch.int32), - token_final_scales=topk_weights, w1_weight=w1, w1_weight_sf=self.w1_sf_mma, w1_alpha=self.g1_alphas, @@ -238,9 +289,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): w2_weight=w2, w2_weight_sf=self.w2_sf_mma, w2_alpha=self.g2_alphas, - num_experts=global_num_experts, - top_k=top_k, - num_local_experts=self.num_local_experts, - output_dtype=self.out_dtype, - output=output, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights, ) + output.copy_(wrapper_output) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 047ae46c0d3..5f112380cf9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -38,15 +38,20 @@ from vllm.model_executor.layers.fused_moe.utils import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + kFp8Dynamic128Sym, kFp8DynamicTokenSym, kFp8Static128BlockSym, kFp8StaticChannelSym, + kFp8StaticTensorSym, kInt4Static, + kInt8DynamicTokenSym, kInt8Static, + kInt8StaticChannelSym, kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, kMxfp8Static, + kNvfp4Dynamic, kNvfp4Static, ) from vllm.platforms import current_platform @@ -61,9 +66,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def get_humming_moe_gemm_type() -> str | None: +def get_humming_moe_gemm_type() -> str: env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE - gemm_type = None + gemm_type = "indexed" if env_gemm_type is not None: env_gemm_type = env_gemm_type.lower() if env_gemm_type == "indexed": @@ -87,7 +92,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): num_dispatchers: int | None = None, ): self.layer = layer - self.num_experts = self.layer.num_experts + self.num_experts = self.layer.local_num_experts self.global_num_experts = self.layer.global_num_experts self.init_humming_moe() @@ -186,6 +191,24 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): (kInt4Static, kFp8DynamicTokenSym), (kInt8Static, None), (kInt8Static, kFp8DynamicTokenSym), + # Checkpoint-driven (weight, activation) pairs the dense/MoE oracles + # pass. Humming defers input quant (see expects_unquantized_inputs), + # so the activation key does not constrain support. + # fp8 (compressed-tensors / native / modelopt) + (kFp8StaticChannelSym, kFp8StaticTensorSym), + (kFp8StaticChannelSym, kFp8Dynamic128Sym), + (kFp8StaticTensorSym, None), + (kFp8StaticTensorSym, kFp8DynamicTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8Dynamic128Sym), + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + # int8 (compressed-tensors w8a8 / experts_int8) + (kInt8StaticChannelSym, None), + (kInt8StaticChannelSym, kInt8DynamicTokenSym), + # nvfp4 (compressed-tensors / modelopt / quark) + (kNvfp4Static, kNvfp4Dynamic), + # mxfp8 (compressed-tensors / modelopt / online) + (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A @@ -463,13 +486,12 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): assert hasattr(cls, "humming_gemm_type") gemm_type = cls.humming_gemm_type().value.lower() preferred_gemm_type = get_humming_moe_gemm_type() - if preferred_gemm_type is not None: - supported = preferred_gemm_type.lower() == gemm_type - if not supported: - reason = ( - f"preferred gemm type {preferred_gemm_type} != " - f"supported gemm type {gemm_type}" - ) + supported = preferred_gemm_type.lower() == gemm_type + if not supported: + reason = ( + f"preferred gemm type {preferred_gemm_type} != " + f"supported gemm type {gemm_type}" + ) return supported, reason diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index dd7429d76e1..117f744aeea 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py @@ -51,7 +51,7 @@ class MoELoRAContext: # Events are paired one-per-overlap-pair: events[0,1] for w13, # events[2,3] for w2, so the two pairs do not race on the same event. aux_stream: torch.cuda.Stream | None = None - events: tuple[torch.Event, ...] | None = None + events: tuple[torch.cuda.Event, ...] | None = None # Per-rank token→LoRA mapping after EP dispatch. Set by # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read diff --git a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py index b511e368f4a..9839756880a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py @@ -19,7 +19,6 @@ and the top-k weighted reduction run in PyTorch between/after the two GEMMs. import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( Mxfp8TritonExpertsBase, ) @@ -32,7 +31,26 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -logger = init_logger(__name__) + +def _select_cfg(M, N, K, block_m): + """Pick the launch config from host constants only (M=num_valid_tokens, N, K, + block_m) — graph-capture safe (no GPU-scalar branch).""" + # Per-regime winners (measured, isolated cuda-event A/B on gfx950, GPU 3): + # BLOCK_K=256 (fewer K-iters + bigger MX scale-load coalesced with the dot), + # num_stages=2 (software-pipeline overlaps the E8M0 scale-load with the scaled + # MFMA), GROUP_SIZE_M=4 (XCD-friendly swizzle that keeps each touched expert's + # A rows + B column-tiles L2/MALL-resident across its N-tiles -> kills the + # redundant A re-read). num_warps stays 8 (wave64 here did NOT spill: VGPR fits + # and 4 warps was neutral/worse in measurement). BLOCK_K must be a power of two + # and divide K (K=6144 and K=768 are both multiples of 256). + BLOCK_K = 256 if K % 256 == 0 else 128 + return { + "BLOCK_N": 128, + "BLOCK_K": BLOCK_K, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + } @triton.jit @@ -46,6 +64,7 @@ def _mxfp8_grouped_gemm_kernel( sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + EM, N, K, num_valid_tokens, @@ -67,9 +86,33 @@ def _mxfp8_grouped_gemm_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, ): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) + # ---- Grid-swizzle (super-grouping over M) so consecutive program-ids cover a + # GROUP_SIZE_M x grid_n super-block. This keeps a touched expert's A rows + its + # B-weight column-tiles L2/MALL-resident across the N-tiles they share, killing + # most of the redundant A re-read (the single largest redundant-traffic source: + # dot_scaled keeps operands in registers, so A is otherwise re-fetched per N-tile). + # ``num_pid_m`` uses EM (= sorted_token_ids.shape[0]), the SAME bound the grid + # is sized from, NOT the runtime ``num_tokens_post_padded`` (<= EM). This keeps + # the swizzle consistent with the grid: ``group_size_m`` is always >= 1 for every + # launched program, so the ``% group_size_m`` / ``// group_size_m`` below can never + # hit modulo/division-by-zero. ``num_tokens_post_padded`` is loaded afterwards and + # only gates the early-return. Mirrors the reference ``fused_moe_kernel``. + pid = tl.program_id(0) + num_pid_m = tl.cdiv(EM, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + if GROUP_SIZE_M == 1: + pid_m = pid % num_pid_m + pid_n = pid // num_pid_m + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + num_post = tl.load(num_tokens_post_padded_ptr) if pid_m * BLOCK_M >= num_post: return @@ -140,9 +183,6 @@ def _grouped_gemm_mxfp8( a_div: int, mul_weight_by: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, - block_n: int = 128, - num_warps: int = 8, - num_stages: int = 2, ) -> torch.Tensor: M_routed = num_valid_tokens E, N, K = w.shape @@ -152,8 +192,16 @@ def _grouped_gemm_mxfp8( # written — zero them so the downstream reduction ignores their garbage. alloc = torch.zeros if expert_map is not None else torch.empty out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device) - BLOCK_K = 128 - grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, block_n)) + + cfg = _select_cfg(M_routed, N, K, block_m) + BLOCK_N = cfg["BLOCK_N"] + BLOCK_K = cfg["BLOCK_K"] + GROUP_SIZE_M = cfg["GROUP_SIZE_M"] + + n_pid_m = triton.cdiv(sorted_token_ids.shape[0], block_m) + n_pid_n = triton.cdiv(N, BLOCK_N) + grid = (n_pid_m * n_pid_n,) + _mxfp8_grouped_gemm_kernel[grid]( a_q, a_scale, @@ -164,6 +212,7 @@ def _grouped_gemm_mxfp8( sorted_token_ids, expert_ids, num_tokens_post_padded, + sorted_token_ids.shape[0], # EM: sizes both the grid and the swizzle N, K, num_valid_tokens, @@ -183,29 +232,15 @@ def _grouped_gemm_mxfp8( A_DIV=a_div, MUL_WEIGHT=mul_weight_by is not None, BLOCK_M=block_m, - BLOCK_N=block_n, + BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, - num_warps=num_warps, - num_stages=num_stages, + GROUP_SIZE_M=GROUP_SIZE_M, + num_warps=cfg["num_warps"], + num_stages=cfg["num_stages"], ) return out -# Tuned native-MXFP8 launch tiles for gfx950 (CDNA4) at MiniMax-M3 MoE shapes. -# For example, 8k/1k, 1k/1k cases. - -_MXFP8_PREFILL_TILES = dict(block_m=128, block_n=256, num_warps=8, num_stages=2) -_MXFP8_DECODE_TILES = dict(block_m=64, block_n=64, num_warps=4, num_stages=2) -_MXFP8_PREFILL_MIN_TOKENS = 1024 - - -def _mxfp8_moe_tiles(num_tokens: int) -> dict: - """Pick grouped-GEMM launch tiles by regime (token count).""" - if num_tokens >= _MXFP8_PREFILL_MIN_TOKENS: - return _MXFP8_PREFILL_TILES - return _MXFP8_DECODE_TILES - - def fused_moe_mxfp8_native( hidden_states: torch.Tensor, # [T, H] bf16 w13: torch.Tensor, # [E, 2I, H] fp8 @@ -225,8 +260,7 @@ def fused_moe_mxfp8_native( top_k = topk_ids.shape[1] M = T * top_k - tiles = _mxfp8_moe_tiles(T) - block_m = tiles["block_m"] + block_m = 64 # Bin by the actual number of expert weight rows. With fused shared experts # the weight tensor has more rows than ``global_num_experts`` (the routed # count), and their ids fall outside [0, global_num_experts); binning by the @@ -257,9 +291,6 @@ def fused_moe_mxfp8_native( hidden_states.dtype, a_div=top_k, expert_map=expert_map, - block_n=tiles["block_n"], - num_warps=tiles["num_warps"], - num_stages=tiles["num_stages"], ) # [M, 2I] # SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the @@ -288,9 +319,6 @@ def fused_moe_mxfp8_native( a_div=1, mul_weight_by=topk_weights.reshape(-1).to(torch.float32), expert_map=expert_map, - block_n=tiles["block_n"], - num_warps=tiles["num_warps"], - num_stages=tiles["num_stages"], ) # [M, H] == [T*top_k, H] return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 550d6b5341d..bd1b9ecaa8a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, ) @@ -145,4 +146,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, activation_type=activation_to_flashinfer_int(activation), + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index a7faa5f6e17..a4cce79741e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -15,7 +15,10 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights +from vllm.model_executor.layers.fused_moe.utils import ( + fi_moe_largest_bucket, + trtllm_moe_pack_topk_ids_weights, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, ) @@ -249,6 +252,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, output=output, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) @@ -419,6 +423,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit use_shuffled_weight=use_shuffled_weight, weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL: kwargs["activation_type"] = activation_type @@ -475,6 +480,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit use_routing_scales_on_input=apply_router_weight_on_input, routing_method_type=self.routing_method_type, activation_type=activation_type, + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), ) return out 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 518c87ce4df..f046dfeaf26 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 @@ -16,7 +16,10 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights +from vllm.model_executor.layers.fused_moe.utils import ( + fi_moe_largest_bucket, + trtllm_moe_pack_topk_ids_weights, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, ) @@ -319,6 +322,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula do_finalize=True, activation_type=activation_to_flashinfer_int(activation), output=output, + tune_max_num_tokens=min( + fi_moe_largest_bucket(self.moe_config), self._get_chunk_size() + ), ) def apply( @@ -479,4 +485,5 @@ class TrtLlmNvFp4ExpertsMonolithic( routing_method_type=self.routing_method_type, do_finalize=True, activation_type=activation_to_flashinfer_int(activation), + tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), )[0] diff --git a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py index cf49e01e628..b1588c4a2a4 100644 --- a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py @@ -15,7 +15,7 @@ from aiter.ops.flydsl.kernels.moe_gemm_2stage import ( ) from vllm.logger import init_logger -from vllm.platforms import current_platform +from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -106,7 +106,7 @@ def build_routing_buffers( @functools.lru_cache def try_get_optimal_config(num_experts, inter_dim): - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() json_file_name = ( f"E={num_experts},N={inter_dim},device_name={device_name}," "dtype=int4_w4a16,backend=flydsl.json" diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 269b6e3da0b..62cf12ea24e 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.utils import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -1034,7 +1035,7 @@ def zero_experts_compute_triton( def get_config_file_name( E: int, N: int, dtype: str | None, block_shape: list[int] | None = None ) -> str: - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() # Set device_name to H200 if a device from the H200 family is detected if "H200" in device_name.split("_"): device_name = "NVIDIA_H200" diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index a522a39153b..5c8f7010f6d 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -127,6 +127,7 @@ def FusedMoE( num_redundant_experts: int = 0, has_bias: bool = False, is_sequence_parallel: bool = False, + reduce_results: bool = True, ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"), n_shared_experts: int | None = None, router_logits_dtype: torch.dtype | None = None, @@ -184,6 +185,9 @@ def FusedMoE( num_redundant_experts: Number of redundant experts for EPLB has_bias: Whether expert layers have bias terms is_sequence_parallel: Whether sequence parallelism is enabled + reduce_results: Whether to all-reduce the final output. Setting this + to False (to fuse the all-reduce downstream) is only honored on the + late-AR path. expert_mapping: Expert parameter mapping for weight loading n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) @@ -220,6 +224,14 @@ def FusedMoE( parallel_config=vllm_config.parallel_config, ) + # Resolve the deferred all-reduce request against the parallel config. + skip_final_all_reduce = ( + not reduce_results + and not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.is_sequence_parallel + and zero_expert_type is None + ) + global_num_experts, logical_num_experts, num_fused_shared_experts = ( determine_expert_counts( num_experts, @@ -340,6 +352,7 @@ def FusedMoE( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, + skip_final_all_reduce=skip_final_all_reduce, ) logger.debug("FusedMoEConfig = %s", moe_config) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 862f292009c..c60e68232d4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch @@ -45,6 +46,7 @@ class Fp8MoeBackend(Enum): DEEPGEMM = "DEEPGEMM" BATCHED_DEEPGEMM = "BATCHED_DEEPGEMM" MARLIN = "MARLIN" + HUMMING = "HUMMING" TRITON = "TRITON" BATCHED_TRITON = "BATCHED_TRITON" AITER = "AITER" @@ -83,6 +85,7 @@ def _get_priority_backends( Fp8MoeBackend.VLLM_CUTLASS, Fp8MoeBackend.TRITON, Fp8MoeBackend.MARLIN, + Fp8MoeBackend.HUMMING, Fp8MoeBackend.BATCHED_DEEPGEMM, Fp8MoeBackend.BATCHED_VLLM_CUTLASS, Fp8MoeBackend.BATCHED_TRITON, @@ -162,6 +165,19 @@ def backend_to_kernel_cls( return [BatchedDeepGemmExperts] + elif backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == Fp8MoeBackend.MARLIN: from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( MarlinExperts, @@ -240,6 +256,7 @@ def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend: "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": Fp8MoeBackend.FLASHINFER_CUTLASS, "marlin": Fp8MoeBackend.MARLIN, + "humming": Fp8MoeBackend.HUMMING, "aiter": Fp8MoeBackend.AITER, "hpc": Fp8MoeBackend.HPC, } @@ -402,6 +419,41 @@ def select_fp8_moe_backend( return Fp8MoeBackend.NONE, None +def _humming_fp8_weight_schema( + layer: RoutedExperts, weight: torch.Tensor, weight_scale: torch.Tensor +) -> dict[str, Any]: + """Build the humming weight schema from the canonical on-device fp8/mxfp8 + tensors (scale dtype/shape, block size), not the producing quant method.""" + # mxfp8: e8m0 group-32 scales (stored as uint8 bytes or e8m0). humming has + # no compressed-tensors mxfp8 loader; its modelopt schema fits both sources. + if weight_scale.dtype in (torch.uint8, torch.float8_e8m0fnu): + return {"quant_method": "modelopt", "quant_algo": "mxfp8"} + + if hasattr(layer, "w13_weight_scale_inv"): + assert hasattr(layer, "weight_block_size") + return {"quant_method": "fp8", "weight_block_size": layer.weight_block_size} + + # fp8 (e4m3): recover the strategy from the scale layout (block from + # weight_block_size, else channel vs tensor by per-expert scale count). + config: dict[str, Any] = { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "type": "float", + "num_bits": 8, + "symmetric": True, + } + weight_block_size = getattr(layer, "weight_block_size", None) + num_experts, num_output = weight.shape[0], weight.shape[-2] + if weight_block_size is not None: + config["strategy"] = "block" + config["block_structure"] = list(weight_block_size) + elif weight_scale.numel() >= num_experts * num_output: + config["strategy"] = "channel" + else: + config["strategy"] = "tensor" + return config + + def convert_to_fp8_moe_kernel_format( fp8_backend: Fp8MoeBackend, # TODO(bnell): replace layer with weight_block_size @@ -429,6 +481,18 @@ def convert_to_fp8_moe_kernel_format( w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights( w13, w2, w13_scale, w2_scale ) + elif fp8_backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale) + ) + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale elif fp8_backend == Fp8MoeBackend.MARLIN: weight_block_size = getattr(layer, "weight_block_size", None) if weight_block_size == [1, 32]: @@ -511,6 +575,7 @@ def make_fp8_moe_quant_config( swiglu_limit: float | None = None, gemm1_alpha: float | None = None, gemm1_beta: float | None = None, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: """ Create FusedMoEQuantConfig for the specified FP8 Backend. @@ -537,6 +602,14 @@ def make_fp8_moe_quant_config( gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, ) + elif fp8_backend == Fp8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) # Flashinfer CUTLASS or HPC per-tensor uses single dq scale # (alpha = w_scale * a_scale) and inverse a2 scale. @@ -600,6 +673,7 @@ def make_fp8_moe_kernel( experts_cls: type[mk.FusedMoEExperts], fp8_backend: Fp8MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -613,6 +687,11 @@ def make_fp8_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if fp8_backend == Fp8MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -622,11 +701,13 @@ def make_fp8_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index e31a3ca07ee..5a2b4c3a75b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch @@ -22,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8DynamicTokenSym, kInt8StaticChannelSym, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform logger = init_logger(__name__) @@ -29,6 +31,7 @@ logger = init_logger(__name__) class Int8MoeBackend(Enum): TRITON = "TRITON" + HUMMING = "HUMMING" CPU = "CPU" @@ -40,6 +43,7 @@ def _get_priority_backends( """ _AVAILABLE_BACKENDS = [ Int8MoeBackend.TRITON, + Int8MoeBackend.HUMMING, Int8MoeBackend.CPU, ] @@ -62,6 +66,19 @@ def backend_to_kernel_cls( return [TritonExperts] + elif backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUExpertsInt8, @@ -77,6 +94,7 @@ def map_int8_backend(runner_backend: MoEBackend) -> Int8MoeBackend: """Map user's MoEBackend to Int8MoeBackend.""" mapping = { "triton": Int8MoeBackend.TRITON, + "humming": Int8MoeBackend.HUMMING, } if backend := mapping.get(runner_backend): return backend @@ -163,6 +181,7 @@ def select_int8_moe_backend( def make_int8_moe_quant_config( + int8_backend: Int8MoeBackend, w1_scale: torch.Tensor, w2_scale: torch.Tensor, a1_scale: torch.Tensor | None = None, @@ -170,11 +189,21 @@ def make_int8_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, per_act_token_quant: bool = False, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: assert (a1_scale is None and a2_scale is None) or ( a1_scale is not None and a2_scale is not None ), "a1_scale and a2_scale must both be provided or both be None" + if int8_backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) + if a1_scale is None or a2_scale is None: return int8_w8a16_moe_quant_config( w1_scale=w1_scale, @@ -196,13 +225,56 @@ def make_int8_moe_quant_config( ) +def _humming_int8_weight_schema( + weight: torch.Tensor, weight_scale: torch.Tensor +) -> dict[str, Any]: + """Build the humming compressed-tensors int8 schema from the canonical + on-device tensors; humming does the signed-int8 -> native conversion.""" + config: dict[str, Any] = { + "quant_method": "compressed-tensors", + "format": "int-quantized", + "type": "int", + "num_bits": 8, + "symmetric": True, + "strategy": "channel", + } + num_experts, num_output = weight.shape[0], weight.shape[-2] + if weight_scale.numel() < num_experts * num_output: + config["strategy"] = "tensor" + return config + + def convert_to_int8_moe_kernel_format( int8_backend: Int8MoeBackend, w13: torch.Tensor, w2: torch.Tensor, + layer: torch.nn.Module | None = None, + w13_scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Convert INT8 MoE weights to backend-specific kernel format.""" - if int8_backend == Int8MoeBackend.CPU: + if int8_backend == Int8MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + assert layer is not None + # Humming reads canonical CT scales (w*_weight_scale) from the layer. + # Online int8 produces per-channel (E, N) w*_scale; expose them as the + # (E, N, 1) w*_weight_scale humming's loader expects. + for sub in ("w13", "w2"): + if hasattr(layer, f"{sub}_weight_scale"): + continue + scale = getattr(layer, f"{sub}_scale").data + if scale.dim() < 3: + scale = scale.unsqueeze(-1) + replace_parameter(layer, f"{sub}_weight_scale", scale) + delattr(layer, f"{sub}_scale") + convert_to_humming_moe_kernel_format( + layer, + quant_config=_humming_int8_weight_schema(w13, layer.w13_weight_scale), + ) + return layer.w13_weight, layer.w2_weight + elif int8_backend == Int8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( prepare_int8_moe_layer_for_cpu, ) @@ -215,10 +287,12 @@ def convert_to_int8_moe_kernel_format( def make_int8_moe_kernel( + int8_backend: Int8MoeBackend, moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -232,6 +306,11 @@ def make_int8_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if int8_backend == Int8MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -241,11 +320,13 @@ def make_int8_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( 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 0cf1382d406..d0cc08ea141 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -11,6 +11,7 @@ from compressed_tensors.quantization import ( import vllm._custom_ops as ops import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -46,6 +47,7 @@ logger = init_logger(__name__) class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" + HUMMING = "HUMMING" CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" @@ -55,7 +57,19 @@ def backend_to_kernel_cls( backend: WNA16MoEBackend, ) -> list[type[mk.FusedMoEExperts]]: """Return the experts class for the given backend, or None for NONE.""" - if backend == WNA16MoEBackend.MARLIN: + if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] + elif backend == WNA16MoEBackend.MARLIN: return [MarlinExperts] elif backend == WNA16MoEBackend.BATCHED_MARLIN: return [BatchedMarlinExperts] @@ -90,10 +104,26 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.HUMMING, ] return _AVAILABLE_BACKENDS +def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: + """Map user's MoEBackend to WNA16MoEBackend.""" + mapping = { + "marlin": WNA16MoEBackend.MARLIN, + "humming": WNA16MoEBackend.HUMMING, + "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, + } + if backend := mapping.get(runner_backend): + return backend + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for WNA16 MoE. " + f"Expected one of {list(mapping.keys())}." + ) + + def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, @@ -146,6 +176,14 @@ def select_wna16_moe_backend( return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) + # Handle explicit moe_backend from user. + runner_backend = config.moe_backend + if runner_backend != "auto": + requested_backend = map_wna16_backend(runner_backend) + return _return_or_raise( + requested_backend, config, weight_key, None, activation_format + ) + # Select kernels in order of backend. AVAILABLE_BACKENDS = _get_priority_backends() @@ -210,6 +248,8 @@ def make_wna16_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], + backend: WNA16MoEBackend = WNA16MoEBackend.MARLIN, + layer: torch.nn.Module | None = None, is_k_full: bool = False, w13_g_idx: torch.Tensor | None = None, w2_g_idx: torch.Tensor | None = None, @@ -228,14 +268,18 @@ def make_wna16_moe_kernel( ) # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, - # BatchedMarlinExperts, XPUExpertsWNA16, and CPUExpertsInt4 - assert experts_cls in ( + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming + # grouped/indexed experts. + allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, ) + if backend == WNA16MoEBackend.HUMMING: + allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING)) + assert experts_cls in allowed_experts is_monolithic = experts_cls.is_monolithic() @@ -251,7 +295,10 @@ def make_wna16_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") extra_args: dict[str, Any] = {} - if issubclass(experts_cls, MarlinExpertsBase): + if backend == WNA16MoEBackend.HUMMING: + assert layer is not None + extra_args = {"layer": layer} + elif issubclass(experts_cls, MarlinExpertsBase): extra_args = { "w13_g_idx": w13_g_idx, "w2_g_idx": w2_g_idx, @@ -941,6 +988,35 @@ def _process_weights_xpu( ) +def _humming_wna16_weight_schema( + quant_config: QuantizationConfig | QuantizationArgs | None, +) -> dict[str, Any]: + """Humming weight schema for a WNA16 checkpoint, derived from the quant + config rather than the running kernel.""" + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig + + if isinstance(quant_config, AutoAWQConfig): + return { + "quant_method": "awq", + "bits": quant_config.weight_bits, + "group_size": quant_config.group_size, + "zero_point": quant_config.zero_point, + } + if isinstance(quant_config, AutoGPTQConfig): + return { + "quant_method": "gptq", + "bits": quant_config.weight_bits, + "group_size": quant_config.group_size, + "desc_act": quant_config.desc_act, + "sym": quant_config.is_sym, + } + raise TypeError( + "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, " + f"got {type(quant_config).__name__}." + ) + + def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -956,26 +1032,30 @@ def convert_to_wna16_moe_kernel_format( 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 -]: +) -> ( + 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 + ] + | None +): """Dispatch weight post-processing to the appropriate per-backend handler. To add a new backend, implement a ``_process_weights_`` helper and - add a branch here. + add a branch here. Backends that rewrite the layer's parameters in place + (e.g. Humming) return ``None``; the caller then skips the param scatter. Args: backend: the selected ``WNA16MoEBackend``. @@ -983,6 +1063,16 @@ def convert_to_wna16_moe_kernel_format( quant_config: the ``QuantizationConfig`` for this layer. input_dtype: optional activation dtype, usually should be 16 bit. """ + if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + convert_to_humming_moe_kernel_format( + layer, quant_config=_humming_wna16_weight_schema(quant_config) + ) + return None + if backend in ( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 06b622a6c4b..b9086cfa48a 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -25,6 +25,7 @@ _SUPPORTED_BACKENDS = ( # is_supported_config passes (gfx950 + flydsl installed + not EP). On other # devices / no flydsl / EP it is skipped and native is used. Fp8MoeBackend.AITER_MXFP8, + Fp8MoeBackend.HUMMING, ) _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { @@ -34,6 +35,7 @@ _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { "xpu": Fp8MoeBackend.XPU, "aiter": Fp8MoeBackend.AITER_MXFP8, "triton": Fp8MoeBackend.TRITON_MXFP8, + "humming": Fp8MoeBackend.HUMMING, } diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 408c69fea09..f295163568d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -43,6 +43,7 @@ class NvFp4MoeBackend(Enum): FLASHINFER_B12X = "FLASHINFER_B12X" VLLM_CUTLASS = "VLLM_CUTLASS" MARLIN = "MARLIN" + HUMMING = "HUMMING" EMULATION = "EMULATION" @@ -119,6 +120,18 @@ def backend_to_kernel_cls( ) return [MarlinExperts] + elif backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ) + + return [ + BatchedHummingGroupedExperts, + HummingGroupedExperts, + HummingIndexedExperts, + ] elif backend == NvFp4MoeBackend.EMULATION: from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import ( Nvfp4QuantizationEmulationTritonExperts, @@ -138,6 +151,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, "marlin": NvFp4MoeBackend.MARLIN, + "humming": NvFp4MoeBackend.HUMMING, "emulation": NvFp4MoeBackend.EMULATION, } if backend := mapping.get(runner_backend): @@ -169,6 +183,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, + NvFp4MoeBackend.HUMMING, NvFp4MoeBackend.EMULATION, ] @@ -346,6 +361,41 @@ def convert_to_nvfp4_moe_kernel_format( a2_scale=a2_scale, is_act_and_mul=is_act_and_mul, ) + elif nvfp4_backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + convert_to_humming_moe_kernel_format, + ) + + # Discriminate the source checkpoint layout by its global-scale param: + # compressed-tensors uses *_weight_global_scale, modelopt *_weight_scale_2. + # The logical schema is identical (nvfp4 group-16); only the on-layer + # param names differ. TODO: normalize both methods to a single canonical + # layout upstream so the oracle needs neither the probe nor the re-alias. + if hasattr(layer, "w13_weight_global_scale"): + quant_config = { + "quant_method": "compressed-tensors", + "format": "nvfp4-pack-quantized", + "type": "float", + "num_bits": 4, + "strategy": "group", + "group_size": 16, + } + # CT pack-quantized reads `weight_packed`; the method renamed it to + # `weight`. Re-alias (convert replaces all params anyway). + layer.w13_weight_packed = layer.w13_weight + layer.w2_weight_packed = layer.w2_weight + else: + quant_config = {"quant_method": "modelopt", "quant_algo": "nvfp4"} + + convert_to_humming_moe_kernel_format(layer, quant_config=quant_config) + a13_scale = None + a2_scale = None + w13 = layer.w13_weight + w13_scale = layer.w13_weight_scale + w13_scale_2 = getattr(layer, "w13_global_scale", None) + w2 = layer.w2_weight + w2_scale = layer.w2_weight_scale + w2_scale_2 = getattr(layer, "w2_global_scale", None) elif nvfp4_backend == NvFp4MoeBackend.MARLIN: a13_scale = None a2_scale = None @@ -418,8 +468,17 @@ def make_nvfp4_moe_quant_config( a13_scale: torch.Tensor, a2_scale: torch.Tensor, swiglu_limit: float | None = None, + layer: torch.nn.Module | None = None, ) -> FusedMoEQuantConfig: - if backend == NvFp4MoeBackend.MARLIN: + if backend == NvFp4MoeBackend.HUMMING: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + assert isinstance(layer, RoutedExperts) + return get_humming_moe_quant_config(layer) + elif backend == NvFp4MoeBackend.MARLIN: return nvfp4_w4a16_moe_quant_config( g1_alphas=w13_scale_2, g2_alphas=w2_scale_2, @@ -467,7 +526,9 @@ def make_nvfp4_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts], + backend: NvFp4MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + layer: torch.nn.Module | None = None, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -481,6 +542,11 @@ def make_nvfp4_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) + extra_kwargs = {} + if backend == NvFp4MoeBackend.HUMMING: + assert layer is not None + extra_kwargs = {"layer": layer} + # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() @@ -490,11 +556,13 @@ def make_nvfp4_moe_kernel( quant_config=moe_quant_config, max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), + **extra_kwargs, ) else: experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, + **extra_kwargs, ) kernel = mk.FusedMoEKernel( diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index cda0eaf7300..639fce5234c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -232,7 +232,9 @@ def select_unquantized_moe_backend( raise ValueError(_make_log_unsupported(backend, reason)) runner_backend = moe_config.moe_backend - if runner_backend != "auto": + # 'humming' is quantization-only; an unquantized layer (e.g. excluded via + # modules_to_not_convert) falls through to auto instead of erroring. + if runner_backend not in ["auto", "humming"]: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py index e4f4d497568..c66e43ccba8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py @@ -270,7 +270,7 @@ def convert_to_w4a8_int8_moe_format( """ # Derive dimensions from tensor shapes E = w13_weight.shape[0] # num_experts - I2 = w13_weight.shape[1] # w13_out_features (2*IN) + w13_out_features = w13_weight.shape[1] # 2 * intermediate_size H = w13_weight.shape[2] # w13_in_features (hidden_size) IN = w2_weight.shape[2] # w2_in_features (intermediate_size) w2_out_features = w2_weight.shape[1] # Should equal H @@ -286,7 +286,7 @@ def convert_to_w4a8_int8_moe_format( w13_weight_scale[e], # [2I, H/g or 1] w13_bias[e] if w13_bias is not None else None, # [2I] H, - I2, + w13_out_features, group_size, ) ) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index d505c5ce4b7..058a96ed6b5 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -43,6 +43,7 @@ def vllm_topk_sigmoid( gating_output: torch.Tensor, renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float = 1.0, ) -> tuple[torch.Tensor, ...]: ops.topk_sigmoid( topk_weights, @@ -51,6 +52,7 @@ def vllm_topk_sigmoid( gating_output, renormalize, e_score_correction_bias, + routed_scaling_factor, ) return topk_weights, topk_indices @@ -216,9 +218,8 @@ def fused_topk_bias( gating_output, renormalize, e_score_correction_bias, + routed_scaling_factor, ) - if routed_scaling_factor != 1.0: - topk_weights *= routed_scaling_factor return topk_weights, topk_ids elif scoring_func == "sqrtsoftplus": return vllm_topk_softplus_sqrt( diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index d8deb73e9b4..cf00fa452a9 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -428,7 +428,6 @@ class MoERunner(MoERunnerInterface): if ( shared_output is not None and not self.moe_config.is_sequence_parallel - and not self.moe_config.skip_final_all_reduce and self._fused_output_is_reduced ): shared_output = tensor_model_parallel_all_reduce(shared_output) @@ -446,6 +445,13 @@ class MoERunner(MoERunnerInterface): here. Skipped when sequence-parallel is active (SP handles its own reduction) or when the early path already reduced both outputs. """ + # skip_final_all_reduce must not coexist with a pre-reduced fused + # output. This should be enforced by MoE config initialization. + if self.moe_config.skip_final_all_reduce: + assert not self._fused_output_is_reduced, ( + "skip_final_all_reduce requires an un-reduced fused output" + ) + # We don't need to reduce the final output if: # - We are not running with TP or DP # - The MK already reduced the fused output itself. @@ -725,8 +731,8 @@ class MoERunner(MoERunnerInterface): @property def do_naive_dispatch_combine(self) -> bool: return ( - self.moe_config.dp_size > 1 and not self._quant_method.supports_internal_mk - ) + self.moe_config.dp_size > 1 or self.moe_config.is_sequence_parallel + ) and not self._quant_method.supports_internal_mk def _maybe_dispatch( self, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index fce74346d62..ed512e1ff8f 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools from math import prod +from typing import TYPE_CHECKING import torch import torch.nn.functional as F @@ -33,6 +34,9 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig + @triton.jit def _count_expert_num_tokens( @@ -403,6 +407,23 @@ def _pack_topk_ids_weights_kernel( tl.store(output_ptr + offsets, packed, mask=mask) +def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int: + """Estimate FlashInfer's MoE autotuning maximum token count. + + All DP ranks may contribute `max_num_tokens` to one invocation. + Keep FlashInfer's default moe `tune_max_num_tokens=8192` + floor to avoid over-underestimation. + DeepEP, SP, or PCP may make this underestimate, however overestimation + may be dangerous, increasing tuning- cost and memory use. + + NOTE: The DP factor applies even when EP is disabled: + > Without `--enable-expert-parallel`, MoE layers would use tensor parallelism. + + For a detailed explanation, see: `docs/serving/data_parallel_deployment.md` + """ + return max(moe_config.max_num_tokens * moe_config.dp_size, 8192) + + def trtllm_moe_pack_topk_ids_weights( topk_ids: torch.Tensor, topk_weights: torch.Tensor, diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index 7eee2a6eb30..2b07d949c02 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -7,6 +7,7 @@ Decoupled from HpcAttentionImpl; extra params are passed via layer attrs. from __future__ import annotations +import importlib.util from enum import IntEnum from typing import Any @@ -122,6 +123,12 @@ class HpcRopeNorm(CustomOp, HpcModule): qk_norm_policy: QkNormPolicy = QkNormPolicy.ROPE_THEN_NORM, ) -> None: super().__init__() + if importlib.util.find_spec("hpc") is None: + raise ImportError( + "HPCRopeNorm requires the hpc module to be installed. " + "Please install it from https://github.com/Tencent/hpc-ops" + ) + self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim @@ -172,6 +179,15 @@ class HpcRopeNorm(CustomOp, HpcModule): self.layer_name: str | None = None self.register_layer_name(layer_name) + import hpc + + if self.use_fp8: + self._quant_type = ( + hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value + ) + else: + self._quant_type = None + @classmethod def support( cls, @@ -304,12 +320,6 @@ class HpcRopeNorm(CustomOp, HpcModule): self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None ) - # Dynamic per-token-per-head Q quant + per-tensor K/V (dqskv). - # rope_norm_store_kv_fp8 is registered as a torch op whose ``quant_policy`` - # argument is typed as ``int``; pybind cannot cast the hpc.QuantType enum - # automatically, so pass its integer ``.value``. - QUANT_POLICY_DQSKV = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR.value - # --- Prefill --- if num_prefill_reqs > 0: seq_lens_prefill = attn_metadata.seq_lens[num_decode_reqs:] @@ -333,7 +343,7 @@ class HpcRopeNorm(CustomOp, HpcModule): is_prefill=True, k_scale=k_scale, v_scale=v_scale, - quant_policy=QUANT_POLICY_DQSKV, + quant_policy=self._quant_type, max_seqlens=max_seqlens, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, @@ -364,9 +374,7 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv_decode = qkv[:num_decode_tokens] # Single-token decode: q_index is the per-request prefix sum # [0, 1, ..., num_decode_reqs]. - qo_indptr_decode = torch.arange( - num_decode_reqs + 1, dtype=torch.int32, device=qkv.device - ) + decode_query_len = attn_metadata.decode_query_len out_q_decode = output[:num_decode_tokens] if self.use_fp8: @@ -376,13 +384,13 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv=qkv_decode, cos_sin=self.cos_sin_cache, num_seqlen_per_req=num_seq_kvcache, - q_index=qo_indptr_decode, + q_index=attn_metadata.qo_indptr_decode, kvcache_indices=block_table_decode, is_prefill=False, k_scale=k_scale, v_scale=v_scale, - quant_policy=QUANT_POLICY_DQSKV, - max_seqlens=1, + quant_policy=self._quant_type, + max_seqlens=decode_query_len, q_norm_weight=q_norm_weight, k_norm_weight=k_norm_weight, qk_norm_policy=self.qk_norm_policy, @@ -398,7 +406,7 @@ class HpcRopeNorm(CustomOp, HpcModule): qkv_decode, self.cos_sin_cache, num_seq_kvcache, - qo_indptr_decode, + attn_metadata.qo_indptr_decode, block_table_decode, False, # is_prefill q_norm_weight=q_norm_weight, diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index e487b91e989..d92b9fc7d00 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1000,16 +1000,12 @@ class QKVParallelLinear(ColumnParallelLinear): self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) self.num_kv_head_replicas = 1 input_size = self.hidden_size - output_size = ( - self.num_heads * self.head_size - + self.num_kv_heads * self.head_size - + self.num_kv_heads * self.v_head_size - ) * tp_size self.output_sizes = [ self.num_heads * self.head_size * tp_size, # q_proj self.num_kv_heads * self.head_size * tp_size, # k_proj self.num_kv_heads * self.v_head_size * tp_size, # v_proj ] + output_size = sum(self.output_sizes) super().__init__( input_size=input_size, diff --git a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py index dd963f829d8..a00fbc74bf8 100644 --- a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py @@ -30,11 +30,10 @@ from vllm.model_executor.layers.mamba.linear.base import LinearAttention from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import ( MiniMaxText01LinearAttention, MiniMaxText01LinearKernel, - clear_linear_attention_cache_for_new_sequences, linear_attention_decode, - linear_attention_prefill_and_mix, ) from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata @@ -57,6 +56,290 @@ def _build_rope_parameters(config: PretrainedConfig) -> dict | None: return rope_parameters or None +def clear_linear_attention_cache_for_new_sequences( + kv_cache: torch.Tensor, + state_indices_tensor: torch.Tensor, + attn_metadata: LinearAttentionMetadata, +) -> None: + num_prefills = getattr(attn_metadata, "num_prefills", 0) + if num_prefills <= 0: + return + + num_decodes = getattr(attn_metadata, "num_decodes", 0) + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None) + for prefill_idx in range(num_prefills): + if num_decodes + prefill_idx + 1 >= len(attn_metadata.query_start_loc): + break + q_start = attn_metadata.query_start_loc[num_decodes + prefill_idx] + q_end = attn_metadata.query_start_loc[num_decodes + prefill_idx + 1] + query_len = q_end - q_start + context_len = attn_metadata.seq_lens[num_decodes + prefill_idx] - query_len + if context_len == 0: + if prefill_state_indices is not None: + block_to_clear = prefill_state_indices[prefill_idx] + else: + block_to_clear = state_indices_tensor[num_decodes + prefill_idx] + kv_cache[block_to_clear, ...] = 0 + + +def linear_attention_prefill_and_mix( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kv_cache: torch.Tensor, + state_indices_tensor: torch.Tensor, + attn_metadata: LinearAttentionMetadata, + slope_rate: torch.Tensor, + block_size: int, + decode_fn, + prefix_fn, + layer_idx: int | None = None, +) -> torch.Tensor: + hidden = [] + req_offset = getattr(attn_metadata, "num_decodes", 0) + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None) + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)): + if req_offset + _prefill_idx + 1 >= len(attn_metadata.query_start_loc): + break + if prefill_state_indices is not None and _prefill_idx >= len( + prefill_state_indices + ): + break + if prefill_state_indices is None and _prefill_idx >= len(state_indices_tensor): + break + _start = attn_metadata.query_start_loc[req_offset + _prefill_idx] + _end = attn_metadata.query_start_loc[req_offset + _prefill_idx + 1] + if prefill_state_indices is not None: + slot_id = prefill_state_indices[_prefill_idx] + else: + slot_id = state_indices_tensor[req_offset + _prefill_idx] + qs = q[_start:_end].transpose(0, 1).contiguous() + ks = k[_start:_end].transpose(0, 1).contiguous() + vs = v[_start:_end].transpose(0, 1).contiguous() + slice_layer_cache = kv_cache[slot_id, ...] + out_slice = prefix_fn( + qs, + ks, + vs, + slice_layer_cache, + slope_rate, + block_size, + layer_idx=layer_idx, + ) + hidden.append(out_slice.contiguous()) + + if attn_metadata.num_decode_tokens > 0: + hidden_decode = decode_fn( + q, k, v, kv_cache, state_indices_tensor, attn_metadata + ) + hidden.insert(0, hidden_decode) + + if not hidden: + return torch.empty((0, q.size(1) * q.size(2)), device=q.device, dtype=q.dtype) + + hidden = torch.concat(hidden, dim=0).contiguous() + return hidden + + +@triton.jit +def _bailing_linear_attn_decode_spec_step_kernel( + q_ptr, + k_ptr, + v_ptr, + kv_cache_ptr, + slope_rate_ptr, + state_indices_ptr, + query_start_loc_ptr, + num_accepted_tokens_ptr, + output_ptr, + q_start: tl.constexpr, + D: tl.constexpr, + q_b_stride, + q_h_stride, + q_d_stride, + k_b_stride, + k_h_stride, + k_d_stride, + v_b_stride, + v_h_stride, + v_d_stride, + cache_b_stride, + cache_h_stride, + cache_d0_stride, + cache_d1_stride, + state_indices_b_stride, + state_indices_t_stride, + output_b_stride, + output_d_stride, + DRAFT_IDX: tl.constexpr, + STATE_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_id = tl.program_id(0) + head_id = tl.program_id(1) + block_id = tl.program_id(2) + + req_start = tl.load(query_start_loc_ptr + req_id).to(tl.int64) + req_end = tl.load(query_start_loc_ptr + req_id + 1).to(tl.int64) + query_len = req_end - req_start + if query_len <= DRAFT_IDX: + return + + dst_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + DRAFT_IDX * state_indices_t_stride + ).to(tl.int64) + if dst_slot == -1: + return + + if DRAFT_IDX == 0: + accepted_offset = tl.load(num_accepted_tokens_ptr + req_id).to(tl.int64) - 1 + accepted_offset = tl.maximum(accepted_offset, 0) + accepted_offset = tl.minimum(accepted_offset, STATE_WIDTH - 1) + src_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + accepted_offset * state_indices_t_stride + ).to(tl.int64) + else: + src_slot = tl.load( + state_indices_ptr + + req_id * state_indices_b_stride + + (DRAFT_IDX - 1) * state_indices_t_stride + ).to(tl.int64) + if src_slot == -1: + return + + token_idx = req_start - q_start + DRAFT_IDX + qk_offsets = tl.arange(0, D) + v_offsets = tl.arange(0, BLOCK_SIZE) + block_id * BLOCK_SIZE + qk_mask = qk_offsets < D + v_mask = v_offsets < D + kv_mask = qk_mask[:, None] & v_mask[None, :] + + q = tl.load( + q_ptr + token_idx * q_b_stride + head_id * q_h_stride + qk_offsets * q_d_stride, + mask=qk_mask, + other=0.0, + ) + k = tl.load( + k_ptr + token_idx * k_b_stride + head_id * k_h_stride + qk_offsets * k_d_stride, + mask=qk_mask, + other=0.0, + ) + v = tl.load( + v_ptr + token_idx * v_b_stride + head_id * v_h_stride + v_offsets * v_d_stride, + mask=v_mask, + other=0.0, + ) + + cache_offsets = ( + qk_offsets[:, None] * cache_d0_stride + v_offsets[None, :] * cache_d1_stride + ) + src_cache_ptr = ( + kv_cache_ptr + + src_slot * cache_b_stride + + head_id * cache_h_stride + + cache_offsets + ) + dst_cache_ptr = ( + kv_cache_ptr + + dst_slot * cache_b_stride + + head_id * cache_h_stride + + cache_offsets + ) + + slope = tl.load(slope_rate_ptr + head_id) + decay = tl.exp(-slope) + kv_old = tl.load(src_cache_ptr, mask=kv_mask, other=0.0) + kv_new = k[:, None] * v[None, :] + decay * kv_old + + output = tl.sum(q[:, None].to(tl.float32) * kv_new, axis=0) + tl.store(dst_cache_ptr, kv_new, mask=kv_mask) + tl.store( + output_ptr + + token_idx * output_b_stride + + (head_id * D + v_offsets) * output_d_stride, + output, + mask=v_mask, + ) + + +def bailing_linear_attention_decode_spec( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kv_cache: torch.Tensor, + slope_rate: torch.Tensor, + state_indices_tensor: torch.Tensor, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + q_start: int, + q_end: int | None, + slot_start: int, + slot_end: int | None, + block_size: int, +) -> torch.Tensor: + q_decode = q[q_start:q_end] + k_decode = k[q_start:q_end] + v_decode = v[q_start:q_end] + hidden = torch.empty( + (q_decode.shape[0], q.shape[1] * q.shape[2]), + device=q.device, + dtype=q.dtype, + ) + hidden.zero_() + + state_indices_tensor = state_indices_tensor[slot_start:slot_end] + query_start_loc = query_start_loc.to(device=q.device) + + batch_size = state_indices_tensor.shape[0] + num_heads = q_decode.shape[1] + head_dim = q_decode.shape[2] + assert k_decode.shape == (q_decode.shape[0], num_heads, head_dim) + assert v_decode.shape == (q_decode.shape[0], num_heads, head_dim) + state_width = state_indices_tensor.shape[1] + + grid = (batch_size, num_heads, triton.cdiv(head_dim, block_size)) + for draft_idx in range(state_width): + _bailing_linear_attn_decode_spec_step_kernel[grid]( + q_decode, + k_decode, + v_decode, + kv_cache, + slope_rate, + state_indices_tensor, + query_start_loc, + num_accepted_tokens[:batch_size], + hidden, + q_start, + head_dim, + q_decode.stride(0), + q_decode.stride(1), + q_decode.stride(2), + k_decode.stride(0), + k_decode.stride(1), + k_decode.stride(2), + v_decode.stride(0), + v_decode.stride(1), + v_decode.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + kv_cache.stride(3), + state_indices_tensor.stride(0), + state_indices_tensor.stride(1), + hidden.stride(0), + hidden.stride(1), + DRAFT_IDX=draft_idx, + STATE_WIDTH=state_width, + BLOCK_SIZE=block_size, + ) + + return hidden + + class BailingGroupRMSNormGate(RMSNormGated): def __init__( self, @@ -208,6 +491,13 @@ class BailingMoELinearAttention(LinearAttention): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def get_attn_backend(self): + from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionBackend, + ) + + return BailingLinearAttentionBackend + @staticmethod def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Load weight for linear attention layers. @@ -368,13 +658,42 @@ class BailingMoELinearAttention(LinearAttention): def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata): """Handle decode (single token per sequence).""" + decode_state_indices = getattr(attn_metadata, "state_indices_tensor_d", None) + num_accepted_tokens = getattr(attn_metadata, "num_accepted_tokens", None) + query_start_loc = getattr(attn_metadata, "query_start_loc_d", None) + if ( + decode_state_indices is not None + and decode_state_indices.dim() > 1 + and num_accepted_tokens is not None + and query_start_loc is not None + ): + return bailing_linear_attention_decode_spec( + q, + k, + v, + kv_cache, + self.tp_slope, + decode_state_indices, + query_start_loc, + num_accepted_tokens, + q_start=0, + q_end=attn_metadata.num_decode_tokens, + slot_start=0, + slot_end=attn_metadata.num_decodes, + block_size=32, + ) + decode_state_indices = ( + state_indices_tensor + if decode_state_indices is None + else decode_state_indices + ) hidden = linear_attention_decode( q, k, v, kv_cache, self.tp_slope, - state_indices_tensor, + decode_state_indices, q_start=0, q_end=attn_metadata.num_decode_tokens, slot_start=0, diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float16.json new file mode 100644 index 00000000000..0156d54c2ca --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.4.0", + "128": { + "BLOCK_SIZE_M": 32, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "num_warps": 2 + }, + "4096": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "16384": { + "BLOCK_SIZE_M": 32, + "num_warps": 2 + }, + "32768": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "131072": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "196608": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "262144": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float32.json new file mode 100644 index 00000000000..1232717b802 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float32.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.4.0", + "128": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "32768": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "131072": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + }, + "196608": { + "BLOCK_SIZE_M": 64, + "num_warps": 1 + }, + "262144": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float16.json new file mode 100644 index 00000000000..8a78ec1a4b6 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float16.json @@ -0,0 +1,35 @@ +{ + "triton_version": "3.6.0", + "10": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "80": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "160": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "320": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "640": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "1280": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "2560": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "5120": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float32.json new file mode 100644 index 00000000000..d9f535e320a --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI355_OAM,cache_dtype=float32.json @@ -0,0 +1,35 @@ +{ + "triton_version": "3.6.0", + "10": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "80": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "160": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "320": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "640": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "1280": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "2560": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "5120": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 8c5a6355803..d348defcc76 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -19,6 +19,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.v1.attention.backends.utils import NULL_BLOCK_ID if current_platform.is_xpu(): @@ -53,7 +54,7 @@ def get_ssm_config_file_name( def get_ssm_device_name() -> str: - return current_platform.get_device_name().replace(" ", "_") + return get_device_name_as_file_name() def _canonical_cache_dtype(cache_dtype: str) -> str: diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 79976dfff14..e7e36f2fc53 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) +from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum @@ -90,7 +91,94 @@ class ShortConv(MambaBase, CustomOp): hidden_states: torch.Tensor, output: torch.Tensor, ): - return + # Reference torch causal conv1d; runs on all CPU platforms. AMX kernels + # for causal conv can be plugged in here later. + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_torch, + causal_conv1d_update_torch, + ) + + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + attn_metadata: AttentionMetadata | None = None + if attn_metadata_raw is not None: + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata, ShortConvAttentionMetadata) + + BCx, _ = self.in_proj(hidden_states) + B, C, x = BCx.chunk(3, dim=-1) + + # (dim, kernel_size) — same reshape as forward_cuda + conv_weights = self.conv.weight.view( + self.conv.weight.size(0), self.conv.weight.size(2) + ) + + if attn_metadata is None: + # Profile run — output value doesn't matter + Bx = (B * x).contiguous() + output_tensor, _ = self.out_proj(C * Bx) + output[: hidden_states.shape[0]] = output_tensor + return + + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) # (num_blocks, dim, state_len) + + num_prefills = attn_metadata.num_prefills + num_decodes = attn_metadata.num_decode_tokens + num_prefill_tokens = attn_metadata.num_prefill_tokens + has_prefill = num_prefills > 0 + has_decode = num_decodes > 0 + num_actual_tokens = num_decodes + num_prefill_tokens + + B_d, B_p = torch.split( + B[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + C_d, C_p = torch.split( + C[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + x_d, x_p = torch.split( + x[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0 + ) + + conv_output_list = [] + + if has_prefill: + assert attn_metadata.state_indices_tensor_p is not None + Bx_p = (B_p * x_p).transpose(0, 1) # (dim, num_prefill_tokens) + out_p = causal_conv1d_torch( + Bx_p, + conv_weights, + self.conv.bias, + conv_state, + attn_metadata.query_start_loc_p, + attn_metadata.state_indices_tensor_p.flatten(), + attn_metadata.has_initial_states_p, + activation=None, + ).transpose(0, 1)[:num_prefill_tokens] # (num_prefill_tokens, dim) + conv_output_list.append(C_p * out_p) + + if has_decode: + assert attn_metadata.state_indices_tensor_d is not None + state_indices_d = attn_metadata.state_indices_tensor_d.flatten() + Bx_d = (B_d * x_d).unsqueeze(-1) # (num_decodes, dim, 1) + # Advanced indexing returns a copy; update in-place then scatter back + gathered = conv_state[state_indices_d] # (num_decodes, dim, state_len) + out_d = causal_conv1d_update_torch( + Bx_d, + gathered, + conv_weights, + self.conv.bias, + activation=None, + ).squeeze(-1) # (num_decodes, dim) + conv_state[state_indices_d] = gathered + conv_output_list.insert(0, C_d * out_d) + + hidden_states_out = torch.vstack(conv_output_list) + output[:num_actual_tokens], _ = self.out_proj(hidden_states_out) def forward( self, @@ -235,7 +323,10 @@ def short_conv( ) -> None: forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] - self.forward_cuda(hidden_states=hidden_states, output=output) + if not current_platform.is_cpu(): + self.forward_cuda(hidden_states=hidden_states, output=output) + else: + self.forward_native(hidden_states=hidden_states, output=output) def short_conv_fake( diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index cebfad7e596..58104fa7d25 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe import ( UnquantizedFusedMoEMethod, ) from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, convert_to_wna16_moe_kernel_format, make_wna16_moe_kernel, make_wna16_moe_quant_config, @@ -663,6 +664,26 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_moe_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=self.input_dtype, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_qzeros=layer.w13_qzeros, + w2_qzeros=layer.w2_qzeros, + w13_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + ( w13, w2, @@ -678,20 +699,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): w2_input_global_scale, w13_bias, w2_bias, - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_moe_backend, - layer=layer, - quant_config=self.quant_config, - input_dtype=self.input_dtype, - w13=layer.w13_qweight, - w2=layer.w2_qweight, - w13_scale=layer.w13_scales, - w2_scale=layer.w2_scales, - w13_qzeros=layer.w13_qzeros, - w2_qzeros=layer.w2_qzeros, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), - ) + ) = converted replace_parameter(layer, "w13_qweight", w13) replace_parameter(layer, "w2_qweight", w2) @@ -734,6 +742,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_moe_backend, + layer=layer, 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), @@ -743,6 +753,12 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): ) def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, @@ -783,8 +799,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -806,8 +822,8 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index aca76162f9a..b056f5222af 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -658,6 +658,26 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): "W8A8-INT8 is not supported by marlin kernel." ) + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_moe_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=self.input_dtype, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_g_idx=layer.w13_g_idx, + w2_g_idx=layer.w2_g_idx, + w13_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + ) + + if converted is None: + # Backend rewrote the layer's params in place (e.g. Humming). + self._setup_kernel(layer) + return + ( w13, w2, @@ -673,20 +693,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): w2_input_global_scale, w13_bias, w2_bias, - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_moe_backend, - layer=layer, - quant_config=self.quant_config, - input_dtype=self.input_dtype, - w13=layer.w13_qweight, - w2=layer.w2_qweight, - w13_scale=layer.w13_scales, - w2_scale=layer.w2_scales, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), - ) + ) = converted replace_parameter(layer, "w13_qweight", w13) replace_parameter(layer, "w2_qweight", w2) @@ -733,6 +740,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False) ) + # The modular kernel reads w13_weight/w2_weight; marlin keeps *_qweight. + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight + self._setup_kernel(layer) def _setup_kernel(self, layer: RoutedExperts) -> None: @@ -743,15 +754,24 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_moe_backend, + layer=layer, is_k_full=self.is_k_full, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, + w13_g_idx=getattr(layer, "w13_g_idx", None), + w2_g_idx=getattr(layer, "w2_g_idx", None), 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(), ) def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) + from vllm.model_executor.layers.fused_moe.config import ( gptq_marlin_moe_quant_config, ) @@ -795,8 +815,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -818,8 +838,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_qweight, - w2=layer.w2_qweight, + w1=layer.w13_weight, + w2=layer.w2_weight, router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 9a051c038f9..f74ed2d7b38 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -236,7 +236,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -259,6 +261,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py index a64104d3ffd..2d2190e614e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py @@ -5,7 +5,6 @@ import torch from compressed_tensors.quantization import ( QuantizationArgs, - QuantizationStrategy, ) from vllm.logger import init_logger @@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ScaleDesc, ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs -from vllm.platforms import CpuArchEnum, current_platform logger = init_logger(__name__) @@ -64,32 +62,17 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): weight_quant.group_size if (weight_quant.group_size is not None) else -1 ) - # Validate scheme: weights=W4 (channel or group), - # activations=dynamic TOKEN (A8) - - # Must be dynamic per-token activations - if ( - input_quant.strategy != QuantizationStrategy.TOKEN - or not input_quant.dynamic + # make sure group size is valid + if self.group_size != -1 and ( + moe.hidden_dim % self.group_size != 0 + or moe.intermediate_size_per_partition % self.group_size != 0 ): raise ValueError( - "W4A8-int MoE needs dynamic per-token activation quantization." + f"Group size ({self.group_size}) must evenly divide both " + f"hidden size ({moe.hidden_dim}) and intermediate size per " + f"partition ({moe.intermediate_size_per_partition})." ) - if weight_quant.num_bits != 4: - raise ValueError("This method only supports 4-bit weights (num_bits=4).") - - # Arm: check _dyn ops availability - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - try: - _ = torch.ops.aten._dyn_quant_matmul_4bit - _ = torch.ops.aten._dyn_quant_pack_4bit_weight - except AttributeError as err: - raise RuntimeError( - f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops; - install a newer build.""" - ) from err - # Construct QuantKey for weights from QuantizationArgs # W4A8 INT4: 4-bit weights (stored as int8), static quantization if self.group_size == -1: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index 14ef8bf614c..35d2a378424 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -338,6 +338,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def maybe_make_prepare_finalize( @@ -355,12 +356,13 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - a1_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, + a1_scale=getattr(layer, "w13_input_scale", None), + a2_scale=getattr(layer, "w2_input_scale", None), per_act_token_quant=is_per_token, per_out_ch_quant=is_per_token, block_shape=self.weight_block_size, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py index c29472cfc6b..d304ca56bf1 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py @@ -146,6 +146,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): int8_backend=self.int8_backend, w13=layer.w13_weight, w2=layer.w2_weight, + layer=layer, + w13_scale=layer.w13_weight_scale, ) replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) @@ -153,10 +155,12 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.experts_cls is not None self.moe_kernel = make_int8_moe_kernel( + int8_backend=self.int8_backend, moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def maybe_make_prepare_finalize( @@ -170,11 +174,13 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_int8_moe_quant_config( + int8_backend=self.int8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, a1_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, per_act_token_quant=True, + layer=layer, ) def apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py index 2e6e01ca766..468aed29a9e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -140,6 +140,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -155,6 +156,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def maybe_make_prepare_finalize( 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 46fa36180d9..8af36bcb102 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 @@ -416,6 +416,27 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Process weights using the shared oracle infrastructure is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM + converted = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.weight_quant, + input_dtype=self.marlin_input_dtype, + w13=layer.w13_weight_packed, + w2=layer.w2_weight_packed, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_qzeros=getattr(layer, "w13_weight_zero_point", None), + w2_qzeros=getattr(layer, "w2_weight_zero_point", None), + ) + if converted is None: + # In-place backends (e.g. Humming) are not wired through this + # marlin-only method; fail clearly rather than unpacking None. + raise NotImplementedError( + f"{type(self).__name__} does not support the " + f"{self.wna16_backend.value} MoE backend." + ) ( w13_qweight, w2_qweight, @@ -431,20 +452,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_input_global_scale, _, # w13_bias _, # w2_bias - ) = convert_to_wna16_moe_kernel_format( - backend=self.wna16_backend, - layer=layer, - quant_config=self.weight_quant, - input_dtype=self.marlin_input_dtype, - w13=layer.w13_weight_packed, - w2=layer.w2_weight_packed, - w13_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_qzeros=getattr(layer, "w13_weight_zero_point", None), - w2_qzeros=getattr(layer, "w2_weight_zero_point", None), - ) + ) = converted # Replace common parameters replace_parameter(layer, "w13_weight_packed", w13_qweight) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 1301c98f45b..23e3510614b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -137,8 +137,16 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): "weight_scale", convert_to_channelwise(layer.weight_scale, layer.logical_widths), ) + self.strategy = QuantizationStrategy.CHANNEL + self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy] + self.linear_kernel.config.weight_quant_key = self.weight_quant_key + # Canonicalize to (K, N) for the kernel. replace_parameter(layer, "weight", layer.weight.t()) + # Preserve the dim tags dropped by the transpose so layout-aware + # kernels see (K, N). + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 self.linear_kernel.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index 1a240f6540d..e50778f0acb 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -181,6 +181,10 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): # required by torch.compile to be torch.nn.Parameter layer.weight = Parameter(weight.data, requires_grad=False) + # Preserve the dim tags dropped by the transpose so layout-aware + # kernels (humming) see (K, N) instead of assuming (N, K). + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 layer.weight_scale = Parameter(weight_scale.data, requires_grad=False) if input_scale is not None: layer.input_scale = Parameter(input_scale.data, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 86818ed4b7e..626fc83cdff 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -714,6 +714,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: @@ -786,6 +787,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) # Inject biases into the quant config if the model has them diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py new file mode 100644 index 00000000000..cf54d1b531e --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from functools import lru_cache +from typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +logger = init_logger(__name__) + +_OPS_REGISTERED = False + + +@lru_cache(maxsize=1) +def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]: + """Return ARK availability, error details, cached module, and QuantLinear.""" + try: + import auto_round_kernel as ark + from auto_round_kernel.qlinear import QuantLinear + + logger.info("Successfully imported auto_round_kernel.") + except ImportError as error: + return False, str(error), None, None + + if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None: + return ( + False, + "No ARK backend library is available.", + None, + None, + ) + + logger.info("Successfully loaded auto_round_kernel backend library.") + return True, None, ark, QuantLinear + + +def _inc_ark_woq_linear_impl( + x: torch.Tensor, + qweight: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + in_features: int, + group_size: int, + compute_type: str, + weight_type: str, + scale_type: str, + asym: bool, +) -> torch.Tensor: + ark = get_ark_state()[2] + assert ark is not None + + return ark.woqgemm_linear( + x, + qweight, + bias, + out_features, + in_features, + group_size, + compute_type, + weight_type, + scale_type, + asym, + ) + + +def _inc_ark_woq_linear_fake( + x: torch.Tensor, + qweight: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + in_features: int, + group_size: int, + compute_type: str, + weight_type: str, + scale_type: str, + asym: bool, +) -> torch.Tensor: + del qweight + del bias + del in_features + del group_size + del compute_type + del weight_type + del scale_type + del asym + return torch.empty( + (*x.shape[:-1], out_features), + dtype=x.dtype, + device=x.device, + ) + + +class ark_ops: + @staticmethod + def register_ops_once() -> None: + global _OPS_REGISTERED + if _OPS_REGISTERED: + return + + is_available, error_str, _, _ = get_ark_state() + if not is_available: + logger.debug( + "Skip registering ark op because ARK is unavailable: %s", + error_str or "unknown error", + ) + return + + direct_register_custom_op( + op_name="inc_ark_woq_linear", + op_func=_inc_ark_woq_linear_impl, + fake_impl=_inc_ark_woq_linear_fake, + dispatch_key=current_platform.dispatch_key, + ) + _OPS_REGISTERED = True + + +ark_ops.register_ops_once() + +__all__ = ["get_ark_state"] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index a212e4d3050..dd6c2fa2eaa 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import lru_cache from typing import TYPE_CHECKING, Any import torch from torch.nn.parameter import Parameter -from vllm.logger import init_logger from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -22,35 +20,10 @@ from vllm.scalar_type import scalar_types from .inc_scheme import INCLinearScheme -logger = init_logger(__name__) - if TYPE_CHECKING: from ..config_parser import INCLayerConfig -@lru_cache(maxsize=1) -def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]: - """Return ARK availability, error details, cached module, and QuantLinear.""" - try: - import auto_round_kernel as ark - from auto_round_kernel.qlinear import QuantLinear - - logger.info("Successfully imported auto_round_kernel.") - except ImportError as error: - return False, str(error), None, None - - if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None: - return ( - False, - "No ARK backend library is available.", - None, - None, - ) - logger.info("Successfully loaded auto_round_kernel backend library.") - - return True, None, ark, QuantLinear - - class INCWNA16LinearScheme(INCLinearScheme): def __init__(self, layer_config: "INCLayerConfig") -> None: self.layer_config = layer_config @@ -380,6 +353,8 @@ class INCARKLinearMethod(INCXPULinearBase): def __init__(self, layer_config: "INCLayerConfig") -> None: super().__init__(layer_config) + from .inc_ark_ops import get_ark_state + is_available, error_str, _, quant_linear_cls = get_ark_state() if not is_available or quant_linear_cls is None: reason = error_str or "unknown error" @@ -453,9 +428,13 @@ class INCARKLinearMethod(INCXPULinearBase): ark_linear.bias.copy_(layer.bias.detach()) ark_linear.post_init() - layer.ark_linear = ark_linear - del layer.qweight + layer.qweight = Parameter(ark_linear.qweight.detach(), requires_grad=False) + layer.ark_bias = ark_linear.bias + layer.ark_compute_type = ark_linear.cdt + layer.ark_weight_type = ark_linear.wdt + layer.ark_scale_type = ark_linear.sdt + if hasattr(layer, "qzeros"): del layer.qzeros del layer.scales @@ -466,8 +445,18 @@ class INCARKLinearMethod(INCXPULinearBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - del bias - return layer.ark_linear.forward(x) + return torch.ops.vllm.inc_ark_woq_linear.default( + x, + layer.qweight, + layer.ark_bias, + layer.out_features, + layer.in_features, + self.group_size, + layer.ark_compute_type, + layer.ark_weight_type, + layer.ark_scale_type, + not self.sym, + ) class INCXPUW4A16LinearScheme(INCXPULinearMethod): diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py index e994b034944..80310358619 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py @@ -36,10 +36,10 @@ class INCWna16Scheme(INCScheme): del config, layer if current_platform.is_xpu(): if layer_config.bits == 4 and layer_config.sym: + from .inc_ark_ops import get_ark_state from .inc_wna16_linear import ( INCARKLinearMethod, INCXPULinearMethod, - get_ark_state, ) is_ark_available, ark_error, _, _ = get_ark_state() @@ -57,10 +57,10 @@ class INCWna16Scheme(INCScheme): if current_platform.is_cpu() and layer_config.is_gptq: if layer_config.bits == 4 and layer_config.sym: + from .inc_ark_ops import get_ark_state from .inc_wna16_linear import ( INCARKLinearMethod, INCWNA16LinearScheme, - get_ark_state, ) is_ark_available, ark_error, _, _ = get_ark_state() diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8fa1cb4d544..7df8178ca71 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -63,6 +63,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( ) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( process_fp8_input_tensor_strategy_moe, + process_fp8_weight_channel_strategy, process_fp8_weight_tensor_strategy_moe, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -604,8 +605,11 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase): ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.weight = Parameter(layer.weight.t(), requires_grad=False) - layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + weight, weight_scale, _ = process_fp8_weight_channel_strategy( + layer.weight, layer.weight_scale.data + ) + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) self.fp8_linear.process_weights_after_loading(layer) def apply( @@ -905,6 +909,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: @@ -951,6 +956,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): a1_scale=a1_scale, a2_scale=a2_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) def apply_monolithic( @@ -1602,7 +1608,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) @@ -1616,6 +1624,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, ) @property @@ -2162,6 +2171,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): fp8_backend=self.mxfp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) # No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation @@ -2207,6 +2217,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def apply_monolithic( @@ -2450,16 +2461,29 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): if key.startswith(parent_dot): return info["quant_algo"].upper() - # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj). - for candidate in self._quantized_layer_prefix_candidates(prefix): - parent_dot = candidate.rsplit(".", 1)[0] + "." - algos = { - info["quant_algo"].upper() - for key, info in self.quantized_layers.items() - if key.startswith(parent_dot) and "." not in key[len(parent_dot) :] - } - if len(algos) == 1: - return algos.pop() + # 4. Parent-prefix fallback for fused projections whose config lists + # shard names instead of vLLM's packed module name. + fused_projection_shards = { + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), + } + shard_names = fused_projection_shards.get(proj_name) + if shard_names is not None: + for candidate in self._quantized_layer_prefix_candidates(prefix): + parent_dot = candidate.rsplit(".", 1)[0] + "." + shard_algos: set[str] = set() + for shard_name in shard_names: + shard_prefix = f"{parent_dot}{shard_name}" + if shard_prefix in self.quantized_layers: + algo = self.quantized_layers[shard_prefix]["quant_algo"].upper() + shard_algos.add(algo) + if len(shard_algos) == 1: + return shard_algos.pop() + if len(shard_algos) > 1: + raise ValueError( + f"Mixed quant_algo within fused layer {prefix}: " + f"{shard_algos}. All shards must use the same quantization." + ) return None diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 4d3a3158791..e63c285aa30 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -458,6 +458,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -486,6 +487,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index 4274b16d55a..ef76594164a 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + convert_to_int8_moe_kernel_format, make_int8_moe_kernel, make_int8_moe_quant_config, select_int8_moe_backend, @@ -92,22 +93,36 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): replace_parameter(layer, "w2_scale", w2_scale) def _setup_kernel(self, layer: RoutedExperts) -> None: + w13, w2 = convert_to_int8_moe_kernel_format( + int8_backend=self.int8_backend, + w13=layer.w13_weight, + w2=layer.w2_weight, + layer=layer, + w13_scale=layer.w13_scale, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) assert self.moe_quant_config is not None assert self.experts_cls is not None self.moe_kernel = make_int8_moe_kernel( + int8_backend=self.int8_backend, moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> "FusedMoEQuantConfig | None": return make_int8_moe_quant_config( - w1_scale=layer.w13_scale, - w2_scale=layer.w2_scale, + int8_backend=self.int8_backend, + w1_scale=getattr(layer, "w13_scale", None), + w2_scale=getattr(layer, "w2_scale", None), w1_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), + layer=layer, ) diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 09d581a0734..84a81bd9064 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -200,6 +200,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): fp8_backend=self.fp8_backend, experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -226,6 +227,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): swiglu_limit=getattr(layer, "swiglu_limit", None), gemm1_alpha=getattr(layer, "swiglu_alpha", None), gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, ) def process_weights_after_loading(self, layer: Module) -> None: diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index bce888415ce..7bdf963b512 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1562,7 +1562,9 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + layer=layer, ) def get_fused_moe_quant_config( @@ -1576,6 +1578,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod): w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale_2, a2_scale=layer.w2_input_scale_2, + layer=layer, ) def apply( diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 5b77ac39225..632cca1fda2 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from enum import Enum from typing import TYPE_CHECKING import torch @@ -15,12 +14,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class FlashinferMoeBackend(Enum): - TENSORRT_LLM = "TensorRT-LLM" - CUTLASS = "CUTLASS" - CUTEDSL = "CUTEDSL" - - def activation_to_flashinfer_int(activation: MoEActivation) -> int: return activation_to_flashinfer_type(activation).value @@ -97,16 +90,6 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( ) -def is_flashinfer_supporting_global_sf(backend: FlashinferMoeBackend | None) -> bool: - # TODO(shuw@nvidia): Update when new backends are added. - backends_supporting_global_sf = ( - FlashinferMoeBackend.CUTLASS, - FlashinferMoeBackend.TENSORRT_LLM, - FlashinferMoeBackend.CUTEDSL, - ) - return backend in backends_supporting_global_sf - - def convert_moe_weights_to_flashinfer_trtllm_block_layout( cache_permute_indices: dict[torch.Size, torch.Tensor], w13_weight: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 32a2d86899c..49eadc5a152 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -33,6 +33,7 @@ from vllm.utils.deep_gemm import ( is_deep_gemm_e8m0_used, transform_sf_into_required_layout, ) +from vllm.utils.platform_utils import get_device_name_as_file_name from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -864,7 +865,7 @@ def get_w8a8_block_fp8_configs( # First look up if an optimized configuration is available in the configs # directory - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() json_file_name = f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8,block_shape=[{block_n},{block_k}].json" # noqa: E501 config_file_path = os.path.join( diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index d84a2e12f54..2d9e7aca766 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -764,7 +764,7 @@ def _convert_sublayer_to_humming( Returns: Tuple of (converted_weight_schema, converted_input_schema) """ - from humming.schema import HummingWeightSchema + from vllm.utils.humming import HummingWeightSchema if isinstance(weight_schema, HummingWeightSchema): # Already in Humming format @@ -814,7 +814,7 @@ def _prepare_and_transform_sublayer( This calls Humming's prepare_layer_meta and transform_humming_layer. """ - from humming.layer import HummingMethod + from vllm.utils.humming import HummingMethod HummingMethod.prepare_layer_meta( layer=layer, @@ -866,7 +866,7 @@ def _process_single_sublayer( Returns: Tuple of (final_weight_schema, final_input_schema) """ - from humming.schema import HummingWeightSchema + from vllm.utils.humming import HummingWeightSchema # Step 1: Convert from checkpoint format to humming format if needed current_weight_schema, current_input_schema = _convert_sublayer_to_humming( @@ -958,12 +958,10 @@ def convert_to_humming_moe_kernel_format( "Must provide either weight_schema/input_schema or quant_config" ) - from humming.layer import HummingInputSchema - from humming.schema import BaseWeightSchema - from vllm.model_executor.layers.quantization.utils.humming_utils import ( humming_is_layer_skipped, ) + from vllm.utils.humming import BaseWeightSchema, HummingInputSchema if weight_schema is None: weight_schema = BaseWeightSchema.from_config(quant_config) diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index eac6b11b219..4f624cf4963 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -12,47 +12,11 @@ import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.platform_utils import get_device_name_as_file_name logger = logging.getLogger(__name__) -def apply_w8a8_block_int8_linear( - input: torch.Tensor, - weight: torch.Tensor, - block_size: list[int], - weight_scale: torch.Tensor, - input_scale: torch.Tensor | None = None, - bias: torch.Tensor | None = None, -) -> torch.Tensor: - assert input_scale is None - # View input as 2D matrix for fp8 methods - input_2d = input.view(-1, input.shape[-1]) - output_shape = [*input.shape[:-1], weight.shape[0]] - - q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) - output = w8a8_block_int8_matmul( - q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype - ) - - if bias is not None: - output = output + bias - return output.to(dtype=input.dtype).view(*output_shape) - - -def input_to_int8( - x: torch.Tensor, dtype: torch.dtype = torch.int8 -) -> tuple[torch.Tensor, torch.Tensor]: - """This function quantizes input values to int8 values with - tensor-wise quantization.""" - iinfo = torch.iinfo(dtype) - min_val, max_val = x.aminmax() - amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) - int8_min, int8_max = iinfo.min, iinfo.max - scale = int8_max / amax - x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) - return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() - - def block_dequant( x_q_block: torch.Tensor, x_s: torch.Tensor, @@ -366,7 +330,7 @@ def get_w8a8_block_int8_configs( # First look up if an optimized configuration is available in the configs # directory - device_name = current_platform.get_device_name().replace(" ", "_") + device_name = get_device_name_as_file_name() json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501 config_file_path = os.path.join( diff --git a/vllm/model_executor/layers/quantization/utils/machete_utils.py b/vllm/model_executor/layers/quantization/utils/machete_utils.py index 95d8102ea50..7bf71ca99c5 100644 --- a/vllm/model_executor/layers/quantization/utils/machete_utils.py +++ b/vllm/model_executor/layers/quantization/utils/machete_utils.py @@ -16,10 +16,6 @@ def query_machete_supported_quant_types(zero_points: bool) -> list[ScalarType]: return [scalar_types.uint4b8, scalar_types.uint8b128] -def query_machete_supported_act_types(zero_points: bool) -> list[ScalarType]: - return [torch.float16, torch.bfloat16] - - def query_machete_supported_group_sizes(act_type: torch.dtype) -> list[int]: """ Queries the supported group sizes for Machete based on the activation type. diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index ea47ed06cbf..f6d96f574d8 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -738,71 +738,3 @@ def apply_gptq_marlin_linear( output = marlin_unpad_output(output, output_size_per_partition, padded_n) return output.reshape(out_shape) - - -def apply_awq_marlin_linear( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - weight_zp: torch.Tensor, - g_idx: torch.Tensor, - g_idx_sort_indices: torch.Tensor, - workspace: torch.Tensor, - quant_type: ScalarType, - output_size_per_partition: int, - input_size_per_partition: int, - input_global_scale: torch.Tensor | None = None, - bias: torch.Tensor | None = None, - use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, - input_dtype: torch.dtype | None = None, -) -> torch.Tensor: - reshaped_x = input.reshape(-1, input.shape[-1]) - out_shape = input.shape[:-1] + (output_size_per_partition,) - - padded_n, padded_k = marlin_repacked_nk(weight, quant_type.size_bits) - reshaped_x = marlin_pad_dim(reshaped_x, input_size_per_partition, padded_k) - - use_atomic_add = should_use_atomic_add_reduce( - m=reshaped_x.size(0), - n=padded_n, - k=padded_k, - device=input.device, - dtype=input.dtype, - ) - - a_scales = None - if input_dtype == torch.int8: - assert quant_type == scalar_types.uint4, ( - "W8A8-INT8 is not supported by marlin kernel." - ) - reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - a_scales = a_scales * input_global_scale - elif input_dtype == torch.float8_e4m3fn: - assert quant_type == scalar_types.uint4, ( - "INT8 weight + FP8 activation is not supported." - ) - reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype) - - output = ops.marlin_gemm( - reshaped_x, - None, - weight, - bias, - weight_scale, - a_scales, - None, - weight_zp, - g_idx, - g_idx_sort_indices, - workspace, - quant_type, - size_m=reshaped_x.shape[0], - size_n=padded_n, - size_k=padded_k, - use_atomic_add=use_atomic_add, - use_fp32_reduce=use_fp32_reduce, - is_zp_float=False, - ) - - output = marlin_unpad_output(output, output_size_per_partition, padded_n) - return output.reshape(out_shape) diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index db88ba273cd..0d3bf88cbf1 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -87,6 +87,12 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8): "split_k": 1, } opt_flags.update_opt_flags_constraints(constraints) + # Patches #47303: pad K (num scale groups) to 0 mod 4 + # TODO: Remove once we upgrade to Triton 3.8.0+ kernels + if scale.numel() > 0: + K = scale.shape[-1] + pad_k = -K % 4 + scale = torch.nn.functional.pad(scale, (0, pad_k)) elif current_platform.is_device_capability_family(100): constraints = { "is_persistent": True, diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 0c5cbae2a4f..705da43c373 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -216,6 +216,12 @@ kInt4W4A8StaticGroupSym = QuantKey( torch.int8, kInt4W4A8StaticGroupScale, symmetric=True ) +kInt4W4A8StaticChannelSym = QuantKey( + torch.int8, + ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), + symmetric=True, +) + def create_fp8_quant_key( static: bool, diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 3ed6eab6767..3f30a42fe40 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -47,21 +47,29 @@ class RunaiModelStreamerLoader(BaseModelLoader): # Validate every value before mutating os.environ, so a later # invalid key cannot leave an earlier one partially applied. env_updates: dict[str, str] = {} - for key, env_var in ( - ("concurrency", "RUNAI_STREAMER_CONCURRENCY"), - ("memory_limit", "RUNAI_STREAMER_MEMORY_LIMIT"), - ): - if key in extra_config: - value = extra_config[key] - if ( - isinstance(value, bool) - or not isinstance(value, int) - or value <= 0 - ): - raise ValueError( - f"{key} must be a positive integer, got {value!r}" - ) - env_updates[env_var] = str(value) + if "concurrency" in extra_config: + concurrency = extra_config["concurrency"] + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency <= 0 + ): + raise ValueError( + f"concurrency must be a positive integer, got {concurrency!r}" + ) + env_updates["RUNAI_STREAMER_CONCURRENCY"] = str(concurrency) + + if "memory_limit" in extra_config: + memory_limit = extra_config["memory_limit"] + if ( + isinstance(memory_limit, bool) + or not isinstance(memory_limit, int) + or memory_limit < -1 + ): + raise ValueError( + f"memory_limit must be an integer >= -1, got {memory_limit!r}" + ) + env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit) os.environ.update(env_updates) runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT") diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 7bf5fb04077..5cc057c3b45 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -533,6 +533,10 @@ class BailingMoeV25Model(nn.Module): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.word_embeddings(input_ids) + @property + def embed_tokens(self) -> nn.Module: + return self.word_embeddings + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py new file mode 100644 index 00000000000..da6b1ddb8b6 --- /dev/null +++ b/vllm/model_executor/models/bailing_moe_mtp.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Bailing MoE v2.5 MTP model.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +from transformers.configuration_utils import PretrainedConfig + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +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, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.bailing_moe_linear import ( + BailingMoeV25, + BailingMoeV25MLAAttention, +) +from vllm.sequence import IntermediateTensors + +from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix + + +def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig: + speculative_config = vllm_config.speculative_config + if speculative_config is not None: + draft_model_config = speculative_config.draft_model_config + if draft_model_config is not None: + return draft_model_config.hf_config + return vllm_config.model_config.hf_config + + +class BailingMTPSharedHead(nn.Module): + def __init__( + self, + config: PretrainedConfig, + prefix: str, + vllm_config: VllmConfig, + ) -> None: + super().__init__() + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +class BailingMoeV25MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + layer_id: int, + ) -> None: + super().__init__() + config = _get_draft_hf_config(vllm_config) + self.config = config + self.layer_id = layer_id + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = BailingMoeV25MLAAttention( + config, + quant_config=vllm_config.quant_config, + layer_id=layer_id, + prefix=maybe_prefix(prefix, "self_attn"), + cache_config=vllm_config.cache_config, + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp = BailingMoeV25( + config, + quant_config=vllm_config.quant_config, + layer_id=layer_id, + prefix=maybe_prefix(prefix, "mlp"), + ) + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.shared_head = BailingMTPSharedHead( + config, + maybe_prefix(prefix, "shared_head"), + vllm_config, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + inputs_embeds = self.enorm(inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + ) + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.self_attn(hidden_states, positions) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states.to(residual.device) + return self.final_layernorm(hidden_states) + + +class BailingMoeV25MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = _get_draft_hf_config(vllm_config) + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.layers = nn.ModuleDict( + { + str(idx): BailingMoeV25MultiTokenPredictorLayer( + vllm_config, + f"{prefix}.layers.{idx}", + idx, + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + else: + self.embed_tokens = PPMissingLayer() + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + lm_head: nn.Module | None = None, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + head = lm_head if lm_head is not None else mtp_layer.shared_head.head + return self.logits_processor( + head, + mtp_layer.shared_head(hidden_states), + ) + + +@support_torch_compile +class BailingMoeV25MTPModel(nn.Module): + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = _get_draft_hf_config(vllm_config) + self.lm_head: nn.Module | None = None + self.model = BailingMoeV25MultiTokenPredictor( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model.compute_logits(hidden_states, spec_step_idx, self.lm_head) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=0, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + (".fused_qkv_a_proj", ".q_a_proj", 0), + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + expert_params_mapping = list(self.get_expert_mapping()) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + loaded_mtp_layers: set[int] = set() + + def load_param( + name: str, + loaded_weight: torch.Tensor, + shard_id=None, + ) -> bool: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + return False + if name not in params_dict or is_pp_missing_parameter(name, self): + return False + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + weight_loader(param, loaded_weight) + elif isinstance(shard_id, int): + weight_loader(param, loaded_weight, shard_id) + else: + weight_loader( + param, + loaded_weight, + name, + expert_id=shard_id[0], + shard_id=shard_id[1], + ) + loaded_params.add(name) + return True + + def get_spec_layer_idx(name: str) -> int | None: + if not name.startswith("model.layers."): + return None + try: + layer_idx = int(name.split("model.layers.", 1)[1].split(".", 1)[0]) + except (IndexError, ValueError): + return None + mtp_idx = layer_idx - self.config.num_hidden_layers + if 0 <= mtp_idx < self.config.num_nextn_predict_layers: + return layer_idx + return None + + def normalize_name(name: str) -> str: + name = name.replace(".attention.dense", ".self_attn.o_proj") + name = name.replace(".attention.", ".self_attn.") + return name.replace( + "mlp.gate.e_score_correction_bias", + "mlp.gate.expert_bias", + ) + + def load_lm_head(loaded_weight: torch.Tensor) -> None: + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + name = f"model.layers.{layer_idx}.shared_head.head.weight" + load_param(name, loaded_weight) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if name == "model.word_embeddings.weight": + load_param("model.embed_tokens.weight", loaded_weight) + continue + if name == "lm_head.weight": + load_lm_head(loaded_weight) + continue + + spec_layer = get_spec_layer_idx(name) + if spec_layer is None: + continue + name = normalize_name(name) + + loaded = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name and name not in params_dict: + continue + mapped_name = name.replace(weight_name, param_name) + if load_param(mapped_name, loaded_weight, shard_id): + loaded = True + break + if loaded: + loaded_mtp_layers.add(spec_layer) + continue + + if "mlp.experts" in name: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + mapped_name = name.replace(weight_name, param_name) + if load_param( + mapped_name, + loaded_weight, + (expert_id, shard_id), + ): + loaded = True + break + if loaded: + loaded_mtp_layers.add(spec_layer) + continue + + if load_param(name, loaded_weight): + loaded_mtp_layers.add(spec_layer) + + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if layer_idx not in loaded_mtp_layers: + raise ValueError( + f"Bailing MTP speculative decoding layer {layer_idx} " + "weights are missing from checkpoint. Use a checkpoint " + "that includes MTP layer weights, or disable speculative " + "decoding." + ) + return loaded_params diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index b20db76692f..4290609cbd0 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -279,7 +279,9 @@ class DeepseekV2MoE(nn.Module): config: DeepseekV2Config | DeepseekV3Config, parallel_config: ParallelConfig, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", + apply_routed_scale_to_output: bool = False, ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() @@ -305,9 +307,7 @@ class DeepseekV2MoE(nn.Module): self.gate = GateLinear( config.hidden_size, config.n_routed_experts, - params_dtype=self.router_dtype, out_dtype=self.router_dtype, - force_fp32_compute=self.router_dtype == torch.float32, prefix=f"{prefix}.gate", ) if getattr(config, "topk_method", None) == "noaux_tc": @@ -372,13 +372,13 @@ class DeepseekV2MoE(nn.Module): topk_group=getattr(config, "topk_group", 1), prefix=f"{prefix}.experts", scoring_func=getattr(config, "scoring_func", "softmax"), - # aiter applies routed_scaling_factor internally routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled, + apply_routed_scale_to_output=apply_routed_scale_to_output, e_score_correction_bias=self.gate.e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, + reduce_results=reduce_results, n_shared_experts=config.n_shared_experts if self.is_fusion_moe_shared_experts_enabled else None, @@ -713,6 +713,7 @@ class Indexer(nn.Module): ) self.is_inplace_rope = is_inplace_rope + self.n_head_scale = self.n_head**-0.5 self.use_fused_indexer_q = ( current_platform.is_cuda() and self.quant_block_size == self.head_dim @@ -761,15 +762,16 @@ class Indexer(nn.Module): rotary_emb.cos_sin_cache, weights, self.softmax_scale, - self.n_head**-0.5, + self.n_head_scale, rotary_emb.is_neox_style, ) # rotate only the MQA K - q_dummy = torch.empty_like(k_pe.unsqueeze(1)) - _, k_pe = rotary_emb(positions, q_dummy, k_pe.unsqueeze(1)) - k_pe = k_pe.reshape(-1, 1, self.rope_dim) - k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + k_pe = k_pe.unsqueeze(1) + q_dummy = torch.empty_like(k_pe) + _, k_pe = rotary_emb(positions, q_dummy, k_pe) + k_pe = k_pe.reshape(-1, self.rope_dim) + k = torch.cat([k_pe, k_nope], dim=-1) return self.indexer_op(hidden_states, q_fp8, k, weights) else: @@ -790,13 +792,13 @@ class Indexer(nn.Module): # Note: RoPE (NeoX) can introduce extra leading dimensions during # compilation so we need to reshape back to token-flattened shapes q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) - k_pe = k_pe.reshape(-1, 1, self.rope_dim) + k_pe = k_pe.reshape(-1, self.rope_dim) # `rotary_emb` is shape-preserving; `q_pe` is already # [num_tokens, n_head, rope_dim]. q = torch.cat([q_pe, q_nope], dim=-1) - # `k_pe` is [num_tokens, 1, rope_dim] (MQA). - k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + # `k_pe` is [num_tokens, rope_dim] (MQA). + k = torch.cat([k_pe, k_nope], dim=-1) # we only quant q here since k quant is fused with cache insertion q = q.view(-1, self.head_dim) @@ -807,12 +809,9 @@ class Indexer(nn.Module): use_ue8m0=self.scale_fmt is not None, ) q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) - q_scale = q_scale.view(-1, self.n_head, 1) + q_scale = q_scale.view(-1, self.n_head) - weights = ( - weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 - ) - weights = weights.squeeze(-1) + weights = weights * q_scale * self.softmax_scale * self.n_head_scale return self.indexer_op(hidden_states, q_fp8, k, weights) @@ -830,7 +829,7 @@ def _try_load_fp8_indexer_wk( if "indexer.wk." not in name or "wk_weights" in name: return False # Weight is not an isolated WK weight for the indexer, ignore. is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn - is_scale = "weight_scale_inv" in name + is_scale = "weight_scale" in name if not is_weight and not is_scale: return False # WK is not in FP8 format, ignore. # Buffer this tensor (weight or scale) until both have arrived. @@ -1247,6 +1246,8 @@ class DeepseekV2DecoderLayer(nn.Module): parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.mlp", + # aiter applies routed_scaling_factor internally + apply_routed_scale_to_output=not rocm_aiter_ops.is_fused_moe_enabled(), ) else: self.mlp = DeepseekV2MLP( @@ -1287,13 +1288,10 @@ class DeepseekV2DecoderLayer(nn.Module): hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) hidden_states = hidden_states[:full_num_tokens] - attn_kwargs = { - "positions": positions, - "hidden_states": hidden_states, - } - if not self.use_mha: - attn_kwargs["llama_4_scaling"] = llama_4_scaling - hidden_states = self.self_attn(**attn_kwargs) + if self.use_mha: + hidden_states = self.self_attn(positions, hidden_states) + else: + hidden_states = self.self_attn(positions, hidden_states, llama_4_scaling) if ( not isinstance(self.self_attn, DeepseekAttention) @@ -1309,15 +1307,11 @@ class DeepseekV2DecoderLayer(nn.Module): residual *= 1.0 / self.routed_scaling_factor if self.use_sequence_parallel_moe: - sp_remainder = ( - hidden_states.shape[0] % get_tensor_model_parallel_world_size() - ) + tp_world_size = get_tensor_model_parallel_world_size() + # small trick using minus, eg. -17 % 8 = 7 + sp_pad = (-hidden_states.shape[0]) % tp_world_size # pad if not divisible by world size - if sp_remainder: - sp_pad = get_tensor_model_parallel_world_size() - sp_remainder - hidden_states = torch.nn.functional.pad( - hidden_states, (0, 0, 0, sp_pad) - ) + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) if not input_is_sequence_parallel: residual = sequence_parallel_chunk(residual) @@ -1519,7 +1513,10 @@ class DeepseekV2Model(nn.Module): ("qkv_proj", "v_proj", "v"), ] # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - _pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion + _pending_wk_fp8 = getattr(self, "_pending_indexer_wk_fp8", None) + if _pending_wk_fp8 is None: + self._pending_indexer_wk_fp8 = _pending_wk_fp8 = {} + indexer_fused_mapping = [ ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index eebb5ef148e..11a10131df1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -47,6 +47,7 @@ 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.platforms import current_platform 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 @@ -516,11 +517,6 @@ def _compiled_sample_step( 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) @@ -800,7 +796,10 @@ class DiffusionGemmaModelState(ModelState): max_denoising_steps=max_denoising_steps, device=device, hidden_size=text_config.hidden_size, - stability_threshold=self.gen_config["stability_threshold"], + # In Transformers, `stability_threshold=1` (the default) means the current + # step must match the previous step. In vLLM, the history buffer includes + # the current step, so we add 1 to match the same behavior. + stability_threshold=self.gen_config["stability_threshold"] + 1, ) self._req_id_to_index: dict[str, int] = {} @@ -1273,9 +1272,12 @@ class DiffusionSampler: 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. + # Clear once: the tiled loop below only scatters its own decode slots, + # so it must not re-clear earlier tiles' writes. sampled = self._sampled[:num_reqs] num_sampled = self._num_sampled[:num_reqs] + sampled.zero_() + num_sampled.zero_() all_slots = input_batch.idx_mapping[:num_reqs] @@ -1283,94 +1285,109 @@ class DiffusionSampler: # 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, - sc_vocab_start=self.sc_vocab_start, - sc_vocab_end=self.sc_vocab_end, - tp_size=self.tp_size, - tp_group_name=self.tp_group_name, - ) - - # --- 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, - ) + # Sample over the [num_decode * CL, vocab] logits. The fp32 pipeline in + # _compiled_sample_step keeps several live [group * CL, vocab] copies, so + # size each tile to a fraction of free memory to bound the transient at + # high concurrency. Tiling is bit-identical to a single pass. + group = max(num_decode, 1) + if num_decode > 0: + free, _ = current_platform.mem_get_info() + # ~10 transient fp32 copies of [group * CL, vocab] inside the step + # (eager peaks at ~8; pad for allocator overhead and small tensors). + bytes_per_req = CL * self.vocab_size * 4 * 10 + budget = int(free * 0.5) // max(bytes_per_req, 1) + group = max(1, min(num_decode, budget)) + + for start_req in range(0, num_decode, group): + end_req = min(start_req + group, num_decode) + tile = slice(start_req, end_req) + tile_slots = decode_slots[tile] + + scaled = _compiled_sample_step( + logits[start_req * CL : end_req * CL], + tile_slots, + decode_idx[tile], + all_slots, + valid_canvas_len[tile], + # 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=CL, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + sc_vocab_start=self.sc_vocab_start, + sc_vocab_end=self.sc_vocab_end, + tp_size=self.tp_size, + tp_group_name=self.tp_group_name, + ) + + # Logprobs for denoise steps that just converged (is_encoder_phase + # flipped False→True), stashed per tile so `scaled` is freed each tile. + if max_num_logprobs >= 0: + converged_mask = states.is_encoder_phase[tile_slots] + just_converged = converged_mask & ~is_committing[tile] + 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 = tile_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[start_req + li]) + pos = li * CL + self._pending_logprobs[slot.item()] = compute_topk_logprobs( + flat_logits[pos : pos + 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. + logprobs_tensors = None + if max_num_logprobs >= 0 and 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, diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 9557ca68020..7e7cdcd822c 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -20,7 +20,7 @@ from typing import Any import numpy as np import torch import torch.nn as nn -from transformers import PreTrainedTokenizerFast, WhisperFeatureExtractor +from transformers import TokenizersBackend, WhisperFeatureExtractor from transformers.activations import get_activation from transformers.feature_extraction_utils import BatchFeature from transformers.modeling_outputs import BaseModelOutput @@ -556,15 +556,15 @@ class FunAudioChatProcessingInfo(BaseProcessingInfo): return WhisperFeatureExtractor.from_pretrained(self.model_id) @cached_property - def speech_tokenizer(self) -> PreTrainedTokenizerFast: - return PreTrainedTokenizerFast.from_pretrained( + def speech_tokenizer(self) -> TokenizersBackend: + return TokenizersBackend.from_pretrained( self.model_id, subfolder="speech_tokenizer" ) def get_feature_extractor(self) -> WhisperFeatureExtractor: return self.feature_extractor - def get_speech_tokenizer(self) -> PreTrainedTokenizerFast: + def get_speech_tokenizer(self) -> TokenizersBackend: return self.speech_tokenizer def get_data_parser(self): diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index b30ab4b7aef..f4d200d11d9 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -51,6 +51,7 @@ from .utils import ( AutoWeightsLoader, WeightsMapper, extract_layer_index, + get_draft_quant_config, maybe_prefix, ) @@ -182,14 +183,14 @@ class Gemma4MTPAttention(nn.Module): hidden_size, self.total_num_heads * self.head_dim, bias=config.attention_bias, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.q_proj", ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, hidden_size, bias=config.attention_bias, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.o_proj", ) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) @@ -304,7 +305,7 @@ class Gemma4MTPDecoderLayer(nn.Module): hidden_size=self.hidden_size, intermediate_size=text_config.intermediate_size, hidden_activation=text_config.hidden_activation, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.mlp", ) @@ -357,7 +358,9 @@ class Gemma4MultiTokenPredictor(nn.Module): config = vllm_config.speculative_config.draft_model_config.hf_config text_config = _get_text_config(config) + quant_config = get_draft_quant_config(vllm_config) self.config = text_config + self.quant_config = quant_config self.hidden_size = text_config.hidden_size self.backbone_hidden_size = getattr( @@ -369,6 +372,8 @@ class Gemma4MultiTokenPredictor(nn.Module): self.embed_tokens = VocabParallelEmbedding( self.vocab_size, self.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", ) self.pre_projection = ColumnParallelLinear( @@ -376,6 +381,7 @@ class Gemma4MultiTokenPredictor(nn.Module): self.hidden_size, bias=False, gather_output=True, + quant_config=quant_config, prefix=f"{prefix}.pre_projection", ) @@ -384,6 +390,7 @@ class Gemma4MultiTokenPredictor(nn.Module): self.backbone_hidden_size, bias=False, input_is_parallel=False, + quant_config=quant_config, prefix=f"{prefix}.post_projection", ) @@ -391,7 +398,7 @@ class Gemma4MultiTokenPredictor(nn.Module): Gemma4MTPDecoderLayer( text_config, cache_config=vllm_config.cache_config, - quant_config=vllm_config.quant_config, + quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", ) for idx in range(self.num_mtp_layers) @@ -473,6 +480,7 @@ class Gemma4MTP(nn.Module): super().__init__() config = vllm_config.speculative_config.draft_model_config.hf_config text_config = _get_text_config(config) + self.quant_config = get_draft_quant_config(vllm_config) self.config = config self._stable_full_lm_head_weight: torch.Tensor | None = None @@ -488,6 +496,7 @@ class Gemma4MTP(nn.Module): self.lm_head = ParallelLMHead( text_config.vocab_size, text_config.hidden_size, + quant_config=self.quant_config, prefix=maybe_prefix(prefix, "lm_head"), ) if getattr(config, "tie_word_embeddings", True): diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 01f2752ac54..5fc7e62640f 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -192,7 +192,10 @@ class MLPBlock(torch.nn.Module): quant_config = vllm_config.quant_config parallel_config = vllm_config.parallel_config - self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + self.is_sequence_parallel = ( + parallel_config.use_sequence_parallel_moe + and vllm_config.lora_config is None + ) self.layer_idx = layer_idx self.num_experts = config.num_local_experts diff --git a/vllm/model_executor/models/hunyuan_vision.py b/vllm/model_executor/models/hunyuan_vision.py index 313e13c915c..87520fe29ee 100644 --- a/vllm/model_executor/models/hunyuan_vision.py +++ b/vllm/model_executor/models/hunyuan_vision.py @@ -31,7 +31,11 @@ from typing import Annotated, Any, Literal, TypeAlias import torch import torch.nn as nn import torch.nn.functional as F -from transformers import BatchFeature +from transformers import BatchFeature, HunYuanVLProcessor +from transformers.models.hunyuan_vl.image_processing_hunyuan_vl import ( + HunYuanVLImageProcessor, + smart_resize, +) from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions @@ -75,11 +79,6 @@ from vllm.transformers_utils.configs.hunyuan_vl import ( HunYuanVLConfig, HunYuanVLVisionConfig, ) -from vllm.transformers_utils.processors.hunyuan_vl import HunYuanVLProcessor -from vllm.transformers_utils.processors.hunyuan_vl_image import ( - HunYuanVLImageProcessor, - smart_resize, -) from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -574,9 +573,12 @@ class HunYuanVLProcessingInfo(BaseProcessingInfo): self, **kwargs: object, ) -> HunYuanVLProcessor: + # transformers>=5.13 replaced `use_fast` with `backend`; pin the + # PIL backend to match the released HunyuanOCR checkpoint packing. + kwargs.pop("use_fast", None) + kwargs.setdefault("backend", "pil") return self.ctx.get_hf_processor( HunYuanVLProcessor, - use_fast=kwargs.pop("use_fast", True), **kwargs, ) @@ -724,8 +726,18 @@ class HunYuanVLMultiModalProcessor(BaseMultiModalProcessor[HunYuanVLProcessingIn mm_kwargs: Mapping[str, object], tok_kwargs: Mapping[str, object], ) -> BatchFeature: + hf_processor = self.info.get_hf_processor(**mm_kwargs) + # HunYuanVLProcessor requires image placeholders wrapped with start/end tokens. + if mm_data.get("images") is not None and prompt: + img_tok = hf_processor.image_token + wrapped = ( + f"{hf_processor.image_start_token}{img_tok}" + f"{hf_processor.image_end_token}" + ) + if img_tok in prompt and wrapped not in prompt: + prompt = prompt.replace(img_tok, wrapped) return self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**mm_kwargs), + hf_processor, dict(text=prompt, **mm_data), dict(**mm_kwargs, **tok_kwargs), ) diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index cb2c96fa96a..62bdd028bb5 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -177,7 +177,9 @@ class HYV3MoEFused(nn.Module): else: self.shared_mlp = None - self.expert_bias = nn.Parameter(torch.empty(config.num_experts)) + self.expert_bias = nn.Parameter( + torch.empty(config.num_experts, dtype=torch.float32) + ) scoring_func = "sigmoid" e_score_correction_bias = self.expert_bias diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 7321b913605..e2b74744b41 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -56,6 +56,10 @@ from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config from vllm.transformers_utils.processor import cached_get_image_processor from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor +from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( + KimiK25FusedVisionProcessor, +) +from vllm.utils.import_utils import is_numba_available from vllm.utils.tensor_schema import TensorSchema, TensorShape from .utils import ( @@ -108,10 +112,16 @@ class KimiK25ProcessingInfo(BaseProcessingInfo): self.hf_config = hf_config = self.get_hf_config() tokenizer = self.get_tokenizer() + processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None + logger.info_once( + "Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.", + "fused CPU" if processor_cls is not None else "remote HF", + ) image_processor = cached_get_image_processor( self.ctx.model_config.model, revision=self.ctx.model_config.revision, trust_remote_code=self.ctx.model_config.trust_remote_code, + processor_cls_overrides=processor_cls, ) # Resolve token ID from the tokenizer because transformers v5 diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 8ce86c77807..2d859bd4918 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -233,7 +233,10 @@ class LlamaModel(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor]: if input_embeds is None: input_embeds = self.embed_input_ids(input_ids) - assert hidden_states.shape[-1] == input_embeds.shape[-1] + torch._assert( + hidden_states.shape[-1] == input_embeds.shape[-1], + "hidden_states and input_embeds must have the same last dimension", + ) residual = None for layer in self.layers: diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index fa32b31560c..60b11129821 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -552,6 +552,7 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo): vendored_processor = MiniCPMVProcessor( image_processor=hf_processor.image_processor, tokenizer=hf_processor.tokenizer, + version=self.get_model_version(), ) hf_processor = vendored_processor diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 6cdc3624151..dc75928eb6c 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -35,6 +35,8 @@ from vllm.multimodal.parse import ImageProcessorItems, ImageSize, VideoProcessor from vllm.multimodal.processing.processor import ( PromptReplacement, PromptUpdateDetails, + ResolvedPromptUpdate, + _seq2text, ) from vllm.sequence import IntermediateTensors @@ -408,6 +410,35 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): for modality, pattern in placeholders ] + def _recompute_cached_prompt_update( + self, cached_update: ResolvedPromptUpdate, new_item_idx: int + ) -> ResolvedPromptUpdate: + new_update = super()._recompute_cached_prompt_update( + cached_update, new_item_idx + ) + # MiniCPM-V 4.6 prefixes video placeholders with `{idx}` + # (the base class only rewrites the image modality). + if cached_update.modality == "video": + tokenizer = self.info.get_tokenizer() + id_start = getattr(tokenizer, "image_id_start_token", "") + id_end = getattr(tokenizer, "image_id_end_token", "") + video_token = getattr(tokenizer, "video_token", "<|video_pad|>") + + text = _seq2text(tokenizer, cached_update.content.full) + prev_item_idx = cached_update.item_idx + + new_update = new_update.with_content( + PromptUpdateDetails.select_text( + text.replace( + f"{id_start}{prev_item_idx}{id_end}", + f"{id_start}{new_item_idx}{id_end}", + 1, + ), + video_token, + ) + ) + return new_update + def _get_mm_fields_config( self, hf_inputs, diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index fdfe982fca3..bef8057d2fb 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -862,6 +862,10 @@ class MossQwen3ForCausalLM(Qwen3ForCausalLM): batch_size, dtype, device ) for layer_idx in self.deepstack_inject_layer_indices: + # Non-first PP ranks only receive DeepStack payloads for layers + # at or after their local start layer. + if layer_idx < self.model.start_layer: + continue intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros( (batch_size, self.config.hidden_size), dtype=dtype, diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py new file mode 100644 index 00000000000..5236f5c61be --- /dev/null +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -0,0 +1,801 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Inference-only MOSS-Transcribe-Diarize ASR model. + +The checkpoint layout is: + +* ``model.whisper_encoder.*``: Whisper-medium encoder weights. +* ``model.vq_adaptor.*``: 4x time-merge projector. +* ``model.language_model.*``: Qwen3-0.6B decoder weights. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from typing import Annotated, Any, Literal, TypeAlias + +import torch +from torch import nn +from transformers import BatchFeature + +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.config.speech_to_text import SpeechToTextParams +from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + SupportsTranscription, + _require_is_multimodal, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + _merge_multimodal_embeddings, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.model_executor.models.whisper import ( + WhisperEncoder, + _create_fake_bias_for_k_proj, +) +from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + AudioItem, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + DictEmbeddingItems, + ModalityDataItems, + MultiModalDataItems, + MultiModalDataParser, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.multimodal.processing.processor import ProcessorInputs +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +WHISPER_ENCODER_STRIDE = 2 + +AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" + +DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT = ( + "请将音频转写为文本，每一段需以起始时间戳和说话人编号" + "（[S01]、[S02]、[S03]…）开头，正文为对应的语音内容，" + "并在段末标注结束时间戳，以清晰标明该段语音范围。" +) + + +class MossTranscribeDiarizeAudioInputs(TensorSchema): + """ + Dimensions: + - c: Audio chunks + - m: Mel bins + - f: Mel frames + - n: Number of audio items + """ + + type: Literal["audio_features"] = "audio_features" + + input_features: Annotated[ + torch.Tensor | None, + TensorShape("c", "m", "f"), + ] + audio_feature_lengths: Annotated[ + torch.Tensor | None, + TensorShape("c"), + ] + audio_chunk_counts: Annotated[ + torch.Tensor | None, + TensorShape("n"), + ] + + +class MossTranscribeDiarizeEmbeddingInputs(TensorSchema): + """ + Dimensions: + - n: Number of audio items + - t: Number of audio tokens + - h: Hidden size + """ + + type: Literal["audio_embeds"] = "audio_embeds" + + audio_embeds: Annotated[ + list[torch.Tensor], + TensorShape("n", "t", "h", dynamic_dims={"t"}), + ] + + +MossTranscribeDiarizeInputs: TypeAlias = ( + MossTranscribeDiarizeAudioInputs | MossTranscribeDiarizeEmbeddingInputs +) + + +def _compute_total_audio_tokens( + num_samples: int, + feature_extractor: Any, + audio_merge_size: int, +) -> int: + if num_samples <= 0: + return 0 + + n_samples = int(feature_extractor.n_samples) + stride = ( + int(feature_extractor.hop_length) + * WHISPER_ENCODER_STRIDE + * int(audio_merge_size) + ) + total = 0 + for start in range(0, num_samples, n_samples): + chunk_samples = min(n_samples, num_samples - start) + total += (chunk_samples - 1) // stride + 1 + return total + + +def _get_max_audio_samples(feature_extractor: Any) -> int: + if hasattr(feature_extractor, "chunk_length"): + return int(feature_extractor.chunk_length * feature_extractor.sampling_rate) + return int(feature_extractor.n_samples) + + +def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]: + if isinstance(audio_embeds, torch.Tensor): + if audio_embeds.ndim == 2: + return [audio_embeds] + if audio_embeds.ndim == 3: + return list(audio_embeds.unbind(dim=0)) + raise ValueError( + f"`audio_embeds` must be a 2D or 3D tensor, got {audio_embeds.ndim}D." + ) + + if isinstance(audio_embeds, (list, tuple)) and all( + isinstance(audio_embed, torch.Tensor) for audio_embed in audio_embeds + ): + return list(audio_embeds) + + raise TypeError( + "`audio_embeds` must be a torch.Tensor or a list of torch.Tensor " + f"objects, got {type(audio_embeds)!r}." + ) + + +def _get_required_token_id(tokenizer: Any, token: str) -> int: + token_id = tokenizer.convert_tokens_to_ids(token) + if token_id is None or (isinstance(token_id, int) and token_id < 0): + raise ValueError(f"Tokenizer is missing required token {token!r}.") + return int(token_id) + + +def _get_audios_from_mm_data(mm_data: Mapping[str, object]) -> list[Any]: + mm_data_dict = dict(mm_data) + audio_data = mm_data_dict.pop("audios", mm_data_dict.pop("audio", [])) + if isinstance(audio_data, list): + audios = audio_data + elif isinstance(audio_data, tuple) and not ( + len(audio_data) == 2 and isinstance(audio_data[1], (int, float)) + ): + audios = list(audio_data) + elif audio_data is None: + audios = [] + else: + audios = [audio_data] + + audio_arrays: list[Any] = [] + for audio in audios: + if isinstance(audio, (tuple, list)) and len(audio) == 2: + audio = audio[0] + audio_arrays.append(audio) + return audio_arrays + + +def _add_vllm_audio_metadata( + processed: BatchFeature, + num_audios: int, +) -> BatchFeature: + audio_feature_lengths = processed["audio_feature_lengths"] + if not isinstance(audio_feature_lengths, torch.Tensor): + audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long) + audio_feature_lengths = audio_feature_lengths.to(dtype=torch.long) + + audio_chunk_mapping = processed.get("audio_chunk_mapping") + if audio_chunk_mapping is None: + if num_audios == 1: + audio_chunk_mapping = torch.zeros_like(audio_feature_lengths) + elif audio_feature_lengths.numel() == num_audios: + audio_chunk_mapping = torch.arange(num_audios, dtype=torch.long) + else: + raise ValueError( + "The MOSS processor did not return `audio_chunk_mapping`, and " + "the chunk-to-audio mapping cannot be inferred." + ) + elif not isinstance(audio_chunk_mapping, torch.Tensor): + audio_chunk_mapping = torch.tensor(audio_chunk_mapping, dtype=torch.long) + audio_chunk_mapping = audio_chunk_mapping.to(dtype=torch.long) + if audio_chunk_mapping.numel() != audio_feature_lengths.numel(): + raise ValueError( + "`audio_chunk_mapping` must contain one item per audio chunk: got " + f"{audio_chunk_mapping.numel()} mappings for " + f"{audio_feature_lengths.numel()} chunks." + ) + if audio_chunk_mapping.numel() > 0 and ( + audio_chunk_mapping.min().item() < 0 + or audio_chunk_mapping.max().item() >= num_audios + ): + raise ValueError("`audio_chunk_mapping` contains an out-of-range audio index.") + + audio_chunk_counts = torch.bincount( + audio_chunk_mapping.cpu(), + minlength=num_audios, + ).to(dtype=torch.long) + audio_token_lengths = torch.zeros( + num_audios, + dtype=torch.long, + device=audio_feature_lengths.device, + ) + audio_token_lengths.scatter_add_( + 0, + audio_chunk_mapping.to(device=audio_feature_lengths.device), + audio_feature_lengths, + ) + + processed["audio_feature_lengths"] = audio_feature_lengths + processed["audio_chunk_counts"] = audio_chunk_counts + processed["audio_token_lengths"] = audio_token_lengths + return processed + + +class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder): + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}, + orig_to_new_stacked={ + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + }, + ) + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__( + vllm_config=vllm_config.with_hf_config( + vllm_config.model_config.hf_config.audio_config + ), + prefix=prefix, + ) + self.audio_merge_size = int(vllm_config.model_config.hf_config.audio_merge_size) + self.max_encoder_batch: int | None = None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight") + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def forward( + self, + input_features: torch.Tensor, + audio_feature_lengths: torch.Tensor, + ) -> torch.Tensor: + if input_features.numel() == 0: + return input_features.new_empty((1, 0, self.conv1.out_channels)) + device = self.conv1.weight.device + dtype = self.conv1.weight.dtype + input_features = input_features.to(device=device, dtype=dtype) + audio_feature_lengths = audio_feature_lengths.to(device=device) + batch_size = self.max_encoder_batch or input_features.shape[0] + encoded_parts: list[torch.Tensor] = [] + + for start in range(0, input_features.shape[0], batch_size): + encoded_parts.append( + super().forward([input_features[start : start + batch_size]]) + ) + + hidden = torch.cat(encoded_parts, dim=0) + chunks = [ + hidden[idx : idx + 1, : int(token_len.item()) * self.audio_merge_size] + for idx, token_len in enumerate(audio_feature_lengths) + ] + return torch.cat(chunks, dim=1) + + +class MossTranscribeDiarizeVQAdaptor(nn.Module): + def __init__(self, input_dim: int, hidden_size: int, eps: float) -> None: + super().__init__() + self.layers = nn.Sequential( + nn.Linear(input_dim, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + nn.LayerNorm(hidden_size, eps=eps, bias=True), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layers(x) + + +def _mtd_field_config( + hf_inputs: Mapping[str, torch.Tensor], +) -> Mapping[str, MultiModalFieldConfig]: + fields: dict[str, MultiModalFieldConfig] = {} + if "audio_embeds" in hf_inputs: + fields["audio_embeds"] = MultiModalFieldConfig.batched("audio") + if "audio_chunk_counts" in hf_inputs: + audio_chunk_counts = hf_inputs["audio_chunk_counts"] + fields.update( + input_features=MultiModalFieldConfig.flat_from_sizes( + "audio", + audio_chunk_counts, + ), + audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes( + "audio", + audio_chunk_counts, + ), + audio_chunk_counts=MultiModalFieldConfig.batched("audio"), + audio_token_lengths=MultiModalFieldConfig.batched("audio"), + ) + return fields + + +class MossTranscribeDiarizeMultiModalDataParser(MultiModalDataParser): + def _parse_audio_data( + self, + data: dict[str, torch.Tensor] | ModalityData[AudioItem], + ) -> ModalityDataItems[Any, Any] | None: + if isinstance(data, dict): + return DictEmbeddingItems( + data, + modality="audio", + required_fields={"audio_embeds"}, + fields_factory=_mtd_field_config, + ) + + return super()._parse_audio_data(data) + + +class MossTranscribeDiarizeProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + def get_hf_processor(self, **kwargs: object) -> Any: + return self.ctx.get_hf_processor(**kwargs) + + def get_feature_extractor(self, **kwargs: object) -> Any: + return self.get_hf_processor(**kwargs).feature_extractor + + def get_data_parser(self) -> MultiModalDataParser: + feature_extractor = self.get_feature_extractor() + return MossTranscribeDiarizeMultiModalDataParser( + target_sr=feature_extractor.sampling_rate, + target_channels=1, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + if mm_counts.get("audio", 0) <= 0: + return {} + + feature_extractor = self.get_feature_extractor() + max_audio_samples = _get_max_audio_samples(feature_extractor) + max_audio_tokens = _compute_total_audio_tokens( + max_audio_samples, + feature_extractor, + self.get_hf_processor().audio_merge_size, + ) + return {"audio": min(seq_len, max_audio_tokens)} + + +class MossTranscribeDiarizeDummyInputsBuilder( + BaseDummyInputsBuilder[MossTranscribeDiarizeProcessingInfo] +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio", 0) + if num_audios == 0: + return {} + + feature_extractor = self.info.get_feature_extractor() + return { + "audio": self._get_dummy_audios( + length=_get_max_audio_samples(feature_extractor), + num_audios=num_audios, + overrides=mm_options.get("audio"), + ) + } + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> ProcessorInputs: + dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data) + num_audios = mm_counts.get("audio", 0) + tokenizer = self.info.get_tokenizer() + prompt = tokenizer.encode( + AUDIO_PLACEHOLDER * num_audios, + add_special_tokens=False, + ) or tokenizer.encode( + "\n", + add_special_tokens=False, + ) + return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items) + + +class MossTranscribeDiarizeMultiModalProcessor( + BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo] +): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + tokenizer = self.info.get_tokenizer() + audios = _get_audios_from_mm_data(mm_data) + if not audios: + input_ids = tokenizer.encode( + prompt, + add_special_tokens=tok_kwargs.get("add_special_tokens", False), + ) + return BatchFeature({"input_ids": [input_ids]}, tensor_type="pt") + + processed = self.info.ctx.call_hf_processor( + self.info.get_hf_processor(**mm_kwargs), + dict(text=prompt, audio=audios), + dict(**mm_kwargs, **tok_kwargs), + ) + return _add_vllm_audio_metadata(processed, len(audios)) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return mm_items.get_count("audio", strict=False) > 0 + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return _mtd_field_config(hf_inputs) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + out_mm_data = out_mm_kwargs.get_data() + audio_token_lengths_tensor = out_mm_data.get("audio_token_lengths") + if audio_token_lengths_tensor is None: + audio_embeds = out_mm_data.get("audio_embeds") + if audio_embeds is None: + audio_token_lengths: list[int] = [] + else: + audio_token_lengths = [ + int(audio_embed.shape[0]) + for audio_embed in _as_audio_embedding_list(audio_embeds) + ] + else: + if not isinstance(audio_token_lengths_tensor, torch.Tensor): + raise TypeError( + "`audio_token_lengths` must be a torch.Tensor, got " + f"{type(audio_token_lengths_tensor)!r}." + ) + audio_token_lengths = [ + int(length) for length in audio_token_lengths_tensor.tolist() + ] + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + audio_start_id = _get_required_token_id(tokenizer, processor.audio_start_token) + audio_token_id = int(processor.audio_token_id) + audio_end_id = _get_required_token_id(tokenizer, processor.audio_end_token) + + def get_num_tokens(item_idx: int) -> int: + if item_idx >= len(audio_token_lengths): + raise ValueError( + "Cannot determine the number of audio tokens for audio item " + f"{item_idx}." + ) + num_tokens = audio_token_lengths[item_idx] + if num_tokens <= 0: + raise ValueError("Audio input is too short to produce any tokens.") + return num_tokens + + def get_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]: + num_tokens = get_num_tokens(item_idx) + audio_tokens = processor._audio_span_ids(num_tokens) + return PromptUpdateDetails.select_token_id( + [audio_start_id] + audio_tokens + [audio_end_id], + embed_token_id=audio_token_id, + ) + + return [ + PromptReplacement( + modality="audio", + target=AUDIO_PLACEHOLDER, + replacement=get_replacement, + ), + ] + + +@MULTIMODAL_REGISTRY.register_processor( + MossTranscribeDiarizeMultiModalProcessor, + info=MossTranscribeDiarizeProcessingInfo, + dummy_inputs=MossTranscribeDiarizeDummyInputsBuilder, +) +class MossTranscribeDiarizeForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsTranscription, +): + supports_transcription = True + supports_transcription_only = True + supports_segment_timestamp = False + supported_languages = ISO639_1_SUPPORTED_LANGS + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "language_model.layers.": "language_model.model.layers.", + "language_model.embed_tokens.": "language_model.model.embed_tokens.", + "language_model.norm.": "language_model.model.norm.", + "model.language_model.model.": "language_model.model.", + "model.language_model.lm_head.": "language_model.lm_head.", + "model.language_model.": "language_model.model.", + "model.whisper_encoder.": "whisper_encoder.", + "model.vq_adaptor.": "vq_adaptor.", + "model.": None, + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + return AUDIO_PLACEHOLDER if modality.startswith("audio") else None + + @classmethod + def get_speech_to_text_config( + cls, + model_config: ModelConfig, + task_type: str, + ) -> SpeechToTextConfig: + processor = cached_processor_from_config(model_config) + return SpeechToTextConfig( + max_audio_clip_s=None, + sample_rate=processor.feature_extractor.sampling_rate, + min_energy_split_window_size=None, + ) + + @classmethod + def get_num_audio_tokens( + cls, + audio_duration_s: float, + stt_config: SpeechToTextConfig, + model_config: ModelConfig, + ) -> int | None: + processor = cached_processor_from_config(model_config) + num_samples = math.ceil(audio_duration_s * stt_config.sample_rate) + return _compute_total_audio_tokens( + num_samples, + processor.feature_extractor, + processor.audio_merge_size, + ) + + @classmethod + def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: + stt_config = stt_params.stt_config + question = stt_params.request_prompt or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT + question = question.strip() or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT + prompt = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + f"<|im_start|>user\n{AUDIO_PLACEHOLDER}\n" + f"{question}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + return TextPrompt( + prompt=prompt, + multi_modal_data={"audio": (stt_params.audio, stt_config.sample_rate)}, + ) + + @classmethod + def post_process_output(cls, text: str) -> str: + return text.strip() + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = vllm_config.model_config.hf_config + self.dtype = vllm_config.model_config.dtype + + with self._mark_tower_model(vllm_config, "audio"): + self.whisper_encoder = MossTranscribeDiarizeWhisperEncoder( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "whisper_encoder"), + ) + self.vq_adaptor = MossTranscribeDiarizeVQAdaptor( + input_dim=int(self.config.adaptor_input_dim), + hidden_size=int(self.config.text_config.hidden_size), + eps=float(self.config.text_config.rms_norm_eps), + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=self.config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["Qwen3ForCausalLM"], + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def _time_merge(self, features: torch.Tensor) -> torch.Tensor: + batch, seq_len, dim = features.shape + merge_size = int(self.config.audio_merge_size) + seq_len_trim = (seq_len // merge_size) * merge_size + return features[:, :seq_len_trim, :].reshape( + batch, + seq_len_trim // merge_size, + dim * merge_size, + ) + + def _parse_and_validate_audio_input( + self, + **kwargs: object, + ) -> MossTranscribeDiarizeInputs | None: + input_features = kwargs.pop("input_features", None) + audio_embeds = kwargs.pop("audio_embeds", None) + audio_feature_lengths = kwargs.pop("audio_feature_lengths", None) + audio_chunk_counts = kwargs.pop("audio_chunk_counts", None) + if input_features is None and audio_embeds is None: + return None + if audio_embeds is not None: + return MossTranscribeDiarizeEmbeddingInputs( + type="audio_embeds", + audio_embeds=_as_audio_embedding_list(audio_embeds), + ) + return MossTranscribeDiarizeAudioInputs( + type="audio_features", + input_features=input_features, + audio_feature_lengths=audio_feature_lengths, + audio_chunk_counts=audio_chunk_counts, + ) + + def _process_audio_input( + self, audio_input: MossTranscribeDiarizeInputs + ) -> list[torch.Tensor]: + if audio_input["type"] == "audio_embeds": + return list(audio_input["audio_embeds"]) + + input_features = audio_input["input_features"] + audio_feature_lengths = audio_input["audio_feature_lengths"] + audio_chunk_counts = audio_input["audio_chunk_counts"] + if input_features is None or audio_feature_lengths is None: + raise ValueError( + "MOSS-Transcribe-Diarize audio inputs require both " + "`input_features` and `audio_feature_lengths`." + ) + if audio_feature_lengths.numel() != input_features.shape[0]: + raise ValueError( + "`audio_feature_lengths` must contain one length per " + "`input_features` chunk: got " + f"{audio_feature_lengths.numel()} lengths for " + f"{input_features.shape[0]} chunks." + ) + if audio_chunk_counts is None: + audio_chunk_counts = audio_feature_lengths.new_tensor( + [input_features.shape[0]], + dtype=torch.long, + ) + else: + audio_chunk_counts = audio_chunk_counts.to(dtype=torch.long) + if audio_chunk_counts.numel() == 0: + raise ValueError("`audio_chunk_counts` must contain at least one item.") + if torch.any(audio_chunk_counts <= 0): + raise ValueError("`audio_chunk_counts` must contain positive counts.") + num_audio_chunks = int(audio_chunk_counts.sum().item()) + if num_audio_chunks != input_features.shape[0]: + raise ValueError( + "`audio_chunk_counts` must sum to the number of input feature chunks: " + f"got {num_audio_chunks} chunks for " + f"{input_features.shape[0]} input chunks." + ) + + features = self.whisper_encoder(input_features, audio_feature_lengths) + merged = self._time_merge(features.to(dtype=self.dtype)) + projected = self.vq_adaptor(merged).squeeze(0) + + audio_chunk_offsets = torch.cumsum(audio_chunk_counts, dim=0) + audio_chunk_offsets = torch.cat( + [audio_chunk_offsets.new_zeros(1), audio_chunk_offsets] + ) + tokens_per_item = [ + int(audio_feature_lengths[start:end].sum().item()) + for start, end in zip( + audio_chunk_offsets[:-1].tolist(), + audio_chunk_offsets[1:].tolist(), + ) + ] + if any(num_tokens <= 0 for num_tokens in tokens_per_item): + raise ValueError("Audio input is too short to produce any tokens.") + return list(projected.split(tokens_per_item, dim=0)) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + audio_input = self._parse_and_validate_audio_input(**kwargs) + if audio_input is None: + return [] + return self._process_audio_input(audio_input) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.language_model.embed_input_ids(input_ids) + if not multimodal_embeddings: + return inputs_embeds + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=_require_is_multimodal(is_multimodal), + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + return self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, + mapper=self.hf_to_vllm_mapper, + ) diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index fe737438c30..bd4b908dc61 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -415,11 +415,7 @@ class NemotronHMTP(nn.Module, SupportsPP): for name, loaded_weight in weights: # Only process MTP weights - skip all non-MTP weights - if ( - not name.startswith("mtp.") - and "embeddings" not in name - and "lm_head" not in name - ): + if not name.startswith("mtp.") and "embeddings" not in name: continue # Skip rotary embeddings (computed, not loaded) if "rotary_emb.inv_freq" in name: diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 6261fc50fac..12bc1430726 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -86,14 +86,15 @@ def _resolve_layer_attention( if layer_types is not None: num_sliding = sum(lt == SLIDING_ATTENTION for lt in layer_types) any_sliding = num_sliding > 0 - all_sliding = num_sliding == len(layer_types) - if any_sliding and not all_sliding: - # Mixed sliding/full attention needs per-layer causal metadata and - # multiple KV-cache groups, which DFlash does not yet support. + # Mixed sliding/full attention needs multiple KV groups (V2 runner only). + if ( + 0 < num_sliding < len(layer_types) + and not get_current_vllm_config().use_v2_model_runner + ): raise NotImplementedError( - "DFlash does not yet support mixed sliding/full attention via " - "layer_types; see " - "https://github.com/vllm-project/vllm/issues/40898." + "DFlash drafters with mixed sliding/full attention require " + "the V2 model runner; relaunch with " + "VLLM_USE_V2_MODEL_RUNNER=1." ) default_causal = False @@ -207,8 +208,7 @@ class DFlashQwen3Attention(nn.Module): attn_type=attn_type, sinks=self.attention_sink_bias, ) - # NOTE: `causal` is currently unused here, but will be needed in the future - # to support models with different causality per-layer. + self.causal = causal self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) @@ -709,6 +709,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): ) -> torch.Tensor: return self.model(input_ids, positions, inputs_embeds) + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.attn.layer_name for layer in self.model.layers] + + def get_draft_attn_causal(self) -> list[bool]: + """Per-layer attention causality, aligned with + get_draft_kv_cache_layer_names.""" + return [layer.self_attn.causal for layer in self.model.layers] + def compute_logits( self, hidden_states: torch.Tensor, diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index a85286164ea..5252afd78e4 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -33,7 +33,7 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F -from transformers import BatchFeature +from transformers import BatchFeature, ProcessorMixin from transformers.models.qwen2_vl import Qwen2VLImageProcessor from transformers.models.qwen2_vl.image_processing_qwen2_vl import ( smart_resize as image_smart_resize, @@ -1191,34 +1191,31 @@ def _replace_video_token_placeholders( target: list[int], replacements: list[list[int]], ) -> list[int]: - """Replace each 3-token video placeholder with its expanded sequence. + """Replace each video placeholder with its expanded sequence. Args: prompt_ids: Token IDs of the original (unexpanded) prompt. - target: 3-element list ``[vision_start_id, video_pad_id, - vision_end_id]`` to search for. + target: Token-ID sequence to search for. New-style + (transformers>=5.10) processors expand only the bare + `[video_pad_id]`; older ones expand the full + `[vision_start_id, video_pad_id, vision_end_id]` triplet. replacements: Per-video expanded token sequences, in prompt order. Returns: - Token IDs with every placeholder triplet replaced. + Token IDs with every placeholder replaced. """ result: list[int] = [] repl_idx = 0 i = 0 n = len(prompt_ids) - t0, t1, t2 = target + k = len(target) num_repl = len(replacements) while i < n: - if ( - i + 2 < n - and prompt_ids[i] == t0 - and prompt_ids[i + 1] == t1 - and prompt_ids[i + 2] == t2 - ): + if prompt_ids[i : i + k] == target: result.extend(replacements[repl_idx]) repl_idx += 1 - i += 3 + i += k else: result.append(prompt_ids[i]) i += 1 @@ -1230,6 +1227,15 @@ def _replace_video_token_placeholders( class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]): + @staticmethod + def _expands_only_video_token(hf_processor: ProcessorMixin) -> bool: + """Transformers>=5.10 processors override `replace_video_token` + to expand only the bare video token, keeping the prompt's outer + `<|vision_start|>`/`<|vision_end|>` markers.""" + mixin_impl = getattr(ProcessorMixin, "replace_video_token", None) + proc_impl = getattr(type(hf_processor), "replace_video_token", None) + return proc_impl is not None and proc_impl is not mixin_impl + def _call_hf_processor( self, prompt: str, @@ -1365,14 +1371,17 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) tok_kwargs=tok_kwargs, ) - # Replace each placeholder triplet with pre-computed video tokens. + # Replace each placeholder with pre-computed video tokens. if video_input_ids_lst: hf_config = self.info.get_hf_config() - video_target = [ - hf_config.vision_start_token_id, - hf_config.video_token_id, - hf_config.vision_end_token_id, - ] + if self._expands_only_video_token(self.info.get_hf_processor()): + video_target = [hf_config.video_token_id] + else: + video_target = [ + hf_config.vision_start_token_id, + hf_config.video_token_id, + hf_config.vision_end_token_id, + ] input_ids = processed_outputs.pop("input_ids") if not isinstance(input_ids, list): input_ids = input_ids.tolist() @@ -1464,17 +1473,24 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) select_token_id=select_token_id, ) + if self._expands_only_video_token(hf_processor): + # transformers>=5.10 expands only the bare video_token + video_target = hf_processor.video_token + else: + # Old-style processors expand the full placeholder + # NOTE: We match string on purpose since searching sequence of + # token ids takes more time. + video_target = "<|vision_start|><|video_pad|><|vision_end|>" + return [ PromptReplacement( modality="image", target=hf_processor.image_token, replacement=get_image_replacement_qwen3vl, ), - # NOTE: We match string on purpose since searching sequence of - # token ids takes more time. PromptReplacement( modality="video", - target="<|vision_start|><|video_pad|><|vision_end|>", + target=video_target, replacement=get_video_replacement_qwen3vl, ), ] diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 44ce9a88173..c9e43ed6355 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -447,6 +447,10 @@ _MULTIMODAL_MODELS = { "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), + "MossTranscribeDiarizeForConditionalGeneration": ( + "moss_transcribe_diarize", + "MossTranscribeDiarizeForConditionalGeneration", + ), "LightOnOCRForConditionalGeneration": ( "lightonocr", "LightOnOCRForConditionalGeneration", @@ -612,6 +616,7 @@ _SPECULATIVE_DECODING_MODELS = { "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"), + "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"), "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index cb224e5cbc0..78a12876e66 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -16,6 +16,10 @@ # limitations under the License. """Wrapper around `transformers` models""" +from typing import TYPE_CHECKING + +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + from vllm.model_executor.models.transformers.base import Base from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin @@ -32,6 +36,36 @@ from vllm.model_executor.models.transformers.pooling import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY +if TYPE_CHECKING: + import torch + + from vllm.model_executor.layers.attention import Attention + + +def vllm_attention_forward( + # Transformers args + module: "torch.nn.Module", + query: "torch.Tensor", + key: "torch.Tensor", + value: "torch.Tensor", + attention_mask: "torch.Tensor", + # Transformers kwargs + scaling: float | None = None, + # vLLM kwargs + attention_instances: dict[int, "Attention"] | None = None, + **kwargs, +): + self_attn = attention_instances[module.layer_idx] + if scaling is not None: + self_attn.impl.scale = float(scaling) + hidden = query.shape[-2] + query, key, value = (x.transpose(1, 2) for x in (query, key, value)) + query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) + return self_attn.forward(query, key, value), None + + +ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward + # Text only models class TransformersForCausalLM(CausalMixin, Base): ... diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index bcda62918f3..510abe82ce7 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -31,7 +31,6 @@ from transformers.conversion_mapping import ( WeightRenaming, get_model_conversion_mapping, ) -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from vllm.compilation.decorators import support_torch_compile from vllm.config.utils import getattr_iter @@ -42,6 +41,7 @@ from vllm.model_executor.layers.attention import ( Attention, EncoderOnlyAttention, ) +from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.models.interfaces import ( SupportsEagle, @@ -51,6 +51,7 @@ from vllm.model_executor.models.interfaces import ( SupportsQuant, ) from vllm.model_executor.models.interfaces_base import VllmModel +from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers from vllm.model_executor.models.transformers.utils import ( can_enable_torch_compile, get_feature_request_tip, @@ -58,7 +59,6 @@ from vllm.model_executor.models.transformers.utils import ( log_replacement, replace_conv_class, replace_linear_class, - replace_rms_norm_class, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -74,37 +74,10 @@ if TYPE_CHECKING: from transformers import PreTrainedModel from vllm.config import VllmConfig -else: - PreTrainedModel = object logger = init_logger(__name__) -def vllm_flash_attention_forward( - # Transformers args - module: torch.nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor, - # Transformers kwargs - scaling: float | None = None, - # vLLM kwargs - attention_instances: dict[int, Attention] | None = None, - **kwargs, -): - self_attn = attention_instances[module.layer_idx] - if scaling is not None: - self_attn.impl.scale = float(scaling) - hidden = query.shape[-2] - query, key, value = (x.transpose(1, 2) for x in (query, key, value)) - query, key, value = (x.reshape(hidden, -1) for x in (query, key, value)) - return self_attn.forward(query, key, value), None - - -ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward - - class Base( nn.Module, VllmModel, @@ -141,6 +114,9 @@ class Base( """Ignore unexpected weights whose qualname starts with these prefixes.""" self.ignore_unexpected_suffixes: list[str] = [] """Ignore unexpected weights whose qualname ends with these suffixes.""" + self.packed_modules_mapping: dict[str, list[str]] = {} + """Fused module -> constituent projections, populated by `recursive_replace` + for the quantization machinery and loaders (e.g. bitsandbytes).""" # Attrs for Eagle3 (see self.set_aux_hidden_state_layers) self._target_class: type[nn.Module] = nn.Module @@ -217,7 +193,7 @@ class Base( self.text_config._attn_implementation = "vllm" self.config.dtype = torch.get_default_dtype() - def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]: + def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]: """ Get the decoder class from the model. @@ -236,7 +212,7 @@ class Base( def _decorate_cls_for_torch_compile( self, - cls: type[PreTrainedModel], + cls: type["PreTrainedModel"], dynamic_arg_dims: dict[str, int] | None, enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, @@ -300,13 +276,13 @@ class Base( - Any quantization config specific mappings """ self.hf_to_vllm_mapper = WeightsMapper() - orig_to_new_renamings = self.hf_to_vllm_mapper.orig_to_new_renamings + orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex for mapping in get_model_conversion_mapping(self.model): # Handle weights which have been renamed in Transformers if isinstance(mapping, WeightRenaming): - orig_to_new_renamings.append(mapping) + orig_to_new_renaming.append(mapping) # TODO: Handle WeightConverter to enable layer merging # Handle unexpected weights which should be ignored @@ -356,12 +332,24 @@ class Base( if self.pp_group.world_size <= 1: return - if not self.model.supports_pp_plan: + if self.model.supports_pp_plan: + module = self.model + names = list(module._pp_plan.keys()) + else: + module = self.model.get_decoder() + has_parameters = lambda m: next(m.parameters(), None) is not None + names = [n for n, c in module.named_children() if has_parameters(c)] tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - raise ValueError( - f"{type(self.model)} does not support pipeline parallel. {tip}" + logger.warning( + "%s does not define a pipeline parallel plan. The Transformers " + "modeling backend will infer the split from the layers of %s in order " + "of declaration and keep parameter-free modules on every rank. This " + "may fail if the model's structure is non-standard. %s", + type(self.model), + type(module), + tip, ) def attrsetter(attr: str) -> Callable[[object, object], None]: @@ -376,10 +364,9 @@ class Base( module_lists = [] module_list_idx = None - pp_plan = list(self.model._pp_plan.keys()) - for i, name in enumerate(pp_plan): + for i, name in enumerate(names): # attrgetter in case the module is nested (e.g. "text_model.layers") - if isinstance(attrgetter(name)(self.model), nn.ModuleList): + if isinstance(attrgetter(name)(module), nn.ModuleList): module_lists.append(name) module_list_idx = i @@ -389,16 +376,16 @@ class Base( "in the base model are not supported yet!" ) if module_list_idx is None: - raise ValueError(f"Could not find `ModuleList` in {type(self.model)}") + raise ValueError(f"Could not find `ModuleList` in {type(module)}") # Layers before module list - for name in pp_plan[:module_list_idx]: + for name in names[:module_list_idx]: if self.pp_group.is_first_rank or ( self._get_tie_word_embeddings() and self.pp_group.is_last_rank ): continue # attrsetter in case the module is nested (e.g. "text_model.embed_tokens") - attrsetter(name)(self.model, PPMissingLayer()) + attrsetter(name)(module, PPMissingLayer()) # Module list start_layer, end_layer = get_pp_indices( @@ -406,41 +393,61 @@ class Base( self.pp_group.rank_in_group, self.pp_group.world_size, ) - layers_name = pp_plan[module_list_idx] + layers_name = names[module_list_idx] # attrgetter in case the module is nested (e.g. "text_model.layers") - layers = attrgetter(layers_name)(self.model) + layers = attrgetter(layers_name)(module) for i in range(len(layers)): if start_layer <= i and i < end_layer: continue layers[i] = PPMissingLayer() # Layers after module list - for name in pp_plan[module_list_idx + 1 :]: + for name in names[module_list_idx + 1 :]: # Modules that should be on last rank if not self.pp_group.is_last_rank: # attrsetter in case the module is nested (e.g. "text_model.norm") - attrsetter(name)(self.model, PPMissingLayer()) + attrsetter(name)(module, PPMissingLayer()) def recursive_replace(self): """Recursively replace modules in the model as needed. Currently, this replaces: + - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul` + - Attention QKV projections with a fused `QKVParallelLinear` + split - `nn.Linear` with vLLM's tensor parallel linear classes - - `*RMSNorm` with vLLM's `RMSNorm` + - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d` + - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm` """ - tp_plan = self.model.tp_plan + tp_plan = self.model.tp_plan or {} if not tp_plan and self.tp_group.world_size > 1: tip = get_feature_request_tip( self.model_config.model, self.model_config.trust_remote_code ) - raise ValueError( - f"{type(self.model)} does not support tensor parallel. {tip}" + logger.warning_once( + "%s does not define a tensor parallel plan. The Transformers modeling " + "backend will shard the model the best it can during graph fusion and " + "replicate the rest. This may be suboptimal or fail if the model does " + "not fuse cleanly. %s", + type(self.model), + tip, ) # Prefix the patterns because we always start from `self.model` tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()} + # Detect fusable patterns once per module class (cached, so this is cheap) + fusers = Fusers(self.model, self.model_config) + + def register_fusion(fuser: BaseFuser, prefix: str): + """Register a fused layer's mappings just before it is built.""" + orig_to_new_stacked = fuser.orig_to_new_stacked(prefix) + self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked) + + packed_modules_mapping = fuser.packed_modules_mapping + self.packed_modules_mapping.update(packed_modules_mapping) + if self.quant_config is not None: + self.quant_config.packed_modules_mapping.update(packed_modules_mapping) def _recursive_replace(module: nn.Module, prefix: str): for child_name, child_module in module.named_children(): @@ -479,11 +486,16 @@ class Base( ) elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)): new_module = replace_conv_class(child_module) - elif child_module.__class__.__name__.endswith("RMSNorm"): - new_module = replace_rms_norm_class( - child_module, self.text_config.hidden_size + elif (fuser := fusers[child_module]) is not None: + register_fusion(fuser, qual_name) + new_module = fuser.fuse( + child_module, qual_name, self.model_config, self.quant_config ) - else: + logger.info_once(fuser.info(child_name)) + _recursive_replace(new_module, prefix=qual_name) + elif not isinstance(child_module, MoERunner): + # MoERunner can contain aliases of shared experts and gates, + # so we don't want to recurse into it and break weight loading. _recursive_replace(child_module, prefix=qual_name) if new_module is not child_module: @@ -538,7 +550,7 @@ class Base( num_heads=num_heads, head_size=head_size, # NOTE: We use Llama scale as default, if it's set by - # Transformers, it's updated in vllm_flash_attention_forward + # Transformers, it's updated in vllm_attention_forward scale=head_size**-0.5, num_kv_heads=num_kv_heads, cache_config=self.cache_config, diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py new file mode 100644 index 00000000000..0aa7c419ceb --- /dev/null +++ b/vllm/model_executor/models/transformers/fuser.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fuser detection for the Transformers modeling backend. + +`get_fuser` traces a module class once (see `fx_utils`) and matches it against +each concrete fuser in `fusers`; `Fusers` caches the result per class for a +whole model. `base.recursive_replace` then applies the matched fuser per +instance. RMSNorm-shaped modules the tracer cannot match are warned about. +""" + +from collections import UserDict +from typing import TYPE_CHECKING + +from cachetools import cached +from torch import nn + +from vllm.logger import init_logger +from vllm.model_executor.models.transformers.fusers import ( + BaseFuser, + GLUFuser, + QKVFuser, + RMSNormFuser, + StackedFuser, +) +from vllm.model_executor.models.transformers.fx_utils import trace + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + +logger = init_logger(__name__) + + +@cached(cache={}, key=type) +def get_fuser(module: nn.Module) -> BaseFuser | None: + """The fuser for `type(module)` (cached per class), or `None` if no match.""" + # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a + # leaf module (raw tensor math, no submodules). Nothing else can match, and + # tracing is skipped for it. + n_linear = sum(isinstance(c, nn.Linear) for c in module.children()) + is_leaf = next(module.children(), None) is None + if n_linear < 2 and not is_leaf: + return None + if (graph := trace(module)) is None: + return None + for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser): + if (fuser := fuser_cls.match(graph, module)) is not None: + if isinstance(fuser, StackedFuser): + try: + fuser.update_forward(module) + except Exception as exc: + # An unrecognised source just means we cannot fuse here. + logger.debug( + "Could not rewrite %s for fusion: %s", type(module), exc + ) + return None + return fuser + # A norm we could not match structurally is left unfused; flag likely misses. + if module.__class__.__name__.endswith("RMSNorm"): + logger.warning_once( + "%s looks like an RMSNorm but its computation did not match the " + "expected pattern, so it was left unfused.", + module.__class__.__name__, + ) + return None + + +class Fusers(UserDict): + """Mapping from module class to fuser, for all fusable classes in a model.""" + + def __init__(self, model: nn.Module, model_config: "ModelConfig"): + self.model_config = model_config + super().__init__({type(m): get_fuser(m) for m in model.modules()}) + + def __getitem__(self, m: nn.Module) -> BaseFuser | None: + fuser = self.data.get(type(m)) + if fuser is not None and fuser.validate(m, self.model_config): + return fuser + return None diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py new file mode 100644 index 00000000000..58910b0ecc3 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Concrete fusers for the Transformers modeling backend.""" + +from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser +from vllm.model_executor.models.transformers.fusers.glu import GLUFuser +from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser +from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser +from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser + +__all__ = [ + "BaseFuser", + "StackedFuser", + "GLUFuser", + "MoEBlockFuser", + "QKVFuser", + "RMSNormFuser", +] diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py new file mode 100644 index 00000000000..54fb2d09165 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base classes for the Transformers backend fusers.""" + +import types +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn + +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + + +@dataclass +class BaseFuser(ABC): + """A detected fusion and how to apply it. + + `match` analyses the module *class* once (cached, see `get_fuser`); `fuse` + then applies the fusion to an instance in `recursive_replace`, returning the + module to install in its place. + """ + + @abstractmethod + def info(self, name: str) -> str: + """A human-readable description of the fusion at `name`, for logging.""" + + @classmethod + @abstractmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None": + """Match the pattern in `graph`, returning a fuser if found.""" + + @abstractmethod + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + """Whether this fuser can be applied to this `module` instance.""" + + @abstractmethod + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Apply the fusion to an already-validated `module`, returning the + module to install in its place (mutated in place, or freshly built).""" + + def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: + """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes + (none unless it stacks weights).""" + return {} + + @property + def packed_modules_mapping(self) -> dict[str, list[str]]: + """`packed_modules_mapping` entries this fuser contributes (none unless + it stacks weights).""" + return {} + + +@dataclass +class StackedFuser(BaseFuser): + """A fuser that merges sibling projections into one stacked linear and + rewrites the forward to call it. + + `match` and `update_forward` analyse the class once; `fuse` builds the merged + submodule and binds the compiled forward on an instance in place, so it keeps + its class and any attribute the fusion does not consume. + """ + + merged_name: ClassVar[str] + """Attribute name of the merged module created by `update_attrs`.""" + merged_cls: ClassVar[str] + """Name of the vLLM class the merged projection becomes (for logging).""" + + source_cls: str + """Class of the HF module the fused projections belonged to (for logging).""" + + fused_forward: Callable = field(init=False, repr=False) + """The compiled rewritten forward, set by `update_forward`.""" + + def info(self, name: str) -> str: + sources = " + ".join(shard for shard, _ in self.shards) + return ( + f"Fused: {sources} ({name}: {self.source_cls}) -> " + f"{self.merged_name} ({self.merged_cls})" + ) + + @property + @abstractmethod + def shards(self) -> list[tuple[str, ShardId]]: + """Each projection's original name and its shard id in the merged module. + + Source for both `orig_to_new_stacked` and `packed_modules_mapping`.""" + + def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: + """`WeightsMapper.orig_to_new_stacked` entries for one fused instance. + + Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname + so only this exact layer is remapped, never a same-named projection + elsewhere (e.g. an unfused MoE expert's `gate_proj`).""" + merged = maybe_prefix(prefix, self.merged_name) + return { + maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards + } + + @property + def packed_modules_mapping(self) -> dict[str, list[str]]: + """`{merged_name: [projection names]}` so quantization can unpack the + fused layer into its per-shard configs.""" + return {self.merged_name: [name for name, _ in self.shards]} + + @abstractmethod + def update_forward(self, module: nn.Module) -> None: + """Rewrite and compile `type(module)`'s forward source. + + Raises if the source does not admit the rewrite (fusion is then skipped). + """ + + @abstractmethod + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + """Replace `module`'s submodules with the merged module.""" + + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). + + Builds the merged submodule and binds the compiled forward.""" + self.update_attrs(module, prefix, model_config, quant_config) + module.forward = types.MethodType(self.fused_forward, module) + return module diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py new file mode 100644 index 00000000000..951eb35777a --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/glu.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear.""" + +import ast +import operator +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn +from transformers.activations import ACT2CLS + +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import ( + _ACTIVATION_AND_MUL_REGISTRY, + get_act_and_mul_fn, +) +from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + find_node, + is_linear, + peel, + recover_forward, + replace_expr, + single_self_call, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + +logger = init_logger(__name__) + + +CLS2ACT: dict[type, list[str]] = {} +for _act_name, _act_cls in ACT2CLS.items(): + if isinstance(_act_cls, tuple): + _act_cls = _act_cls[0] + CLS2ACT.setdefault(_act_cls, []).append(_act_name) + +ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys()) + + +@dataclass +class GLUFuser(StackedFuser): + """Fuser for the GLU pattern `act(gate(x)) * up(x)`.""" + + act_name: str + gate_name: str + up_name: str + down_name: str | None + merged_name: ClassVar[str] = "gate_up_proj" + merged_cls: ClassVar[str] = "MergedColumnParallelLinear" + + @property + def shards(self) -> list[tuple[str, ShardId]]: + return [(self.gate_name, 0), (self.up_name, 1)] + + @classmethod + def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool: + """Is node `act(gate(x))` where `gate` is linear and `act` is not linear.""" + return ( + node.op == "call_module" + and not is_linear(node, module) + and len(node.args) == 1 + and isinstance(node.args[0], fx.Node) + and is_linear(node.args[0], module) + ) + + @classmethod + def _get_glu_nodes( + cls, graph: fx.Graph, module: nn.Module + ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None: + """Search graph for the GLU pattern `act(gate(x)) * up(x)`.""" + for mul in graph.nodes: + if ( + mul.op == "call_function" + and mul.target == operator.mul + and len(mul.args) == 2 + and all(isinstance(arg, fx.Node) for arg in mul.args) + ): + a, b = mul.args + if cls._is_act_of_gate(a, module) and is_linear(b, module): + act, gate, up = a, a.args[0], b + elif cls._is_act_of_gate(b, module) and is_linear(a, module): + act, gate, up = b, b.args[0], a + else: + continue + if ( + all(len(args) == 1 for args in (gate.args, up.args)) + and isinstance(x := gate.args[0], fx.Node) + and x is up.args[0] + ): + return act, gate, up, mul + return None + + @staticmethod + def _get_act_and_mul_name(act: nn.Module) -> str | None: + """Get the name of `act` if it has an `...AndMul` equivalent.""" + for name in CLS2ACT.get(type(act), []): + if name in ACT_AND_MUL_NAMES: + return name + # nn.GELU is not in ACT2CLS, but could be in model code + if type(act) is nn.GELU: + return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu" + return None + + @classmethod + def _get_act_and_mul(cls, act: nn.Module) -> nn.Module: + """Get the `...AndMul` equivalent of a Transformers activation module.""" + if name := cls._get_act_and_mul_name(act): + return get_act_and_mul_fn(name) + raise ValueError(f"No AndMul equivalent for {type(act)}") + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None": + if (glu_nodes := cls._get_glu_nodes(graph, module)) is None: + return None + act_node, gate_node, up_node, mul_node = glu_nodes + + gate = module.get_submodule(gate_node.target) + up = module.get_submodule(up_node.target) + # Shapes must be compatible for a single merged GEMM. + if gate.in_features == up.in_features and (gate.bias is None) == ( + up.bias is None + ): + predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node + down_node = find_node(graph, predicate) + return cls( + source_cls=type(module).__name__, + act_name=act_node.target, + gate_name=gate_node.target, + up_name=up_node.target, + down_name=down_node.target if down_node is not None else None, + ) + return None + + def update_forward(self, module: nn.Module) -> None: + """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source.""" + funcdef, fn = recover_forward(type(module)) + act_call = single_self_call(funcdef, self.act_name) + gate_call = single_self_call(funcdef, self.gate_name) + up_call = single_self_call(funcdef, self.up_name) + if act_call.args[0] is not gate_call: + raise ValueError("activation does not directly wrap the gate") + if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]): + raise ValueError("gate and up inputs are written differently") + muls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.BinOp) + and isinstance(node.op, ast.Mult) + and {id(node.left), id(node.right)} == {id(act_call), id(up_call)} + ] + if len(muls) != 1: + raise ValueError("no multiply of the activation and up projection") + + # act(gate(x)) * up(x) -> act(gate_up(x)) + assert isinstance(gate_call.func, ast.Attribute) + gate_call.func.attr = self.merged_name + replace_expr(funcdef, muls[0], act_call) + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + act = module.get_submodule(self.act_name) + if self._get_act_and_mul_name(act) is None: + logger.debug("No AndMul equivalent for %s; skipping fusion", type(act)) + return False + return True + + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + act_fn = self._get_act_and_mul(module.get_submodule(self.act_name)) + gate = module.get_submodule(self.gate_name) + up = module.get_submodule(self.up_name) + merged = MergedColumnParallelLinear( + input_size=gate.in_features, + output_sizes=[gate.out_features, up.out_features], + bias=gate.bias is not None, + quant_config=quant_config, + prefix=maybe_prefix(prefix, self.merged_name), + return_bias=False, + ) + logger.debug( + "%s: %s, %s: %s -> %s: %s", + self.gate_name, + gate, + self.up_name, + up, + self.merged_name, + merged, + ) + setattr(module, self.merged_name, merged) + setattr(module, self.act_name, act_fn) + # Drop the consumed submodules so their (meta) params are not expected. + delattr(module, self.gate_name) + delattr(module, self.up_name) + # If there is a down projection, we know it must be rowwise. + if self.down_name is not None: + down_prefix = maybe_prefix(prefix, self.down_name) + down = module.get_submodule(self.down_name) + new_down = replace_linear_class( + down, "rowwise", quant_config, prefix=down_prefix + ) + setattr(module, self.down_name, new_down) + log_replacement(down_prefix, down, new_down) diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py new file mode 100644 index 00000000000..6a3c7e85d6e --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing.""" + +import ast +import inspect +import textwrap +import types +from collections.abc import Iterator +from dataclasses import dataclass +from itertools import chain + +import torch +from torch import fx, nn + +from vllm.distributed import tensor_model_parallel_all_gather +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.transformers.fx_utils import ( + find_node, + is_op, + peel, + trace, +) +from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk + + +def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: + """`module`'s own state (i.e. named parameters and buffers).""" + return chain(module.named_parameters(), module.named_buffers()) + + +def _own_returns(node: ast.AST) -> Iterator[ast.Return]: + """`return` statements in `node`'s own scope, not in nested functions.""" + stack = list(ast.iter_child_nodes(node)) + while stack: + child = stack.pop() + if isinstance(child, ast.Return): + yield child + elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + stack.extend(ast.iter_child_nodes(child)) + + +def _returns_tuple(cls: type[nn.Module]) -> bool: + """Does `cls.forward()` return a tuple?""" + try: + source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward))) + forward = ast.parse(source).body[0] + except (OSError, SyntaxError, TypeError, IndexError): + return True + # Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`. + tuple_names = { + target.id + for node in ast.walk(forward) + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple) + for target in node.targets + if isinstance(target, ast.Name) + } + + def yields_tuple(value: ast.expr | None) -> bool: + if isinstance(value, ast.Tuple): + return True + if isinstance(value, ast.Name): + return value.id in tuple_names + if isinstance(value, ast.IfExp): + return yields_tuple(value.body) or yields_tuple(value.orelse) + return False + + return any(yields_tuple(ret.value) for ret in _own_returns(forward)) + + +def _is_scalar_gate(module: nn.Module) -> bool: + """A linear projecting to a single logit (the shared-expert sigmoid gate).""" + weight = getattr(module, "weight", None) + return ( + isinstance(module, nn.Linear) + and weight is not None + and weight.ndim == 2 + and weight.shape[0] == 1 + ) + + +def _reaches(node: fx.Node, key: str) -> set[fx.Node]: + """Returns the set of nodes reachable from `node` by following `key` edges.""" + seen: set[fx.Node] = set() + stack = [node] + while stack: + n = stack.pop() + if n in seen: + continue + seen.add(n) + stack.extend(getattr(n, key)) + return seen + + +class SharedExpertMLP(nn.Module): + """Wraps an HF shared expert, applying the output gating it is paired with.""" + + def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None): + super().__init__() + self.shared_experts = shared_experts + self.gate = gate + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + out = self.shared_experts(hidden_states) + if self.gate is not None: + out = torch.sigmoid(self.gate(hidden_states)[0]) * out + return out + + +def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + """Standard MoE block forward. + + Routing and any shared experts are handled inside `self.experts: MoERunner`.""" + orig_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, orig_shape[-1]) + num_tokens = hidden_states.shape[0] + is_sequence_parallel = self.experts.moe_config.is_sequence_parallel + if is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + out = self.experts(hidden_states, router_logits=hidden_states) + if is_sequence_parallel: + out = tensor_model_parallel_all_gather(out, 0)[:num_tokens] + return out.reshape(orig_shape) + + +@dataclass +class MoEBlockFuser: + """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional).""" + + gate_name: str + scoring_func: str + shared_name: str | None + shared_gate_name: str | None + + @staticmethod + def _match_router(gate: nn.Module) -> str | None: + """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`.""" + if [name for name, _ in named_state(gate)] != ["weight"]: + return None + graph = trace(gate) + if graph is None: + return None + topk = find_node(graph, lambda n: is_op(n, "topk")) + if topk is None: + return None + # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. + scorers = [ + n + for n in _reaches(topk, "all_input_nodes") + if is_op(n, "softmax") or is_op(n, "sigmoid") + ] + if len(scorers) != 1: + return None + scorer = scorers[0] + if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")): + return None + return "softmax" if is_op(scorer, "softmax") else "sigmoid" + + @staticmethod + def _match_shared_experts( + graph: fx.Graph, experts: str + ) -> tuple[str | None, str | None]: + """Detects the shared expert and its optional gate by dataflow.""" + experts_predicate = lambda n: n.op == "call_module" and n.target == experts + if (experts_node := find_node(graph, experts_predicate)) is None: + return None, None + from_experts = _reaches(experts_node, "users") + for add in graph.nodes: + if not is_op(add, "add"): + continue + operands = [a for a in add.args if isinstance(a, fx.Node)] + # Exactly one side is the experts' output; the other is the shared path. + sides = [a in from_experts for a in operands] + if len(operands) != 2 or sides.count(True) != 1: + continue + cone = _reaches(operands[sides.index(False)], "all_input_nodes") + modules = [n for n in cone if n.op == "call_module" and n.target != experts] + # A sigmoid wrapping one of those modules marks the shared-expert gate. + gate = next( + ( + src + for n in cone + if is_op(n, "sigmoid") + and isinstance(src := peel(n.args[0]), fx.Node) + and src in modules + ), + None, + ) + shared = [n for n in modules if n is not gate] + if len(shared) != 1: + return None, None + return shared[0].target, (gate.target if gate is not None else None) + return None, None + + @classmethod + def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None": + # Standard MoE block returns a single tensor. + if _returns_tuple(type(moe_block)): + return None + # Router: the child that scores + top-k selects. + gate_name = scoring_func = None + for name, child in moe_block.named_children(): + if name != experts_name and (func := cls._match_router(child)) is not None: + gate_name, scoring_func = name, func + break + if gate_name is None or scoring_func is None: + return None + # Shared expert: a child the block adds to the experts' output. + shared_name = shared_gate_name = None + others = [ + n + for n, _ in moe_block.named_children() + if n not in {experts_name, gate_name} + ] + if others: + graph = trace(moe_block) + if graph is None: + return None + shared_name, shared_gate_name = cls._match_shared_experts( + graph, experts_name + ) + if shared_gate_name is not None and not _is_scalar_gate( + getattr(moe_block, shared_gate_name) + ): + return None + # Fail closed: `rewrite_forward` runs only the experts and the detected + # shared expert, so any other stateful child would be dropped. + accounted = {experts_name, gate_name, shared_name, shared_gate_name} + for name, child in moe_block.named_children(): + if name not in accounted and next(named_state(child), None) is not None: + return None + return cls(gate_name, scoring_func, shared_name, shared_gate_name) + + def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear: + """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE.""" + num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape + gate = ReplicatedLinear( + hidden_size, + num_experts, + bias=False, + prefix=maybe_prefix(prefix, self.gate_name), + ) + setattr(moe_block, self.gate_name, gate) + return gate + + def shared_experts( + self, moe_block: nn.Module, prefix: str + ) -> SharedExpertMLP | None: + """Build the HF shared expert (and its optional gate) + as a `SharedExpertMLP` for vLLM's fused MoE.""" + if self.shared_name is None: + return None + shared_experts = getattr(moe_block, self.shared_name) + gate = None + if self.shared_gate_name is not None: + hf_gate = getattr(moe_block, self.shared_gate_name) + gate = ReplicatedLinear( + hf_gate.in_features, + hf_gate.out_features, + bias=hf_gate.bias is not None, + prefix=maybe_prefix(prefix, self.shared_gate_name), + ) + setattr(moe_block, self.shared_gate_name, gate) + return SharedExpertMLP(shared_experts, gate) + + def rewrite_forward(self, moe_block: nn.Module) -> None: + """Rewrite `moe_block.forward` to route through vLLM's fused MoE.""" + moe_block.forward = types.MethodType(_moe_block_forward, moe_block) diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py new file mode 100644 index 00000000000..010a0018acc --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split.""" + +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +from torch import fx, nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fx_utils import ( + compile_forward, + innermost_block, + is_linear, + recover_forward, + replace_expr, + single_self_call, +) +from vllm.model_executor.models.transformers.utils import ( + log_replacement, + replace_linear_class, +) +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + +logger = init_logger(__name__) + + +@dataclass +class QKVFuser(StackedFuser): + """Fuser for the attention QKV pattern `q(x), k(x), v(x)`.""" + + q_name: str + k_name: str + v_name: str + o_name: str | None + merged_name: ClassVar[str] = "qkv_proj" + merged_cls: ClassVar[str] = "QKVParallelLinear" + + @property + def shards(self) -> list[tuple[str, ShardId]]: + return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")] + + @classmethod + def _get_qkv_nodes( + cls, graph: fx.Graph, module: nn.Module + ) -> tuple[fx.Node, fx.Node, fx.Node] | None: + """Search `graph` for the QKV pattern `q(x), k(x), v(x)`.""" + by_input: dict[fx.Node, list[fx.Node]] = {} + for node in graph.nodes: + if ( + is_linear(node, module) + and len(node.args) == 1 + and not node.kwargs + and isinstance(node.args[0], fx.Node) + and node.args[0].op == "placeholder" + ): + by_input.setdefault(node.args[0], []).append(node) + triples = [nodes for nodes in by_input.values() if len(nodes) == 3] + if len(triples) != 1: + return None + + q_node, k_node, v_node = nodes = triples[0] + outs = [module.get_submodule(node.target).out_features for node in nodes] + if len(set(outs)) == 2: + # q is identified as the larger projection (GQA) + (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1) + k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2) + if module.get_submodule(q_node.target).out_features != max(outs): + return None + elif len(set(outs)) != 1: + return None + return q_node, k_node, v_node + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None": + if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None: + return None + q, k, v = qkv_nodes + names = dict(q_name=q.target, k_name=k.target, v_name=v.target) + attn_width = module.get_submodule(q.target).out_features + candidates = [ + name + for name, child in module.named_children() + if isinstance(child, nn.Linear) + and name not in names.values() + and child.in_features == attn_width + ] + names["o_name"] = candidates[0] if len(candidates) == 1 else None + return cls(source_cls=type(module).__name__, **names) + + def update_forward(self, module: nn.Module) -> None: + """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source.""" + funcdef, fn = recover_forward(type(module)) + calls = [ + single_self_call(funcdef, name) + for name in (self.q_name, self.k_name, self.v_name) + ] + arg_dumps = {ast.dump(call.args[0]) for call in calls} + if len(arg_dumps) != 1: + raise ValueError("projection inputs are written differently") + # The trace may be partial, so prove projection exclusivity in source: + # no other linear child may consume the same input (else the matched + # three may not be q, k and v) + other_linears = { + name + for name, child in module.named_children() + if isinstance(child, nn.Linear) + } - {self.q_name, self.k_name, self.v_name} + for node in ast.walk(funcdef): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in other_linears + and any(ast.dump(arg) in arg_dumps for arg in node.args) + ): + raise ValueError("another linear consumes the same input") + blocks = [innermost_block(funcdef.body, call) for call in calls] + if any(found is None for found in blocks): + raise ValueError("projection calls not found in the function body") + if len({id(block) for block, _ in blocks}) != 1: + raise ValueError("projection calls are in different blocks") + + # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1) + names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)} + temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)] + if names & set(temps): + raise ValueError("fused temporaries would shadow existing names") + merged = f"self.{self.merged_name}" + sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]" + template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)" + assign = ast.parse(template).body[0] + arg = next( + node + for node in ast.walk(assign) + if isinstance(node, ast.Name) and node.id == "__arg__" + ) + replace_expr(assign, arg, calls[0].args[0]) + block, index = blocks[0] + ast.copy_location(assign, block[index]) + block.insert(min(index for _, index in blocks), assign) + for call, temp in zip(calls, temps): + replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load())) + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + """Shapes must be compatible for a single merged, head-sharded GEMM.""" + q = module.get_submodule(self.q_name) + k = module.get_submodule(self.k_name) + v = module.get_submodule(self.v_name) + head_size = model_config.get_head_size() + compatible = ( + q.in_features == k.in_features == v.in_features + and len({proj.bias is None for proj in (q, k, v)}) == 1 + and k.out_features == v.out_features + and q.out_features % head_size == 0 + and k.out_features % head_size == 0 + ) + if not compatible: + logger.debug("%s is not compatible with QKV fusion", type(module)) + return compatible + + def update_attrs( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> None: + head_size = model_config.get_head_size() + q = module.get_submodule(self.q_name) + k = module.get_submodule(self.k_name) + merged = QKVParallelLinear( + hidden_size=q.in_features, + head_size=head_size, + total_num_heads=q.out_features // head_size, + total_num_kv_heads=k.out_features // head_size, + bias=q.bias is not None, + quant_config=quant_config, + prefix=maybe_prefix(prefix, self.merged_name), + return_bias=False, + ) + logger.debug( + "%s: %s, %s: %s, %s: %s -> %s: %s", + self.q_name, + q, + self.k_name, + k, + self.v_name, + module.get_submodule(self.v_name), + self.merged_name, + merged, + ) + setattr(module, self.merged_name, merged) + # Drop the consumed submodules so their (meta) params are not expected. + for name in (self.q_name, self.k_name, self.v_name): + delattr(module, name) + # If there is an output projection, we know it must be rowwise. + if self.o_name is not None: + o_proj_prefix = maybe_prefix(prefix, self.o_name) + o_proj = module.get_submodule(self.o_name) + new_o = replace_linear_class( + o_proj, "rowwise", quant_config, prefix=o_proj_prefix + ) + setattr(module, self.o_name, new_o) + log_replacement(o_proj_prefix, o_proj, new_o) diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py new file mode 100644 index 00000000000..829fd9a541c --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm.""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import fx, nn + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.distributed.parallel_state import model_parallel_is_initialized +from vllm.distributed.utils import split_tensor_along_last_dim +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm +from vllm.model_executor.models.transformers.fusers.base import BaseFuser +from vllm.model_executor.models.transformers.fx_utils import ( + find_node, + forward_input_count, + is_op, + peel, + trace, +) + +if TYPE_CHECKING: + from vllm.config.model import ModelConfig + from vllm.model_executor.layers.quantization import QuantizationConfig + + +def _is_squared(node: object, x: fx.Node) -> bool: + """`x**2`, `x.square()` or `x * x`, through any dtype casts.""" + node = peel(node) + if is_op(node, "pow"): + base, exp = node.args + return peel(base) is x and exp == 2 + if is_op(node, "square"): + return peel(node.args[0]) is x + if is_op(node, "mul"): + a, b = node.args + return peel(a) is x and peel(b) is x + return False + + +def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None: + """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape.""" + add = peel(rsqrt.args[0]) + if not is_op(add, "add"): + return None + consts = [a for a in add.args if isinstance(a, (int, float))] + nodes = [a for a in add.args if isinstance(a, fx.Node)] + if len(consts) != 1 or len(nodes) != 1: + return None + mean = peel(nodes[0]) + if not is_op(mean, "mean"): + return None + if not _is_squared(mean.args[0], x): + return None + return float(consts[0]) + + +def _is_one_plus(node: object) -> bool: + """`1 + weight` in either operand order (marks a zero-centered weight).""" + node = peel(node) + if not is_op(node, "add"): + return False + return any(isinstance(a, (int, float)) and a == 1 for a in node.args) + + +def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool: + """Does the forward compute anything after `node` before returning?""" + output = find_node(graph, lambda n: n.op == "output") + if output is None or not output.args: + return False + return peel(output.args[0]) is not node + + +class TPAwareNormMixin(nn.Module): + """Mixin for RMSNorms that reconstructs a TP-sharded input before normalizing.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if model_parallel_is_initialized(): + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + else: + self.tp_size, self.tp_rank = 1, 0 + + def forward( + self, x: torch.Tensor, residual: torch.Tensor | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if self.tp_size > 1 and x.shape[-1] < (full := self.weight.shape[0]): + if x.shape[-1] * self.tp_size != full: + raise ValueError( + f"Cannot gather norm of width {full}: a TP-sharded input of " + f"width {x.shape[-1]} does not tile it evenly across " + f"{self.tp_size} ranks (replicated or uneven sharding)." + ) + x = tensor_model_parallel_all_gather(x.contiguous()) + x = super().forward(x) + splits = split_tensor_along_last_dim(x, num_partitions=self.tp_size) + return splits[self.tp_rank] + return super().forward(x, residual) + + +class TPAwareRMSNorm(TPAwareNormMixin, RMSNorm): + """`RMSNorm` that reconstructs a TP-sharded input before normalizing.""" + + +class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm): + """`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing.""" + + +@dataclass +class RMSNormFuser(BaseFuser): + """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights.""" + + zero_centered: bool + """Gemma-style `(1 + weight)` scaling (weight initialised at zero).""" + source_cls: str + """Class name of the norm this was matched from (for logging).""" + + def info(self, name: str) -> str: + norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm" + return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)" + + @classmethod + def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None": + """Match a graph to the RMSNorm pattern, returning a fuser if found.""" + if forward_input_count(type(module)) != 1: + return None + x = find_node(graph, lambda n: n.op == "placeholder") + if x is None: + return None + # Handle native torch `rms_norm` op. + rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm")) + if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x: + if _has_trailing_compute(graph, rms_norm): + return None + return cls(zero_centered=False, source_cls=type(module).__name__) + # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern. + # The rsqrt over the mean-square variance is the spine of the norm. + rsqrt = None + for node in graph.nodes: + if is_op(node, "rsqrt") and _variance_eps(node, x) is not None: + rsqrt = node + break + if rsqrt is None: + return None + # The `x * rsqrt(...)` normalize multiply. + normalize = find_node( + graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args) + ) + if normalize is None: + return None + # An optional later `weight * normalized` (or `(1 + weight) * normalized`). + tail, zero_centered = normalize, False + for node in graph.nodes: + if not is_op(node, "mul") or node is normalize: + continue + operands = [peel(a) for a in node.args if isinstance(a, fx.Node)] + if len(operands) == 2 and normalize in operands: + weight = next(o for o in operands if o is not normalize) + tail, zero_centered = node, _is_one_plus(weight) + break + # The norm must be the last compute in forward, or it is not a pure norm. + if _has_trailing_compute(graph, tail): + return None + return cls(zero_centered=zero_centered, source_cls=type(module).__name__) + + @staticmethod + def _eps_from_graph(graph: fx.Graph) -> float | None: + """Extract the `eps` constant from the graph, if present.""" + if (x := find_node(graph, lambda n: n.op == "placeholder")) is None: + return None + fused = find_node(graph, lambda n: is_op(n, "rms_norm")) + if fused is not None and fused.args and peel(fused.args[0]) is x: + args, kwargs = fused.args, fused.kwargs + eps = args[3] if len(args) > 3 else kwargs.get("eps") + return eps if isinstance(eps, (int, float)) else None + for node in graph.nodes: + if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None: + return eps + return None + + def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool: + return True + + def fuse( + self, + module: nn.Module, + prefix: str, + model_config: "ModelConfig", + quant_config: "QuantizationConfig", + ) -> nn.Module: + """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" + weight = getattr(module, "weight", None) + hidden_size = ( + weight.size(0) if weight is not None else model_config.get_hidden_size() + ) + graph = trace(module) + eps = self._eps_from_graph(graph) if graph is not None else None + if eps is None: + # If eps not in graph, match torch behaviour. + dtype = weight.dtype if weight is not None else model_config.dtype + eps = torch.finfo(dtype).eps + if self.zero_centered: + return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps) + has_weight = weight is not None + return TPAwareRMSNorm( + hidden_size=hidden_size, + eps=eps, + has_weight=has_weight, + dtype=weight.dtype if has_weight else None, + ) diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py new file mode 100644 index 00000000000..0e043941d8a --- /dev/null +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""fx tracing and forward-source rewriting for the Transformers backend fusers. + +A small engine, independent of any particular pattern: trace a module's forward +with `torch.fx` (tolerating a partial graph), inspect the resulting nodes, and +rewrite the forward's *source* (AST) so only matched calls change while the rest +stays live Python. `fusion.py` builds the concrete fusion patterns on top. +""" + +import ast +import inspect +import operator +import textwrap +from collections.abc import Callable + +import torch +from torch import fx, nn +from torch.nn import functional as F + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _infer_len(node: fx.Node) -> int | None: + """Concrete length of a proxy's value, inferred from its node chain. + + Lets tracing pass through the shape unpacks and `*`-splats (e.g. + `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention. + """ + # `x.shape` has the rank of `x`, when known + if ( + node.op == "call_function" + and node.target is getattr + and node.args[1] == "shape" + and (rank := _rank(node.args[0])) is not None + ): + return rank + # Slices of known-length values + if node.op == "call_function" and node.target is operator.getitem: + src_len = _infer_len(node.args[0]) + index = node.args[1] + if src_len is not None and isinstance(index, slice): + return len(range(*index.indices(src_len))) + return None + + +def _rank(node: fx.Node) -> int | None: + """The tensor rank of `node`'s value, if known.""" + # vLLM always feeds the model [1, seq_len, hidden_size] hidden states + if node.op == "placeholder" and node.target == "hidden_states": + return 3 + return None + + +class _SizedProxy(fx.Proxy): + """Proxy whose `len` is inferred from the graph (see `_infer_len`).""" + + def __len__(self) -> int: + length = _infer_len(self.node) + if length is None: + return super().__len__() + return length + + +class _AllLeafTracer(fx.Tracer): + """Tracer that treats every submodule as a leaf. + + Each child stays one `call_module` node, so matching sees the module's own + forward structure (activations aren't decomposed into e.g. `sigmoid * x`). + `iter` traces through the leading shape unpacks (see `_infer_len`); anything + else untraceable ends the trace early and the partial graph is matched. + """ + + def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: + return True + + def proxy(self, node: fx.Node) -> fx.Proxy: + return _SizedProxy(node, self) + + def iter(self, obj: fx.Proxy): + length = _infer_len(obj.node) + if length is None: + return super().iter(obj) + return iter([obj[i] for i in range(length)]) + + +def trace(module: nn.Module) -> fx.Graph | None: + """Trace `module.forward`, returning the partial graph on failure. + + The graph is only evidence for matching, and the patterns sit at the top of + their forwards, so a trace that fails partway can still be matched.""" + tracer = _AllLeafTracer() + try: + return tracer.trace(module) + except Exception as exc: + logger.debug("Could not fully trace %s: %s", type(module), exc) + return getattr(tracer, "graph", None) + + +def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]: + """Parse the source of `cls.forward`, ready for rewriting.""" + fn = inspect.unwrap(cls.forward) + if fn.__code__.co_freevars: + raise ValueError("forward is a closure") + tree = ast.parse(textwrap.dedent(inspect.getsource(fn))) + funcdef = tree.body[0] + if not isinstance(funcdef, ast.FunctionDef): + raise ValueError("source is not a plain function definition") + # `fn` is already unwrapped; don't re-apply its decorators + funcdef.decorator_list.clear() + # Annotations may not evaluate outside the defining module (e.g. with + # postponed evaluation); they're not needed at runtime + funcdef.returns = None + args = funcdef.args + for arg in ( + *args.posonlyargs, + *args.args, + *args.kwonlyargs, + *filter(None, (args.vararg, args.kwarg)), + ): + arg.annotation = None + # Recompiling outside the class body would break name mangling + for node in ast.walk(funcdef): + name = getattr(node, "attr", None) or getattr(node, "id", None) + if name and name.startswith("__") and not name.endswith("__"): + raise ValueError(f"{name} would be name mangled") + return funcdef, fn + + +def forward_input_count(cls: type[nn.Module]) -> int: + """The number of tensor inputs `cls.forward` declares, excluding `self` and + any `*args`/`**kwargs`. Read from the signature, so it is independent of + whether the trace completes (unlike counting placeholders).""" + try: + params = list(inspect.signature(cls.forward).parameters.values())[1:] + except (ValueError, TypeError): + return 1 # uninspectable: assume a single input and let matching decide + fixed = ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + return sum(1 for p in params if p.kind in fixed) + + +def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable: + """Compile `funcdef` in `fn`'s module so tracebacks point at the source.""" + module = ast.Module(body=[funcdef], type_ignores=[]) + ast.fix_missing_locations(module) + ast.increment_lineno(module, fn.__code__.co_firstlineno - 1) + code = compile(module, fn.__code__.co_filename, "exec") + namespace: dict = {} + exec(code, fn.__globals__, namespace) + return namespace[funcdef.name] + + +def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call: + """The unique `self.(arg)` call in `funcdef`. + + Raises unless `name` appears exactly once, as such a call, so the source + rewrite agrees with the fx match. + """ + uses = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Attribute) and node.attr == name + ] + if len(uses) != 1: + raise ValueError(f"{name} is referenced {len(uses)} times") + calls = [ + node + for node in ast.walk(funcdef) + if isinstance(node, ast.Call) + and node.func is uses[0] + and len(node.args) == 1 + and not isinstance(node.args[0], ast.Starred) + and not node.keywords + ] + if ( + len(calls) != 1 + or not isinstance(uses[0].value, ast.Name) + or uses[0].value.id != "self" + ): + raise ValueError(f"{name} is not a single-argument call on self") + return calls[0] + + +def innermost_block( + block: list[ast.stmt], node: ast.AST +) -> tuple[list[ast.stmt], int] | None: + """The innermost statement list containing `node`, and the index within.""" + for index, stmt in enumerate(block): + if not any(child is node for child in ast.walk(stmt)): + continue + child_blocks = [ + getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody") + ] + child_blocks += [h.body for h in getattr(stmt, "handlers", [])] + child_blocks += [c.body for c in getattr(stmt, "cases", [])] + for child_block in child_blocks: + if ( + isinstance(child_block, list) + and child_block + and (found := innermost_block(child_block, node)) is not None + ): + return found + return block, index + return None + + +def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None: + """Replace the expression `old` (by identity) with `new` within `module`.""" + + class _Replacer(ast.NodeTransformer): + def visit(self, node: ast.AST) -> ast.AST: + if node is old: + return new + return super().generic_visit(node) + + _Replacer().visit(module) + + +def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: + """The first node in `graph` matching `predicate`, or `None`.""" + return next((n for n in graph.nodes if predicate(n)), None) + + +def is_linear(node: fx.Node, module: nn.Module) -> bool: + """Is node `nn.Linear.__call__()`.""" + return node.op == "call_module" and isinstance( + module.get_submodule(node.target), nn.Linear + ) + + +_DTYPE_CASTS = frozenset({"to", "float", "double", "half", "bfloat16", "type_as"}) + + +def peel(node: object) -> object: + """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`).""" + while ( + isinstance(node, fx.Node) + and node.op == "call_method" + and node.target in _DTYPE_CASTS + ): + node = node.args[0] + return node + + +def is_fn(node: object, target: Callable) -> bool: + """Is node `()`.""" + return ( + isinstance(node, fx.Node) + and node.op == "call_function" + and node.target is target + ) + + +def is_method(node: object, name: str) -> bool: + """Is node `.()`.""" + return ( + isinstance(node, fx.Node) and node.op == "call_method" and node.target == name + ) + + +def is_op(node: object, name: str) -> bool: + """ + Is node `torch.()`, `F.()`, `operator.()`, or `Tensor.()`. + """ + return any( + is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) + ) or (hasattr(torch.Tensor, name) and is_method(node, name)) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index d5267a26179..d1f5dba0373 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -26,9 +26,11 @@ import torch.nn as nn from vllm.config.utils import getattr_iter from vllm.distributed import get_dp_group, get_ep_group from vllm.forward_context import ForwardContext, get_forward_context +from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import maybe_prefix from vllm.utils.torch_utils import direct_register_custom_op @@ -37,6 +39,8 @@ from .utils import log_replacement if TYPE_CHECKING: from vllm.config import VllmConfig +logger = init_logger(__name__) + @dataclass class TransformersMoEState: @@ -142,14 +146,15 @@ class MoEMixin(MixtureOfExperts): self.num_physical_experts = num_physical_experts self.num_local_physical_experts = num_local_physical_experts self.num_redundant_experts = num_physical_experts - self.num_logical_experts - for mlp in self.mlp_layers: - mlp.n_local_physical_experts = num_local_physical_experts - mlp.n_physical_experts = num_physical_experts - mlp.n_redundant_experts = self.num_redundant_experts - mlp.experts.update_expert_map() + for moe_block in self.mlp_layers: + moe_block.n_local_physical_experts = num_local_physical_experts + moe_block.n_physical_experts = num_physical_experts + moe_block.n_redundant_experts = self.num_redundant_experts + moe_block.experts.update_expert_map() def recursive_replace(self): """Initialize the MoE layers.""" + experts_name = "experts" text_config = self.text_config # Positional arguments @@ -217,10 +222,13 @@ class MoEMixin(MixtureOfExperts): # down_proj = (num_experts, intermediate_size, hidden_size) params = list(child_module.parameters()) is_3d = len(params) > 0 and all(p.ndim == 3 for p in params) - if child_name == "experts" and (is_modulelist or is_3d): + if child_name == experts_name and (is_modulelist or is_3d): # Alias for readability - mlp = module + moe_block = module experts = child_module + # Class of the fused block (parent of gate/experts/shared) + moe_block_cls = type(moe_block).__name__ + experts_cls = type(experts).__name__ # Do the experts have biases has_bias = False for experts_param_name, _ in experts.named_parameters(): @@ -230,64 +238,93 @@ class MoEMixin(MixtureOfExperts): # If the config does not specify num_shared_experts, but # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: - for mlp_param_name, _ in mlp.named_parameters(): - if "shared_expert" in mlp_param_name: + for moe_block_param_name, _ in moe_block.named_parameters(): + if "shared_expert" in moe_block_param_name: self.num_shared_experts = 1 break - # Replace experts module with FusedMoE - moe_state = TransformersMoEState() - - def custom_routing_function( - hidden_states: torch.Tensor, - gating_output: torch.Tensor, - topk: int, - renormalize: bool, - moe_state: TransformersMoEState, - ): - """Return `topk_weights` from `gating_output` and the - `topk_ids` we stored in the layer earlier.""" - topk_weights = gating_output - topk_ids = moe_state.topk_ids - assert topk_ids is not None - # Handle all gather in expert parallel - if topk_ids.size(0) != hidden_states.size(0): - dp_metadata = get_forward_context().dp_metadata - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - is_sp = moe_state.is_sequence_parallel - dist_group = get_ep_group() if is_sp else get_dp_group() - assert sizes[dist_group.rank_in_group] == topk_ids.shape[0] - (topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes) - return topk_weights, topk_ids - - fused_experts = FusedMoE( + kwargs: dict[str, Any] = dict( num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, intermediate_size=intermediate_size, renormalize=renormalize, - # Hard coded because topk happens in Transformers use_grouped_topk=False, - num_expert_group=num_expert_group, - topk_group=topk_group, quant_config=self.quant_config, prefix=qual_name, activation=activation, enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, has_bias=has_bias, - custom_routing_function=partial( - custom_routing_function, - moe_state=moe_state, - ), - runner_cls=TransformersMoERunner, routed_experts_cls=TransformersRoutedExperts, - runner_args={"moe_state": moe_state}, ) - mlp.experts = fused_experts + fuser = MoEBlockFuser.match(moe_block, experts_name) + if self.num_expert_groups <= 1 and fuser is not None: + # MoE block forward is fully replaced. + # gate/router and shared expert (if any) runs in FusedMoE. + kwargs |= dict( + scoring_func=fuser.scoring_func, + is_sequence_parallel=( + self.parallel_config.use_sequence_parallel_moe + ), + gate=fuser.gate(moe_block, prefix), + shared_experts=fuser.shared_experts(moe_block, prefix), + ) + fuser.rewrite_forward(moe_block) + routed = "gate + experts" + if fuser.shared_name: + routed += " + shared experts" + logger.info_once( + "Fused: %s (%s) -> FusedMoE (internal routing)", + routed, + moe_block_cls, + ) + else: + # MoE block forward is unmodified. + # gate/router and shared expert (if any) runs in Transformers. + # We then smuggle the topk_ids in using a custom op. + moe_state = TransformersMoEState() + + def custom_routing_function( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + moe_state: TransformersMoEState, + ): + """Return `topk_weights` from `gating_output` and the + `topk_ids` we stored in the layer earlier.""" + topk_weights = gating_output + topk_ids = moe_state.topk_ids + assert topk_ids is not None + # Handle all gather in expert parallel + if topk_ids.size(0) != hidden_states.size(0): + dp_metadata = get_forward_context().dp_metadata + sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + is_sp = moe_state.is_sequence_parallel + group = get_ep_group() if is_sp else get_dp_group() + assert sizes[group.rank_in_group] == topk_ids.shape[0] + (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes) + return topk_weights, topk_ids + + kwargs |= dict( + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=partial( + custom_routing_function, moe_state=moe_state + ), + runner_cls=TransformersMoERunner, + runner_args={"moe_state": moe_state}, + ) + logger.info_once( + "Fused: experts (%s) -> FusedMoE (external routing)", + experts_cls, + ) + fused_experts = FusedMoE(**kwargs) + moe_block.experts = fused_experts log_replacement(qual_name, experts, fused_experts) # Update MixtureOfExperts mixin state - self.mlp_layers.append(mlp) + self.mlp_layers.append(moe_block) self.moe_layers.append(fused_experts) else: _recursive_replace(child_module, prefix=qual_name) diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 0a4ca94c5e9..4d9b01ce393 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -23,10 +23,8 @@ from typing import TYPE_CHECKING, Literal import torch from torch import nn -from vllm.config.utils import getattr_iter from vllm.logger import init_logger from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer -from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, ReplicatedLinear, @@ -183,45 +181,6 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: ) -def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm: - """Replace a Transformers RMSNorm with vLLM's RMSNorm. - - This method assumes: - - Weight is stored as `weight`. - - Epsilon is stored as `eps` or `variance_epsilon`. - - `with_scale` indicates whether the layer has a weight (Gemma3n only). - - `var_hidden_size` is only ever used for Intern vision encoder in vLLM - and Transformers doesn't appear to have the same concept. - """ - eps = getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6) - kwargs = {"hidden_size": hidden_size, "eps": eps} - # Update hidden size if weight is available - weight_meta = getattr(rms_norm, "weight", None) - if weight_meta is not None: - kwargs["hidden_size"] = weight_meta.size(0) - # Check if weight is all zeros, which indicates GemmaRMSNorm - # We must create a new instance because rms_norm is on meta - try: - with torch.device("cpu"): - weight_test = getattr(rms_norm.__class__(1), "weight", None) - except Exception: - logger.warning( - "Failed to determine if RMSNorm weight is centered on zero or one. " - "Defaulting to one." - ) - weight_test = None - if weight_test is not None and torch.all(weight_test == 0): - return GemmaRMSNorm(**kwargs) - # Otherwise assume it's a regular RMSNorm - kwargs["has_weight"] = getattr(rms_norm, "with_scale", True) - if weight_meta is not None: - kwargs["dtype"] = weight_meta.dtype - else: - # No weight, fall back to weightless RMSNorm - kwargs["has_weight"] = False - return RMSNorm(**kwargs) - - def recursive_replace_linear( model: nn.Module, quant_config: "QuantizationConfig | None", diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 6f4524400c0..86dc72de76c 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -11,7 +11,6 @@ import regex as re import torch import torch.nn as nn from torch.nn.modules.module import register_module_module_registration_hook -from transformers import PretrainedConfig from vllm.config import VllmConfig from vllm.distributed import ( @@ -33,6 +32,9 @@ from vllm.utils.torch_utils import ( ) if TYPE_CHECKING: + from transformers import PretrainedConfig + from transformers.conversion_mapping import WeightRenaming + from vllm.model_executor.layers.quantization import QuantizationConfig logger = init_logger(__name__) @@ -46,7 +48,7 @@ class WeightsMapper: If a key maps to a value of `None`, the corresponding weight is ignored.""" - orig_to_new_renamings: list[Any] = field(default_factory=list) + orig_to_new_renaming: list["WeightRenaming"] = field(default_factory=list) orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) orig_to_new_stacked: Mapping[str, tuple[str, ShardId]] = field(default_factory=dict) @@ -56,9 +58,9 @@ class WeightsMapper: def __or__(self, other: "WeightsMapper") -> "WeightsMapper": """Combine two `WeightsMapper`s by merging their mappings.""" return WeightsMapper( - orig_to_new_renamings=[ - *self.orig_to_new_renamings, - *other.orig_to_new_renamings, + orig_to_new_renaming=[ + *self.orig_to_new_renaming, + *other.orig_to_new_renaming, ], orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex}, orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr}, @@ -92,7 +94,7 @@ class WeightsMapper: "k_scale to v_scale" ) - for renaming in self.orig_to_new_renamings: + for renaming in self.orig_to_new_renaming: key, _ = renaming.rename_source_key(key) for pattern, new_key in self.orig_to_new_regex.items(): @@ -426,7 +428,7 @@ def init_vllm_registered_model( vllm_config: VllmConfig, *, prefix: str = "", - hf_config: PretrainedConfig | None = None, + hf_config: "PretrainedConfig | None" = None, architectures: list[str] | None = None, ) -> nn.Module: """ diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 8c59532e3b6..b54a4dccba4 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -17,7 +17,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.config.speech_to_text import SpeechToTextParams from vllm.engine.protocol import StreamingInput -from vllm.envs import VLLM_ENGINE_ITERATION_TIMEOUT_S from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsRealtime @@ -265,13 +264,11 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim await buffer.append_audio(right_pad.audio_array) await buffer.append_audio(None) # signal end - # Feed output tokens back into buffer in background + # Feed output tokens back into the buffer. Idle waits are normal here; + # request cleanup still cancels this task through the finally block. async def feed_tokens(): while True: - all_outputs = await asyncio.wait_for( - input_stream.get(), - timeout=VLLM_ENGINE_ITERATION_TIMEOUT_S, - ) + all_outputs = await input_stream.get() await buffer.append_tokens(all_outputs[-1:]) audio_task = asyncio.create_task(feed_audio()) diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 2c860632650..466d8c13ce7 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -388,7 +388,7 @@ class _ModuleOffloader: # Event to signal when H2D copy to static buffer is complete. # Used for per-layer synchronization (both eager and capture modes). - self._copy_done_event = torch.Event() + self._copy_done_event = torch.cuda.Event() # Track whether _copy_done_event is valid for eager-mode wait_event. # False when: (1) never recorded, or (2) last recorded during a @@ -518,7 +518,7 @@ class _ModuleOffloader: # Fork: record event on compute stream, copy_stream waits on it # This joins copy_stream to any active CUDA graph capture - fork_event = torch.Event() + fork_event = torch.cuda.Event() torch.cuda.current_stream().record_event(fork_event) self.copy_stream.wait_event(fork_event) diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index d41604fc7a6..cfff491ab2b 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -11,6 +11,9 @@ from tqdm import tqdm import vllm.envs as envs from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank +from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( + DeepGemmFp8BlockScaledMMKernel, +) from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( compute_aligned_M_and_alignment, @@ -147,6 +150,12 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool: ): return False + if not isinstance( + getattr(module.quant_method, "fp8_linear", None), + DeepGemmFp8BlockScaledMMKernel, + ): + return False + w, _, block_sizes = _extract_data_from_linear_base_module(module) return ( block_sizes == get_mk_alignment_for_contiguous_layout() diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index e31a14db663..b7f3c265704 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -148,15 +148,8 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: use_persistent_cache = True - deepep_a2a_backends = { - "deepep_high_throughput", - "deepep_low_latency", - "deepep_v2", - } - if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends: - # DeepEP dispatch/combine can timeout when only rank 0 - # performs autotune and falls behind other ranks. - # Thus we skip persistent cache in this case. + # When distributed, tune on every rank so the collectives stay synchronized. + if get_world_group().world_size > 1: use_persistent_cache = False if not use_persistent_cache: diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index b6ed0aa4d4f..de1a985c6f7 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -206,8 +206,6 @@ def _warm_zero_kv_blocks_kernel( N_SEGS=config.n_segs, PAGE_SIZE_EL=config.page_size_el, BLOCK_SIZE=config.block_size, - num_warps=4, - num_stages=3, ) @@ -372,11 +370,10 @@ def qwen_triton_warmup( logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type) zero_config = _zero_kv_warmup_config(runner) - if _warm_zero_kv_blocks_with_runner_zeroer(runner): - pass - elif zero_config is not None: + warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner) + if zero_config is not None: _warm_zero_kv_blocks_kernel(device, zero_config) - else: + elif not warmed_zeroer: logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.") _warm_compute_slot_mapping_kernel(device) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index b139b8ba23f..22ced9fa94d 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -72,16 +72,17 @@ class DeepseekV32DecoderLayer(torch.nn.Module): and layer_idx >= config.first_k_dense_replace and layer_idx % moe_layer_freq == 0 ): + # Defer the MoE cross-rank all-reduce; it is fused into the next + # layer's input_layernorm (or the final norm) via + # fused_allreduce_rms_norm. self.mlp = DeepseekV2MoE( config=config, parallel_config=parallel_config, quant_config=quant_config, + reduce_results=False, prefix=f"{prefix}.mlp", + apply_routed_scale_to_output=False, ) - # Defer the MoE cross-rank all-reduce; it is fused into the next - # layer's input_layernorm (or the final norm) via - # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner. - self.mlp.experts.moe_config.skip_final_all_reduce = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 118f27459bb..060dfbf7bdd 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -89,14 +89,22 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module): hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) - # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the + # mtp_block's MoE output is left un-reduced (reduce_results=False); the # main model fuses that all-reduce into the next norm, but here the # recycle hidden is consumed directly, so reduce it now. hidden_states = tensor_model_parallel_all_reduce(hidden_states) - # Return the pre-final-norm recycle hidden (re-fed as the next spec - # step's previous_hidden_states); shared_head norm is applied in - # compute_logits. Matches the V2-runner / deepseek_v4 MTP contract. - return residual + hidden_states + # Recycle the POST-final-norm hidden into the next draft step. The + # residual-add is fused into the final RMSNorm so it is computed + # exactly once, and the result is returned for both tuple positions: + # the draft-logits hidden (compute_logits applies the LM head only) and + # the recycled previous_hidden_states. Recycling the pre-final-norm + # hidden mismatches the draft model's hnorm and lowers MTP acceptance; + # post-norm recycle matches deepseek_mtp.py (PR #45895). The tuple form + # is understood by both the V2 speculator (isinstance-tuple check) and + # the legacy proposer (model_returns_tuple is True for the + # DeepSeekMTPModel architecture). + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + return hidden_states, hidden_states class DeepseekV32MultiTokenPredictor(nn.Module): @@ -168,9 +176,10 @@ class DeepseekV32MultiTokenPredictor(nn.Module): ) -> torch.Tensor: current_step_idx = spec_step_idx % self.num_mtp_layers mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] - return self.logits_processor( - mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states) - ) + # hidden_states is already post-final-norm (produced in the layer + # forward and recycled as-is); apply the LM head only, without a + # second RMSNorm. + return self.logits_processor(mtp_layer.shared_head.head, hidden_states) class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 519f5f9a144..5628a6d0d72 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -56,7 +56,11 @@ from vllm.v1.attention.backends.mla.indexer import ( get_max_prefill_buffer_size, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache -from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + get_kv_quant_mode, +) logger = init_logger(__name__) @@ -272,7 +276,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins # before post-GEMM starts. - self.ln_events = [torch.Event() for _ in range(4)] + self.ln_events = [torch.cuda.Event() for _ in range(4)] assert cache_config is not None, "DeepseekV4 attention requires cache_config" # ---- Attention / KV-cache setup ---- @@ -614,8 +618,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, compress_ratio=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else None, + alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) @@ -643,15 +648,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, compress_ratio=self.compress_ratio, - # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with - # the indexer's compressor state cache. V3.2 keeps the legacy layout. - alignment=576, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, ) def forward(self): ... @@ -760,7 +765,10 @@ class DeepseekV4Indexer(nn.Module): # None on ROCm — maybe_execute_in_parallel falls back to sequential. self.aux_stream = aux_stream - self.ln_events: list[torch.Event] = [torch.Event(), torch.Event()] + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] def forward( self, diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index ffaec528aa8..55106d06af9 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -654,6 +654,8 @@ def build_flashinfer_mixed_sparse_indices( topk: int, decode_compressed_indices_are_local: bool = False, decode_is_valid_token: torch.Tensor | None = None, + swa_block_span: int | None = None, + compressed_block_span: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -730,6 +732,13 @@ def build_flashinfer_mixed_sparse_indices( max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 + # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV. + swa_span = swa_block_size if swa_block_span is None else swa_block_span + compressed_span = ( + compressed_block_size + if compressed_block_span is None + else compressed_block_span + ) _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)]( sparse_indices, sparse_indices.stride(0), @@ -748,9 +757,11 @@ def build_flashinfer_mixed_sparse_indices( swa_block_table, swa_block_table.stride(0), swa_block_size, + swa_span, compressed_block_table, compressed_block_table.stride(0), compressed_block_size, + compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, COMPRESS_RATIO=compress_ratio, @@ -767,6 +778,18 @@ def build_flashinfer_mixed_sparse_indices( return sparse_indices, sparse_topk_lens +@triton.jit +def _remap_flashinfer_index(values, block_size, block_span): + # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so + # packed pages (#44577) need block*block_size+off -> block*block_span+off. + # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed. + is_valid = values >= 0 + safe_values = tl.where(is_valid, values, 0) + values = (safe_values // block_size) * block_span + values += safe_values % block_size + return tl.where(is_valid, values, -1) + + @triton.jit( do_not_specialize=[ "sparse_indices_stride", @@ -775,8 +798,10 @@ def build_flashinfer_mixed_sparse_indices( "prefill_topk_stride", "swa_block_table_stride", "swa_block_size", + "swa_block_span", "compressed_block_table_stride", "compressed_block_size", + "compressed_block_span", "NUM_DECODE_TOKENS", "PREFILL_TOPK_STRIDE", ] @@ -799,9 +824,11 @@ def _build_flashinfer_mixed_sparse_indices_kernel( swa_block_table_ptr, swa_block_table_stride, swa_block_size, + swa_block_span, compressed_block_table_ptr, compressed_block_table_stride, compressed_block_size, + compressed_block_span, NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, COMPRESS_RATIO: tl.constexpr, @@ -825,6 +852,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( mask=mask, other=-1, ) + values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride + offset, values, @@ -858,6 +886,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( values = block_numbers * compressed_block_size + block_offsets values = tl.where(is_valid, values, -1) compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0) + values = _remap_flashinfer_index( + values, compressed_block_size, compressed_block_span + ) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride @@ -904,6 +935,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( block_offsets = pos_offset % swa_block_size slot_ids = block_numbers * swa_block_size + block_offsets slot_ids = tl.where(offset < swa_len, slot_ids, -1) + slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride + offset, slot_ids, @@ -930,6 +962,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( block_offsets = local_idx % compressed_block_size slot_ids = block_numbers * compressed_block_size + block_offsets slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1) + slot_ids = _remap_flashinfer_index( + slot_ids, compressed_block_size, compressed_block_span + ) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 1efa987fe7b..24838c237ce 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -104,12 +104,9 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CompressorMetadata: - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu - num_reqs = common_attn_metadata.num_reqs - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + token_to_req_indices = common_attn_metadata.token_to_req_indices( + self.token_to_req_indices + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, @@ -165,7 +162,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase): head_size=self.state_dim, dtype=self.dtype, sliding_window=self.sliding_window, - alignment=576 if uses_fp8_ds_mla_layout else None, + alignment=576 if uses_fp8_ds_mla_layout else 512, ) def forward(self): ... diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index be4a87b323f..e4e258372a2 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -269,6 +269,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): # load_dspark_model always aliases the target's. has_own_embed_tokens = False has_own_lm_head = False + # Full-vocab draft: draft ids are target ids, no remapping needed. + draft_id_to_target_id = None def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index f35fa03252d..1848c1930db 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -45,6 +45,21 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: return workspace +def _packed_block_span(pool: torch.Tensor) -> int: + """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): == + block_size for unpacked KV, larger when packed (#44577). Raises if not + token-aligned.""" + block_stride = pool.stride(0) + token_stride = pool.stride(-2) + if block_stride % token_stride != 0: + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride " + f"({block_stride}) to be a multiple of the per-token stride " + f"({token_stride}); this layout is not supported yet." + ) + return block_stride // token_stride + + class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend): """FlashInfer backend using the DSv4 sparse metadata/cache layout. @@ -368,6 +383,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): ) cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None) if cached_sparse is None: + swa_block_span = _packed_block_span(swa_k_cache) + compressed_block_span = _packed_block_span(compressed_kv_cache) sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices( decode_swa_indices, decode_compressed_indices, @@ -385,6 +402,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): top_k, decode_compressed_indices_are_local=decode_compressed_indices_are_local, decode_is_valid_token=decode_is_valid_token, + swa_block_span=swa_block_span, + compressed_block_span=compressed_block_span, ) if cache_key != "c4a": swa_metadata.flashinfer_sparse_index_cache[cache_key] = ( diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 136a96a45da..1aaf3f1a141 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -5,7 +5,6 @@ from dataclasses import dataclass from typing import Any, ClassVar -import numpy as np import torch from vllm.config import VllmConfig @@ -13,7 +12,6 @@ from vllm.config.cache import CacheDType from vllm.platforms.interface import DeviceCapability from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -203,18 +201,7 @@ class DeepseekV4FlashMLAMetadataBuilder( fast_build: bool = False, ) -> DeepseekV4FlashMLAMetadata: cm = common_attn_metadata - num_tokens = cm.num_actual_tokens - starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32) - seg_lengths = np.diff(starts) - req_id_per_token = np.repeat( - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths - ) - # Zero-fill for cudagraphs - self.req_id_per_token_buffer.fill_(0) - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( - np_to_pinned_tensor(req_id_per_token), non_blocking=True - ) - req_id_per_token = self.req_id_per_token_buffer[:num_tokens] + req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer) slot_mapping = cm.slot_mapping if self.compress_ratio > 1: diff --git a/vllm/models/deepseek_v4/xpu/xpu_sparse.py b/vllm/models/deepseek_v4/xpu/xpu_sparse.py index 74d27d7bc41..77cc35cf492 100644 --- a/vllm/models/deepseek_v4/xpu/xpu_sparse.py +++ b/vllm/models/deepseek_v4/xpu/xpu_sparse.py @@ -44,6 +44,17 @@ class DeepseekV4XPUAttention(DeepseekV4Attention): backend_cls = DeepseekV4XPUSparseBackend use_flashmla_fp8_layout = True + def __init__(self, *args, **kwargs) -> None: + # torch.cuda.Event() raises RuntimeError on XPU ("dummy base class"). + # The Base and DeepseekV4Indexer both create cuda Events in __init__, so + # we temporarily redirect torch.cuda.Event → torch.xpu.Event. + _orig_event = torch.cuda.Event + torch.cuda.Event = torch.xpu.Event # type: ignore[misc] + try: + super().__init__(*args, **kwargs) + finally: + torch.cuda.Event = _orig_event # type: ignore[misc] + def _fused_qnorm_rope_kv_insert(self, q, kv, positions, attn_metadata): from typing import cast diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index bb1ed619320..c66e7ce5267 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -256,6 +256,13 @@ class MiniMaxM3IndexerMetadataBuilder( dtype=torch.int32, device=device, ) + # Stable per-token causal page-count buffer for decode cudagraph replays + # (consumed by the MSA top-k path's sparse_topk_select num_valid_pages). + self.num_valid_pages_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + dtype=torch.int32, + device=device, + ) class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): diff --git a/vllm/models/minimax_m3/common/ops/__init__.py b/vllm/models/minimax_m3/common/ops/__init__.py index b3a7c2d9f6e..1fb6fa22255 100644 --- a/vllm/models/minimax_m3/common/ops/__init__.py +++ b/vllm/models/minimax_m3/common/ops/__init__.py @@ -4,6 +4,7 @@ from .index_topk import ( minimax_m3_index_decode, + minimax_m3_index_decode_score, minimax_m3_index_score, minimax_m3_index_topk, ) @@ -11,6 +12,7 @@ from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode __all__ = [ "minimax_m3_index_decode", + "minimax_m3_index_decode_score", "minimax_m3_index_score", "minimax_m3_index_topk", "minimax_m3_sparse_attn", diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 28becf7bfec..ed677fd871e 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -757,25 +757,26 @@ def minimax_m3_index_topk( @torch.no_grad() -def minimax_m3_index_decode( +def minimax_m3_index_decode_score( idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] block_table: torch.Tensor, # [num_reqs, max_blocks] seq_lens: torch.Tensor, # [num_reqs] int32 max_seq_len: int, - topk: int, init_blocks: int, local_blocks: int, num_kv_heads: int, decode_query_len: int, max_decode_query_len: int, - out: torch.Tensor | None = None, + score_out: torch.Tensor | None = None, ) -> torch.Tensor: - """Decode index block-score + top-k, both split-K (cudagraph-safe). + """Decode index block-score (split-K, cudagraph-safe); no top-k. - Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). - When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into - ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. + Returns score [num_kv_heads, total_q, >=max_block] (fp32; init/local blocks + forced to 1e30/1e29). When ``score_out`` is given the scores are written into + it (read/written by strides, so a transposed view of a unified buffer is + accepted) instead of a fresh tensor -- used to share a unified score buffer + with the prefill side and run a single top-k over both. """ total_q, num_idx_heads, head_dim = idx_q.shape assert num_idx_heads == num_kv_heads, ( @@ -783,7 +784,6 @@ def minimax_m3_index_decode( ) assert decode_query_len <= max_decode_query_len assert total_q == seq_lens.shape[0] * decode_query_len - batch = total_q max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) use_pdl = current_platform.is_arch_support_pdl() # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA @@ -800,13 +800,16 @@ def minimax_m3_index_decode( if num_idx_heads > 1 and max_decode_query_len > 1: score_kwargs.update({"num_warps": 4, "num_stages": 2}) - # Keep score strides 16-divisible to avoid Triton recompiles. - score_block_stride = round_up(max_block, 16) - score = torch.empty( - (num_idx_heads, total_q, score_block_stride), - dtype=torch.float32, - device=idx_q.device, - ) + if score_out is not None: + score = score_out + else: + # Keep score strides 16-divisible to avoid Triton recompiles. + score_block_stride = round_up(max_block, 16) + score = torch.empty( + (num_idx_heads, total_q, score_block_stride), + dtype=torch.float32, + device=idx_q.device, + ) # split-K over seq blocks; chunk count depends only on shape constants so # the grid is fixed within a cuda graph. TARGET_GRID = 512 @@ -848,6 +851,55 @@ def minimax_m3_index_decode( USE_PDL=use_pdl, **score_kwargs, ) + return score + + +@torch.no_grad() +def minimax_m3_index_decode( + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + max_seq_len: int, + topk: int, + init_blocks: int, + local_blocks: int, + num_kv_heads: int, + decode_query_len: int, + max_decode_query_len: int, + out: torch.Tensor | None = None, + score_out: torch.Tensor | None = None, +) -> torch.Tensor: + """Decode index block-score + top-k, both split-K (cudagraph-safe). + + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. + When ``score_out`` ([num_kv_heads, total_q, >=max_block]) is given, the block + scores are written into it (read back by the top-k) instead of a fresh + tensor -- used to share a unified score buffer with the prefill side. Reads + via strides, so a transposed view of a block-major buffer is accepted. + """ + total_q, num_idx_heads, _ = idx_q.shape + batch = total_q + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + use_pdl = current_platform.is_arch_support_pdl() + pdl_kwargs: dict[str, bool | int] = {} + if use_pdl: + pdl_kwargs.update({"launch_pdl": True}) + score = minimax_m3_index_decode_score( + idx_q, + index_kv_cache, + block_table, + seq_lens, + max_seq_len, + init_blocks, + local_blocks, + num_kv_heads, + decode_query_len, + max_decode_query_len, + score_out=score_out, + ) if out is not None: topk_idx = out[:, :total_q, :] diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index f04652c89a3..08d375dc610 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -43,7 +43,6 @@ _FP8_DTYPES = ( { "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]), "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] * triton.next_power_of_2(args["gqa_group_size"]), } @@ -84,7 +83,6 @@ def _gqa_sparse_fwd_kernel( BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) BLOCK_SIZE_D: tl.constexpr, BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_QH: tl.constexpr, USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load ): @@ -109,9 +107,10 @@ def _gqa_sparse_fwd_kernel( for j in range(real_q_loop): pid_q_j = pid_q * num_q_loop + j t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th - off_t = tl.arange(0, BLOCK_SIZE_T) - topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1) - real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0) + # Valid block count from seq position (no sentinel): block_size_q == 1. + q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q + valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, valid_blocks) q_ptrs = tl.make_block_ptr( base=q_ptr + q_start * stride_qn + pid_h * stride_qh, shape=(q_len, gqa_group_size, head_dim), @@ -202,7 +201,6 @@ def _gqa_sparse_fwd_kernel( 16, triton.next_power_of_2(args["gqa_group_size"]) ), "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]), } ) @triton.jit(do_not_specialize=["decode_query_len"]) @@ -243,7 +241,6 @@ def _gqa_sparse_decode_kernel( NUM_TOPK_CHUNKS: tl.constexpr, BLOCK_SIZE_H: tl.constexpr, BLOCK_SIZE_D: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load USE_PDL: tl.constexpr, ): @@ -268,11 +265,10 @@ def _gqa_sparse_decode_kernel( # attention range instead of letting padded rows produce negative lengths. kv_len = tl.maximum(query_pos + 1, 0) - # number of valid (non-padded) selected blocks for this query token - off_t = tl.arange(0, BLOCK_SIZE_T) + # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)). idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn - topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1) - real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0) + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk) off_n = tl.arange(0, BLOCK_SIZE_K) diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py index 432c8bb790d..16a9277d11a 100644 --- a/vllm/models/minimax_m3/nvidia/indexer_msa.py +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -2,16 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MSA (SM100/Blackwell) indexer impl for MiniMax M3. -Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects -top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much -faster than Triton for the wide prefill score (benchmarked ~3-5x). +Both sides write block scores into one unified token-major buffer +``[total_q, H, max_k_tiles]``, then a single ``fmha_sm100.sparse_topk_select`` +selects the top-k blocks for the whole batch (decode ``[:nd]`` + prefill +``[nd:]``) into the shared ``topk_indices_buffer``. It bounds each row by its +causal page count and force-includes the init/local blocks, so the unwritten +tail of the buffer is pre-filled with ``-inf``. -Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the -Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix -score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K -top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap) -by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids) -and writes the shared ``topk_indices_buffer`` via ``out=``. +Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faster +than Triton for the wide prefill score, benchmarked ~3-5x), writing its +``max_score`` straight into the buffer's prefill region (stride-aware, no copy). + +Decode scores with the Triton split-K ``minimax_m3_index_decode_score`` (a +purpose-built vector x matrix score, no wasted tensor-core tiles, 256-way +split-K, cudagraph-safe by shape-constant grids), writing into the decode +region. Its tuning heuristics are kept; only the top-k is shared with prefill. ``fmha_sm100`` imports are function-local so this module is import-safe on AMD / non-SM100. @@ -22,6 +27,7 @@ from typing import ClassVar import torch +from vllm.config import VllmConfig from vllm.forward_context import get_forward_context from vllm.models.minimax_m3.common.indexer import ( MiniMaxM3IndexerBackend, @@ -31,8 +37,7 @@ from vllm.models.minimax_m3.common.indexer import ( MiniMaxM3IndexerMetadataBuilder, ) from vllm.models.minimax_m3.common.ops.index_topk import ( - minimax_m3_index_decode, - minimax_m3_index_topk, + minimax_m3_index_decode_score, ) from vllm.v1.attention.backend import ( AttentionBackend, @@ -40,10 +45,22 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.kv_cache_interface import AttentionSpec # Page size == sparse block size == index-K block; fmha tile id == M3 block id. PAGE_SIZE = 128 +# Fill for unwritten score tiles: -inf so they never win the top-k (score kernels +# only write causally-valid blocks). +_SCORE_SENTINEL = float("-inf") + +# Tile (KV-block) dim of the unified score buffer, hardcoded as a cudagraph +# capture-time constant so the decode score kernel's buffer shape is frozen +# across replays. 8192 tiles == 1M tokens of context; -inf padding + +# num_valid_pages bound each row to its causal range, so shorter replays reuse +# the same buffer safely. +MAX_K_TILES = 8192 + class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend): """Indexer side-cache backend selecting the MSA builder.""" @@ -71,6 +88,25 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): (the base ``prefill`` field is unused on this path).""" prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None + # Per-forward view (``[:num_tokens]``) of the builder's persistent unified + # score buffer ``[total_q, H, MAX_K_TILES]``, shared by decode and prefill + # and reused across all layers. Pre-filled with the -inf sentinel in + # ``build()``; each layer overwrites its valid tiles before its own top-k. + unified_scores: torch.Tensor | None = None + # Tile (KV-block) dim of the unified score buffer (== ``MAX_K_TILES``). + # Forced as the fmha plan's ``max_k_tiles`` so prefill writes its max_score + # straight into the shared buffer. + max_k_tiles: int = 0 + # Batch-wide inputs for the single top-k over the unified buffer (decode + + # prefill in one call). ``cu_seqlens_q`` is the per-request query-start + # offsets (== query_start_loc, a stable view for cudagraph); ``prefix_lens`` + # is the per-request context length (== context_lens). + topk_cu_seqlens_q: torch.Tensor | None = None + topk_prefix_lens: torch.Tensor | None = None + topk_max_query_len: int = 0 + # Per-token causal page count cdiv(seq_pos+1, PAGE_SIZE), [total_q] int32: + # drives sparse_topk_select force_end_blocks + -1 out-of-range clamp. + topk_num_valid_pages: torch.Tensor | None = None class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): @@ -79,6 +115,28 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + # Persistent unified score buffer [T, H, MAX_K_TILES] shared by all + # indexer layers and reused across forwards. Stable address (required + # for the captured decode path) + fixed tile dim so the decode score + # kernel's shape is frozen at capture. Filled with -inf per forward in + # build(); the valid/padding partition is a per-forward constant, so + # every layer overwrites the same valid tiles before its own top-k. + self.unified_scores_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.num_index_heads, + MAX_K_TILES, + dtype=torch.float32, + device=device, + ) + def build( self, common_prefix_len: int, @@ -107,6 +165,19 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): common_attn_metadata.compute_num_computed_tokens(), non_blocking=True ) + # Per-token causal page count for the top-k, into the stable cg buffer. + positions = common_attn_metadata.positions + assert positions is not None + num_valid_pages = self.num_valid_pages_buffer[:num_tokens] + num_valid_pages.copy_(positions[:num_tokens] // PAGE_SIZE + 1) + + # Unified score buffer: a per-forward view of the persistent buffer, + # reset to the -inf sentinel once here and shared by every layer (the + # tile dim is the capture-time constant MAX_K_TILES). + max_k_tiles = MAX_K_TILES + unified_scores = self.unified_scores_buffer[:num_tokens] + unified_scores.fill_(_SCORE_SENTINEL) + decode: MiniMaxM3IndexerDecodeMetadata | None = None if num_decodes > 0: qsl_cpu = common_attn_metadata.query_start_loc_cpu @@ -148,6 +219,11 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): causal=True, num_kv_splits=1, ) + # Force the plan's tile dim to the unified buffer's so prefill writes + # its max_score straight into unified[:, :, nd:] (the stride-aware + # binding shape-matches the tile dim exactly). max_k_tiles >= the + # plan's natural value, so the extra tiles are simply never written. + plan["max_k_tiles"] = max_k_tiles cols = torch.arange(block_table.shape[1], device=block_table.device) valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None] prefill = MiniMaxM3IndexerMSAPrefillMetadata( @@ -171,6 +247,12 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): num_prefill_tokens=num_prefill_tokens, decode=decode, prefill_msa=prefill, + unified_scores=unified_scores, + max_k_tiles=max_k_tiles, + topk_cu_seqlens_q=query_start_loc[: num_reqs + 1], + topk_prefix_lens=context_lens, + topk_max_query_len=common_attn_metadata.max_query_len, + topk_num_valid_pages=num_valid_pages, ) @@ -195,36 +277,49 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): -1, self.num_index_heads, self.index_head_dim ) kv = self.index_cache.kv_cache - # Both sides write into the single shared persistent topk_indices_buffer: - # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]). + # Shared persistent top-k output buffer; the unified top-k below writes + # the selected block ids into buf[:num_tokens]. buf = self.topk_indices_buffer + assert buf is not None - decode_topk: torch.Tensor | None = None + # Unified token-major score buffer [total_q, H, MAX_K_TILES]: the tile + # dim is innermost/contiguous, so both fmha writes (native [T,H,K]) and + # the block-iterating top-k reads hit contiguous tiles. Each side gets a + # contiguous slice on dim 0: decode [:nd], prefill [nd:]; the kernels + # read/write by strides. The builder allocates it once (persistent, + # shared by all layers) and resets it to the sentinel each forward, so + # the top-k never picks an unwritten tile. + unified_scores = md.unified_scores + assert unified_scores is not None + + # Decode scores -> unified[:nd] (transposed [H, nd, MK] view; the kernel + # writes by strides). Top-k is deferred to the single unified call below. if md.decode is not None: d = md.decode - decode_topk = minimax_m3_index_decode( + minimax_m3_index_decode_score( index_q[:nd], kv, d.block_table, d.seq_lens, d.max_seq_len, - self.topk_blocks, self.init_blocks, self.local_blocks, self.num_kv_heads, d.decode_query_len, d.max_decode_query_len, - out=buf, + score_out=unified_scores[:nd].transpose(0, 1), ) - prefill_topk: torch.Tensor | None = None if md.prefill_msa is not None: from vllm.third_party.fmha_sm100.api import _fmha_sm100 p = md.prefill_msa # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D). k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim) - _, max_score = _fmha_sm100( + # fmha writes its max_score natively as [nnz_p, H, max_k_tiles] into + # the prefill region (stride-aware; plan max_k_tiles forced to + # md.max_k_tiles so the shape matches exactly -> no copy). + _fmha_sm100( index_q[nd:], k_pages, k_pages, # V placeholder; not read in OnlyScore @@ -233,19 +328,24 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): output_o=False, output_maxscore=True, sm_scale=self.scale, - ) - # Triton top-k wants [num_index_heads, num_tokens, max_block]; the - # transpose is a strided view (the kernel reads via strides). - out = buf[:, nd:, :] if buf is not None else None - prefill_topk = minimax_m3_index_topk( - max_score.transpose(1, 2), - p.cu_seqlens_q, - p.prefix_lens, - p.max_query_len, - self.topk_blocks, - self.init_blocks, - self.local_blocks, - out=out, + max_score=unified_scores[nd:], ) - return decode_topk, prefill_topk + # Single top-k over the unified buffer via fmha_sm100 sparse_topk_select + # (THK, no transpose) into ``buf``. num_valid_pages drives force_end_blocks + # (always keep each token's local block; fmha OnlyScore won't) + the -1 + # out-of-range clamp. + from vllm.third_party.fmha_sm100.api import sparse_topk_select + + sparse_topk_select( + unified_scores, + self.topk_blocks, + num_valid_pages=md.topk_num_valid_pages, + force_begin_blocks=self.init_blocks, + force_end_blocks=self.local_blocks, + output=buf[:num_tokens], + max_score_layout="THK", + ) + + # The attend reads ``buf`` directly; this return is vestigial. + return None, None diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index a30bc335fcf..66a0b67202c 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -196,6 +196,7 @@ class MiniMaxM3MoE(nn.Module): config: PretrainedConfig, layer_id: int, quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -258,10 +259,10 @@ class MiniMaxM3MoE(nn.Module): swiglu_alpha=config.swiglu_alpha, swiglu_beta=config.swiglu_beta, routed_scaling_factor=self.routed_scaling_factor, - apply_routed_scale_to_output=True, router_logits_dtype=self.gate.out_dtype, shared_experts=self.shared_experts, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.experts", ) @@ -659,14 +660,10 @@ class MiniMaxM3DecoderLayer(nn.Module): layer_id = int(prefix.split(sep=".")[-1]) self.layer_id = layer_id - # Complete the preceding dense MLP's deferred all-reduce - # (reduce_results=False), fused into this layer's input_layernorm. - # Disable this fusion when PP is set - self.fuse_input_allreduce = ( - layer_id > 0 - and not _is_moe_layer(config, layer_id - 1) - and vllm_config.parallel_config.pipeline_parallel_size == 1 - ) + # When set, complete the preceding layer's deferred FFN all-reduce + # fused into this layer's input_layernorm. + # Configured by MiniMaxM3Model.__init__ + self.fuse_input_allreduce = False is_sparse_attention_layer = ( force_sparse_attn or layer_id in _sparse_attention_layer_ids(config) @@ -692,12 +689,19 @@ class MiniMaxM3DecoderLayer(nn.Module): # Dense layers store the FFN under `mlp`; MoE layers under # `block_sparse_moe` -- matching the checkpoint's naming. + # Leave the FFN output un-reduced so its all-reduce fuses into the + # next RMSNorm. MTP blocks add the residual directly and PP sends + # hidden states across stages, so both must reduce. + reduce_results = ( + is_mtp_block or vllm_config.parallel_config.pipeline_parallel_size > 1 + ) self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id) if self.is_moe_layer: self.block_sparse_moe = MiniMaxM3MoE( config=config, layer_id=layer_id, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.block_sparse_moe", ) else: @@ -706,7 +710,7 @@ class MiniMaxM3DecoderLayer(nn.Module): intermediate_size=config.dense_intermediate_size, quant_config=quant_config, prefix=f"{prefix}.mlp", - reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1, + reduce_results=reduce_results, ) # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm. @@ -745,6 +749,14 @@ class MiniMaxM3DecoderLayer(nn.Module): hidden_states = ffn(hidden_states) return hidden_states, residual + @property + def ffn_all_reduce_deferred(self) -> bool: + """This layer's FFN output is left un-reduced; the caller fuses the + all-reduce into the next RMSNorm.""" + if self.is_moe_layer: + return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce + return not self.mlp.down_proj.reduce_results + class MiniMaxM3Model(nn.Module, EagleModelMixin): fall_back_to_pt_during_load = False @@ -770,8 +782,9 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): # Reserved top-k indices buffer shared by all sparse-attention indexer # layers (mirrors DeepseekV4); kept at a stable address so the indexer's - # top-k output survives cudagraph capture/replay. Shape matches the - # per-head index top-k output [num_index_heads, total_q, topk]. + # top-k output survives cudagraph capture/replay. Token-major + # [total_q, num_index_heads, topk] so the indexer writes its native + # [token, head, topk] top-k; the attend transposes to [H, tokens, topk]. sparse_cfg = getattr(config, "sparse_attention_config", None) if sparse_cfg is not None: tp_size = get_tensor_model_parallel_world_size() @@ -781,8 +794,8 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4 self.topk_indices_buffer = torch.empty( - num_index_heads, padded_num_tokens, + num_index_heads, sparse_cfg["sparse_topk_blocks"], dtype=torch.int32, ) @@ -807,6 +820,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): ["hidden_states", "residual"], config.hidden_size ) + # Configure cross-layer all-reduce/RMSNorm fusion: a layer whose FFN output + # is left un-reduced has that all-reduce fused into the next layer's + # input_layernorm (or the final norm). + prev_defers = False + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + layer.fuse_input_allreduce = idx > 0 and prev_defers + prev_defers = layer.ffn_all_reduce_deferred + self.fuse_final_norm_allreduce = prev_defers + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -841,7 +863,12 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): {"hidden_states": hidden_states, "residual": residual} ) - hidden_states, _ = self.norm(hidden_states, residual) + if self.fuse_final_norm_allreduce: + hidden_states, _ = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.norm + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states @@ -958,6 +985,11 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin): class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3): """MiniMax M3 (sparse/dense backbone) for causal language modeling.""" + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_text_config @@ -1025,6 +1057,11 @@ class MiniMaxM3SparseForConditionalGeneration( # ranks (see ``_process_image_input`` / ``_process_video_input``). supports_encoder_tp_data = True + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "multi_modal_projector.": "vision_tower.multi_modal_projector.", diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index 0df2ed85bbd..545667c8cfe 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -39,7 +39,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): nd = main_md.num_decode_tokens num_tokens = main_md.num_actual_tokens - # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:]. + # Indexer top-k from the shared token-major buffer [total_q, H, MK]; the + # kernels want [H, tokens, MK], so slice tokens on dim 0 then transpose. topk = layer.topk_indices_buffer # type: ignore[attr-defined] assert topk is not None hd = self.head_size @@ -56,7 +57,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): minimax_m3_sparse_attn_decode( q[:nd], kv_cache, - topk[:, :nd, :], + topk[:nd].transpose(0, 1), d.block_table, d.seq_lens, self.num_kv_heads, @@ -74,7 +75,9 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): p = main_md.prefill assert p is not None - prefill_topk = topk[:, nd:num_tokens, :] + # [H, prefill, MK] transposed view; build_k2q_csr consumes the + # strided view directly (topK stays innermost-contiguous). + prefill_topk = topk[nd:num_tokens].transpose(0, 1) qp = q[nd:] k_cache = kv_cache[:, 0].transpose(1, 2) v_cache = kv_cache[:, 1].transpose(1, 2) diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index 582b6fde565..a440e69ef1c 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -13,14 +13,17 @@ from pathlib import Path from typing import Any, TypeVar from urllib.request import url2pathname +import aiohttp import numpy as np import numpy.typing as npt +import requests import torch from PIL import Image, UnidentifiedImageError from urllib3.util import Url, parse_url import vllm.envs as envs from vllm.connections import HTTPConnection, global_http_connection +from vllm.exceptions import VLLMUnprocessableEntityError from vllm.logger import init_logger from vllm.multimodal.video import get_video_loader_backend_for_processor from vllm.utils.registry import ExtensionManager @@ -48,6 +51,65 @@ MODALITY_IO_MAP: dict[str, type[MediaIO]] = { } +def _wrap_media_fetch_error( + url: str, exc: Exception +) -> VLLMUnprocessableEntityError | Exception: + """Convert media fetch exceptions to VLLMUnprocessableEntityError. + + This handles HTTP errors that indicate the media resource is invalid + (4xx responses except 408/429, malformed URLs) and converts them to a + 422 Unprocessable Entity error instead of 500. + + Transient errors (5xx, 408, 429, DNS failures, connection errors, + timeouts) are returned as-is to allow retry logic to handle them + appropriately. + + Returns: + VLLMUnprocessableEntityError for permanent client errors (4xx except + 408/429, invalid URL) + Original exception for transient errors (5xx, 408, 429, network blips) + or other exceptions + """ + if isinstance(exc, aiohttp.ClientResponseError): + if exc.status in (408, 429): + return exc + if exc.status < 500: + return VLLMUnprocessableEntityError( + f"Failed to fetch media from URL: HTTP {exc.status} error", + parameter="image_url", + value=url, + ) + return exc + + if isinstance(exc, requests.exceptions.HTTPError): + if exc.response is not None: + status_code = exc.response.status_code + if status_code in (408, 429): + return exc + if status_code < 500: + return VLLMUnprocessableEntityError( + f"Failed to fetch media from URL: HTTP {status_code} error", + parameter="image_url", + value=url, + ) + return exc + + if isinstance(exc, requests.exceptions.InvalidURL): + return VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Invalid URL format", + parameter="image_url", + value=url, + ) + + if isinstance(exc, ValueError): + return VLLMUnprocessableEntityError( + "Failed to fetch media from URL: Invalid URL", + parameter="image_url", + value=url, + ) + return exc + + def merge_media_io_kwargs( defaults: dict[str, dict[str, Any]] | None, overrides: dict[str, dict[str, Any]] | None, @@ -303,11 +365,17 @@ class MediaConnector: return media_io.load_bytes(cached) connection = self.connection - data = connection.get_bytes( - url_spec.url, - timeout=fetch_timeout, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) + try: + data = connection.get_bytes( + url_spec.url, + timeout=fetch_timeout, + allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, + ) + except Exception as e: + wrapped = _wrap_media_fetch_error(url, e) + if isinstance(wrapped, VLLMUnprocessableEntityError): + raise wrapped from e + raise self._put_cached_bytes(url, data) return media_io.load_bytes(data) @@ -348,11 +416,17 @@ class MediaConnector: return await future connection = self.connection - data = await connection.async_get_bytes( - url_spec.url, - timeout=fetch_timeout, - allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, - ) + try: + data = await connection.async_get_bytes( + url_spec.url, + timeout=fetch_timeout, + allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS, + ) + except Exception as e: + wrapped = _wrap_media_fetch_error(url, e) + if isinstance(wrapped, VLLMUnprocessableEntityError): + raise wrapped from e + raise await loop.run_in_executor( global_thread_pool, self._put_cached_bytes, url, data diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index 404f5a0e7cf..45ea4c2fdf4 100644 --- a/vllm/multimodal/media/video.py +++ b/vllm/multimodal/media/video.py @@ -10,11 +10,14 @@ import pybase64 from PIL import Image from vllm import envs +from vllm.logger import init_logger from ..video import VIDEO_LOADER_REGISTRY from .base import MediaIO from .image import ImageMediaIO +logger = init_logger(__name__) + class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): """Configuration values can be user-provided either by --media-io-kwargs or @@ -28,6 +31,24 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]): default_kwargs: dict[str, Any] | None, runtime_kwargs: dict[str, Any] | None, ) -> dict[str, Any]: + if runtime_kwargs: + # Block request-level selection of GPU video backends that + # were not configured (and VRAM-reserved) at startup. + for key in ("video_backend", "backend"): + requested = runtime_kwargs.get(key) + if requested and VIDEO_LOADER_REGISTRY.backend_requires_gpu(requested): + static_val = (default_kwargs or {}).get(key) + if static_val != requested: + logger.warning_once( + "Stripping request-level %s=%r: GPU video " + "backend not configured at startup.", + key, + requested, + ) + runtime_kwargs = { + k: v for k, v in runtime_kwargs.items() if k != key + } + merged = super().merge_kwargs(default_kwargs, runtime_kwargs) # fps and num_frames interact with each other, so if either is # overridden at request time, wipe the other from defaults to diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 725e33e3f8b..aab8bfd3b3c 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -15,7 +15,7 @@ import torch from vllm import envs from vllm.logger import init_logger -from vllm.utils.import_utils import PlaceholderModule +from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available from vllm.utils.mem_constants import MiB_bytes from vllm.utils.registry import ExtensionManager @@ -31,6 +31,13 @@ try: except ImportError: av = PlaceholderModule("av") # type: ignore[assignment] +try: + from torchcodec.decoders import VideoDecoder +except (ImportError, RuntimeError): + VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment] + "decoders.VideoDecoder" + ) + logger = init_logger(__name__) @@ -39,6 +46,7 @@ class VideoLoaderRegistry(ExtensionManager): def __init__(self) -> None: super().__init__() self.processor2backend: dict[str, str] = {} + self._requires_gpu: dict[str, bool] = {} @staticmethod def _normalize_registered_video_processors( @@ -62,11 +70,13 @@ class VideoLoaderRegistry(ExtensionManager): name: str, *, video_processor: str | tuple[str, ...] | None = None, + requires_gpu: bool = False, ): processors = self._normalize_registered_video_processors(video_processor) def wrap(cls_to_register): self.name2class[name] = cls_to_register + self._requires_gpu[name] = requires_gpu for processor_name in processors: self.processor2backend[processor_name] = name return cls_to_register @@ -82,6 +92,9 @@ class VideoLoaderRegistry(ExtensionManager): return self.processor2backend.get(video_processor) + def backend_requires_gpu(self, name: str) -> bool: + return self._requires_gpu.get(name, False) + def get_video_loader_backend_for_processor( video_processor: str | None, @@ -562,6 +575,53 @@ class PyAVVideoBackendMixin: return np.stack(frames_list), valid_indices +class TorchCodecVideoBackendMixin: + """TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities. + + Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory + bytes and extracts the sampled indices with a single batched + ``get_frames_at`` call, while releasing the GIL during decode. + """ + + @staticmethod + def make_torchcodec_decoder( + data: bytes, + *, + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", + ) -> "VideoDecoder": + # NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest + # of the pipeline expects, avoiding a transpose. + return VideoDecoder( + data, + dimension_order="NHWC", + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + + @staticmethod + def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata: + md = decoder.metadata + total_frames = md.num_frames or 0 + fps = float(md.average_fps) if md.average_fps else 0.0 + duration = float(md.duration_seconds) if md.duration_seconds else 0.0 + if total_frames == 0 and duration > 0 and fps > 0: + total_frames = int(duration * fps) + return VideoSourceMetadata(total_frames, fps, duration) + + @staticmethod + def decode_torchcodec_frames( + decoder: "VideoDecoder", + frame_indices: list[int], + ) -> tuple[npt.NDArray, list[int]]: + """Decode the requested indices in one batched, index-exact call.""" + if not frame_indices: + return np.empty((0,), dtype=np.uint8), [] + # Note: torchcodec releases the GIL for the entire call + batch = decoder.get_frames_at(frame_indices) + return batch.data.numpy(), list(frame_indices) + + class PyNvVideoCodecVideoBackendMixin: """PyNvVideoCodec utilities for GPU-backed frame decode.""" @@ -771,14 +831,15 @@ class VideoBackend( VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin, + TorchCodecVideoBackendMixin, PyNvVideoCodecVideoBackendMixin, ): """Uniform-sampling video backend. Samples ``num_frames`` uniformly across the video (or one frame every ``1/fps`` seconds, whichever produces fewer frames). The decoding codec - is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, or - ``"pynvvideocodec"``), which can be passed through + is selected via the ``backend`` kwarg (``"opencv"``, ``"pyav"``, + ``"torchcodec"`` or ``"pynvvideocodec"``), which can be passed through ``--media-io-kwargs``. Defaults to ``"opencv"``. """ @@ -824,7 +885,9 @@ class VideoBackend( max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", + num_ffmpeg_threads: int = 0, + seek_mode: Literal["exact", "approximate"] = "exact", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: """Load sampled frames from raw video bytes. @@ -837,8 +900,20 @@ class VideoBackend( dynamic subclass; ignored here. frame_recovery: Enable forward-scan recovery for failed frames. Only honored by the OpenCV codec. - backend: Decoding codec — ``"opencv"``, ``"pyav"``, or - ``"pynvvideocodec"``. + backend: Decoding codec — ``"opencv"``, ``"pyav"``, + ``"torchcodec"`` or ``"pynvvideocodec"``. + num_ffmpeg_threads: Number of FFmpeg decoding threads, only used by + TorchCodec: ``0`` (default) relies on the FFmpeg default value + which is ``min(cpu_count + 1, 16)``. + OpenCV will always use ``min(cpu_count, 16)`` while pyav will + always use ``min(cpu_count, (height + 15) / 16)``. + seek_mode: Seek mode for the TorchCodec decoder, only used by + TorchCodec: ``"exact"`` (default) guarantees frame-accurate + sampling by scanning the file on creation, while + ``"approximate"`` skips that scan for faster decoder creation + at the cost of relying on the file's metadata. See + https://meta-pytorch.org/torchcodec/stable/generated_examples/decoding/approximate_mode.html + for details. Returns: Tuple of ``(frames_array, metadata_dict)``. @@ -877,6 +952,25 @@ class VideoBackend( frames, valid = cls.decode_frames( container, frame_idx, source.original_fps, source.duration ) + elif backend == "torchcodec": + assert not frame_recovery, ( + "frame_recovery is only available for `opencv` backend" + ) + check_torchcodec_available() + decoder = cls.make_torchcodec_decoder( + data, + num_ffmpeg_threads=num_ffmpeg_threads, + seek_mode=seek_mode, + ) + _check_frame_pixel_limit( + decoder.metadata.width or 0, + decoder.metadata.height or 0, + ) + source = cls._prepare_source(cls.get_torchcodec_metadata(decoder)) + frame_idx = cls.compute_frames_index_to_sample( + source=source, target=target, **kwargs + ) + frames, valid = cls.decode_torchcodec_frames(decoder, frame_idx) elif backend == PYNVVIDEOCODEC_VIDEO_BACKEND: if frame_recovery: raise ValueError( @@ -891,7 +985,8 @@ class VideoBackend( else: raise ValueError( f"Unknown video codec backend {backend!r}; " - "valid options: 'opencv', 'pyav', 'pynvvideocodec'." + "valid options: 'opencv', 'pyav', 'torchcodec', " + "'pynvvideocodec'." ) if len(valid) < len(frame_idx): @@ -909,7 +1004,7 @@ class VideoBackend( ) -@VIDEO_LOADER_REGISTRY.register(PYNVVIDEOCODEC_VIDEO_BACKEND) +@VIDEO_LOADER_REGISTRY.register(PYNVVIDEOCODEC_VIDEO_BACKEND, requires_gpu=True) class PyNvVideoCodecVideoBackend(VideoBackend): """Hardware-accelerated video backend using PyNvVideoCodec. @@ -978,7 +1073,7 @@ class Qwen3VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1057,7 +1152,7 @@ class Qwen2VLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1149,7 +1244,7 @@ class DynamicVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1274,7 +1369,7 @@ class GLM46VVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: return super().load_bytes( @@ -1372,7 +1467,7 @@ class GLMGAVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( @@ -1695,7 +1790,7 @@ class NemotronVLVideoBackend(VideoBackend): max_duration: int = 300, frame_recovery: bool = False, *, - backend: Literal["opencv", "pyav", "pynvvideocodec"] = "opencv", + backend: Literal["opencv", "pyav", "torchcodec", "pynvvideocodec"] = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: frames, metadata = super().load_bytes( diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 80fd02e4cec..5043ca191f3 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, NamedTuple -from openai_harmony import HarmonyError +from openai_harmony import HarmonyError, Message, Role from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -91,6 +91,9 @@ class HarmonyParser(DelegatingParser): self._next_tool_call_index = 0 self._num_processed_messages = 0 + # For error recovery + self._current_message_tokens: list[int] = [] + @property def _harmony_parser(self) -> StreamableParser: """Lazily initializes the Harmony parser.""" @@ -107,32 +110,48 @@ class HarmonyParser(DelegatingParser): self._num_processed_messages += 1 return msg - def flush(self) -> Segment | None: + def flush(self) -> list[Segment]: + segments: list[Segment] = [] try: self._harmony_parser.process_eos() msg = self._poll_completed_message() except HarmonyError: logger.warning( "Harmony parser ended in a non-terminal state; returning the " - "raw unparsed output. This usually indicates a malformed " - "assistant turn, e.g. a 'final' channel missing the " - "<|message|> delimiter." + "recovered raw output." ) - raise - finally: - # Reset to the initial assistant-parser state for the next turn. - self._parser = None - self._num_processed_messages = 0 + + final_channel = "final" + text = self.model_tokenizer.decode(self._current_message_tokens) + segments.append( + Segment( + channel=final_channel, + recipient=None, + delta=text, + completed_message=None, + ) + ) + msg = Message.from_role_and_content(Role.ASSISTANT, text).with_channel( + final_channel + ) + + # Reset to the initial assistant-parser state for the next turn. + self._parser = None + self._num_processed_messages = 0 + self._current_message_tokens.clear() if msg is None: - return None + return segments - return Segment( - channel=msg.channel, - recipient=msg.recipient, - delta="", - completed_message=msg, + segments.append( + Segment( + channel=msg.channel, + recipient=msg.recipient, + delta="", + completed_message=msg, + ) ) + return segments def parse( self, @@ -147,12 +166,9 @@ class HarmonyParser(DelegatingParser): Callers must decide whether to surface them. """ result = self.process_chunk(model_output_token_ids) - try: - flushed_segment = self.flush() - except HarmonyError: - return None, model_output, None - if flushed_segment is not None: - result.segments.append(flushed_segment) + flushed_segments = self.flush() + if flushed_segments: + result.segments.extend(flushed_segments) reasoning_parts: list[str] = [] content_parts: list[str] = [] @@ -209,13 +225,9 @@ class HarmonyParser(DelegatingParser): ) result = self.process_chunk(delta_token_ids) if finished: - try: - flushed_segment = self.flush() - except HarmonyError: - self._next_tool_call_index = 0 - return DeltaMessage(content=delta_text) - if flushed_segment is not None: - result.segments.append(flushed_segment) + flushed_segments = self.flush() + if flushed_segments: + result.segments.extend(flushed_segments) combined_content = "" combined_reasoning = "" tool_messages: list[DeltaToolCall] = [] @@ -298,6 +310,11 @@ class HarmonyParser(DelegatingParser): delta = self._harmony_parser.last_content_delta or "" completed_message = self._poll_completed_message() + if completed_message is not None: + self._current_message_tokens.clear() + else: + self._current_message_tokens.append(token_id) + if channel == "analysis" or ( channel == "commentary" and recipient is not None ): diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 605650525af..9eac95e0324 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -727,7 +727,7 @@ class NvmlCudaPlatform(CudaPlatformBase): @with_nvml_context def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: try: - physical_device_id = cls.device_id_to_physical_device_id(device_id) + physical_device_id = cls.visible_device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) return DeviceCapability(major=major, minor=minor) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index a81c34d7a50..7cf009c8071 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -271,6 +271,16 @@ class Platform: f"{cls.device_name}." ) from e + # GPU device IDs can refer to three distinct namespaces: + # - logical: vLLM-local IDs such as local ranks. These index + # assigned_physical_gpu_ids when it is set. + # - visible: torch/CUDA ordinals in the current process after applying + # the device-control env var, e.g. CUDA_VISIBLE_DEVICES. + # - physical: global GPU IDs used by topology and management APIs such as + # NVML, which are not remapped by CUDA_VISIBLE_DEVICES. + # Keep conversions explicit. In particular, torch device indices are + # visible IDs, not vLLM logical IDs. + @classmethod def device_id_to_physical_device_id(cls, device_id: int): """Map a vLLM-local logical device ID to a physical device ID. @@ -411,7 +421,12 @@ class Platform: cls, device_id: int = 0, ) -> DeviceCapability | None: - """Stateless version of [torch.cuda.get_device_capability][].""" + """Stateless version of [torch.cuda.get_device_capability][]. + + Args: + device_id: Device index in the visible device namespace, matching + the argument accepted by torch.cuda. + """ return None @classmethod diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 867833c9d7e..cbfa579313b 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -148,6 +148,15 @@ class XPUPlatform(Platform): if selected_backend == AttentionBackendEnum.TRITON_ATTN: logger.info_once("Using Triton backend.") return AttentionBackendEnum.TRITON_ATTN.get_path() + elif attn_selector_config.use_mm_prefix: + # Flash Attention on XPU has no FA4 kernel, so it cannot apply the + # multimodal prefix-LM bidirectional mask. Fall back to Triton + # Attention, which supports mm_prefix. + logger.warning_once( + "Flash Attention on XPU does not support multimodal prefix-LM " + "attention. Falling back to Triton Attention backend." + ) + return AttentionBackendEnum.TRITON_ATTN.get_path() elif dtype == torch.float32: logger.warning_once( "Flash Attention on XPU does not support float32 dtype. " @@ -286,6 +295,8 @@ class XPUPlatform(Platform): "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", "fuse_rope_kvcache": "RoPE + KV cache fusion", + "fuse_rope_kvcache_cat_mla": "RoPE + KV cache + MLA fusion", + "enable_qk_norm_rope_fusion": "QK Norm + RoPE fusion", } if compilation_config.mode != CompilationMode.NONE: for flag, feature_name in fusion_passes_to_disable.items(): diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py index 89fadad7a8f..95e895c279b 100644 --- a/vllm/plugins/__init__.py +++ b/vllm/plugins/__init__.py @@ -3,10 +3,14 @@ import logging from collections.abc import Callable -from typing import Any +from typing import TYPE_CHECKING, Any import vllm.envs as envs +if TYPE_CHECKING: + from vllm.plugins.endpoint_plugins.interface import EndpointPlugin + from vllm.tasks import SupportedTask + logger = logging.getLogger(__name__) # Default plugins group will be loaded in all processes(process0, engine core @@ -20,6 +24,10 @@ PLATFORM_PLUGINS_GROUP = "vllm.platform_plugins" # Stat logger plugins group will be loaded in process0 only when serve vLLM with # async mode. STAT_LOGGER_PLUGINS_GROUP = "vllm.stat_logger_plugins" +# Endpoint plugins group is loaded in the API server front end process only. +# Each entry point resolves to a factory returning an `EndpointPlugin` +# (see `vllm/plugins/endpoint_plugins/interface.py`). +ENDPOINT_PLUGINS_GROUP = "vllm.endpoint_plugins" # make sure one process only loads plugins once plugins_loaded = False @@ -80,3 +88,71 @@ def load_general_plugins(): # general plugins, we only need to execute the loaded functions for func in plugins.values(): func() + + +def load_endpoint_plugins( + supported_tasks: "tuple[SupportedTask, ...] | None" = None, +) -> "list[EndpointPlugin]": + """Discover, gate and instantiate `vllm.endpoint_plugins` entry points. + + Endpoint plugins add HTTP routes to the API server, so they default to + not loading. Unlike other plugin groups, a plugin here is only + considered when it is explicitly named in `VLLM_PLUGINS`. This is a + stricter posture than `load_plugins_by_group` which "load everything unless + an allowlist says otherwise". This posture is taken to handle potentially + larger exposed network surface. + + A discovered plugin is loaded only if both hold: + - it is named in `VLLM_PLUGINS` (enforced by not calling the loader + at all when `VLLM_PLUGINS` is unset). Note that `VLLM_PLUGINS=""` + parses to `[""]`, not `None`, so it is treated as a (non strict) + allowlist that matches no plugin name, not as "unset". + - its `required_tasks` is `None` or intersects `supported_tasks`. + + Args: + supported_tasks: Tasks the server supports. `None` means no plugin + with a non `None` `required_tasks` will be loaded. + + Returns: + Instantiated plugins that passed gating in discovery order. + """ + from importlib.metadata import entry_points + + if envs.VLLM_PLUGINS is None: + discovered = entry_points(group=ENDPOINT_PLUGINS_GROUP) + if discovered: + logger.warning( + "Found endpoint plugin(s) %s but VLLM_PLUGINS is not set. " + "Endpoint plugins add HTTP routes and must be explicitly " + "allowlisted via VLLM_PLUGINS to be loaded.", + [p.name for p in discovered], + ) + return [] + + factories = load_plugins_by_group(ENDPOINT_PLUGINS_GROUP) + + endpoint_plugins: list[EndpointPlugin] = [] + for name, factory in factories.items(): + try: + plugin = factory() + except Exception: + logger.exception("Failed to instantiate endpoint plugin %s", name) + continue + + required_tasks = plugin.required_tasks + if required_tasks is not None and ( + supported_tasks is None or not set(required_tasks) & set(supported_tasks) + ): + logger.info( + "Skipping endpoint plugin %s: requires one of tasks %s, " + "server supports %s", + name, + required_tasks, + supported_tasks, + ) + continue + + logger.info("Loaded endpoint plugin %s", name) + endpoint_plugins.append(plugin) + + return endpoint_plugins diff --git a/vllm/plugins/endpoint_plugins/__init__.py b/vllm/plugins/endpoint_plugins/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/plugins/endpoint_plugins/interface.py b/vllm/plugins/endpoint_plugins/interface.py new file mode 100644 index 00000000000..99487f57b68 --- /dev/null +++ b/vllm/plugins/endpoint_plugins/interface.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Contract for `vllm.endpoint_plugins` entry points. + +An endpoint plugin adds HTTP routes to the OpenAI compatible API server. +Its scope is HTTP surface only. It registers routes and optionally +per app state used by those routes. It must not open new paths into the +engine by reaching the engine the same way an in-tree serving handler does +via `EngineClient` (e.g. `engine_client.collective_rpc(...)`). + +If a plugin also needs engine side behavior (a new worker side RPC method, +a custom stat, etc.) pair this entry point with one registered under +`vllm.general_plugins` (see `vllm/plugins/__init__.py`). The +`general_plugins` entry installs the engine side method and the +`endpoint_plugins` entry exposes it over HTTP. The two are registered and +loaded independently where neither implies the other. + +Plugins are opt-in. See `load_endpoint_plugins` in `vllm/plugins/__init__.py` +for the loading/gating rules and `docs/usage/security.md` for the security +posture of exposing plugin defined routes. + +The CPU only render server (see `build_and_serve_renderer` in +`vllm/entrypoints/openai/api_server.py`) has no `EngineClient`. A plugin +eligible for the `render` task (`required_tasks` is `None` or includes +`"render"`) still gets `attach_router` called but `init_state` receives +`engine_client=None`. Plugins that cannot function without an engine should +either exclude `"render"` from `required_tasks` or check for `None` in +`init_state`/their route handlers and degrade gracefully. +""" + +from argparse import Namespace +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from fastapi import FastAPI +from starlette.datastructures import State + +if TYPE_CHECKING: + from vllm.engine.protocol import EngineClient + from vllm.tasks import SupportedTask + + +@runtime_checkable +class EndpointPlugin(Protocol): + """Protocol implemented by `vllm.endpoint_plugins` entry point factories. + + An entry point registered under the `vllm.endpoint_plugins` group must + resolve to a zero argument callable (a class or factory function) that + returns an object satisfying this protocol. + """ + + name: str + """Unique plugin name used in logs and for `VLLM_PLUGINS` allowlisting.""" + + required_tasks: "tuple[SupportedTask, ...] | None" + """Tasks the server must support for this plugin to be loaded. + + The plugin is loaded only if this set intersects the server's + `supported_tasks`. `None` means the plugin has no task requirement and + is always eligible (subject to the `VLLM_PLUGINS` allowlist). + """ + + def attach_router(self, app: FastAPI) -> None: + """Register this plugin's routes on `app`. + + Called once during `build_app()` after all core routers have been + attached. Routes attached here can shadow core routes with the same + path. There is currently no conflict enforcement (see RFC #46565 follow ups). + """ + ... + + async def init_state( + self, engine_client: "EngineClient | None", state: State, args: Namespace + ) -> None: + """Initialize per app state consumed by this plugin's routes. + + Called once during `init_app_state()` after core state has been + initialized. Use `engine_client` (e.g. `collective_rpc`) to reach + the engine. Do not open new engine access paths. + + `engine_client` is `None` on the CPU only render server which has + no engine. This only happens for plugins eligible for the `render` + task (`required_tasks` is `None` or includes `"render"`). Handle + `None` explicitly (e.g. skip engine dependent setup, or have route + handlers return an error) if the plugin is loadable for `render` but + cannot function without an engine. + """ + ... diff --git a/vllm/reasoning/abs_reasoning_parsers.py b/vllm/reasoning/abs_reasoning_parsers.py index 4e28e50702e..04d6a937b3f 100644 --- a/vllm/reasoning/abs_reasoning_parsers.py +++ b/vllm/reasoning/abs_reasoning_parsers.py @@ -41,7 +41,7 @@ class ReasoningParser: @cached_property def vocab(self) -> dict[str, int]: - # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab + # NOTE: Only TokenizersBackend is guaranteed to have .vocab # whereas all tokenizers have .get_vocab() return self.model_tokenizer.get_vocab() diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 490f589af45..a7a6693154b 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -228,20 +228,12 @@ def _try_get_processor_chat_template( if cache_key in _PROCESSOR_CHAT_TEMPLATES: return _PROCESSOR_CHAT_TEMPLATES[cache_key] - from transformers import ( - PreTrainedTokenizer, - PreTrainedTokenizerFast, - ProcessorMixin, - ) + from transformers import ProcessorMixin, PythonBackend, TokenizersBackend try: processor = cached_get_processor( tokenizer.name_or_path, - processor_cls=( - PreTrainedTokenizer, - PreTrainedTokenizerFast, - ProcessorMixin, - ), + processor_cls=(PythonBackend, TokenizersBackend, ProcessorMixin), trust_remote_code=trust_remote_code, ) if ( @@ -619,15 +611,15 @@ _cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs) @lru_cache def _get_hf_base_chat_template_params() -> frozenset[str]: - from transformers import PreTrainedTokenizer + from transformers import PythonBackend # Get standard parameters from HuggingFace's base tokenizer class. - # This dynamically extracts parameters from PreTrainedTokenizer's + # This dynamically extracts parameters from PythonBackend's # apply_chat_template method, ensuring compatibility with tokenizers # that use **kwargs to receive standard parameters. # Read signature from HF's base class - the single source of truth - base_sig = inspect.signature(PreTrainedTokenizer.apply_chat_template) + base_sig = inspect.signature(PythonBackend.apply_chat_template) # Exclude VAR_KEYWORD (**kwargs) and VAR_POSITIONAL (*args) placeholders return frozenset( diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index 45a3898d4ac..15a4023fcee 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -12,6 +12,7 @@ from vllm.entrypoints.chat_utils import ( ConversationMessage, ) from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) from vllm.entrypoints.openai.completion.protocol import ( @@ -134,8 +135,12 @@ class OnlineRenderer: ) elif request.tool_choice != "auto": # "required" or named tool requires tool parser + if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): + tool_choice_desc = f'function "{request.tool_choice.function.name}"' + else: + tool_choice_desc = f'"{request.tool_choice}"' return self.create_error_response( - f'tool_choice="{request.tool_choice}" requires ' + f"tool_choice={tool_choice_desc} requires " "--tool-call-parser to be set" ) diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 5df2e8cfc17..f0966902d36 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -745,6 +745,7 @@ class SamplingParams( self._validate_logits_processors(model_config) self._validate_allowed_token_ids(tokenizer) self._validate_spec_decode(speculative_config) + self._validate_diffusion(model_config) self._validate_structured_outputs( model_config, structured_outputs_config, tokenizer ) @@ -878,6 +879,28 @@ class SamplingParams( "are not yet supported with speculative decoding." ) + def _validate_diffusion(self, model_config: ModelConfig) -> None: + if not model_config.is_diffusion: + return + + # Diffusion models denoise a whole canvas per step with a fixed + # temperature schedule, so per-request sampling parameters are not + # supported. Penalties are ignored by the sampler with a warning. + if ( + self.temperature != 1.0 + or self.min_p > _SAMPLING_EPS + or self.seed is not None + or self.min_tokens > 0 + or self.logit_bias + or self.bad_words + or self.allowed_token_ids + ): + raise ValueError( + "The temperature, min_p, seed, min_tokens, logit_bias, " + "bad_words, and allowed_token_ids sampling parameters " + "are not yet supported with diffusion models." + ) + def _validate_structured_outputs( self, model_config: ModelConfig, diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index 51199de5c47..b388f057930 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -3,7 +3,7 @@ import copy from typing import Any -from transformers import PreTrainedTokenizerFast +from transformers import TokenizersBackend from vllm.entrypoints.chat_utils import ChatCompletionMessageParam @@ -85,5 +85,5 @@ def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: class DeepseekV32Tokenizer(TokenizerLike): @classmethod def from_pretrained(cls, *args, **kwargs) -> HfTokenizer: - tokenizer = PreTrainedTokenizerFast.from_pretrained(*args, **kwargs) + tokenizer = TokenizersBackend.from_pretrained(*args, **kwargs) return get_cached_tokenizer(get_deepseek_v32_tokenizer(tokenizer)) diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index 2a6aaaf7397..3897149626f 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -3,7 +3,7 @@ import copy from typing import Any -from transformers import PreTrainedTokenizerFast +from transformers import TokenizersBackend from vllm.entrypoints.chat_utils import ChatCompletionMessageParam @@ -92,5 +92,5 @@ def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: class DeepseekV4Tokenizer(TokenizerLike): @classmethod def from_pretrained(cls, *args, **kwargs) -> HfTokenizer: - tokenizer = PreTrainedTokenizerFast.from_pretrained(*args, **kwargs) + tokenizer = TokenizersBackend.from_pretrained(*args, **kwargs) return get_cached_tokenizer(get_deepseek_v4_tokenizer(tokenizer)) diff --git a/vllm/tokenizers/hf.py b/vllm/tokenizers/hf.py index 45370bbb394..bdc767acd66 100644 --- a/vllm/tokenizers/hf.py +++ b/vllm/tokenizers/hf.py @@ -6,13 +6,13 @@ import queue from pathlib import Path from typing import TypeAlias, TypeVar -from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast +from transformers import AutoTokenizer, PythonBackend, TokenizersBackend from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config from .protocol import TokenizerLike -HfTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast +HfTokenizer: TypeAlias = PythonBackend | TokenizersBackend _T = TypeVar("_T", bound=TokenizerLike) @@ -24,7 +24,7 @@ class ThreadSafeHFTokenizerMixin: def maybe_make_thread_pool(tokenizer: _T, copies: int = 1): """ - If `tokenizer` is a `PreTrainedTokenizerFast`, modify the tokenizer + If `tokenizer` is a `TokenizersBackend`, modify the tokenizer in-place to make the public interface thread-safe by routing calls through a deep-copied tokenizer pool. @@ -34,14 +34,14 @@ def maybe_make_thread_pool(tokenizer: _T, copies: int = 1): methods like ``add_special_tokens`` or ``add_tokens``. - Adjacent method calls could happen on different deep copies. """ - if not isinstance(tokenizer, PreTrainedTokenizerFast) or isinstance( + if not isinstance(tokenizer, TokenizersBackend) or isinstance( tokenizer, ThreadSafeHFTokenizerMixin ): return tokenizer og_tokenizer = copy.copy(tokenizer) - tokenizer_pool: queue.Queue[PreTrainedTokenizerFast] = queue.Queue() + tokenizer_pool: queue.Queue[TokenizersBackend] = queue.Queue() for _ in range(copies): tokenizer_pool.put(copy.deepcopy(og_tokenizer)) @@ -116,6 +116,16 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: tokenizer_all_special_tokens = tokenizer.all_special_tokens tokenizer_vocab = tokenizer.get_vocab() tokenizer_len = len(tokenizer) + # The underlying tokenizer class could be MistralCommonBackend, + # which does not implement is_fast in Transformers + tokenizer_is_fast = getattr(tokenizer, "is_fast", True) + + # MistralCommonBackend is tekken-backed and needs byte-fallback-aware tokenization. + mistral_tekkenizer = None + if getattr(getattr(tokenizer, "tokenizer", None), "instruct_tokenizer", None): + from vllm.tokenizers.mistral import mistral_common_tekkenizer + + mistral_tekkenizer = mistral_common_tekkenizer(tokenizer) max_token_id = max(tokenizer_vocab.values()) max_chars_per_token = max(len(tok) for tok in tokenizer_vocab) @@ -145,6 +155,31 @@ def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: def max_chars_per_token(self) -> int: return max_chars_per_token + @property + def is_fast(self) -> bool: + return tokenizer_is_fast + + def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False): + if mistral_tekkenizer is not None: + from vllm.tokenizers.mistral import tekken_convert_ids_to_tokens + + return tekken_convert_ids_to_tokens(mistral_tekkenizer, ids) + return super().convert_ids_to_tokens( + ids, skip_special_tokens=skip_special_tokens + ) + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + if mistral_tekkenizer is not None: + from vllm.tokenizers.mistral import tekken_convert_tokens_to_string + + return tekken_convert_tokens_to_string(mistral_tekkenizer, tokens) + try: + return super().convert_tokens_to_string(tokens) + except NotImplementedError: + # The underlying tokenizer class could be MistralCommonBackend, + # which does not implement convert_tokens_to_string in Transformers + return "".join(tokens) + def get_vocab(self) -> dict[str, int]: return tokenizer_vocab diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index 8e29e1e5d6c..1164f7c41a7 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -188,6 +188,39 @@ def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int: return tokenizer.unk_id +def mistral_common_tekkenizer(tokenizer: object) -> "Tekkenizer | None": + """Return the underlying `Tekkenizer` for a `MistralCommonBackend`.""" + mistral = getattr(tokenizer, "tokenizer", None) + instruct = getattr(mistral, "instruct_tokenizer", None) + tekken = getattr(instruct, "tokenizer", None) + return tekken if isinstance(tekken, Tekkenizer) else None + + +def tekken_convert_ids_to_tokens( + tokenizer: "Tekkenizer", ids: Sequence[int] +) -> list[str | bytes]: + """Convert ids to pieces, using raw `bytes` for byte-fallback tokens.""" + tokens: list[str | bytes] = [tokenizer.id_to_piece(i) for i in ids] + if any("�" in t for t in tokens): + tokens = [ + tokenizer.id_to_byte_piece(i, SpecialTokenPolicy.KEEP) + if i >= tokenizer.num_special_tokens + else tokenizer.decode([i], SpecialTokenPolicy.KEEP) + for i in ids + ] + return tokens + + +def tekken_convert_tokens_to_string( + tokenizer: "Tekkenizer", tokens: Sequence[str | bytes] +) -> str: + """Reassemble pieces from `tekken_convert_ids_to_tokens` into text.""" + if any(isinstance(t, bytes) for t in tokens): + ids = [_tekken_token_to_id(tokenizer, t) for t in tokens] + return tokenizer.decode(ids, SpecialTokenPolicy.KEEP) + return "".join(cast(Sequence[str], tokens)) + + class MistralTokenizer(TokenizerLike): IS_MISTRAL_TOKENIZER = True # used by vllm.utils.mistral @@ -453,14 +486,8 @@ class MistralTokenizer(TokenizerLike): if (t in to_decode_special_tokens or t not in self._special_tokens_set) ] - if any(isinstance(t, bytes) for t in tokens): - # we need to encode and decode all tokens again - ids = [_tekken_token_to_id(self.tokenizer, t) for t in tokens] - # We filtered unwanted special tokens before - # so we can decode the rest. - decoded = self.tokenizer.decode(ids, SpecialTokenPolicy.KEEP) - else: - decoded = "".join(tokens) + # We filtered unwanted special tokens before so we can decode the rest. + decoded = tekken_convert_tokens_to_string(self.tokenizer, tokens) else: # make sure certain special tokens like Tool calls are # not decoded diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index f90e427aee0..cef2f7645fe 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -32,6 +32,7 @@ logger = init_logger(__name__) # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better) # - Fix tokenizer_class on the hub for the affected models (best) _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = { + "internlm2", "step3_vl", "step3p7", "unlimited-ocr", diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index a1c4cf1ffae..acb96e28f62 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -112,7 +112,7 @@ class ToolParser: @cached_property def vocab(self) -> dict[str, int]: - # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab + # NOTE: Only TokenizersBackend is guaranteed to have .vocab # whereas all tokenizers have .get_vocab() return self.model_tokenizer.get_vocab() diff --git a/vllm/tool_parsers/granite_tool_parser.py b/vllm/tool_parsers/granite_tool_parser.py index d586db32670..174e2884277 100644 --- a/vllm/tool_parsers/granite_tool_parser.py +++ b/vllm/tool_parsers/granite_tool_parser.py @@ -154,9 +154,11 @@ class GraniteToolParser(ToolParser): current_tool_call: dict = tool_call_arr[self.current_tool_id] delta = None - # case: we are starting a new tool in the array - # -> array has > 0 length AND length has moved past cursor - if len(tool_call_arr) > self.current_tool_id + 1: + # Only advance once the current tool name is streamed; granite + # emits arguments before name, so advancing early would drop it. + if len(tool_call_arr) > self.current_tool_id + 1 and ( + self.current_tool_id < 0 or self.current_tool_name_sent + ): # if we're moving on to a new call, first make sure we # haven't missed anything in the previous one that was # auto-generated due to JSON completions, but wasn't @@ -184,7 +186,7 @@ class GraniteToolParser(ToolParser): ) # re-set stuff pertaining to progress in the current tool - self.current_tool_id = len(tool_call_arr) - 1 + self.current_tool_id += 1 self.current_tool_name_sent = False self.streamed_args_for_tool.append("") logger.debug("starting on new tool %d", self.current_tool_id) diff --git a/vllm/tool_parsers/mistral_tool_parser.py b/vllm/tool_parsers/mistral_tool_parser.py index 1d605557b1f..026098a8735 100644 --- a/vllm/tool_parsers/mistral_tool_parser.py +++ b/vllm/tool_parsers/mistral_tool_parser.py @@ -533,6 +533,7 @@ class MistralToolParser(ToolParser): if prefix == "item" and event == "start_map": self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY + self.starting_new_tool = True if prefix == "item" and event == "map_key" and value == "name": self.streaming_state = StreamingState.PARSING_NAME if prefix == "item.name" and event == "string": @@ -640,18 +641,10 @@ class MistralToolParser(ToolParser): # Given the parsed text and the possible streaming state change, # let's add to the tool delta - if ( - (streaming_state_before_parse != self.streaming_state) - and streaming_state_before_parse - in [StreamingState.WAITING_FOR_TOOL_START, StreamingState.TOOL_COMPLETE] - and self.streaming_state - not in [ - StreamingState.ALL_TOOLS_COMPLETE, - StreamingState.TOOL_COMPLETE, - StreamingState.WAITING_FOR_TOOL_START, - ] - ): - # starting a new tool call + # start_map is the authoritative new-tool signal and survives + # batched deltas, unlike comparing pre/post streaming states + if self.starting_new_tool: + self.starting_new_tool = False if current_tool_call_modified: if self.current_tool_mistral_id is not None: current_tool_call.id = self.current_tool_mistral_id diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index a31420cf1cd..95769bafd7f 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -5,12 +5,14 @@ import ast import json import math import warnings +from dataclasses import dataclass from json import JSONDecodeError, JSONDecoder from typing import Any, TypeAlias import partial_json_parser from openai.types.responses import ( FunctionTool, + NamespaceTool, ToolChoiceFunction, ) from openai.types.responses.tool import Tool as ResponsesTool @@ -166,6 +168,91 @@ def consume_space(i: int, s: str) -> int: return i +_NAMESPACE_TOOL_SEPARATOR = "__" + + +@dataclass(frozen=True) +class ResponsesToolCallName: + name: str + namespace: str | None = None + + +def flat_namespace_tool_name(namespace: str, name: str) -> str: + return f"{namespace}{_NAMESPACE_TOOL_SEPARATOR}{name}" + + +def iter_response_function_tool_info( + tool: ResponsesTool, +) -> list[tuple[str, dict[str, Any] | None]]: + if isinstance(tool, FunctionTool): + return [(tool.name, tool.parameters)] + if not isinstance(tool, NamespaceTool): + return [] + + namespace = tool.name + return [ + ( + flat_namespace_tool_name(namespace, namespaced_tool.name), + namespaced_tool.parameters, + ) + for namespaced_tool in tool.tools + if namespaced_tool.type == "function" + ] + + +def iter_response_function_tool_dicts( + tools: list[ResponsesTool], +) -> list[dict[str, Any]]: + function_tools: list[dict[str, Any]] = [] + for tool in tools: + if isinstance(tool, NamespaceTool): + namespace = tool.name + for namespaced_tool in tool.tools: + if namespaced_tool.type != "function": + continue + tool_dict = namespaced_tool.model_dump() + tool_dict["name"] = flat_namespace_tool_name( + namespace, namespaced_tool.name + ) + function_tools.append(tool_dict) + else: + function_tools.append(tool.model_dump()) + return function_tools + + +def build_responses_tool_call_name_map( + tools: list[ResponsesTool] | None, +) -> dict[str, ResponsesToolCallName]: + if not tools: + return {} + + name_map: dict[str, ResponsesToolCallName] = {} + for tool in tools: + if not isinstance(tool, NamespaceTool): + continue + namespace = tool.name + for namespaced_tool in tool.tools: + if namespaced_tool.type != "function": + continue + flat_name = flat_namespace_tool_name(namespace, namespaced_tool.name) + name_map[flat_name] = ResponsesToolCallName( + name=namespaced_tool.name, + namespace=namespace, + ) + return name_map + + +def resolve_responses_tool_call_name( + name: str, + tools: list[ResponsesTool] | None = None, + tool_call_name_map: dict[str, ResponsesToolCallName] | None = None, +) -> ResponsesToolCallName: + name_map = tool_call_name_map + if name_map is None: + name_map = build_responses_tool_call_name_map(tools) + return name_map.get(name, ResponsesToolCallName(name=name)) + + def _is_function_tool(tool: Tool) -> bool: return isinstance(tool, (FunctionTool, ChatCompletionToolsParam)) @@ -189,6 +276,11 @@ def find_tool_properties( if not tools: return {} for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + for name, params in iter_response_function_tool_info(tool): + if name == tool_name: + return (params or {}).get("properties", {}) + continue if not _is_function_tool(tool): continue name, params = _extract_tool_info(tool) @@ -205,6 +297,11 @@ def find_tool_name( if not tools: return False for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + for name, _ in iter_response_function_tool_info(tool): + if name == tool_name: + return True + continue if not _is_function_tool(tool): continue name, _ = _extract_tool_info(tool) @@ -213,8 +310,9 @@ def find_tool_name( return False -def _get_tool_schema_from_tool(tool: Tool) -> dict: - name, params = _extract_tool_info(tool) +def _get_tool_schema_from_name_and_params( + name: str, params: dict[str, Any] | None +) -> dict: params = params if params else {"type": "object", "properties": {}} return { "properties": { @@ -225,6 +323,11 @@ def _get_tool_schema_from_tool(tool: Tool) -> dict: } +def _get_tool_schema_from_tool(tool: Tool) -> dict: + name, params = _extract_tool_info(tool) + return _get_tool_schema_from_name_and_params(name, params) + + def _get_tool_schema_defs( tools: list[Tool], ) -> dict: @@ -247,13 +350,25 @@ def _get_tool_schema_defs( def _get_json_schema_from_tools( tools: list[Tool], ) -> dict: - fn_tools = [t for t in tools if _is_function_tool(t)] + fn_tool_schemas: list[dict[str, Any]] = [] + fn_tools: list[Tool] = [] + for tool in tools: + if isinstance(tool, (FunctionTool, NamespaceTool)): + fn_tool_schemas.extend( + _get_tool_schema_from_name_and_params(name, params) + for name, params in iter_response_function_tool_info(tool) + ) + if isinstance(tool, FunctionTool): + fn_tools.append(tool) + elif _is_function_tool(tool): + fn_tool_schemas.append(_get_tool_schema_from_tool(tool)) + fn_tools.append(tool) json_schema = { "type": "array", "minItems": 1, "items": { "type": "object", - "anyOf": [_get_tool_schema_from_tool(tool) for tool in fn_tools], + "anyOf": fn_tool_schemas, }, } json_schema_defs = _get_tool_schema_defs(fn_tools) @@ -274,23 +389,30 @@ def get_json_schema_from_tools( tool_choice, ToolChoiceFunction ): tool_name = tool_choice.name - tool_map = {tool.name: tool for tool in tools if isinstance(tool, FunctionTool)} - if tool_name not in tool_map: + responses_tool_map: dict[str, dict[str, Any] | None] = {} + for tool in tools: + if not isinstance(tool, (FunctionTool, NamespaceTool)): + continue + for name, params in iter_response_function_tool_info(tool): + responses_tool_map[name] = params + if "__" in name: + responses_tool_map.setdefault(name.rsplit("__", 1)[1], params) + if tool_name not in responses_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") - return tool_map[tool_name].parameters + return responses_tool_map[tool_name] # tool_choice: Forced Function (ChatCompletion) if (not isinstance(tool_choice, str)) and isinstance( tool_choice, ChatCompletionNamedToolChoiceParam ): tool_name = tool_choice.function.name - tool_map = { + chat_tool_map: dict[str, ChatCompletionToolsParam] = { tool.function.name: tool for tool in tools if isinstance(tool, ChatCompletionToolsParam) } - if tool_name not in tool_map: + if tool_name not in chat_tool_map: raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.") - return tool_map[tool_name].function.parameters + return chat_tool_map[tool_name].function.parameters # tool_choice: "required" if tool_choice == "required": return _get_json_schema_from_tools(tools) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index ff1251665af..a848bbe142f 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -71,6 +71,7 @@ class LazyConfigDict(dict): _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( afmoe="AfmoeConfig", + arctic="ArcticConfig", bagel="BagelConfig", umm="CheersConfig", chatglm="ChatGLMConfig", @@ -106,6 +107,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( minimax_m3_vl="MiniMaxM3Config", minimax_m3_mtp="MiniMaxM3MTPConfig", moondream3="Moondream3Config", + moss_transcribe_diarize="MossTranscribeDiarizeConfig", eagle="EAGLEConfig", speculators="SpeculatorsConfig", nemotron="NemotronConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 808a8bf0774..f48b0dd6df6 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -16,6 +16,7 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { "AfmoeConfig": "vllm.transformers_utils.configs.afmoe", + "ArcticConfig": "vllm.transformers_utils.configs.arctic", "AXK1Config": "vllm.transformers_utils.configs.AXK1", "BagelConfig": "vllm.transformers_utils.configs.bagel", "CheersConfig": "vllm.transformers_utils.configs.cheers", @@ -60,6 +61,9 @@ _CLASS_TO_MODULE: dict[str, str] = { "Moondream3Config": "vllm.transformers_utils.configs.moondream3", "Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3", "Moondream3VisionConfig": "vllm.transformers_utils.configs.moondream3", + "MossTranscribeDiarizeConfig": ( + "vllm.transformers_utils.configs.moss_transcribe_diarize" + ), "MoonViTConfig": "vllm.transformers_utils.configs.moonvit", "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", @@ -92,6 +96,7 @@ _CLASS_TO_MODULE: dict[str, str] = { __all__ = [ "AfmoeConfig", + "ArcticConfig", "AXK1Config", "BagelConfig", "CheersConfig", @@ -134,6 +139,7 @@ __all__ = [ "Moondream3Config", "Moondream3TextConfig", "Moondream3VisionConfig", + "MossTranscribeDiarizeConfig", "MoonViTConfig", "KimiLinearConfig", "KimiVLConfig", diff --git a/vllm/transformers_utils/configs/moss_transcribe_diarize.py b/vllm/transformers_utils/configs/moss_transcribe_diarize.py new file mode 100644 index 00000000000..20abbd8feec --- /dev/null +++ b/vllm/transformers_utils/configs/moss_transcribe_diarize.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from transformers import PretrainedConfig, Qwen3Config +from transformers.models.whisper.configuration_whisper import WhisperConfig + + +class MossTranscribeDiarizeConfig(PretrainedConfig): + """Configuration for MOSS-Transcribe-Diarize.""" + + model_type = "moss_transcribe_diarize" + sub_configs = {"text_config": Qwen3Config, "audio_config": WhisperConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: dict[str, Any] | Qwen3Config | None = None, + audio_config: dict[str, Any] | WhisperConfig | None = None, + audio_token_id: int = 151671, + audio_merge_size: int = 4, + adaptor_input_dim: int | None = None, + tie_word_embeddings: bool = True, + **kwargs: Any, + ) -> None: + text_config_obj: Qwen3Config + if text_config is None: + text_config_obj = Qwen3Config( + vocab_size=151936, + hidden_size=1024, + intermediate_size=3072, + num_hidden_layers=28, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=128, + max_position_embeddings=40960, + tie_word_embeddings=tie_word_embeddings, + rope_theta=1_000_000.0, + layer_types=["full_attention"] * 28, + ) + elif isinstance(text_config, dict): + text_config_obj = Qwen3Config(**text_config) + else: + text_config_obj = text_config + + audio_config_obj: WhisperConfig + if audio_config is None: + audio_config_obj = WhisperConfig( + num_mel_bins=80, + d_model=1024, + encoder_layers=24, + encoder_attention_heads=16, + encoder_ffn_dim=4096, + max_source_positions=1500, + dropout=0.0, + attention_dropout=0.0, + activation_dropout=0.0, + activation_function="gelu", + encoder_layerdrop=0.0, + scale_embedding=False, + ) + elif isinstance(audio_config, dict): + audio_config_obj = WhisperConfig(**audio_config) + else: + audio_config_obj = audio_config + + text_config_obj.tie_word_embeddings = tie_word_embeddings + if not getattr(text_config_obj, "layer_types", None): + text_config_obj.layer_types = [ + "full_attention" + ] * text_config_obj.num_hidden_layers + + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + self.text_config = text_config_obj + self.audio_config = audio_config_obj + self.audio_token_id = int(audio_token_id) + self.audio_merge_size = int(audio_merge_size) + self.adaptor_input_dim = ( + int(adaptor_input_dim) + if adaptor_input_dim is not None + else int(audio_config_obj.d_model) * int(audio_merge_size) + ) + + self.vocab_size = int(text_config_obj.vocab_size) + self.hidden_size = int(text_config_obj.hidden_size) + self.intermediate_size = int(text_config_obj.intermediate_size) + self.num_hidden_layers = int(text_config_obj.num_hidden_layers) + self.num_attention_heads = int(text_config_obj.num_attention_heads) + self.num_key_value_heads = int(text_config_obj.num_key_value_heads) + self.head_dim = int(text_config_obj.head_dim) + self.hidden_act = text_config_obj.hidden_act + self.max_position_embeddings = int(text_config_obj.max_position_embeddings) + self.rms_norm_eps = float(text_config_obj.rms_norm_eps) + rope_parameters = getattr(text_config_obj, "rope_parameters", None) + rope_theta = float(getattr(text_config_obj, "rope_theta", 1_000_000.0)) + if rope_parameters is None: + rope_parameters = { + "rope_type": "default", + "rope_theta": rope_theta, + } + text_config_obj.rope_parameters = rope_parameters + self.rope_parameters = rope_parameters + self.rope_theta = float(rope_parameters.get("rope_theta", rope_theta)) + self.attention_bias = bool(text_config_obj.attention_bias) + self.attention_dropout = float(text_config_obj.attention_dropout) + self.is_causal = True diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index f4b0f73f62c..e034cec9745 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -120,6 +120,10 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: "mask_token_id": config_dict["mask_token_id"], "target_layer_ids": [i - 1 for i in aux_layer_ids], } + # Enable causal masking in SWA for vllm-project/speculators models + pre_trained_config["dflash_config"]["causal"] = not config_dict.get( + "sliding_window_non_causal", True + ) @register_speculator("dspark") diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 3b39c911095..7ced083afb5 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -271,6 +271,7 @@ class ModelArchConfigConvertorBase: "pangu_ultra_moe", "pangu_ultra_moe_mtp", "bailing_hybrid", + "bailing_hybrid_mtp", ): # check is deepseek_v4 model if hasattr(self.hf_text_config, "compress_ratios"): @@ -539,6 +540,11 @@ class Qwen3NextMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) +class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): + def get_num_hidden_layers(self) -> int: + return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) + + class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_num_hidden_layers(self) -> int: return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0) @@ -633,6 +639,7 @@ class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase): # hf_config.model_type -> convertor class MODEL_ARCH_CONFIG_CONVERTORS = { + "bailing_hybrid_mtp": BailingHybridMTPModelArchConfigConvertor, "cohere_asr": CohereAsrModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index fa4c558a739..aa33faed916 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -81,6 +81,7 @@ _transformers_v4_compatibility_import() _transformers_v4_compatibility_init() _P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin) +_I = TypeVar("_I", bound=BaseImageProcessor, default=BaseImageProcessor) _V = TypeVar("_V", bound=BaseVideoProcessor, default=BaseVideoProcessor) @@ -440,12 +441,14 @@ def get_image_processor( *args: Any, revision: str | None = None, trust_remote_code: bool = False, + processor_cls_overrides: type[_I] | None = None, **kwargs: Any, ): """Load an image processor for the given model name via HuggingFace.""" try: processor_name = convert_model_repo_to_path(processor_name) - processor = AutoImageProcessor.from_pretrained( + processor_cls = processor_cls_overrides or AutoImageProcessor + processor = processor_cls.from_pretrained( processor_name, *args, revision=revision, diff --git a/vllm/transformers_utils/processors/fireredlid.py b/vllm/transformers_utils/processors/fireredlid.py index cb041397d03..3afdca4c7bb 100644 --- a/vllm/transformers_utils/processors/fireredlid.py +++ b/vllm/transformers_utils/processors/fireredlid.py @@ -232,7 +232,7 @@ class FireRedLIDProcessor(ProcessorMixin): """ feature_extractor_class = "FireRedLIDFeatureExtractor" - tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast") + tokenizer_class = ("PythonBackend", "TokenizersBackend") def __init__(self, feature_extractor, tokenizer): super().__init__(feature_extractor, tokenizer) diff --git a/vllm/transformers_utils/processors/glm4v.py b/vllm/transformers_utils/processors/glm4v.py index 3ecb1bae531..a8da2395525 100644 --- a/vllm/transformers_utils/processors/glm4v.py +++ b/vllm/transformers_utils/processors/glm4v.py @@ -3,7 +3,7 @@ # Adapted from # https://github.com/zai-org/CogAgent -from transformers import PreTrainedTokenizer +from transformers import PythonBackend from transformers.image_processing_utils_fast import BaseImageProcessorFast from transformers.image_utils import PILImageResampling from transformers.processing_utils import ProcessorMixin @@ -30,7 +30,7 @@ class GLM4VProcessor(ProcessorMixin): def __init__( self, image_processor: GLM4VImageProcessorFast, - tokenizer: PreTrainedTokenizer, + tokenizer: PythonBackend, ) -> None: self.image_processor = image_processor self.tokenizer = tokenizer diff --git a/vllm/transformers_utils/processors/hunyuan_vl.py b/vllm/transformers_utils/processors/hunyuan_vl.py deleted file mode 100644 index 2d0e4db97a6..00000000000 --- a/vllm/transformers_utils/processors/hunyuan_vl.py +++ /dev/null @@ -1,226 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# adapted from https://github.com/ManaEstras/transformers/blob/v4.57.1.hyvl/src/transformers/models/hunyuan_vl/processing_hunyuan_vl.py - -import numpy as np -import torch -from transformers.feature_extraction_utils import BatchFeature -from transformers.image_utils import ImageInput -from transformers.processing_utils import ProcessorMixin -from transformers.tokenization_utils_base import PreTokenizedInput, TextInput -from transformers.video_utils import VideoInput - - -class HunYuanVLProcessor(ProcessorMixin): - attributes = ["image_processor", "tokenizer"] - valid_kwargs = ["chat_template"] - image_processor_class = "AutoImageProcessor" - tokenizer_class = "AutoTokenizer" # ("AutoTokenizer", None) - - def __init__( - self, - image_processor=None, - tokenizer=None, - chat_template=None, - **kwargs, - ): - # TODO Fix the init - self.tokenizer = tokenizer - self.image_token_id = 120120 # self.tokenizer.image_token_id - self.image_token = self.tokenizer.convert_ids_to_tokens(self.image_token_id) - self.im_start_token_id = 120118 # self.tokenizer.im_start_id - self.im_start_token = self.tokenizer.convert_ids_to_tokens( - self.im_start_token_id - ) - self.im_end_token_id = 120119 # self.tokenizer.im_end_id - self.im_end_token = self.tokenizer.convert_ids_to_tokens(self.im_end_token_id) - self.placeholder_token = self.tokenizer.convert_ids_to_tokens( - self.tokenizer.vocab_size - 1 - ) - self.pad_id = 120002 # self.tokenizer.pad_token_id - - super().__init__(image_processor, tokenizer, chat_template=chat_template) - - def __call__( - self, - images: ImageInput = None, - text: TextInput - | PreTokenizedInput - | list[TextInput] - | list[PreTokenizedInput] = None, - videos: VideoInput = None, - **kwargs, - ) -> BatchFeature: - image_inputs = {} - if images is not None: - image_inputs = self.image_processor(images=images) - image_grid_thw = image_inputs["image_grid_thw"] - - if not isinstance(text, list): - text = [text] - - text = text.copy() # below lines change text in-place - - image_tokens_cumsum = [0] - if images is not None: - index = 0 - for i in range(len(text)): - while self.image_token in text[i]: - grid_h, grid_w = image_grid_thw[index][-2:] - patch_h = grid_h // self.image_processor.merge_size - patch_w = grid_w // self.image_processor.merge_size - num_image_tokens = patch_h * (patch_w + 1) + 2 - image_tokens_cumsum.append( - image_tokens_cumsum[-1] + num_image_tokens - ) - # text[i] = text[i].replace(self.image_token, self.im_start_token + self.placeholder_token * num_image_tokens + self.im_end_token, 1) # noqa: E501 - text[i] = text[i].replace( - self.image_token, self.placeholder_token * num_image_tokens, 1 - ) - index += 1 - text[i] = text[i].replace(self.placeholder_token, self.image_token) - # text[i] = self.tokenizer.bos_token + text[i] - - text_inputs = self.tokenizer(text, add_special_tokens=False, **kwargs) - self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) - - input_ids = text_inputs["input_ids"] - position_ids = torch.arange(len(input_ids[0])) - position_ids_w = torch.arange(len(input_ids[0])) - position_ids_h = torch.arange(len(input_ids[0])) - position_ids_t = torch.arange(len(input_ids[0])) - - if images is not None: - image_token_pos_indices = torch.where(input_ids[0] == self.image_token_id)[ - 0 - ] - for i in range(len(image_grid_thw)): - grid_h, grid_w = image_grid_thw[i][-2:] - patch_h = grid_h // self.image_processor.merge_size - patch_w = grid_w // self.image_processor.merge_size - start_pos = image_token_pos_indices[image_tokens_cumsum[i]].item() + 1 - replace_num = (patch_w + 1) * patch_h - position_ids_w[start_pos : start_pos + replace_num] = torch.tensor( - list(range(patch_w + 1)) * patch_h, dtype=torch.int64 - ) - patch_h_list = [] - for h in range(patch_h): - patch_h_list += [h] * (patch_w + 1) - position_ids_h[start_pos : start_pos + replace_num] = torch.tensor( - patch_h_list, dtype=torch.int64 - ) - position_ids_t[start_pos : start_pos + replace_num] = 0 - - position_ids = torch.stack( - [position_ids, position_ids_w, position_ids_h, position_ids_t] - ).unsqueeze(0) - text_inputs["position_ids"] = position_ids - - attention_mask = input_ids.ne(self.pad_id) - text_inputs["attention_mask"] = attention_mask - text_inputs["imgs_pos"] = [self.get_imgs_pos(e) for e in input_ids] - # image_inputs["imgs"] = [[image_inputs["pixel_values"]]] - - return_tensors = kwargs.pop("return_tensors", None) - return BatchFeature( - data={**text_inputs, **image_inputs}, - tensor_type=return_tensors, - ) - - def batch_decode(self, *args, **kwargs): - return self.tokenizer.batch_decode(*args, **kwargs) - - def decode(self, *args, **kwargs): - return self.tokenizer.decode(*args, **kwargs) - - def post_process_image_text_to_text( - self, - generated_outputs, - skip_special_tokens=True, - clean_up_tokenization_spaces=False, - **kwargs, - ): - assert 0 - - def apply_chat_template(self, *args, **kwargs): - kwargs["return_dict"] = False - return self.tokenizer.apply_chat_template(*args, **kwargs) - - def get_imgs_pos(self, doc_ids): - doc_ids = np.array(doc_ids, dtype=np.int64) - img_begin_index = np.where(doc_ids == self.im_start_token_id)[0] - img_end_index = np.where(doc_ids == self.im_end_token_id)[0] - imgs_pos = np.concatenate( - ( - np.reshape(img_begin_index + 1, (-1, 1)), - np.reshape(img_end_index, (-1, 1)), - ), - axis=-1, - ).tolist() - return imgs_pos - - @property - def model_input_names(self): - tokenizer_input_names = self.tokenizer.model_input_names - image_processor_input_names = self.image_processor.model_input_names - return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) - - -def split_image_into_patch_blocks( - pixel_values: torch.Tensor, # shape: [batch_size, 3, H, W] - patch_size: int = 16, # e.g. 16 - adaptor_patch_div: int = 4, # e.g. 4 --> each patch_size is cut into 4x4 small regions, i.e. patch_size // 4 # noqa: E501 -) -> torch.Tensor: - """ - Split the input image tensor (supporting batch) into large patches of size `patch_size`, - and then further divide each large patch into smaller regions of size - (patch_size // adaptor_patch_div) x (patch_size // adaptor_patch_div). - Each small region is extracted as a tensor of shape [3, patch_size, patch_size]. - The final output contains all such small region tensors. - - Args: - pixel_values: Input image tensor of shape [batch_size, 3, H, W]. - patch_size: Size of the large patch, e.g., 16. - adaptor_patch_div: Each large patch is divided into - (patch_size // adaptor_patch_div) x (patch_size // adaptor_patch_div) - smaller regions. - - Returns: - patches: A tensor of shape [N, 3, patch_size, patch_size], - where N = batch_size * (H // patch_size) * (W // patch_size) * (patch_size // adaptor_patch_div)^2. - Each element in the batch corresponds to one small image region. - """ # noqa: E501 - batch_size, channels, height, width = pixel_values.shape - assert channels == 3, "Pixel values must have 3 channels in dim=1" - assert height % patch_size == 0 and width % patch_size == 0, ( - "H and W must be divisible by patch_size" - ) - - patch_height_num = height // patch_size - patch_width_num = width // patch_size - - # Reshape to [B, 3, ph, ps, pw, ps] - img = pixel_values.reshape( - batch_size, 3, patch_height_num, patch_size, patch_width_num, patch_size - ) - - # Further split each psxps patch into (ps//aps)x(ps//aps) small regions - img = img.reshape( - batch_size, - 3, - patch_height_num, - patch_size // adaptor_patch_div, # ps // aps - adaptor_patch_div, - patch_width_num, - patch_size // adaptor_patch_div, # ps // aps - adaptor_patch_div, - ) - - # Permute to group the small regions: [B, ph, pw, ps//aps, ps//aps, 3, aps, aps] - img = img.permute(0, 2, 5, 3, 6, 1, 4, 7) - - # Reshape into [B * ph * pw * (ps//aps)^2, 3, patch_size, patch_size] - patches = img.reshape(-1, 3, patch_size, patch_size) - - return patches diff --git a/vllm/transformers_utils/processors/hunyuan_vl_image.py b/vllm/transformers_utils/processors/hunyuan_vl_image.py deleted file mode 100644 index 0b10ae249db..00000000000 --- a/vllm/transformers_utils/processors/hunyuan_vl_image.py +++ /dev/null @@ -1,477 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# adapted from https://github.com/ManaEstras/transformers/blob/v4.57.1.hyvl/src/transformers/models/hunyuan_vl/image_processing_hunyuan_vl.py -"""Image processor class for HunYuanVL.""" - -# isort conflicts with ruff for transformers imports -# isort: skip_file -import math - -import numpy as np -import torchvision.transforms as transforms -from transformers import AutoImageProcessor -from transformers.image_processing_utils import BaseImageProcessor, BatchFeature -from transformers.image_transforms import ( - convert_to_rgb, -) -from transformers.image_utils import ( - OPENAI_CLIP_MEAN, - OPENAI_CLIP_STD, - ChannelDimension, - ImageInput, - PILImageResampling, - make_flat_list_of_images, - make_list_of_images, - valid_images, - validate_preprocess_arguments, -) -from transformers.utils import TensorType, logging -from transformers.video_utils import VideoInput, make_batched_videos - -logger = logging.get_logger(__name__) - - -def smart_resize( - height: int, - width: int, - factor: int = 16, - min_pixels: int = 512 * 512, - max_pixels: int = 2048 * 2048, -): - """Rescales the image so that the following conditions are met: - - 1. Both dimensions (height and width) are divisible by 'factor'. - - 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. - - 3. The aspect ratio of the image is maintained as closely as possible. - - """ - if max(height, width) / min(height, width) > 200: - raise ValueError( - "absolute aspect ratio must be smaller than 200, got " - f"{max(height, width) / min(height, width)}" - ) - h_bar = round(height / factor) * factor - w_bar = round(width / factor) * factor - if h_bar * w_bar > max_pixels: - beta = math.sqrt((height * width) / max_pixels) - h_bar = max(factor, math.floor(height / beta / factor) * factor) - w_bar = max(factor, math.floor(width / beta / factor) * factor) - elif h_bar * w_bar < min_pixels: - beta = math.sqrt(min_pixels / (height * width)) - h_bar = math.ceil(height * beta / factor) * factor - w_bar = math.ceil(width * beta / factor) * factor - return h_bar, w_bar - - -class HunYuanVLImageProcessor(BaseImageProcessor): - model_input_names = [ - "pixel_values", - "image_grid_thw", - "pixel_values_videos", - "video_grid_thw", - ] - - def __init__( - self, - do_resize: bool = True, - size: dict[str, int] | None = None, - resample: PILImageResampling = PILImageResampling.BICUBIC, - do_rescale: bool = True, - rescale_factor: int | float = 1 / 255, - do_normalize: bool = True, - image_mean: float | list[float] | None = None, - image_std: float | list[float] | None = None, - do_convert_rgb: bool = True, - min_pixels: int | None = None, - max_pixels: int | None = None, - patch_size: int = 16, - temporal_patch_size: int = 2, - merge_size: int = 2, - **kwargs, - ) -> None: - super().__init__(**kwargs) - if size is not None and ( - "shortest_edge" not in size or "longest_edge" not in size - ): - raise ValueError( - "size must contain 'shortest_edge' and 'longest_edge' keys." - ) - else: - size = {"shortest_edge": 512 * 512, "longest_edge": 2048 * 2048} - # backward compatibility: override size with min_pixels and max_pixels - # if they are provided. - if min_pixels is not None: - size["shortest_edge"] = min_pixels - if max_pixels is not None: - size["longest_edge"] = max_pixels - self.min_pixels = size["shortest_edge"] - self.max_pixels = size["longest_edge"] - self.size = size - - self.do_resize = do_resize - self.resample = resample - self.do_rescale = do_rescale - self.rescale_factor = rescale_factor - self.do_normalize = do_normalize - self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN - self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD - - self.patch_size = patch_size - self.temporal_patch_size = temporal_patch_size - self.merge_size = merge_size - self.do_convert_rgb = do_convert_rgb - - # hard-code - - def _preprocess( - self, - images: ImageInput | VideoInput, - do_resize: bool | None = None, - size: dict[str, int] | None = None, - resample: PILImageResampling = None, - do_rescale: bool | None = None, - rescale_factor: float | None = None, - do_normalize: bool | None = None, - image_mean: float | list[float] | None = None, - image_std: float | list[float] | None = None, - patch_size: int = 16, - temporal_patch_size: int = 2, - merge_size: int = 2, - do_convert_rgb: bool | None = None, - data_format: ChannelDimension | None = ChannelDimension.FIRST, - input_data_format: str | ChannelDimension | None = None, - ): - """ - Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`. - - Args: - images (`ImageInput`): - Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`. - do_resize (`bool`, *optional*, defaults to `self.do_resize`): - Whether to resize the image. - size (`dict[str, int]`, *optional*, defaults to `self.size`): - Size of the image after resizing. `shortest_edge` and `longest_edge` keys must be present. - resample (`PILImageResampling`, *optional*, defaults to `self.resample`): - Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums. - do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): - Whether to rescale the image. - rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): - Scale factor to use if rescaling the image. - do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): - Whether to normalize the image. - image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`): - Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image. - image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`): - Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image. - patch_size (`int`, *optional*, defaults to `self.patch_size`): - The spatial patch size of the vision encoder. - temporal_patch_size (`int`, *optional*, defaults to `self.temporal_patch_size`): - The temporal patch size of the vision encoder. - merge_size (`int`, *optional*, defaults to `self.merge_size`): - The merge size of the vision encoder to llm encoder. - do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): - Whether to convert the image to RGB. - data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`): - The channel dimension format for the output image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - Unset: Use the channel dimension format of the input image. - input_data_format (`ChannelDimension` or `str`, *optional*): - The channel dimension format for the input image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - """ # noqa: E501 - images = make_list_of_images(images) - - if do_convert_rgb: - images = [convert_to_rgb(image) for image in images] - - width, height = images[0].width, images[0].height - resized_width, resized_height = width, height - processed_images = [] - for image in images: - if do_resize: - resized_height, resized_width = smart_resize( - height=height, - width=width, - factor=patch_size * merge_size, - min_pixels=self.min_pixels, - max_pixels=self.max_pixels, - ) - image = image.resize((resized_width, resized_height)) - - if do_normalize: - image = transforms.Compose( - [ - transforms.ToTensor(), - transforms.Normalize(self.image_mean, self.image_std), - ] - )(image) - processed_images.append(image) - - patches = np.array(processed_images) - channel = patches.shape[1] - grid_t = patches.shape[0] // temporal_patch_size - grid_h, grid_w = resized_height // patch_size, resized_width // patch_size - patches = patches.reshape( - 1, - channel, - grid_h // merge_size, - merge_size, - patch_size, - grid_w // merge_size, - merge_size, - patch_size, - ) - patches = patches.transpose(0, 2, 3, 5, 6, 1, 4, 7) - flatten_patches = patches.reshape( - 1 * grid_h * grid_w, channel * patch_size * patch_size - ) - - return flatten_patches, (grid_t, grid_h, grid_w) - - def preprocess( - self, - images: ImageInput, - videos: VideoInput = None, - do_resize: bool | None = None, - size: dict[str, int] | None = None, - min_pixels: int | None = None, - max_pixels: int | None = None, - resample: PILImageResampling = None, - do_rescale: bool | None = None, - rescale_factor: float | None = None, - do_normalize: bool | None = None, - image_mean: float | list[float] | None = None, - image_std: float | list[float] | None = None, - patch_size: int | None = None, - temporal_patch_size: int | None = None, - merge_size: int | None = None, - do_convert_rgb: bool | None = None, - return_tensors: str | TensorType | None = None, - data_format: ChannelDimension | None = ChannelDimension.FIRST, - input_data_format: str | ChannelDimension | None = None, - ): - """ - Args: - images (`ImageInput`): - Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If - passing in images with pixel values between 0 and 1, set `do_rescale=False`. - videos (`VideoInput`): - Video to preprocess. Expects a single or batch of videos with pixel values ranging from 0 to 255. If - passing in videos with pixel values between 0 and 1, set `do_rescale=False`. - do_resize (`bool`, *optional*, defaults to `self.do_resize`): - Whether to resize the image. - size (`dict[str, int]`, *optional*, defaults to `self.size`): - Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with - the longest edge resized to keep the input aspect ratio. - resample (`int`, *optional*, defaults to `self.resample`): - Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only - has an effect if `do_resize` is set to `True`. - do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): - Whether to rescale the image. - rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): - Rescale factor to rescale the image by if `do_rescale` is set to `True`. - do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): - Whether to normalize the image. - image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`): - Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`. - image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`): - Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to - `True`. - min_pixels (`int`, *optional*, defaults to `self.min_pixels`): - The min pixels of the image to resize the image. - max_pixels (`int`, *optional*, defaults to `self.max_pixels`): - The max pixels of the image to resize the image. - patch_size (`int`, *optional*, defaults to `self.patch_size`): - The spatial patch size of the vision encoder. - temporal_patch_size (`int`, *optional*, defaults to `self.temporal_patch_size`): - The temporal patch size of the vision encoder. - merge_size (`int`, *optional*, defaults to `self.merge_size`): - The merge size of the vision encoder to llm encoder. - do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): - Whether to convert the image to RGB. - return_tensors (`str` or `TensorType`, *optional*): - The type of tensors to return. Can be one of: - - Unset: Return a list of `np.ndarray`. - - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`. - - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`. - - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`. - - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`. - data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`): - The channel dimension format for the output image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - Unset: Use the channel dimension format of the input image. - input_data_format (`ChannelDimension` or `str`, *optional*): - The channel dimension format for the input image. If unset, the channel dimension format is inferred - from the input image. Can be one of: - - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. - - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. - - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - - """ # noqa: E501 - min_pixels = min_pixels if min_pixels is not None else self.min_pixels - max_pixels = max_pixels if max_pixels is not None else self.max_pixels - - if size is not None: - if "shortest_edge" not in size or "longest_edge" not in size: - raise ValueError( - "size must contain 'shortest_edge' and 'longest_edge' keys." - ) - min_pixels = size["shortest_edge"] - elif min_pixels is not None and max_pixels is not None: - # backward compatibility: override size with min_pixels and max_pixels - # if they are provided. - size = {"shortest_edge": min_pixels, "longest_edge": max_pixels} - else: - size = {**self.size} - - do_resize = do_resize if do_resize is not None else self.do_resize - - resample = resample if resample is not None else self.resample - do_rescale = do_rescale if do_rescale is not None else self.do_rescale - rescale_factor = ( - rescale_factor if rescale_factor is not None else self.rescale_factor - ) - do_normalize = do_normalize if do_normalize is not None else self.do_normalize - image_mean = image_mean if image_mean is not None else self.image_mean - image_std = image_std if image_std is not None else self.image_std - patch_size = patch_size if patch_size is not None else self.patch_size - temporal_patch_size = ( - temporal_patch_size - if temporal_patch_size is not None - else self.temporal_patch_size - ) - merge_size = merge_size if merge_size is not None else self.merge_size - do_convert_rgb = ( - do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb - ) - - if images is not None: - images = make_flat_list_of_images(images) - - if images is not None and not valid_images(images): - raise ValueError( - "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " - "torch.Tensor, tf.Tensor or jax.ndarray." - ) - - validate_preprocess_arguments( - rescale_factor=rescale_factor, - do_normalize=do_normalize, - image_mean=image_mean, - image_std=image_std, - do_resize=do_resize, - size=size, - resample=resample, - ) - - data = {} - if images is not None: - pixel_values, vision_grid_thws = [], [] - for image in images: - patches, image_grid_thw = self._preprocess( - image, - do_resize=do_resize, - size=size, - resample=resample, - do_rescale=do_rescale, - rescale_factor=rescale_factor, - do_normalize=do_normalize, - image_mean=image_mean, - image_std=image_std, - patch_size=patch_size, - temporal_patch_size=temporal_patch_size, - merge_size=merge_size, - data_format=data_format, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - ) - pixel_values.extend(patches) - vision_grid_thws.append(image_grid_thw) - pixel_values = np.array(pixel_values) - vision_grid_thws = np.array(vision_grid_thws) - data.update( - {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws} - ) - - # kept for BC only and should be removed after v5.0 - if videos is not None: - logger.warning( - "`HunYuanVLV1ImageProcessor` works only with image inputs " - "and doesn't process videos anymore. " - "This is a deprecated behavior and will be removed in v5.0. " - "Your videos should be forwarded to `HunYuanVLV1VideoProcessor`. " - ) - videos = make_batched_videos(videos) - pixel_values_videos, vision_grid_thws_videos = [], [] - for images in videos: - patches, video_grid_thw = self._preprocess( - images, - do_resize=do_resize, - size=size, - resample=resample, - do_rescale=do_rescale, - rescale_factor=rescale_factor, - do_normalize=do_normalize, - image_mean=image_mean, - image_std=image_std, - patch_size=patch_size, - temporal_patch_size=temporal_patch_size, - merge_size=merge_size, - data_format=data_format, - do_convert_rgb=do_convert_rgb, - input_data_format=input_data_format, - ) - pixel_values_videos.extend(patches) - vision_grid_thws_videos.append(video_grid_thw) - data.update( - { - "pixel_values_videos": np.array(pixel_values_videos), - "video_grid_thw": np.array(vision_grid_thws_videos), - } - ) - - return BatchFeature(data=data, tensor_type=return_tensors) - - def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None): - """ - A utility that returns number of image patches for a given image size. - - Args: - height (`int`): - Height of the input image. - width (`int`): - Width of the input image. - images_kwargs (`dict`, *optional*): - Any kwargs to override defaults of the image processor. - Returns: - `int`: Number of image patches per image. - """ - min_pixels = ( - images_kwargs["min_pixels"] - if "min_pixels" in images_kwargs - else self.size["shortest_edge"] - ) - max_pixels = ( - images_kwargs["max_pixels"] - if "max_pixels" in images_kwargs - else self.size["longest_edge"] - ) - patch_size = images_kwargs.get("patch_size", self.patch_size) - merge_size = images_kwargs.get("merge_size", self.merge_size) - - factor = patch_size * merge_size - resized_height, resized_width = smart_resize( - height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels - ) - grid_h, grid_w = resized_height // patch_size, resized_width // patch_size - return grid_h * (grid_w + 1) + 2 - - -AutoImageProcessor.register("HunYuanVLImageProcessor", HunYuanVLImageProcessor) diff --git a/vllm/transformers_utils/processors/internvl.py b/vllm/transformers_utils/processors/internvl.py index fc582deef97..22e3f5be98a 100644 --- a/vllm/transformers_utils/processors/internvl.py +++ b/vllm/transformers_utils/processors/internvl.py @@ -12,7 +12,12 @@ import numpy.typing as npt import torch import torchvision.transforms as T from PIL import Image -from transformers import BatchFeature, TensorType +from transformers import ( + BaseVideoProcessor, + BatchFeature, + ImageProcessingMixin, + TensorType, +) from transformers.processing_utils import ProcessorMixin from vllm.multimodal.image import convert_image_mode @@ -215,7 +220,7 @@ def video_to_pixel_values_internvl( return pixel_values -class InternVLImageProcessor: +class InternVLImageProcessor(ImageProcessingMixin): def __init__( self, image_size: int, @@ -312,7 +317,7 @@ class InternVLImageProcessor: return BatchFeature(image_inputs, tensor_type=return_tensors) -class InternVLVideoProcessor: +class InternVLVideoProcessor(BaseVideoProcessor): def __init__( self, image_size: int, diff --git a/vllm/transformers_utils/processors/isaac.py b/vllm/transformers_utils/processors/isaac.py index da548a8f1e0..2d791df0167 100644 --- a/vllm/transformers_utils/processors/isaac.py +++ b/vllm/transformers_utils/processors/isaac.py @@ -7,7 +7,12 @@ import numpy as np import torch import torch.nn.functional as F from PIL import Image -from transformers import BatchFeature, ProcessorMixin, TensorType +from transformers import ( + BatchFeature, + ImageProcessingMixin, + ProcessorMixin, + TensorType, +) from transformers.processing_utils import ProcessingKwargs from typing_extensions import Unpack @@ -322,7 +327,7 @@ class IsaacProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call- } -class IsaacImageProcessor: +class IsaacImageProcessor(ImageProcessingMixin): model_input_names = ["pixel_values", "image_grid_thw"] def __init__( diff --git a/vllm/transformers_utils/processors/kimi_k25_vision_fused.py b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py new file mode 100644 index 00000000000..63907898175 --- /dev/null +++ b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks.""" + +import io +import json +import math +from typing import Any + +import numpy as np +import pybase64 as base64 +import torch +from PIL import Image +from transformers.image_processing_utils import BaseImageProcessor, BatchFeature +from transformers.utils import TensorType + +from vllm.utils.import_utils import is_numba_available +from vllm.utils.jit_monitor import numba_workqueue_threading_layer + +if is_numba_available(): + from numba import njit, prange + + @njit(parallel=True, cache=True) + def _write_fused_patches( + frames: np.ndarray, + out: np.ndarray, + out_offset: int, + new_h: int, + new_w: int, + padded_h: int, + padded_w: int, + patch_size: int, + normalize_lut: np.ndarray, + ) -> None: + # frames: [T, new_h, new_w, 3] uint8, without padding. + # out: [total_patches, 3, patch_size, patch_size] float32. + t_size = frames.shape[0] + patch_h = padded_h // patch_size + patch_w = padded_w // patch_size + total = t_size * padded_h * padded_w * 3 + hwc = padded_h * padded_w * 3 + wc = padded_w * 3 + + for linear in prange(total): + t = linear // hwc + rem = linear - t * hwc + y = rem // wc + rem = rem - y * wc + x = rem // 3 + c = rem - x * 3 + + value = frames[t, y, x, c] if y < new_h and x < new_w else 0 + + patch_idx = ( + out_offset + + t * patch_h * patch_w + + (y // patch_size) * patch_w + + (x // patch_size) + ) + out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c] + +else: + + def _write_fused_patches(*args: Any, **kwargs: Any) -> None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +) -> dict[str, int]: + s1 = math.sqrt( + in_patch_limit + / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) + ) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size) + new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size) + + factor = merge_kernel_size * patch_size + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + num_tokens = token_height * token_width + + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def navit_resize_video( + width: int, + height: int, + nframes: int, + avg_fps: float, + sample_fps: float, + patch_size: int, + merge_kernel_size: int, + in_patch_limit_each_frame: int, + patch_limit_on_one_side: int, + in_patch_limit_total: int | None, + max_num_frames_each_video: int | None, + fixed_output_tokens_each_frame: int | None, +) -> dict[str, int]: + sample_fps = min(sample_fps, avg_fps) + sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1) + if max_num_frames_each_video is not None: + sampled_nframes = min(sampled_nframes, max_num_frames_each_video) + + if in_patch_limit_total is not None: + in_patch_limit_each_frame = min( + round(in_patch_limit_total / sampled_nframes), + in_patch_limit_each_frame, + ) + + ret = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit_each_frame, + patch_limit_on_one_side, + fixed_output_tokens_each_frame, + ) + ret["sampled_nframes"] = sampled_nframes + return ret + + +def _to_pil(data: Any) -> Image.Image: + if hasattr(data, "media") and hasattr(data, "original_bytes"): + data = data.media + if isinstance(data, Image.Image): + return data if data.mode == "RGB" else data.convert("RGB") + if isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",", 1)[1] + return Image.open(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB") + return Image.open(data).convert("RGB") + if isinstance(data, bytes): + return Image.open(io.BytesIO(data)).convert("RGB") + raise ValueError(f"Unsupported data type: {type(data)}") + + +def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]: + if media["type"] == "image": + media["image"] = _to_pil(media["image"]) + return media + if media["type"] == "video_chunk": + media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]] + return media + raise ValueError(f"Unsupported media type: {media['type']}") + + +class KimiK25FusedVisionProcessor(BaseImageProcessor): + model_type = "kimi_k25" + + def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None: + super().__init__(**kwargs) + media_proc_cfg = dict(media_proc_cfg) + merge_kernel_size = media_proc_cfg["merge_kernel_size"] + if isinstance(merge_kernel_size, (list, tuple)): + media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0]) + self.media_proc_cfg = media_proc_cfg + self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"] + values = np.arange(256, dtype=np.float32)[:, None] + image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32) + image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32) + self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[ + None, : + ] + + def media_tokens_calculator(self, media: dict[str, Any]) -> int: + media = _ensure_media_type(media) + ret = self.get_resize_config(media) + return ret["num_tokens"] + + def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]: + if media_input["type"] == "image": + width, height = media_input["image"].size + return navit_resize_image( + width, + height, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + self.media_proc_cfg["in_patch_limit"], + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["fixed_output_tokens"], + ) + + if media_input["type"] == "video_chunk": + frame = media_input["video_chunk"][0] + width, height = frame.size + num_frames = len(media_input["video_chunk"]) + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"] + if in_patch_limit_each_frame is None: + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"] + + return navit_resize_video( + width, + height, + num_frames, + 1.0, + math.inf, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + in_patch_limit_each_frame, + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["in_patch_limit_video"], + None, + self.media_proc_cfg["fixed_output_tokens"], + ) + + raise ValueError(f"Unsupported type: {media_input['type']}") + + @staticmethod + def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray: + image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) + return np.asarray(image) + + def preprocess( + self, + medias: list[dict[str, Any]], + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + if not isinstance(medias, list): + medias = [medias] + if not medias: + return BatchFeature(data={}, tensor_type=return_tensors) + + if njit is None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + patch_size = int(self.media_proc_cfg["patch_size"]) + prepared = [] + grid_thws_np = np.empty((len(medias), 3), dtype=np.int64) + total_patches = 0 + + for idx, item in enumerate(medias): + item = _ensure_media_type(item) + resize_config = self.get_resize_config(item) + new_width = resize_config["new_width"] + new_height = resize_config["new_height"] + pad_width = resize_config["pad_width"] + pad_height = resize_config["pad_height"] + padded_width = new_width + pad_width + padded_height = new_height + pad_height + + if item["type"] == "image": + image_np = self.resize_image(item["image"], new_width, new_height) + frames = image_np[np.newaxis, ...] + elif item["type"] == "video_chunk": + frames = np.stack( + [ + self.resize_image(frame, new_width, new_height) + for frame in item["video_chunk"] + ], + axis=0, + ) + else: + raise ValueError(f"Unsupported type: {item['type']}") + + t_size = frames.shape[0] + grid_h = padded_height // patch_size + grid_w = padded_width // patch_size + grid_thws_np[idx, 0] = t_size + grid_thws_np[idx, 1] = grid_h + grid_thws_np[idx, 2] = grid_w + + num_patches = t_size * grid_h * grid_w + prepared.append( + ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) + ) + total_patches += num_patches + + pixel_values_np = np.empty( + (total_patches, 3, patch_size, patch_size), dtype=np.float32 + ) + out_offset = 0 + with numba_workqueue_threading_layer(): + for ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) in prepared: + _write_fused_patches( + frames, + pixel_values_np, + out_offset, + new_height, + new_width, + padded_height, + padded_width, + patch_size, + self.normalize_lut, + ) + out_offset += num_patches + + data = { + "pixel_values": torch.from_numpy(pixel_values_np), + "grid_thws": torch.from_numpy(grid_thws_np), + } + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, "tolist"): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py index d5e5750ca5d..899e0402ba5 100644 --- a/vllm/transformers_utils/processors/minicpmo.py +++ b/vllm/transformers_utils/processors/minicpmo.py @@ -268,7 +268,7 @@ class MiniCPMOProcessor(ProcessorMixin): def batch_decode(self, *args, **kwargs): """ This method forwards all its arguments to LlamaTokenizerFast's - [`~PreTrainedTokenizer.batch_decode`]. Please refer to the + [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ output_ids = args[0] @@ -289,7 +289,7 @@ class MiniCPMOProcessor(ProcessorMixin): def decode(self, *args, **kwargs): """ This method forwards all its arguments to LlamaTokenizerFast's - [`~PreTrainedTokenizer.decode`]. Please refer to the docstring + [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ result = args[0] diff --git a/vllm/transformers_utils/processors/minicpmv.py b/vllm/transformers_utils/processors/minicpmv.py index 91c3a8e479f..712cbbb47a0 100644 --- a/vllm/transformers_utils/processors/minicpmv.py +++ b/vllm/transformers_utils/processors/minicpmv.py @@ -56,14 +56,9 @@ class MiniCPMVProcessor(ProcessorMixin): image_processor_class = "AutoImageProcessor" tokenizer_class = "AutoTokenizer" - def __init__(self, image_processor=None, tokenizer=None): + def __init__(self, image_processor=None, tokenizer=None, version=None): super().__init__(image_processor, tokenizer) - # Newer (transformers v5.7+) MiniCPM-V image processors, e.g. - # MiniCPMV4_6ImageProcessor, no longer carry a `version` attribute. - # Fall back to None instead of hard-crashing: `version` is only used - # to special-case the 2.5 tokenization path in `_convert`, and any - # value other than 2.5 takes the default branch anyway. - self.version = getattr(image_processor, "version", None) + self.version = version def __call__( self, @@ -100,7 +95,7 @@ class MiniCPMVProcessor(ProcessorMixin): def batch_decode(self, *args, **kwargs): """ This method forwards all its arguments to LlamaTokenizerFast's - [`~PreTrainedTokenizer.batch_decode`]. Please refer to the + [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ output_ids = args[0] @@ -133,7 +128,7 @@ class MiniCPMVProcessor(ProcessorMixin): def decode(self, *args, **kwargs): """ This method forwards all its arguments to LlamaTokenizerFast's - [`~PreTrainedTokenizer.decode`]. Please refer to the docstring + [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ result = args[0] @@ -161,7 +156,7 @@ class MiniCPMVProcessor(ProcessorMixin): def _convert(self, input_str, max_inp_length: int | None = None): add_bos = getattr(self.tokenizer, "add_bos_token", False) - if self.version == 2.5 or add_bos: + if self.version == (2, 5) or add_bos: input_ids = self.tokenizer.encode(input_str) else: bos_id = getattr( diff --git a/vllm/transformers_utils/processors/moondream3.py b/vllm/transformers_utils/processors/moondream3.py index 289c40dd175..ae6833b2007 100644 --- a/vllm/transformers_utils/processors/moondream3.py +++ b/vllm/transformers_utils/processors/moondream3.py @@ -195,7 +195,7 @@ class Moondream3Processor(ProcessorMixin): The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1' instead of having tokenizer files in the model repo. """ - from transformers import AutoTokenizer, PreTrainedTokenizerFast + from transformers import AutoTokenizer, TokenizersBackend from transformers.utils import cached_file tokenizer = kwargs.pop("tokenizer", None) @@ -237,7 +237,7 @@ class Moondream3Processor(ProcessorMixin): "tokenizer.json", **cached_file_kwargs, ) - return PreTrainedTokenizerFast( + return TokenizersBackend( tokenizer_file=tokenizer_file, clean_up_tokenization_spaces=False, ) diff --git a/vllm/transformers_utils/processors/openvla.py b/vllm/transformers_utils/processors/openvla.py index 162f4023830..e520f3b1398 100644 --- a/vllm/transformers_utils/processors/openvla.py +++ b/vllm/transformers_utils/processors/openvla.py @@ -7,7 +7,7 @@ from typing import Any import numpy as np import torch from PIL import Image -from transformers.processing_utils import ProcessorMixin +from transformers import ImageProcessingMixin, ProcessorMixin IMAGENET_MEAN = np.array([0.484375, 0.455078125, 0.40625], dtype=np.float32) IMAGENET_STD = np.array([0.228515625, 0.2236328125, 0.224609375], dtype=np.float32) @@ -65,7 +65,7 @@ def preprocess_openvla_image(image: Any, image_size: int) -> torch.Tensor: return torch.from_numpy(pixel_values) -class OpenVLAImageProcessor: +class OpenVLAImageProcessor(ImageProcessingMixin): def __init__(self, *, image_size: int) -> None: self.image_size = image_size diff --git a/vllm/transformers_utils/processors/ovis.py b/vllm/transformers_utils/processors/ovis.py index da80f24e75c..907c68035fc 100644 --- a/vllm/transformers_utils/processors/ovis.py +++ b/vllm/transformers_utils/processors/ovis.py @@ -417,14 +417,14 @@ class OvisProcessor(ProcessorMixin): def batch_decode(self, *args, **kwargs): """ - This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please + This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): """ - This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to + This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to the docstring of this method for more information. """ return self.tokenizer.decode(*args, **kwargs) diff --git a/vllm/transformers_utils/processors/pixtral.py b/vllm/transformers_utils/processors/pixtral.py index c03360a2a56..588ad6b7fad 100644 --- a/vllm/transformers_utils/processors/pixtral.py +++ b/vllm/transformers_utils/processors/pixtral.py @@ -4,13 +4,13 @@ import torch from mistral_common.protocol.instruct.chunk import ImageChunk from mistral_common.tokens.tokenizers.multimodal import ImageEncoder from PIL import Image -from transformers import BatchFeature, ProcessorMixin, TensorType +from transformers import BatchFeature, ImageProcessingMixin, ProcessorMixin, TensorType from transformers.image_utils import ImageInput from vllm.tokenizers.mistral import MistralTokenizer -class MistralCommonImageProcessor: +class MistralCommonImageProcessor(ImageProcessingMixin): """ Provide a HF-compatible interface for `mistral_common.tokens.tokenizers.multimodal.ImageEncoder`. diff --git a/vllm/transformers_utils/processors/step3_vl.py b/vllm/transformers_utils/processors/step3_vl.py index 71540f433fd..8957a7c353c 100644 --- a/vllm/transformers_utils/processors/step3_vl.py +++ b/vllm/transformers_utils/processors/step3_vl.py @@ -8,7 +8,12 @@ import torch from PIL import Image from torchvision import transforms from torchvision.transforms.functional import InterpolationMode -from transformers import BatchFeature, ProcessorMixin, TensorType +from transformers import ( + BatchFeature, + ImageProcessingMixin, + ProcessorMixin, + TensorType, +) from vllm.tokenizers import TokenizerLike @@ -240,7 +245,7 @@ class ImagePatcher: ) -class Step3VLImageProcessor: +class Step3VLImageProcessor(ImageProcessingMixin): def __init__( self, image_size: int = 728, diff --git a/vllm/transformers_utils/processors/voxtral.py b/vllm/transformers_utils/processors/voxtral.py index f67bfe9d2e2..5403a7e47f9 100644 --- a/vllm/transformers_utils/processors/voxtral.py +++ b/vllm/transformers_utils/processors/voxtral.py @@ -6,13 +6,18 @@ from math import ceil import numpy as np import torch from mistral_common.tokens.tokenizers.audio import AudioEncoder -from transformers import BatchFeature, ProcessorMixin, TensorType +from transformers import ( + BatchFeature, + ProcessorMixin, + SequenceFeatureExtractor, + TensorType, +) from transformers.audio_utils import AudioInput from vllm.tokenizers.mistral import MistralTokenizer -class MistralCommonFeatureExtractor: +class MistralCommonFeatureExtractor(SequenceFeatureExtractor): """ Provide a HF-compatible interface for `mistral_common.tokens.tokenizers.multimodal.AudioEncoder`. diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 5506af4cac8..a758f5d535f 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -229,10 +229,11 @@ def get_model_path(model: str | Path, revision: str | None = None): if os.path.exists(model): return model assert huggingface_hub.constants.HF_HUB_OFFLINE - common_kwargs = { - "local_files_only": huggingface_hub.constants.HF_HUB_OFFLINE, - "revision": revision, - } + common_kwargs = dict( + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_patterns="*", + revision=revision, + ) if envs.VLLM_USE_MODELSCOPE: from modelscope.hub.snapshot_download import snapshot_download diff --git a/vllm/utils/async_utils.py b/vllm/utils/async_utils.py index 60c26569751..3faf728384e 100644 --- a/vllm/utils/async_utils.py +++ b/vllm/utils/async_utils.py @@ -103,8 +103,15 @@ async def merge_async_iterators( """ if len(iterators) == 1: # Fast-path single iterator case. - async for item in iterators[0]: - yield 0, item + iterator: AsyncGenerator[T, None] | None = iterators[0] + try: + async for item in iterator: # type: ignore[union-attr] + yield 0, item + iterator = None + finally: + if iterator is not None: + with contextlib.suppress(BaseException): + await iterator.aclose() return loop = asyncio.get_running_loop() diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 043798a584b..812adb59e6c 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -487,6 +487,11 @@ def has_nixl_ep() -> bool: return _has_module("nixl_ep") +def is_numba_available() -> bool: + """Whether the optional `numba` package is available.""" + return _has_module("numba") + + def has_triton_kernels() -> bool: """Whether the optional `triton_kernels` package is available.""" is_available = _has_module("triton_kernels") or _has_module( @@ -547,3 +552,21 @@ def has_cutedsl() -> bool: def has_humming() -> bool: """Whether the optional `humming` package is available.""" return _has_module("humming") + + +def check_torchcodec_available(): + """Whether the optional `torchcodec` package is available.""" + try: + import torchcodec # noqa: F401 + except RuntimeError as e: + # torchcodec will raise RuntimeError during import instead + # of ImportError when system ffmpeg unavailable, with a + # message that can leak sensitive system information. + # Trim it down to avoid it. + marker = ( + "The following exceptions were raised as we tried to load libtorchcodec:" + ) + message = str(e) + if marker in message: + raise RuntimeError(message.split(marker, 1)[0].rstrip()) from None + raise e diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 4565ffdae06..7ba8ecde653 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -20,9 +20,10 @@ Currently monitors: (via ``knobs.runtime.jit_post_compile_hook``) """ +import contextlib import functools import os -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from typing import Literal from vllm.logger import init_logger @@ -287,3 +288,32 @@ def _setup_cutedsl_jit_hook() -> None: cute.compile = _compile_with_monitor _cutedsl_hook_installed = True + + +@contextlib.contextmanager +def numba_workqueue_threading_layer() -> Iterator[None]: + """Force numba's fork-safe `workqueue` threading layer for this block. + + GNU OpenMP (numba's default `omp` threading layer) aborts the process + if a forked child re-enters an OpenMP-active runtime. vLLM forks the + EngineCore subprocess from a process that may already have launched + numba's parallel accelerator, so the first call to any + `@njit(parallel=True)` function must happen under `workqueue` instead. + The threading layer choice is sticky for the life of the process once + launched, so restoring the config on exit does not undo the effect. + """ + import numba + + key = "NUMBA_THREADING_LAYER" + previous_env = os.environ.get(key) + previous_config = numba.config.THREADING_LAYER + os.environ[key] = "workqueue" + numba.config.THREADING_LAYER = "workqueue" + try: + yield + finally: + if previous_env is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous_env + numba.config.THREADING_LAYER = previous_config diff --git a/vllm/utils/multi_stream_utils.py b/vllm/utils/multi_stream_utils.py index fed38ea1a35..2203221c5a1 100644 --- a/vllm/utils/multi_stream_utils.py +++ b/vllm/utils/multi_stream_utils.py @@ -20,8 +20,8 @@ class EventType(Enum): def maybe_execute_in_parallel( fn0: Callable[[], Any], fn1: Callable[[], Any], - event0: torch.Event, - event1: torch.Event, + event0: torch.cuda.Event, + event1: torch.cuda.Event, aux_stream: torch.cuda.Stream | None = None, ) -> tuple[Any, Any]: """Run two functions potentially in parallel on separate CUDA streams. @@ -61,8 +61,8 @@ def maybe_execute_in_parallel( def execute_in_parallel( default_fn: Callable[[], Any], aux_fns: list[Callable[[], Any] | None], - start_event: torch.Event, - done_events: list[torch.Event], + start_event: torch.cuda.Event, + done_events: list[torch.cuda.Event], aux_streams: list[torch.cuda.Stream] | None = None, enable: bool = False, ) -> tuple[Any, list[Any]]: @@ -108,7 +108,7 @@ def execute_in_parallel( ) aux_results = [None] * len(aux_fns) - pending: list[torch.Event] = [] + pending: list[torch.cuda.Event] = [] start_event.record() for i, fn in enumerate(aux_fns): diff --git a/vllm/utils/platform_utils.py b/vllm/utils/platform_utils.py index cc69d9a241c..5d7fed3c990 100644 --- a/vllm/utils/platform_utils.py +++ b/vllm/utils/platform_utils.py @@ -7,6 +7,7 @@ from concurrent.futures.process import ProcessPoolExecutor from functools import cache from typing import Any +import regex as re import torch @@ -62,3 +63,12 @@ def num_compute_units(device_id: int = 0) -> int: from vllm.platforms import current_platform return current_platform.num_compute_units(device_id) + + +@cache +def get_device_name_as_file_name(device_id: int = 0) -> str: + from vllm.platforms import current_platform + + name = current_platform.get_device_name(device_id) + name = re.sub(r"[\s/]+", "_", name) + return name diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 7f56f972a4f..e2ac15a949f 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -308,26 +308,26 @@ def set_ulimit(target_soft_limit: int = 65535): def find_loaded_library(lib_name: str) -> str | None: """ - According to according to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, + According to https://man7.org/linux/man-pages/man5/proc_pid_maps.5.html, the file `/proc/self/maps` contains the memory maps of the process, which includes the shared libraries loaded by the process. We can use this file to find the path of the loaded library. """ # noqa - found_line = None + # Match the mapped file's name, not the whole line: an unrelated library + # whose name merely contains lib_name (e.g. TileLang's libcudart_stub.so + # when looking for libcudart) or a directory component containing it must + # not win. Legitimate filenames are {lib_name}.so[.*] or a name-mangled + # {lib_name}-.so[.*], and /proc/self/maps is ordered by mapping + # address rather than load order, so a substring hit is a + # nondeterministic hijack. with open("/proc/self/maps") as f: for line in f: - if lib_name in line: - found_line = line - break - if found_line is None: - # the library is not loaded in the current process - return None - # if lib_name is libcudart, we need to match a line with: - # address /path/to/libcudart-hash.so.11.0 - start = found_line.index("/") - path = found_line[start:].strip() - filename = path.split("/")[-1] - assert filename.rpartition(".so")[0].startswith(lib_name), ( - f"Unexpected filename: {filename} for library {lib_name}" - ) - return path + start = line.find("/") + if start == -1: + continue + path = line[start:].strip() + filename = path.rsplit("/", maxsplit=1)[-1] + if filename.startswith((f"{lib_name}.", f"{lib_name}-")): + return path + # the library is not loaded in the current process + return None diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 22c6a382287..bfecb3c952e 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kNvfp4Dynamic, ) +from vllm.utils.torch_utils import np_to_pinned_tensor if TYPE_CHECKING: from vllm.config import VllmConfig @@ -467,6 +468,7 @@ class CommonAttentionMetadata: _num_computed_tokens_cpu: torch.Tensor | None = None _num_computed_tokens_cache: torch.Tensor | None = None + _token_to_req_indices_cache: torch.Tensor | None = None def batch_size(self) -> int: return self.seq_lens.shape[0] @@ -515,6 +517,31 @@ class CommonAttentionMetadata: self._num_computed_tokens_cache = self.seq_lens - query_lens return self._num_computed_tokens_cache + def token_to_req_indices(self, buffer: torch.Tensor) -> torch.Tensor: + """Build or reuse the per-token request index mapping.""" + num_tokens = self.num_actual_tokens + if self._token_to_req_indices_cache is not None: + assert self._token_to_req_indices_cache.device == buffer.device + assert self._token_to_req_indices_cache.dtype == torch.int32 + assert self._token_to_req_indices_cache.shape[0] >= num_tokens + return self._token_to_req_indices_cache[:num_tokens] + + starts = np.asarray(self.query_start_loc_cpu, dtype=np.int32) + query_lens = np.diff(starts) + token_to_req_indices = np.repeat( + np.arange(query_lens.shape[0], dtype=np.int32), query_lens + ) + num_mapped_tokens = token_to_req_indices.shape[0] + assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens) + # copy from CPU to GPU + buffer[:num_mapped_tokens].copy_( + np_to_pinned_tensor(token_to_req_indices), non_blocking=True + ) + if num_mapped_tokens < num_tokens: + buffer[num_mapped_tokens:num_tokens].zero_() + self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)] + return self._token_to_req_indices_cache[:num_tokens] + # TODO(lucas): remove once we have FULL-CG spec-decode support def unpadded( self, num_actual_tokens: int, num_actual_reqs: int diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 6b2d202d3f1..a5735bf313f 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -437,27 +437,24 @@ def _riscv_supports_rvv() -> bool: The RVV path is compiled whenever __riscv_v_min_vlen is defined, so we check that at least one supported zvlb is advertised. """ + # The C++ compile-time check is the ground truth: it knows which + # VLEN the binary was actually compiled for. The cpuinfo check + # below is only a fast-path shortcut. + try: + import torch + + if torch.ops._C.cpu_attn_has_isa("rvv"): + return True + except Exception: + pass + + # Fallback: check /proc/cpuinfo for zvl128b/zvl256b. try: with open("/proc/cpuinfo") as f: cpuinfo = f.read() except OSError: return False - # If VLEN >= 512 is detected, the RVV kernel was not compiled. - if any(f"zvl{n}b" in cpuinfo for n in (512, 1024)): - return False - - # zvl128b or zvl256b explicitly advertised -> RVV kernel available. - if any(f"zvl{n}b" in cpuinfo for n in (128, 256)): - return True - - # No zvlb flag at all (e.g. some hardware reports zve* without - # a VLEN hint). Delegate to the C++ compile-time check instead. - try: - import torch - - return torch.ops._C.cpu_attn_has_isa("rvv") - except Exception: - return False + return any(f"zvl{n}b" in cpuinfo for n in (128, 256)) def _get_attn_isa( diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index a7cc6fb08e9..db5b0dda367 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -892,21 +892,28 @@ class FlashAttentionImpl(AttentionImpl): and self.vllm_flash_attn_version == 4 ): max_ranges = mm_prefix_ranges.shape[1] - # Gemma4: clamp the bidirectional block to the sliding - # window (HF (causal OR blockwise) AND sliding_window on - # local layers). Other mm-prefix models pass sliding_window=0 - # -> unclamped, behavior unchanged. sliding_window_size[0]+1 - # is the window (self.sliding_window = (sw-1, 0) on sliding - # layers); causal mm_prefix never hits the symmetric path. + # Sliding window value in Triton convention + # (1 + window_size[0]). Global-attention layers + # store (-1, -1) → sw stays None / 0. + sw_val = ( + 1 + sliding_window_size[0] + if sliding_window_size is not None + and sliding_window_size[0] >= 0 + else None + ) + # Gemma4: also clamp the bidirectional block to the + # sliding window when the layer opts in + # (mm_prefix_clamp_sliding_window flag from PR #47217). mm_clamp_sw = 0 if ( getattr(layer, "mm_prefix_clamp_sliding_window", False) - and sliding_window_size is not None - and sliding_window_size[0] >= 0 + and sw_val is not None ): - mm_clamp_sw = sliding_window_size[0] + 1 + mm_clamp_sw = sw_val mm_mask_mod = _make_mm_prefix_mask_mod( - max_ranges, sliding_window=mm_clamp_sw + max_ranges, + sliding_window=mm_clamp_sw, + sliding_window_left=sw_val, ) mm_aux = [mm_prefix_ranges] @@ -1204,23 +1211,26 @@ class FlashAttentionImpl(AttentionImpl): return output -def _make_mm_prefix_mask_mod(max_ranges: int, sliding_window: int = 0): - """Build a CuTE-DSL mask_mod implementing (causal OR mm_prefix). +def _make_mm_prefix_mask_mod( + max_ranges: int, + sliding_window: int = 0, + sliding_window_left: int | None = None, +): + """Build a CuTE-DSL mask_mod implementing + ``(causal AND sliding_window) OR mm_prefix``. - Returns a @cute.jit callable that evaluates: - keep = (kv_idx <= q_idx) OR - (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end] - [AND (q_idx - kv_idx) < sliding_window]) - for each mm_prefix range stored in aux_tensors[0]. + The FA4 kernel passes *local* ``q_idx`` (0-based within the current + prefill chunk) while ``kv_idx`` is absolute (0-based over the full + KV cache). We recover the absolute Q position via + ``q_abs = q_idx + seqlen_k - seqlen_q`` (the context-length offset) + so that causal, sliding-window, and mm_prefix range comparisons all + use consistent absolute positions. This matches the Triton + reference path (``compute_kv_seq_mask``). - When sliding_window > 0 (Gemma4 sliding layers), the bidirectional block is - clamped to the sliding window, matching HF's - (causal OR blockwise) AND sliding_window (== sliding_window_overlay - kv > q - sw; future kv passes since q - kv < 0). sliding_window <= 0 - (default) leaves the block unclamped, so other mm-prefix models are - unchanged. q_idx/kv_idx are local offsets, but their difference is - frame-invariant and image bidirectional attention only matters during - prefill (where local == absolute). + ``sliding_window_left`` enforces the sliding window on the causal + term (None = full causal, no window). ``sliding_window`` clamps the + bidirectional block to the window (0 = unclamped; >0 = Gemma4 local + layers via ``mm_prefix_clamp_sliding_window``). """ import cutlass import cutlass.cute as cute @@ -1230,29 +1240,59 @@ def _make_mm_prefix_mask_mod(max_ranges: int, sliding_window: int = 0): scalar_to_ssa, ) - @cute.jit - def mm_prefix_mask_mod( - batch_idx: cute.TensorSSA, - head_idx: cute.TensorSSA, - q_idx: cute.TensorSSA, - kv_idx: cute.TensorSSA, - seqlen_info, - aux_tensors, - ): - keep = kv_idx <= q_idx - ranges = aux_tensors[0] - b = batch_idx[0] - for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined] - r_start = scalar_to_ssa(ranges[b, i, 0], Int32) - r_end = scalar_to_ssa(ranges[b, i, 1], Int32) - valid = r_start < r_end - q_in = (q_idx >= r_start) & (q_idx <= r_end) & valid - k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid - mm = q_in & k_in - if sliding_window > 0: - mm = mm & ((q_idx - kv_idx) < sliding_window) - keep = keep | mm - return keep + if sliding_window_left is not None: + + @cute.jit + def mm_prefix_mask_mod( + batch_idx: cute.TensorSSA, + head_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors, + ): + ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32) + q_abs = q_idx + ctx_off + sw = scalar_to_ssa(Int32(sliding_window_left), Int32) + keep = (kv_idx <= q_abs) & ((q_abs - kv_idx) < sw) + ranges = aux_tensors[0] + b = batch_idx[0] + for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined] + r_start = scalar_to_ssa(ranges[b, i, 0], Int32) + r_end = scalar_to_ssa(ranges[b, i, 1], Int32) + valid = r_start < r_end + q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid + k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid + mm = q_in & k_in + if sliding_window > 0: + mm = mm & ((q_abs - kv_idx) < sw) + keep = keep | mm + return keep + + else: + + @cute.jit + def mm_prefix_mask_mod( + batch_idx: cute.TensorSSA, + head_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors, + ): + ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32) + q_abs = q_idx + ctx_off + keep = kv_idx <= q_abs + ranges = aux_tensors[0] + b = batch_idx[0] + for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined] + r_start = scalar_to_ssa(ranges[b, i, 0], Int32) + r_end = scalar_to_ssa(ranges[b, i, 1], Int32) + valid = r_start < r_end + q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid + k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid + keep = keep | (q_in & k_in) + return keep mm_prefix_mask_mod.use_fast_sampling = True return mm_prefix_mask_mod diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 5e2295c0af8..12eab21e3e1 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1584,11 +1584,11 @@ class FlashInferImpl(AttentionImpl): if query.dtype != q_data_type: assert query.dtype in [torch.float16, torch.bfloat16] assert q_data_type in [torch.float8_e4m3fn, torch.float8_e5m2] - assert query.is_contiguous() assert query.dim() == 3 num_tokens = query.shape[0] num_heads = query.shape[1] head_size = query.shape[2] + assert query.stride(2) == 1 and query.stride(1) == head_size query_quantized, _ = custom_ops.scaled_fp8_quant( query.view(num_tokens, num_heads * head_size), scale=scale ) diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py index 4a3a4383e2e..8c6dcfe5168 100644 --- a/vllm/v1/attention/backends/hpc_attn.py +++ b/vllm/v1/attention/backends/hpc_attn.py @@ -31,8 +31,6 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.utils import ( KVCacheLayoutType, - get_per_layer_parameters, - infer_global_hyperparameters, split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -86,9 +84,23 @@ class HpcAttnMetadata(AttentionMetadata): hpc_prefill_q_scale: torch.Tensor | None = None """FP8 per-token-per-head Q scale for prefill (from RopeNorm).""" hpc_decode_q_scale: torch.Tensor | None = None - """FP8 per-token-per-head Q scale for decode (from RopeNorm).""" + """FP8 per-token-per-head Q scale for decode (persistent buffer). + shape = [max_decode_tokens, num_q_heads], contiguous. + Only the first num_decode_q_scale_tokens rows are valid.""" hpc_split_k_flag: torch.Tensor | None = None - """Split-K flag tensor for FP8 decode (from RopeNorm).""" + """Split-K flag tensor for FP8 decode (persistent buffer). + shape = [max_num_seqs, num_kv_heads], int32.""" + + # --- MTP (Multi-Token Prediction) fields --- + decode_query_len: int = 1 + """Number of query tokens per decode request. + 1 for standard decoding, mtp+1 for speculative decoding (2 or 3).""" + qo_indptr_decode: torch.Tensor | None = None + """Cumulative query offsets for decode requests (GPU tensor). + shape = [num_decodes + 1]. Only set when decode_query_len > 1. + e.g. 3 requests with dql=2: [0, 2, 4, 6].""" + task_map: torch.Tensor | None = None + """Used for HPC dynamic schedule attention""" class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): @@ -105,20 +117,40 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): device: torch.device, ): super().__init__(kv_cache_spec, layer_names, vllm_config, device) - self.model_config = vllm_config.model_config - self.cache_config = vllm_config.cache_config + import hpc - self.num_qo_heads = self.model_config.get_num_attention_heads( - vllm_config.parallel_config - ) self.num_kv_heads = kv_cache_spec.num_kv_heads - self.head_dim = kv_cache_spec.head_size - self.page_size = kv_cache_spec.block_size + self.hpc_dynamic_sched_attn_min_split_len = 1024 - self.cache_dtype = self.cache_config.cache_dtype + # MTP constraint: HPC decode kernel only supports mtp in {0, 1, 2, 3} + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 3 + ): + raise ValueError( + f"HPC attention only supports up to 3 speculative tokens " + f"(mtp ∈ {{0, 1, 2, 3}}), got " + f"num_speculative_tokens={spec_config.num_speculative_tokens}. " + f"Please reduce num_speculative_tokens or use a different " + f"attention backend." + ) - self.global_hyperparameters = infer_global_hyperparameters( - get_per_layer_parameters(vllm_config, layer_names, HpcAttentionImpl) + # Dynamic decode threshold for MTP support. + # _init_reorder_batch_threshold computes: + # no spec_config → threshold=1 (unchanged) + # with spec_config → threshold=1+num_speculative_tokens + self._init_reorder_batch_threshold( + reorder_batch_threshold=1, + supports_spec_as_decode=True, + ) + + self.task_map = hpc.get_attention_decode_task_workspace( + vllm_config.scheduler_config.max_num_seqs, + vllm_config.model_config.max_model_len or 4096, + self.num_kv_heads, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, ) @override # type: ignore[misc] @@ -128,6 +160,13 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): vllm_config: VllmConfig, kv_cache_spec: AttentionSpec, ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if ( + spec_config is not None + and spec_config.num_speculative_tokens is not None + and spec_config.num_speculative_tokens > 0 + ): + return AttentionCGSupport.UNIFORM_BATCH return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE def build( @@ -143,7 +182,8 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=False, + # MTP requires uniform query lengths across decode requests + require_uniform=(self.reorder_batch_threshold > 1), ) ) @@ -152,7 +192,16 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): slot_mapping = common_attn_metadata.slot_mapping max_query_len = common_attn_metadata.max_query_len + # Compute decode_query_len (tokens per decode request). + # Non-MTP: 1, MTP: mtp+1 (2 or 3). + if num_decodes > 0 and num_decode_tokens > num_decodes: + decode_query_len = num_decode_tokens // num_decodes + else: + decode_query_len = 1 + + seq_lens_decode = None qo_indptr = None + qo_indptr_decode = None if num_prefills > 0: qo_indptr_cpu = common_attn_metadata.query_start_loc_cpu prefill_start = num_decodes @@ -161,6 +210,21 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): ) qo_indptr = qo_indptr_prefill_cpu.to(self.device, non_blocking=True) + if num_decodes > 0: + seq_lens_decode = seq_lens[:num_decodes] + # block_table is per-request, indexed by num_decodes (not tokens) + qo_indptr_decode = common_attn_metadata.query_start_loc[: num_decodes + 1] + import hpc + + hpc.assign_attention_decode_task( + seq_lens_decode, + self.task_map, + self.num_kv_heads, + decode_query_len, + new_kv_included=True, + min_process_len=self.hpc_dynamic_sched_attn_min_split_len, + ) + return HpcAttnMetadata( num_actual_tokens=num_actual_tokens, num_decodes=num_decodes, @@ -176,6 +240,9 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]): hpc_prefill_q_scale=None, hpc_decode_q_scale=None, hpc_split_k_flag=None, + decode_query_len=decode_query_len, + qo_indptr_decode=qo_indptr_decode, + task_map=self.task_map, ) @@ -192,6 +259,7 @@ class HpcAttentionBackend(AttentionBackend): ] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", + "bfloat16", "fp8_e4m3", ] @@ -323,6 +391,13 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): self.supports_quant_query_input = False self.splitk = True + import hpc + + if self.use_fp8: + self._quant_type = hpc.QuantType.QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR + else: + self._quant_type = None + def forward( self, layer: torch.nn.Module, @@ -417,6 +492,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): block_table_prefill, seq_lens_prefill, max_seqlens, + quant_type=self._quant_type, output=output_prefill, ) else: @@ -439,6 +515,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): q_decode = query[:num_decode_tokens] output_decode = output[:num_decode_tokens] + mtp = attn_metadata.decode_query_len - 1 + if self.use_fp8: hpc.attention_decode_fp8( q_decode, @@ -449,8 +527,14 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): hpc_decode_q_scale, k_scale, v_scale, + mtp=mtp, + # MTP: split_flag from rope_norm is unavailable + # when using prefill-mode kernel; let HPC decide. + # splitk=(self.splitk if mtp == 0 else True), new_kv_included=True, + quant_type=self._quant_type, splitk=self.splitk, + task_map=attn_metadata.task_map, split_flag=hpc_split_k_flag, output=output_decode, ) @@ -461,6 +545,7 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]): kv_cache[:, 1], block_table_decode, num_seq_kvcache, + mtp=mtp, output=output_decode, new_kv_included=True, splitk=self.splitk, diff --git a/vllm/v1/attention/backends/linear_attn.py b/vllm/v1/attention/backends/linear_attn.py index b2ca151986c..9cdcf0e30e7 100644 --- a/vllm/v1/attention/backends/linear_attn.py +++ b/vllm/v1/attention/backends/linear_attn.py @@ -4,7 +4,7 @@ from dataclasses import dataclass import torch -from vllm.config import VllmConfig +from vllm.config import CompilationConfig, VllmConfig from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -12,6 +12,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.utils import ( + PAD_SLOT_ID, mamba_get_block_table_tensor, split_decodes_and_prefills, ) @@ -91,3 +92,213 @@ class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMet state_indices_tensor=state_indices_tensor, ) return attn_metadata + + +class BailingLinearAttentionBackend(LinearAttentionBackend): + @staticmethod + def get_name() -> str: + return "BAILING_LINEAR_ATTN" + + @staticmethod + def get_builder_cls() -> type["BailingLinearAttentionMetadataBuilder"]: + return BailingLinearAttentionMetadataBuilder + + +@dataclass +class BailingLinearAttentionMetadata(LinearAttentionMetadata): + state_indices_tensor_d: torch.Tensor | None = None + state_indices_tensor_p: torch.Tensor | None = None + num_accepted_tokens: torch.Tensor | None = None + query_start_loc_d: torch.Tensor | None = None + + +class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder): + supports_spec_decode_metadata = True + supports_update_block_table: bool = False + + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: AttentionSpec, + ) -> AttentionCGSupport: + return AttentionCGSupport.UNIFORM_BATCH + + 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) + self.compilation_config: CompilationConfig = vllm_config.compilation_config + self.num_spec_tokens: int = vllm_config.num_speculative_tokens + self.use_spec_decode: bool = self.num_spec_tokens > 0 + self.decode_cudagraph_max_bs: int = vllm_config.scheduler_config.max_num_seqs + if self.compilation_config.max_cudagraph_capture_size is not None: + self.decode_cudagraph_max_bs = min( + self.decode_cudagraph_max_bs, + self.compilation_config.max_cudagraph_capture_size, + ) + self.decode_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens), + dtype=torch.int32, + device=device, + ) + self.decode_legacy_state_indices_tensor: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self.decode_query_start_loc: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs + 1,), + dtype=torch.int32, + device=device, + ) + self.decode_num_accepted_tokens: torch.Tensor = torch.empty( + (self.decode_cudagraph_max_bs,), + dtype=torch.int32, + device=device, + ) + self._init_reorder_batch_threshold(1, self.use_spec_decode) + + def build_for_cudagraph_capture( + self, + common_attn_metadata: CommonAttentionMetadata, + ) -> BailingLinearAttentionMetadata: + num_accepted_tokens = None + if self.use_spec_decode: + assert common_attn_metadata.max_query_len <= 1 + self.num_spec_tokens, ( + "Bailing linear attention only supports speculative decoding " + "with query length <= 1 + number of speculative tokens." + ) + num_accepted_tokens = torch.diff(common_attn_metadata.query_start_loc) + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + num_accepted_tokens=num_accepted_tokens, + ) + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + *, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + ) -> BailingLinearAttentionMetadata: + query_start_loc = common_attn_metadata.query_start_loc + seq_lens = common_attn_metadata.seq_lens + num_reqs = common_attn_metadata.num_reqs + use_spec_decode = self.use_spec_decode and num_accepted_tokens is not None + + state_indices_tensor = mamba_get_block_table_tensor( + common_attn_metadata.block_table_tensor, + common_attn_metadata.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + if state_indices_tensor.dim() == 1: + state_indices_tensor = state_indices_tensor.unsqueeze(-1) + + decode_threshold = self.reorder_batch_threshold if use_spec_decode else 1 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=decode_threshold, + ) + ) + state_indices_tensor_d, state_indices_tensor_p = torch.split( + state_indices_tensor, + [num_decodes, num_prefills], + dim=0, + ) + state_indices_tensor_p = state_indices_tensor_p[:, 0] + + query_start_loc_d = None + if use_spec_decode: + assert num_accepted_tokens is not None + state_indices_tensor_d = state_indices_tensor_d[ + :, : 1 + self.num_spec_tokens + ] + query_start_loc_d = query_start_loc[: num_decodes + 1] + num_accepted_tokens = num_accepted_tokens[:num_decodes] + else: + state_indices_tensor_d = state_indices_tensor_d[:, 0] + num_accepted_tokens = None + + legacy_state_indices_tensor = state_indices_tensor[:, 0] + cudagraph_mode = self.compilation_config.cudagraph_mode + use_full_cudagraph = ( + cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() + ) + if ( + num_prefills == 0 + and num_decodes <= self.decode_cudagraph_max_bs + and use_full_cudagraph + ): + padded_bs = num_reqs + is_padded_decode = seq_lens[:num_decodes] == 0 + if state_indices_tensor_d.dim() > 1: + state_indices_tensor_d = torch.where( + is_padded_decode.unsqueeze(1), + torch.full_like(state_indices_tensor_d, PAD_SLOT_ID), + state_indices_tensor_d, + ) + self.decode_state_indices_tensor[:num_decodes].copy_( + state_indices_tensor_d, + non_blocking=True, + ) + state_indices_tensor_d = self.decode_state_indices_tensor[:padded_bs] + state_indices_tensor_d[num_decodes:] = PAD_SLOT_ID + + self.decode_legacy_state_indices_tensor[:num_decodes].copy_( + torch.where( + is_padded_decode, + torch.full_like( + legacy_state_indices_tensor[:num_decodes], + PAD_SLOT_ID, + ), + legacy_state_indices_tensor[:num_decodes], + ), + non_blocking=True, + ) + legacy_state_indices_tensor = self.decode_legacy_state_indices_tensor[ + :padded_bs + ] + legacy_state_indices_tensor[num_decodes:] = PAD_SLOT_ID + if state_indices_tensor_d.dim() == 1: + state_indices_tensor_d = legacy_state_indices_tensor + + if use_spec_decode and num_accepted_tokens is not None: + assert query_start_loc_d is not None + self.decode_query_start_loc[: num_decodes + 1].copy_( + query_start_loc_d, + non_blocking=True, + ) + decode_num_query_tokens = query_start_loc_d[-1] + query_start_loc_d = self.decode_query_start_loc[: padded_bs + 1] + query_start_loc_d[num_decodes + 1 :] = decode_num_query_tokens + + self.decode_num_accepted_tokens[:num_decodes].copy_( + num_accepted_tokens, + non_blocking=True, + ) + num_accepted_tokens = self.decode_num_accepted_tokens[:padded_bs] + num_accepted_tokens[num_decodes:] = 1 + + return BailingLinearAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + state_indices_tensor=legacy_state_indices_tensor, + state_indices_tensor_d=state_indices_tensor_d, + state_indices_tensor_p=state_indices_tensor_p, + num_accepted_tokens=num_accepted_tokens, + query_start_loc_d=query_start_loc_d, + ) diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index a38b274dfcb..e0d54eee101 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -134,7 +134,7 @@ def get_mla_prefill_backend( f"Reason: {invalid_reasons}" ) assert backend_cls is not None - logger.info("Using %s MLA prefill backend.", selected_backend.name) + logger.info_once("Using %s MLA prefill backend.", selected_backend.name) return backend_cls return _auto_select_mla_prefill_backend( diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index e6a64ee85f8..977a42c6fef 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -196,6 +196,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): kv_cache_dtype_str = "fp8" + q_dtype = dtypes.fp8 else: kv_cache_dtype_str = "bf16" kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 55bba1a7ee8..1e4a2683063 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -522,31 +522,28 @@ class ROCMAiterMLASparseMetadataBuilder( # treated as its own batch entry), so persistent metadata can always # be precomputed here. The kernel switches to the persistent # work-stealing path automatically when work_meta_data is non-None. - # The output is a deterministic function of (num_tokens, max_query_len, - # num_heads, min(seq_lens, topk_tokens)); fingerprint those CPU-side - # and skip the launch when nothing changed. + # The output is a deterministic function of the per-request query and + # context lengths (both clamped to topk_tokens, past which per-token KV + # length saturates) and num_heads; fingerprint those CPU-side and skip + # the launch when nothing changed. num_reqs = common_attn_metadata.num_reqs clamped_seq_lens = np.minimum( common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(), self.topk_tokens, ) + clamped_context_lens = np.minimum( + common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths, + self.topk_tokens, + ) metadata_key = ( num_tokens, int(common_attn_metadata.max_query_len), self._num_attention_heads, clamped_seq_lens.tobytes(), + clamped_context_lens.tobytes(), + seg_lengths.tobytes(), ) - # The persistent MLA kernel is numerically wrong for multi-token prefill - # batches; errors compound across chunked prefill and break long-context - # decode (vllm#47042). Use it only for decode and single-chunk prefills, - # not chunked-prefill continuations (>1 query token, seq_len > query_len). - step_query_lens = seg_lengths - total_seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - is_chunked_continuation = (step_query_lens > 1) & ( - total_seq_lens > step_query_lens - ) - use_persistent = not is_chunked_continuation.any() - if use_persistent and metadata_key != self._prev_metadata_key: + if metadata_key != self._prev_metadata_key: from aiter import get_mla_metadata_v1 get_mla_metadata_v1( @@ -586,7 +583,7 @@ class ROCMAiterMLASparseMetadataBuilder( paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, paged_kv_indptr=paged_kv_indptr, - work_meta_data=self._mla_work_meta_data if use_persistent else None, + work_meta_data=self._mla_work_meta_data, work_indptr=self._mla_work_indptr, work_info_set=self._mla_work_info_set, reduce_indptr=self._mla_reduce_indptr, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index ac722dca9fc..9e54b62a00b 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, SlidingWindowMLASpec, + get_kv_quant_mode, ) # DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct @@ -89,8 +90,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else None, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) def forward(self): ... @@ -399,7 +402,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): For prefill, we use chunked prefill to align with the indexer's chunking. """ - num_reqs = common_attn_metadata.num_reqs seq_lens = common_attn_metadata.seq_lens seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound query_start_loc = common_attn_metadata.query_start_loc @@ -416,10 +418,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): # NOTE: Ensure all metadata tensors maintain fixed memory addresses # for CUDA graph compatibility. - query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + token_to_req_indices = common_attn_metadata.token_to_req_indices( + self.token_to_req_indices + ) is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] is_valid_token.copy_(slot_mapping >= 0) diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 66d57917fa1..57b64cc93df 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -27,12 +27,14 @@ logger = init_logger(__name__) class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend): - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", + "float16", "bfloat16", "fp8", "fp8_e4m3", + "fp8_e5m2", ] @staticmethod diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 500cd3fdf9e..aa81e44e635 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/vllm/v1/attention/backends/rocm_attn.py @@ -421,13 +421,8 @@ class RocmAttentionImpl(AttentionImpl): if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) - # chunked_prefill_paged_decode runs attention with a full-precision - # query (it does not quantize Q to fp8 and does not consume - # q_scale), so q_scale only matters when the query itself is fp8. - # For a non-fp8 query, q_scale is not applicable and is ignored - # (mirrors TritonAttentionImpl). This avoids spuriously failing on - # checkpoints that carry a non-1.0 q_scale while keeping the query - # in full precision. + # q_scale only applies to an fp8 query; this path keeps the query + # in full precision, so a non-1.0 q_scale is not applicable here. if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0: raise NotImplementedError( "A non 1.0 q_scale with an fp8 query is not currently " diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py index dbe3c5705de..6aec1db2e59 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -133,7 +133,7 @@ def _fwd_kernel_stage1( + offs_n // PAGE_SIZE, mask=offs_n < split_kv_end, other=0, - ) + ).to(tl.int64) # page_number * page stride overflows int32 kv_in_page = offs_n % PAGE_SIZE offs_buf_k = ( (kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[ @@ -375,7 +375,7 @@ def _fwd_grouped_kernel_stage1( mask=offs_n < split_kv_end, other=0, cache_modifier=".ca", - ) + ).to(tl.int64) # page_number * page stride overflows int32 kv_off_k = ( kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs ) diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index a13ae96a7a9..93622957b55 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -254,12 +254,19 @@ def kernel_unified_attention( # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}. k_scale_cache_ptr=None, v_scale_cache_ptr=None, - stride_ks_blk: tl.int64 = None, - stride_ks_slot: tl.int64 = None, - stride_ks_head: tl.int64 = None, - stride_vs_blk: tl.int64 = None, - stride_vs_slot: tl.int64 = None, - stride_vs_head: tl.int64 = None, + # ``tl.int64`` cannot be combined with a ``None`` default — Triton's JIT + # rejects ``Optional[tl.int64]`` / ``tl.int64 | None`` at trace time, and + # plain ``tl.int64 = None`` raises ``TypeError: 'NoneType' object cannot + # be interpreted as an integer`` when callers omit these arguments. + # ``int | None`` is the only annotation that lets the wrapper pass + # ``None`` here so Triton can skip materialising the strides when the + # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead. + stride_ks_blk: int | None = None, + stride_ks_slot: int | None = None, + stride_ks_head: int | None = None, + stride_vs_blk: int | None = None, + stride_vs_slot: int | None = None, + stride_vs_head: int | None = None, # KV cache quantization mode handled inside this kernel via constexpr # branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2), # FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own @@ -283,7 +290,10 @@ def kernel_unified_attention( # original (causal AND SW) OR mm_prefix behavior for all other models. MM_PREFIX_CLAMP_SW: tl.constexpr = False, ): - USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2 + # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}. + USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and ( + KV_QUANT_MODE <= 3 + ) USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8 if USE_TD: @@ -1041,9 +1051,9 @@ def unified_attention( # The kernel signature is the same for 2D and 3D — only the launch # grid + a handful of constexpr toggles differ. Per-token-head scale - # caches and their strides are required arguments; non-per-token-head - # modes pass dummy zeros (the code path is dead-code eliminated by - # the ``USE_PER_TOKEN_HEAD_SCALES`` constexpr branch in the kernel). + # caches and their strides are passed as ``None`` when the + # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead so Triton can skip + # materialising those arguments and the associated registers. if use_per_token_head_scales: ks_strides = k_scale_cache.stride() vs_strides = v_scale_cache.stride() @@ -1052,16 +1062,15 @@ def unified_attention( k_scale_ptr = k_scale_cache v_scale_ptr = v_scale_cache else: - ks_blk = ks_slot = ks_head = 0 - vs_blk = vs_slot = vs_head = 0 - # Pass the K cache as a stand-in pointer; never dereferenced. - k_scale_ptr = k - v_scale_ptr = v - # 3D needs real segm tensors; 2D never touches them but Triton wants - # a non-null pointer. Reuse ``out`` as the placeholder. - segm_output_ptr = softmax_segm_output if use_3d else out - segm_max_ptr = softmax_segm_max if use_3d else out - segm_expsum_ptr = softmax_segm_expsum if use_3d else out + ks_blk = ks_slot = ks_head = None + vs_blk = vs_slot = vs_head = None + k_scale_ptr = None + v_scale_ptr = None + # 3D needs real segm tensors; 2D never touches them. Pass ``None`` in + # 2D mode so Triton can skip materialising these pointer arguments. + segm_output_ptr = softmax_segm_output if use_3d else None + segm_max_ptr = softmax_segm_max if use_3d else None + segm_expsum_ptr = softmax_segm_expsum if use_3d else None num_segments = num_par_softmax_segments if use_3d else 1 grid: tuple[Any, ...] diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index a759d7a80ad..4756136f03c 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -67,7 +67,7 @@ class KVCacheCoordinator(ABC): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -107,7 +107,7 @@ class KVCacheCoordinator(ABC): self.single_type_managers = tuple( get_manager_for_kv_cache_spec( kv_cache_spec=kv_cache_group.kv_cache_spec, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, block_pool=self.block_pool, enable_caching=enable_caching, @@ -115,6 +115,7 @@ class KVCacheCoordinator(ABC): dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, scheduler_block_size=self.scheduler_block_size, + needs_kv_cache_zeroing=self.kv_cache_config.needs_kv_cache_zeroing, ) for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) ) @@ -331,7 +332,7 @@ class KVCacheCoordinator(ABC): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """ @@ -340,15 +341,15 @@ class KVCacheCoordinator(ABC): Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length. R-SWA managers use this to free gap blocks between the prefill tail and decode window; other manager types ignore it. """ for manager in self.single_type_managers: manager.remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) def get_blocks(self, request_id: str) -> tuple[list[KVCacheBlock], ...]: @@ -386,7 +387,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_kv_cache_events: bool, dcp_world_size: int, @@ -398,7 +399,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, False, enable_kv_cache_events, @@ -435,7 +436,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -448,7 +449,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -521,7 +522,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): self, kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -534,7 +535,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): super().__init__( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -782,7 +783,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): def get_kv_cache_coordinator( kv_cache_config: KVCacheConfig, max_model_len: int, - max_num_batched_tokens: int, + max_in_flight_tokens: int, use_eagle: bool, enable_caching: bool, enable_kv_cache_events: bool, @@ -796,7 +797,7 @@ def get_kv_cache_coordinator( return KVCacheCoordinatorNoPrefixCache( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_kv_cache_events, dcp_world_size=dcp_world_size, @@ -809,7 +810,7 @@ def get_kv_cache_coordinator( return UnitaryKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, @@ -822,7 +823,7 @@ def get_kv_cache_coordinator( return HybridKVCacheCoordinator( kv_cache_config, max_model_len, - max_num_batched_tokens, + max_in_flight_tokens, use_eagle, enable_caching, enable_kv_cache_events, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 57cd1490e81..89061c6bc17 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -114,7 +114,7 @@ class KVCacheManager: max_model_len: int, scheduler_block_size: int, hash_block_size: int, - max_num_batched_tokens: int | None = None, + max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, log_stats: bool = False, @@ -128,8 +128,8 @@ class KVCacheManager: # When unset, fall back to `max_model_len` so the recycling-aware cap # collapses to the prior (uncapped) admission behavior. The scheduler # always supplies the real value at runtime. - if max_num_batched_tokens is None: - max_num_batched_tokens = max_model_len + if max_in_flight_tokens is None: + max_in_flight_tokens = max_model_len self.enable_caching = enable_caching self.use_eagle = use_eagle @@ -143,7 +143,7 @@ class KVCacheManager: self.coordinator = get_kv_cache_coordinator( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, use_eagle=self.use_eagle, enable_caching=self.enable_caching, enable_kv_cache_events=enable_kv_cache_events, @@ -397,9 +397,12 @@ class KVCacheManager: # insufficient free blocks. # Should call this function before allocating new blocks to reduce # the number of evicted blocks. + # Free on the processed-token basis: in-flight steps' attention windows + # still read blocks below the optimistic boundary, and rejected spec + # tokens can roll it back. self.coordinator.remove_skipped_blocks( request.request_id, - total_computed_tokens, + max(0, total_computed_tokens - request.num_in_flight_tokens), num_prompt_tokens=request.num_prompt_tokens, ) @@ -472,7 +475,7 @@ class KVCacheManager: def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Remove the blocks that are no longer needed from `blocks` and replace @@ -480,12 +483,12 @@ class KVCacheManager: Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for R-SWA gap eviction. """ self.coordinator.remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index b13c23d8040..88e60987baf 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1051,11 +1051,12 @@ def unify_kv_cache_spec_page_size( ) -> dict[str, KVCacheSpec]: """ Unify the page size of the given KVCacheSpec. If the page size of all layers - are the same, return the original KVCacheSpec. If not same, first try to - unify page size by increasing the block size of layers with smaller page - size. If a smaller attention page does not evenly divide the maximum page - size, keep its logical block size and pad its physical page instead --- but - only for attention layers whose backend opts in via + are the same, return the original KVCacheSpec. If not same, unify the page + size by increasing the block size of layers with smaller page size. Two + cases cannot be unified by block size alone and pad their physical page to + the maximum instead: Mamba layers, whose page size comes from state shapes + and is independent of block size; and attention layers whose page does not + evenly divide the maximum and whose backend opts in via ``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through a strided view, which not every backend handles). Raise NotImplementedError if failed to unify the page size. @@ -1076,6 +1077,16 @@ def unify_kv_cache_spec_page_size( for layer_name, layer_spec in kv_cache_spec.items(): if layer_spec.page_size_bytes == max_page_size: new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, MambaSpec): + # MambaSpec's page size is determined by its state shapes and does + # not scale with block_size, so pad the page instead. This is the + # same padding mechanism the platform uses to align Mamba pages + # with the main model's attention page size; it is needed here + # when another layer (e.g. from a draft model) has a larger page + # than the already-aligned Mamba page. + new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec else: layer_page_size = layer_spec.page_size_bytes if max_page_size % layer_page_size == 0: @@ -1312,9 +1323,6 @@ def _get_kv_cache_config_packed( return num_blocks, kv_cache_tensors -_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed - - def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9ce1d94ef3c..483cc2f543b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -262,7 +262,7 @@ class Scheduler(SchedulerInterface): self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, - max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, + max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, log_stats=self.log_stats, @@ -811,8 +811,10 @@ class Scheduler(SchedulerInterface): # Pad new decode requests to uniform spec decoding size to # preserve full cudagraph for this step. + # Not for diffusion where draft tokens can't be padded. if ( (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None) + and self.num_sampled_tokens_per_step > 0 and num_new_tokens == 1 and (scheduled_running_reqs and not prefill_scheduled) ): @@ -1078,10 +1080,11 @@ class Scheduler(SchedulerInterface): self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. + new_attn_block_ids = self.kv_cache_manager.take_new_block_ids() new_block_ids_to_zero = ( - (self.kv_cache_manager.take_new_block_ids() or None) - if self.needs_kv_cache_zeroing - else None + (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None ) # Dynamic speculative decoding: compute optimal K @@ -1177,6 +1180,7 @@ class Scheduler(SchedulerInterface): for req_id, num_scheduled_token in num_scheduled_tokens.items(): request = self.requests[req_id] request.num_computed_tokens += num_scheduled_token + request.num_in_flight_tokens += num_scheduled_token if self.defer_block_free: # Record the in-flight step, to fence deferred block freeing. request.last_sched_seq = self.sched_step_seq @@ -1563,10 +1567,12 @@ class Scheduler(SchedulerInterface): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is not None: + request.num_in_flight_tokens -= num_tokens_scheduled if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: # skip failed or rescheduled requests from KV load failure continue - request = self.requests.get(req_id) if request is None or request.is_finished(): # The request is already finished. This can happen if the # request is aborted while the model is executing it (e.g., @@ -1585,8 +1591,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 or self.num_sampled_tokens_per_step == 0 + # Skip a stale frame still pending discard (async_tokens_to_discard + # > 0): its pre-reset rejection count would underflow the counters. + if ( + scheduled_spec_token_ids + and (generated_token_ids or self.num_sampled_tokens_per_step == 0) + and request.async_tokens_to_discard == 0 ): num_draft_tokens = len(scheduled_spec_token_ids) num_sampled = self.num_sampled_tokens_per_step @@ -2360,10 +2370,12 @@ class Scheduler(SchedulerInterface): return False, None # Free any out-of-window prefix blocks before we hand the block table to - # the connector. + # the connector, on the processed-token basis (see `allocate_slots`). self.kv_cache_manager.remove_skipped_blocks( request_id=request.request_id, - total_computed_tokens=request.num_computed_tokens, + processed_computed_tokens=max( + 0, request.num_computed_tokens - request.num_in_flight_tokens + ), num_prompt_tokens=request.num_prompt_tokens, ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 642fe3e6a08..87c3f8feb72 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -45,6 +45,7 @@ class SingleTypeKVCacheManager(ABC): scheduler_block_size: int, dcp_world_size: int = 1, pcp_world_size: int = 1, + needs_kv_cache_zeroing: bool = False, max_admission_blocks_per_request: int | None = None, ) -> None: """ @@ -55,6 +56,8 @@ class SingleTypeKVCacheManager(ABC): kv_cache_group_id: The id of the kv cache group of this manager. scheduler_block_size: The scheduling granularity (LCM of all group block sizes); a multiple of this manager's ``block_size``. + needs_kv_cache_zeroing: Whether worker-side KV cache zeroing needs + newly allocated block IDs from this manager. max_admission_blocks_per_request: Recycling-aware per-request block cap used by `get_num_blocks_to_allocate`. Only set for spec types that recycle blocks across chunks (SWA, @@ -73,6 +76,14 @@ class SingleTypeKVCacheManager(ABC): self.block_pool = block_pool self.enable_caching = enable_caching self._max_admission_blocks_per_request = max_admission_blocks_per_request + # Record newly allocated block ids only when worker-side zeroing will + # consume them and this manager holds a spec type that gets zeroed. + self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + HiddenStateCacheSpec, + ) self.new_block_ids: list[int] = [] # Mapping from request ID to blocks to track the blocks allocated @@ -268,12 +279,7 @@ class SingleTypeKVCacheManager(ABC): cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) ) req_blocks.extend(allocated_blocks) - if type(self.kv_cache_spec) in ( - FullAttentionSpec, - TQFullAttentionSpec, - MLAAttentionSpec, - HiddenStateCacheSpec, - ): + if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -301,12 +307,7 @@ class SingleTypeKVCacheManager(ABC): else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) - if type(self.kv_cache_spec) in ( - FullAttentionSpec, - TQFullAttentionSpec, - MLAAttentionSpec, - HiddenStateCacheSpec, - ): + if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in new_blocks) return new_blocks @@ -507,7 +508,7 @@ class SingleTypeKVCacheManager(ABC): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """ @@ -519,15 +520,15 @@ class SingleTypeKVCacheManager(ABC): Args: request_id: The request ID. - total_computed_tokens: The total number of computed tokens, including - local computed tokens and external computed tokens. + processed_computed_tokens: Computed-token prefix length covering + fully processed and committed tokens only (safe to free). num_prompt_tokens: Optional prompt length for attention types (e.g. R-SWA) that evict a middle gap rather than a head prefix. Ignored by the default implementation. """ del num_prompt_tokens # Remove the blocks that will be skipped during attention computation. - num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) + num_skipped_tokens = self.get_num_skipped_tokens(processed_computed_tokens) if num_skipped_tokens <= 0: # This indicates that ALL tokens are inside attention window. # Thus we do not need to free any blocks outside attention window. @@ -638,14 +639,14 @@ class RSWAManager(FullAttentionManager): def remove_skipped_blocks( self, request_id: str, - total_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Free gap blocks that are no longer needed for attention. Gap = blocks entirely within [ceil(prefix_len / block_size) * block_size, - max(prefix_len, total_computed_tokens - rswa_window)) + max(prefix_len, processed_computed_tokens - rswa_window)) Freed blocks are replaced with null_block in req_to_blocks so the block_table passed to FA4 is valid (null_block KV is all-zero; @@ -653,7 +654,7 @@ class RSWAManager(FullAttentionManager): """ if num_prompt_tokens is None: super().remove_skipped_blocks( - request_id, total_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) return @@ -661,7 +662,9 @@ class RSWAManager(FullAttentionManager): # First block fully after the prefill boundary. first_gap_block = cdiv(num_prompt_tokens, bs) # Decode window start position; blocks before this are evictable. - window_start = max(num_prompt_tokens, total_computed_tokens - self.rswa_window) + window_start = max( + num_prompt_tokens, processed_computed_tokens - self.rswa_window + ) last_gap_block = window_start // bs # exclusive upper bound self._remove_blocks_in_range(request_id, first_gap_block, last_gap_block) @@ -1143,20 +1146,13 @@ class MambaManager(SingleTypeKVCacheManager): def remove_skipped_blocks( self, request_id: str, - num_computed_tokens: int, + processed_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: assert isinstance(self.kv_cache_spec, MambaSpec) - # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain - # draft tokens from the previous step that may or may not be rejected later. - # This can make us think we are further ahead in the sequence than we actually - # are, so let's assume that all tokens are rejected so we don't free blocks - # that we might actually need. - num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks) - super().remove_skipped_blocks( - request_id, num_computed_tokens, num_prompt_tokens + request_id, processed_computed_tokens, num_prompt_tokens ) if self.mamba_cache_mode == "align": # `last_state_block_idx` refers to the block index allocated two steps ago. @@ -1170,7 +1166,7 @@ class MambaManager(SingleTypeKVCacheManager): if ( last_state_block_idx is not None and last_state_block_idx - < cdiv(num_computed_tokens, self.block_size) - 1 + < cdiv(processed_computed_tokens, self.block_size) - 1 ): blocks = self.req_to_blocks[request_id] if blocks[last_state_block_idx] != self._null_block: @@ -1454,7 +1450,7 @@ class SinkFullAttentionManager(FullAttentionManager): def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, - max_num_batched_tokens: int, + max_in_flight_tokens: int, max_model_len: int, **kwargs, ) -> SingleTypeKVCacheManager: @@ -1467,7 +1463,8 @@ def get_manager_for_kv_cache_spec( Args: kv_cache_spec: The KVCacheSpec instance - max_num_batched_tokens: The maximum number of tokens in a batch + max_in_flight_tokens: The max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens` max_model_len: The maximum context length the model could serve Returns: An instance of the appropriate SingleTypeKVCacheManager subclass @@ -1488,7 +1485,7 @@ def get_manager_for_kv_cache_spec( ): kwargs["max_admission_blocks_per_request"] = ( kv_cache_spec.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, + max_in_flight_tokens=max_in_flight_tokens, max_model_len=max_model_len, ) ) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index d5cf1050ca4..bcb441e7564 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1434,6 +1434,12 @@ class DPLBAsyncMPClient(DPAsyncMPClient): # Increment local waiting count for better balancing between stats # updates from the coordinator (which happen every 100ms). current_counts[eng_index][0] += self.client_count + # Rotate the scan start so that ties (equal scores, e.g. right + # after a coordinator stats reset when engines look equally loaded) + # don't systematically favor the same engine. This removes the + # fixed tie-break bias without affecting load-aware decisions when + # scores actually differ. + self.eng_start_index = (self.eng_start_index + 1) % num_engines chosen_engine = self.core_engines[eng_index] # Record which engine is chosen for this request, to handle aborts. diff --git a/vllm/v1/engine/detokenizer.py b/vllm/v1/engine/detokenizer.py index 4700eecb59a..50f14b9f96a 100644 --- a/vllm/v1/engine/detokenizer.py +++ b/vllm/v1/engine/detokenizer.py @@ -6,7 +6,7 @@ import tokenizers import tokenizers.decoders from packaging import version from tokenizers import Tokenizer -from transformers import PreTrainedTokenizerFast +from transformers import TokenizersBackend from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike @@ -57,7 +57,7 @@ class IncrementalDetokenizer: # No tokenizer => skipping detokenization. return IncrementalDetokenizer() - if USE_FAST_DETOKENIZER and isinstance(tokenizer, PreTrainedTokenizerFast): + if USE_FAST_DETOKENIZER and isinstance(tokenizer, TokenizersBackend): # Fast tokenizer => use tokenizers library DecodeStream. return FastIncrementalDetokenizer(tokenizer, request) @@ -165,7 +165,7 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC): class FastIncrementalDetokenizer(BaseIncrementalDetokenizer): - def __init__(self, tokenizer: PreTrainedTokenizerFast, request: EngineCoreRequest): + def __init__(self, tokenizer: TokenizersBackend, request: EngineCoreRequest): super().__init__(request) sampling_params = request.sampling_params diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 0937fad8f1f..7633ca89cdf 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -280,9 +280,12 @@ class MultiprocExecutor(Executor): logger.debug("MultiprocWorkerMonitor: shutdown already initiated") return _self.is_failed = True - proc_name = next(h.proc.name for h in workers if h.proc.sentinel == died[0]) + proc = next(h.proc for h in workers if h.proc.sentinel == died[0]) logger.error( - "Worker proc %s died unexpectedly, shutting down executor.", proc_name + "Worker proc %s died unexpectedly (exit code: %s), " + "shutting down executor.", + proc.name, + proc.exitcode, ) _self.shutdown() callback = _self.failure_callback diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 323b1e763a5..4204c31be58 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -482,25 +482,28 @@ class ChunkedLocalAttentionSpec(AttentionSpec): attention_chunk_size: int def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. Single source of truth for both startup pool sizing (`max_memory_usage_bytes`) and the runtime admission gate, so requests admitted by startup can also be admitted at runtime. + + `max_in_flight_tokens` is the max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. """ - # During chunked prefill, we hold KV for at most one chunk window. + # During chunked prefill, we hold KV for at most one chunk window plus + # the in-flight tokens, since frees happen on the processed-token basis. num_tokens = min( - self.attention_chunk_size + max_num_batched_tokens, max_model_len + self.attention_chunk_size + max_in_flight_tokens, max_model_len ) return cdiv(num_tokens, self.block_size) def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_model_len=vllm_config.model_config.max_model_len, ) return max_blocks * self.page_size_bytes @@ -544,7 +547,7 @@ class SlidingWindowSpec(AttentionSpec): ) def max_admission_blocks_per_request( - self, max_num_batched_tokens: int, max_model_len: int + self, max_in_flight_tokens: int, max_model_len: int ) -> int: """Per-request admission cap, in blocks. @@ -553,13 +556,14 @@ class SlidingWindowSpec(AttentionSpec): real-held blocks plateau at this bound because `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots` before each chunk's `get_num_blocks_to_allocate`. + + `max_in_flight_tokens` is the max tokens scheduled but not yet settled + (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`. """ # During chunked prefill, we hold KV for the last `sliding_window-1` - # computed tokens plus the newly scheduled tokens, and never more - # than `max_model_len`. - num_tokens = min( - self.sliding_window - 1 + max_num_batched_tokens, max_model_len - ) + # computed tokens plus the in-flight tokens (frees happen on the + # processed-token basis); never more than `max_model_len`. + num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. @@ -569,10 +573,9 @@ class SlidingWindowSpec(AttentionSpec): assert vllm_config.parallel_config.decode_context_parallel_size == 1, ( "DCP not support sliding window." ) - max_model_len = vllm_config.model_config.max_model_len - max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_blocks = self.max_admission_blocks_per_request( - max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len + max_in_flight_tokens=vllm_config.max_in_flight_tokens, + max_model_len=vllm_config.model_config.max_model_len, ) return max_blocks * self.page_size_bytes diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 48838599d6a..5a2e3c184d3 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -12,7 +12,6 @@ from typing import TYPE_CHECKING, Any, NamedTuple, NewType import numpy as np import torch -from typing_extensions import override from vllm.logger import init_logger from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes @@ -85,21 +84,12 @@ class ScheduleEndContext(NamedTuple): preempted_req_ids: Collection[str] -class LoadStoreSpec(ABC): +class LoadStoreSpec: """ - Abstract metadata that encapsulates information allowing a worker + Metadata that encapsulates information allowing a worker to load, and optionally also to store, blocks of KV data. """ - @staticmethod - @abstractmethod - def medium() -> str: - """ - Returns a string representation of the medium type - this store/load targets. - """ - pass - @dataclass class PrepareStoreOutput: @@ -313,6 +303,10 @@ class OffloadingManager(ABC): """ Take the offloading events from the manager. + A tier manager emits only events for storage state it owns. A + composing manager may aggregate child event streams, but should not + synthesize events on behalf of a child tier. + Yields: New OffloadingEvents collected since the last call. """ @@ -392,11 +386,6 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec): self.group_sizes: Sequence[int] = group_sizes self.block_indices: Sequence[int] = block_indices - @staticmethod - @override - def medium() -> str: - return "GPU" - @dataclass class CanonicalKVCacheTensor: diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index 14c96680fd0..29f95e1b407 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -1,21 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing_extensions import override - from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec class CPUOffloadingMetrics: STORES_SKIPPED = "vllm:kv_offload_stores_skipped" CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc" + CPU_ALLOCATION_SIZE = "vllm:kv_offload_cpu_allocation_size" class CPULoadStoreSpec(BlockIDsLoadStoreSpec): """ Spec for loading/storing a KV block to CPU memory. """ - - @staticmethod - @override - def medium() -> str: - return "CPU" diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 0424196c9fd..e7416bf63bd 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -6,6 +6,7 @@ from typing import Literal from typing_extensions import override +from vllm.distributed.kv_events import MEDIUM_CPU from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) @@ -51,7 +52,7 @@ class CPUOffloadingManager(OffloadingManager): store_threshold: int = 1, max_tracker_size: int = 64_000, ): - self.medium: str = CPULoadStoreSpec.medium() + self.medium: str = MEDIUM_CPU self._num_blocks: int = num_blocks self._num_allocated_blocks: int = 0 self._free_list: list[int] = [] @@ -69,6 +70,7 @@ class CPUOffloadingManager(OffloadingManager): self.store_threshold: int = store_threshold self.max_tracker_size: int = max_tracker_size self.stores_skipped_in_current_batch: int = 0 + self.allocation_sizes_in_current_batch: list[int] = [] # Number of block references. It is ordered so can evict the LRU entry in O(1). self.counts: OrderedDict[OffloadKey, int] | None = ( @@ -150,7 +152,7 @@ class CPUOffloadingManager(OffloadingManager): @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: - self._policy.touch(keys) + self._policy.touch(keys, req_context) @override def complete_load( @@ -185,6 +187,7 @@ class CPUOffloadingManager(OffloadingManager): evicted_keys=[], ) + self.allocation_sizes_in_current_batch.append(len(keys_to_store)) num_blocks_to_evict = len(keys_to_store) - self._get_num_free_blocks() to_evict: list[OffloadKey] = [] @@ -300,10 +303,17 @@ class CPUOffloadingManager(OffloadingManager): usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0 stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage) + for allocation_size in self.allocation_sizes_in_current_batch: + stats.observe_histogram( + CPUOffloadingMetrics.CPU_ALLOCATION_SIZE, allocation_size + ) + self.allocation_sizes_in_current_batch.clear() + if self.store_threshold >= 2: stats.increase_counter( CPUOffloadingMetrics.STORES_SKIPPED, self.stores_skipped_in_current_batch, ) self.stores_skipped_in_current_batch = 0 + return stats diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index 7d22e518654..f682a47e45f 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -72,7 +72,7 @@ class ARCCachePolicy(CachePolicy): self.t2.pop(key, None) @override - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): if key in self.t1: block = self.t1.pop(key) diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index f898a60b0f6..2b6681e4992 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -4,7 +4,7 @@ import ctypes from abc import ABC, abstractmethod from collections.abc import Iterable -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext class BlockStatus(ctypes.Structure): @@ -57,8 +57,14 @@ class CachePolicy(ABC): """Remove a block (used to clean up after a failed store).""" @abstractmethod - def touch(self, keys: Iterable[OffloadKey]) -> None: - """Mark blocks as recently used.""" + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + """ + Mark blocks as recently used. + + Args: + keys: Blocks to mark as recently used. + req_context: Per-request context for the request touching these blocks. + """ @abstractmethod def evict( diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index 47e18f5565f..efa24fe9033 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -39,7 +39,7 @@ class LRUCachePolicy(CachePolicy): self.evictable_blocks.pop(key, None) @override - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): if key in self.evictable_blocks: self.evictable_blocks.move_to_end(key) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 16729a9dbb4..26ea3728191 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -12,6 +12,7 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCaches, OffloadingCounterMetadata, OffloadingGaugeMetadata, + OffloadingHistogramMetadata, OffloadingManager, OffloadingMetricMetadata, OffloadingSpec, @@ -37,7 +38,14 @@ class CPUOffloadingSpec(OffloadingSpec): "values indicate transfers (stores or promotions) may be " "dropped due to insufficient capacity." ), - ) + ), + CPUOffloadingMetrics.CPU_ALLOCATION_SIZE: OffloadingHistogramMetadata( + documentation=( + "Histogram of the number of CPU blocks requested by each " + "KV offload prepare_store call." + ), + buckets=(1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144), + ), } store_threshold = int(extra_config.get("store_threshold", 0)) if store_threshold >= 2: diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index b022e1f2f8a..f83113e137c 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -13,6 +13,7 @@ import numpy as np from vllm.v1.kv_offload.base import ( LookupResult, + OffloadingEvent, OffloadingMetricMetadata, OffloadKey, ReqContext, @@ -49,6 +50,43 @@ class JobResult: success: bool +class ParentManager(ABC): + """Interface for secondary tiers to call back into the tiering manager. + + Passed to secondary tiers via serve_external_requests() each step. + The _SecondaryTierFacingParent wrapper implements this, automatically + excluding the calling tier from fan-out operations. + + Required call sequence for each remote request: + 1. on_new_request(req_context) — set up per-request state + 2. lookup(key, req_context) — check block availability + (repeat per block) + 3. create_store_job(keys, req_context) — pin blocks and get a + job handle + 4. on_request_finished(req_context) — clean up per-request state + + Steps 2-3 may be interleaved. Step 4 must be called even if no + blocks were found, to avoid leaking async lookup state (e.g. in + the fs tier's AsyncLookupManager). + """ + + @abstractmethod + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: ... + + @abstractmethod + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: ... + + @abstractmethod + def create_store_job( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + ) -> JobMetadata: ... + + @abstractmethod + def on_request_finished(self, req_context: ReqContext) -> None: ... + + class SecondaryTierManager(ABC): """ Abstract interface for managing a single non-primary offloading tier. @@ -171,6 +209,10 @@ class SecondaryTierManager(ABC): """ return False + def take_events(self) -> Iterable[OffloadingEvent]: + """Take KV events for storage state owned by this tier.""" + return () + def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as recently used for eviction policy. @@ -210,11 +252,20 @@ class SecondaryTierManager(ABC): """ return + def serve_external_requests(self, parent: ParentManager) -> None: + """Process remotely-originated requests using the parent manager. + + Called once per scheduler step, BEFORE _flush_pending_promotions(). + The parent handle is valid only for the duration of this call. + Tiers that don't serve external requests leave this as a no-op. + """ + return + def on_schedule_end(self, context: ScheduleEndContext) -> None: """Called once at the end of each scheduler step. - Secondary tiers may override this for per-step cleanup or - deferred work submission. + Args: + context: Per-step context from the scheduler. """ return diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 81151a7c0c9..728ac8dc86e 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -48,6 +48,7 @@ from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.tiering.base import ( JobId, JobMetadata, + ParentManager, SecondaryTierManager, ) @@ -120,6 +121,35 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): self._mmap_region.cleanup() +class _SecondaryTierFacingParent(ParentManager): + """Wrapper that implements ParentManager by delegating to the + TieringOffloadingManager with exclude_tier set to the origin tier.""" + + __slots__ = ("_m", "_origin") + + def __init__( + self, + manager: "TieringOffloadingManager", + tier: SecondaryTierManager, + ): + self._m = manager + self._origin = tier + + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return self._m.on_new_request(req_context, exclude_tier=self._origin) + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + return self._m.lookup(key, req_context, exclude_tier=self._origin) + + def create_store_job( + self, keys: Collection[OffloadKey], req_context: ReqContext + ) -> JobMetadata: + return self._m.create_store_job(keys, req_context) + + def on_request_finished(self, req_context: ReqContext) -> None: + return self._m.on_request_finished(req_context, exclude_tier=self._origin) + + class TieringOffloadingManager(OffloadingManager): """ Orchestrates multi-tier KV cache offloading. @@ -140,7 +170,6 @@ class TieringOffloadingManager(OffloadingManager): self, primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, - enable_events: bool = False, ): """ Initialize the TieringOffloadingManager. @@ -149,14 +178,11 @@ class TieringOffloadingManager(OffloadingManager): primary_tier: The primary tier manager (CPU-based). secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. - enable_events: Whether to track offloading events """ self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] self._job_id_counter: int = 0 - self.events: list[OffloadingEvent] | None = [] if enable_events else None - # Job tracking: maps job_id to metadata for all in-flight transfers. # JobMetadata.is_promotion distinguishes direction: # True: secondary → primary (promotion) @@ -180,6 +206,12 @@ class TieringOffloadingManager(OffloadingManager): # complete_store(), since complete_store() can still submit cascades. self._req_state: dict[str, RequestState] = {} + # Cached ParentManager wrappers for each secondary tier. + self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { + tier: _SecondaryTierFacingParent(self, tier) + for tier in self.secondary_tiers + } + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -235,7 +267,13 @@ class TieringOffloadingManager(OffloadingManager): ) @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + def lookup( + self, + key: OffloadKey, + req_context: ReqContext, + *, + exclude_tier: SecondaryTierManager | None = None, + ) -> LookupResult: """ Check whether a single block is offloaded and ready. @@ -257,6 +295,10 @@ class TieringOffloadingManager(OffloadingManager): MISS — block not found in any tier, or primary is full and cannot accept a promotion. """ + # Poll first so a promotion that finished since the last call is + # already reflected as HIT (not stale HIT_PENDING/MISS) below, and + # so blocks freed by cascade or promotion completions are evictable + # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() primary_hit = self.primary_tier.lookup(key, req_context) @@ -267,6 +309,8 @@ class TieringOffloadingManager(OffloadingManager): any_retry = False for tier in self.secondary_tiers: + if tier is exclude_tier: + continue result = tier.lookup(key, req_context) if result is LookupResult.HIT: if not self._initiate_promotion(tier, key, req_context): @@ -359,8 +403,8 @@ class TieringOffloadingManager(OffloadingManager): """ Prepare blocks to be loaded from primary tier to GPU. - CRITICAL: This method calls _maybe_process_finished_jobs() FIRST to ensure - that any completed promotions have been finalized and blocks are ready. + Callers only pass keys already confirmed HIT by lookup() earlier this + step. This increments ref_cnt on the blocks in the primary tier, protecting them from eviction during the transfer. @@ -372,9 +416,6 @@ class TieringOffloadingManager(OffloadingManager): Returns: LoadStoreSpec for reading from primary tier. """ - # Process completed promotions to ensure blocks are ready - self._maybe_process_finished_jobs() - return self.primary_tier.prepare_load(keys, req_context) @override @@ -427,8 +468,15 @@ class TieringOffloadingManager(OffloadingManager): evicted, or None if store cannot proceed. """ # Step 1: Poll for completed async jobs FIRST - # This decrements ref_cnt on primary blocks that have been - # successfully transferred to secondary tiers. + # _process_finished_jobs() handles two kinds of completions here: + # - Cascade completions (store to a secondary tier, either a local + # cascade or a store job created for a remote requester via + # create_store_job()): decrements ref_cnt on the primary blocks + # that were read, making them evictable again once ref_cnt hits 0. + # - Promotion completions (secondary->primary loads): sets a + # not-yet-ready block's ref_cnt from -1 to 0 via complete_write(), + # making it evictable for the first time. + # Both must be accounted for before the eviction decision below. self._maybe_process_finished_jobs() # Step 2: Store to primary tier (new blocks only). @@ -478,20 +526,7 @@ class TieringOffloadingManager(OffloadingManager): return for tier in request_level_tiers: - primary_blocks_spec = self.primary_tier.prepare_read( - ready_keys, req_context - ) - - job_id = self._next_job_id() - assert isinstance(primary_blocks_spec, CPULoadStoreSpec) - job_metadata = JobMetadata( - job_id=job_id, - keys=ready_keys, - block_ids=primary_blocks_spec.block_ids, - is_promotion=False, - req_context=req_context, - ) - self._transfer_jobs[job_id] = job_metadata + job_metadata = self.create_store_job(ready_keys, req_context) tier.submit_store(job_metadata) @override @@ -500,7 +535,7 @@ class TieringOffloadingManager(OffloadingManager): keys: Collection[OffloadKey], req_context: ReqContext, success: bool = True, - ): + ) -> None: """ Mark blocks as done storing from GPU to primary tier. @@ -529,22 +564,7 @@ class TieringOffloadingManager(OffloadingManager): # eviction during the async transfer). One prepare_read() call per # secondary tier. for tier in self.secondary_tiers: - primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context) - - # Submit async store job: primary→secondary - job_id = self._next_job_id() - - # Track this store job - assert isinstance(primary_blocks_spec, CPULoadStoreSpec) - job_metadata = JobMetadata( - job_id=job_id, - keys=keys, - block_ids=primary_blocks_spec.block_ids, - is_promotion=False, - req_context=req_context, - ) - self._transfer_jobs[job_id] = job_metadata - + job_metadata = self.create_store_job(keys, req_context) tier.submit_store(job_metadata) # Note: The async transfers are now in flight. Their completion is @@ -555,8 +575,40 @@ class TieringOffloadingManager(OffloadingManager): state.pending_primary_stores -= 1 self._maybe_finalize_request(req_id) + def create_store_job( + self, + keys: Collection[OffloadKey], + req_context: ReqContext, + ) -> JobMetadata: + """Pin blocks in the primary tier and create a tracked store job. + + Calls prepare_read() to increment ref_cnt (protecting blocks + from eviction during the async transfer), allocates a job ID, + and registers the job in _transfer_jobs. + + The caller is responsible for the actual data transfer and + reporting completion via get_finished_jobs(). + """ + primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context) + assert isinstance(primary_blocks_spec, CPULoadStoreSpec) + job_id = self._next_job_id() + job_metadata = JobMetadata( + job_id=job_id, + keys=keys, + block_ids=primary_blocks_spec.block_ids, + is_promotion=False, + req_context=req_context, + ) + self._transfer_jobs[job_id] = job_metadata + return job_metadata + @override - def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + def on_new_request( + self, + req_context: ReqContext, + *, + exclude_tier: SecondaryTierManager | None = None, + ) -> RequestOffloadingContext: """ Query each secondary tier for its offload policy preference. @@ -565,6 +617,8 @@ class TieringOffloadingManager(OffloadingManager): """ state = RequestState(req_context=req_context) for tier in self.secondary_tiers: + if tier is exclude_tier: + continue tier_ctx = tier.on_new_request(req_context) if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL: if state.request_level_tiers is None: @@ -580,13 +634,22 @@ class TieringOffloadingManager(OffloadingManager): return RequestOffloadingContext(policy=policy) @override - def on_request_finished(self, req_context: ReqContext) -> None: + def on_request_finished( + self, + req_context: ReqContext, + *, + exclude_tier: SecondaryTierManager | None = None, + ) -> None: self.primary_tier.on_request_finished(req_context) state = self._req_state[req_context.req_id] state.is_finished = True - self._maybe_finalize_request(req_context.req_id) + self._maybe_finalize_request(req_context.req_id, exclude_tier) - def _maybe_finalize_request(self, req_id: str) -> None: + def _maybe_finalize_request( + self, + req_id: str, + exclude_tier: SecondaryTierManager | None = None, + ) -> None: """Finalize secondary tiers once no more store cascades can be submitted. Finalization means forwarding on_request_finished() to secondary tiers. @@ -600,6 +663,8 @@ class TieringOffloadingManager(OffloadingManager): return for tier in self.secondary_tiers: + if tier is exclude_tier: + continue tier.on_request_finished(state.req_context) del self._req_state[req_id] @@ -611,8 +676,18 @@ class TieringOffloadingManager(OffloadingManager): Called once per scheduler step from OffloadingConnectorScheduler.build_connector_meta(). """ + # Catch-all poll: guarantees jobs are processed even on steps where + # lookup()/prepare_store() were never called (e.g. no requests + # scheduled but a tier still has_pending_work()). self._maybe_process_finished_jobs() + + for tier in self.secondary_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + + # Reset the per-step gate AFTER serve_external_requests so that + # lookup() calls within it skip redundant _process_finished_jobs(). self._processed_jobs_this_step = False + self._flush_pending_promotions() for tier in self.secondary_tiers: tier.on_schedule_end(context) @@ -628,16 +703,14 @@ class TieringOffloadingManager(OffloadingManager): @override def take_events(self) -> Iterable[OffloadingEvent]: - """Yield offloading events collected since the last call. + """Yield events owned by the primary and secondary tiers. Yields: - New OffloadingEvents collected since the last call. + New OffloadingEvents collected by each tier since the last call. """ - if self.events is not None: - yield from self.events - self.events.clear() - yield from self.primary_tier.take_events() + for tier in self.secondary_tiers: + yield from tier.take_events() @override def reset_cache(self) -> None: diff --git a/vllm/v1/kv_offload/tiering/obj/config.py b/vllm/v1/kv_offload/tiering/obj/config.py index 5507c6a198e..76a1ae29e55 100644 --- a/vllm/v1/kv_offload/tiering/obj/config.py +++ b/vllm/v1/kv_offload/tiering/obj/config.py @@ -2,29 +2,45 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Connection configuration for the object store secondary tier.""" -from dataclasses import dataclass +from dataclasses import dataclass, field @dataclass class ObjStoreConfig: - """Connection parameters for an object store backend.""" + """Connection parameters for an object store backend. + + When ``access_key`` and ``secret_key`` are left empty the NIXL OBJ + plugin falls back to the AWS SDK default credential provider chain + (IAM roles, environment variables, credential files, etc.), which + enables workload-identity based auth on Kubernetes. + """ bucket: str endpoint_override: str - access_key: str - secret_key: str + access_key: str = field(default="", repr=False) + secret_key: str = field(default="", repr=False) + session_token: str = field(default="", repr=False) + region: str = "" scheme: str = "http" ca_bundle: str = "" def to_nixl_params(self) -> dict[str, str]: - """Build the NIXL backend params dict.""" + """Build the NIXL backend params dict. + + Credential and optional fields are only included when non-empty + so that the AWS SDK default credential chain can activate. + """ params: dict[str, str] = { "bucket": self.bucket, "endpoint_override": self.endpoint_override, "scheme": self.scheme, - "access_key": self.access_key, - "secret_key": self.secret_key, } - if self.ca_bundle: - params["ca_bundle"] = self.ca_bundle + # Omit empty optional fields so the NIXL OBJ plugin's underlying + # AWS SDK can fall back to its default credential provider chain + # (IAM roles, env vars, credential files, etc.). + # https://github.com/ai-dynamo/nixl/blob/main/src/plugins/obj/README.md + for key in ("access_key", "secret_key", "session_token", "region", "ca_bundle"): + value = getattr(self, key) + if value: + params[key] = value return params diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 4c3fa754bf6..6060370ea06 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -157,8 +157,10 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): except Exception as e: raise RuntimeError( f"Object store tier connectivity probe failed — check bucket, " - f"endpoint_override, access_key, secret_key, and scheme. " - f"Error: {e}" + f"endpoint_override, and scheme. If using explicit credentials " + f"verify access_key and secret_key; otherwise ensure the AWS " + f"SDK default credential chain is configured (IAM role, env " + f"vars, credential file). Error: {e}" ) from e def _exists(self, obj_key: str) -> bool: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 406c94e3d79..3dc31a3622e 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -168,7 +168,6 @@ class TieringOffloadingSpec(CPUOffloadingSpec): tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, - enable_events=self.kv_events_config.enable_kv_cache_events, ) if int(self.extra_config.get("store_threshold", 0)) >= 2: raise ValueError( diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 0e8d4ee006f..e9946a7f76b 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -141,6 +141,11 @@ class Request: self.num_output_placeholders = 0 self.async_tokens_to_discard = 0 + # Tokens of steps whose output is not yet processed (async scheduling + # and PP run ahead of the GPU); `num_computed_tokens` counts them + # optimistically. + self.num_in_flight_tokens = 0 + # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps # so the worker's broadcast slot ring stays consistent. self.next_decode_eligible_step = 0 diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index dfaa2234eb9..44ba902f450 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -122,9 +122,7 @@ class SimpleCPUOffloadScheduler: self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, - max_num_batched_tokens=( - vllm_config.scheduler_config.max_num_batched_tokens - ), + max_in_flight_tokens=vllm_config.max_in_flight_tokens, use_eagle=False, enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 4d27e308e88..756c5f3b371 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1468,6 +1468,26 @@ class SpecDecodeBaseProposer: "Sharing target model embedding weights with the draft model." ) + if share_embeddings: + draft_embed = self.model.model.embed_tokens + # Only share when both models use the same embedding width. + # Guard with isinstance so non-Tensor weights (e.g. in tests) + # are not affected — mirrors the weight-equality check above. + if isinstance(target_embed_tokens.weight, torch.Tensor) and isinstance( + draft_embed.weight, torch.Tensor + ): + target_dim = target_embed_tokens.weight.shape[-1] + draft_dim = draft_embed.weight.shape[-1] + if target_dim != draft_dim: + share_embeddings = False + logger.info( + "Target embedding dim (%d) differs from draft " + "embedding dim (%d). Keeping separate embedding " + "weights.", + target_dim, + draft_dim, + ) + if share_embeddings: if hasattr(self.model.model, "embed_tokens"): del self.model.model.embed_tokens diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index a08b341e80a..ed544bb27c1 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -468,7 +468,7 @@ class NgramProposerGPU: def update_scheduler_for_invalid_drafts( - num_valid_draft_tokens_event: torch.Event, + num_valid_draft_tokens_event: torch.cuda.Event, num_valid_draft_tokens_cpu: torch.Tensor, scheduler_output: "SchedulerOutput", req_id_to_index: dict[str, int], @@ -643,7 +643,7 @@ def _sync_num_tokens( def copy_num_valid_draft_tokens( num_valid_draft_tokens_cpu: torch.Tensor, num_valid_draft_tokens_copy_stream: torch.cuda.Stream, - num_valid_draft_tokens_event: torch.Event, + num_valid_draft_tokens_event: torch.cuda.Event, num_valid_draft_tokens: torch.Tensor | None, batch_size: int, ) -> None: diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 31178e9f246..19b2c76e9d3 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch +from transformers import MistralCommonBackend from vllm.logger import init_logger from vllm.sampling_params import SamplingParams @@ -95,6 +96,10 @@ class GuidanceBackend(StructuredOutputBackend): if is_mistral_tokenizer(self.tokenizer): self.ll_tokenizer = self.tokenizer.llg_tokenizer + elif isinstance(self.tokenizer, MistralCommonBackend): + from mistral_common.guidance.tokenizer import from_mistral_tokenizer + + self.ll_tokenizer = from_mistral_tokenizer(self.tokenizer.tokenizer) else: self.ll_tokenizer = llguidance_hf.from_tokenizer( self.tokenizer, max(self.vocab_size, len(self.tokenizer)) diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index deec52e44ba..bd1f96c71ed 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -48,6 +48,7 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]: torch.Event = _EventPlaceholder +torch.cuda.Event = _EventPlaceholder torch.cuda.Stream = _StreamPlaceholder torch.cuda.set_stream = noop torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder() diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index a9ad16b1520..e4659104f49 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -24,7 +24,8 @@ class AsyncOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.copy_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -81,7 +82,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.pooler_output = pooler_output self.is_valid = is_valid - self.copy_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 3906717b5b7..228fd08a751 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, replace from math import prod from typing import Any, cast @@ -12,8 +12,10 @@ from vllm.config import ( get_layers_from_vllm_config, set_current_vllm_config, ) +from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import ( AttentionCGSupport, @@ -25,6 +27,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, KVQuantMode, MambaSpec, + TQFullAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata @@ -35,6 +38,8 @@ from vllm.v1.worker.utils import ( prepare_kernel_block_sizes, ) +logger = init_logger(__name__) + @dataclass(frozen=True) class AttentionCGSupportInfo: @@ -307,6 +312,7 @@ def _reshape_kv_cache( layer_cache_dtype = ( "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE + and not isinstance(kv_cache_spec, TQFullAttentionSpec) else cache_dtype ) kv_cache_shape = group.backend.get_kv_cache_shape( @@ -365,6 +371,14 @@ def _reshape_kv_cache( kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, ) + elif has_attn and kv_cache_config is not None: + _align_mixed_attention_kv_cache_views( + attn_groups=attn_groups, + kv_caches=kv_caches, + kernel_block_sizes=kernel_block_sizes, + cache_dtype=cache_dtype, + kv_cache_config=kv_cache_config, + ) # Map any sharing layers to their target layer's KV cache. for layer_name, target_layer_name in shared_kv_cache_layers.items(): @@ -373,6 +387,77 @@ def _reshape_kv_cache( return kv_caches +def _align_mixed_attention_kv_cache_views( + attn_groups: Iterable[AttentionGroup], + kv_caches: dict[str, Any], + kernel_block_sizes: list[int], + cache_dtype: str, + kv_cache_config: KVCacheConfig, +) -> None: + """Align shared attention KV views when backends disagree on layout. + + Encoder-decoder models can share one raw allocation between decoder + self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention + (blocks-first backends, block dim 0). Keep the physical storage in the + K/V-first layout expected by ROCM_ATTN, and restride the blocks-first + logical views so block IDs address the same bytes. + """ + block_dims_by_layer: dict[str, int] = {} + for group in attn_groups: + kv_cache_spec = group.kv_cache_spec + if not isinstance(kv_cache_spec, AttentionSpec): + continue + if group.kv_cache_group_id >= len(kernel_block_sizes): + continue + block_dim = group.backend.get_kv_cache_block_dim( + kernel_block_sizes[group.kv_cache_group_id], + kv_cache_spec.num_kv_heads, + kv_cache_spec.head_size, + cache_dtype_str=cache_dtype, + ) + for layer_name in group.layer_names: + if layer_name in kv_caches: + block_dims_by_layer[layer_name] = block_dim + + for kv_tensor in kv_cache_config.kv_cache_tensors: + if kv_tensor.block_stride > 0: + continue + shared_block_dims = { + block_dims_by_layer[layer_name] + for layer_name in kv_tensor.shared_by + if layer_name in block_dims_by_layer + } + if 0 not in shared_block_dims or 1 not in shared_block_dims: + continue + + for layer_name in kv_tensor.shared_by: + if block_dims_by_layer.get(layer_name) == 0: + _restride_blocks_first_kv_cache_to_kv_first_storage( + kv_caches[layer_name] + ) + + +def _restride_blocks_first_kv_cache_to_kv_first_storage( + kv_cache: torch.Tensor, +) -> None: + assert kv_cache.ndim >= 3 + assert kv_cache.shape[1] == 2 + page_size = kv_cache.shape[2:].numel() + num_blocks = kv_cache.shape[0] + expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride() + if kv_cache.stride()[2:] != expected_tail_stride: + logger.warning_once( + "Skipping mixed KV-cache layout alignment for a non-NHD " + "blocks-first attention view with stride %s.", + kv_cache.stride(), + ) + return + kv_cache.as_strided_( + size=kv_cache.shape, + stride=(page_size, num_blocks * page_size, *expected_tail_stride), + ) + + def _update_hybrid_attention_layout( attn_groups: Iterable[AttentionGroup], kv_caches: dict[str, Any], @@ -391,7 +476,10 @@ def _update_hybrid_attention_layout( # (quantization only changes the last dim), so this is a no-op today, # but it keeps both call sites consistent for skip layers. layer_cache_dtype = ( - "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE else cache_dtype + "auto" + if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE + and not isinstance(kv_cache_spec, TQFullAttentionSpec) + else cache_dtype ) block_dim = group.backend.get_kv_cache_block_dim( kernel_block_sizes[group.kv_cache_group_id], @@ -482,9 +570,10 @@ def build_attn_metadata( seq_lens_cpu_upper_bound: torch.Tensor | None = None, dcp_local_seq_lens: torch.Tensor | None = None, positions: torch.Tensor | None = None, + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None, model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, - causal: bool = True, + causal: bool | Mapping[int, bool] = True, rswa_prefix_lens: torch.Tensor | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] @@ -498,6 +587,8 @@ def build_attn_metadata( for i in range(num_kv_cache_groups): block_table = block_tables[i] slot_mapping = slot_mappings[i] + # Per-group causal for hybrid drafters (mixed SWA/full attention). + group_causal = causal if isinstance(causal, bool) else causal.get(i, True) common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) @@ -515,9 +606,10 @@ def build_attn_metadata( max_query_len=max_query_len, block_table_tensor=block_table, slot_mapping=slot_mapping, - causal=causal, + causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, + mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, **common_attn_metadata_extra_kwargs, ) @@ -545,3 +637,27 @@ def build_attn_metadata( for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata return attn_metadata + + +def compute_mm_prefix_ranges( + req_ids: list[str], + mm_features: dict[str, list[MultiModalFeatureSpec]], + sliding_window: int | None = None, +) -> dict[int, list[tuple[int, int]]]: + """Compute PrefixLM bidirectional ranges for multimodal tokens. + + Ranges exceeding sliding_window are skipped to prevent early tokens + from attending across the entire image span. + """ + req_doc_ranges: dict[int, list[tuple[int, int]]] = {} + for req_idx, req_id in enumerate(req_ids): + image_doc_ranges = [] + for mm_feature in mm_features.get(req_id, ()): + if mm_feature.modality not in ("image", "video"): + continue + for r in mm_feature.mm_position.extract_embeds_range(): + if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window: + continue + image_doc_ranges.append(r) + req_doc_ranges[req_idx] = image_doc_ranges + return req_doc_ranges diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 006e11e4500..64c1096dfbe 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -340,10 +340,12 @@ def _combine_sampled_and_draft_tokens_kernel( # Handling prefill tokens. No sampled or draft tokens. return - if NUM_NEW_SAMPLED_TOKENS > 0: + # Keep prompt-tail slots intact; only rewrite generated-token slots. + first_logit_seq_pos = seq_len - num_logits + if NUM_NEW_SAMPLED_TOKENS > 0 and first_logit_seq_pos >= prefill_len: # 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) + tl.store(input_ids_ptr + logits_start, last_token_id) # Write the draft tokens (if any) to input_ids. if num_draft_tokens > 0: diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index e5e89da2b2e..854b71b69fc 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -9,7 +9,10 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.attn_utils import ( + build_attn_metadata, + compute_mm_prefix_ranges, +) from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.rope import get_rope_state @@ -152,6 +155,17 @@ class DefaultModelState(ModelState): max_seq_len = self.max_model_len else: max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() + req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None + if ( + self.supports_mm_inputs + and self.encoder_cache is not None + and self.model_config.is_mm_prefix_lm + ): + req_doc_ranges = compute_mm_prefix_ranges( + req_ids=input_batch.req_ids, + mm_features=self.encoder_cache.mm_features, + sliding_window=self.model_config.get_sliding_window(), + ) attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -167,6 +181,7 @@ class DefaultModelState(ModelState): seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, positions=input_batch.positions, + mm_req_doc_ranges=req_doc_ranges, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, ) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index cf6ab0ac06e..26439a98583 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -255,11 +255,16 @@ class MambaHybridModelState(DefaultModelState): # GDN uses >= 0 to select spec-decode rows, so non-decode rows # need the -1 sentinel rather than a raw zero draft count. num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) - if input_batch.num_draft_tokens_per_req is not None: - has_draft_tokens = input_batch.num_draft_tokens_per_req > 0 - spec_decode_mask = has_draft_tokens & ~input_batch.is_prefilling_np + num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req + if num_draft_tokens_per_req is not None: + # A row is a spec-decode row only when its whole prompt is already + # computed, i.e. exactly one non-draft (decode) token is scheduled. + is_decode = ( + input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1 + ) + spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( - spec_decode_mask, input_batch.num_draft_tokens_per_req, -1 + spec_decode_mask, num_draft_tokens_per_req, -1 ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 00ff95b6dac..9b1786cbe4c 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -9,6 +9,7 @@ import numpy as np import torch from vllm.distributed.parallel_state import get_pp_group +from vllm.platforms import current_platform from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import InputBatch @@ -17,7 +18,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch class PendingRecv: """Per-step slot data for a deferred postprocess on the main stream.""" - event: torch.Event + event: torch.cuda.Event sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] num_sampled: torch.Tensor # [num_reqs] @@ -179,6 +180,10 @@ class PPHandler: return assert sampled_token_ids.dtype == torch.int64 + + if current_platform.is_xpu(): + self.main_stream.synchronize() + with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) torch.distributed.broadcast( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index 3e4b2b7e7f0..a4b4033cafa 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable +from collections.abc import Callable, Mapping import torch @@ -29,7 +29,7 @@ def _prepare_dflash_inputs_to_capture( kv_cache_config: KVCacheConfig, max_model_len: int, skip_attn: bool, - causal: bool, + causal: bool | Mapping[int, bool], ) -> AttentionState: input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) @@ -63,10 +63,6 @@ class DFlashCudaGraphManager(CudaGraphManager): """DFlash CudaGraphManager for the parallel-drafting query forward, building its own attention metadata from scratch.""" - def __init__(self, *args, causal: bool = False, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.causal = causal - def capture( self, forward_fn: Callable, @@ -75,6 +71,7 @@ class DFlashCudaGraphManager(CudaGraphManager): attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, max_model_len: int, + causal: bool | Mapping[int, bool], progress_bar_desc: str = "Capturing CUDA graphs", ) -> None: def create_forward_fn( @@ -97,7 +94,7 @@ class DFlashCudaGraphManager(CudaGraphManager): kv_cache_config, max_model_len, skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), - causal=self.causal, + causal=causal, ) attn_metadata, slot_mappings = attn_state diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index d5d68a01460..f9ffd3135d9 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Mapping from typing import Any import torch @@ -94,7 +95,6 @@ class DFlashSpeculator(DraftModelSpeculator): self.device, cudagraph_mode, decode_query_len=self.num_query_per_req, - causal=self.dflash_causal, ) def capture(self, attn_states: dict | None = None) -> None: @@ -112,6 +112,7 @@ class DFlashSpeculator(DraftModelSpeculator): self.attn_groups, self.kv_cache_config, self.max_model_len, + causal=self._group_causal, progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) @@ -135,9 +136,6 @@ class DFlashSpeculator(DraftModelSpeculator): ] assert self.draft_kv_cache_group_ids, "No draft attention groups found." self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] - self.draft_block_size = self.block_tables.block_sizes[ - self.draft_kv_cache_group_id - ] # Per-group context slot buffers for the precompute (one row per group). self._context_slot_mappings = torch.zeros( @@ -151,7 +149,10 @@ class DFlashSpeculator(DraftModelSpeculator): # of the kv-cache group its cache belongs to. Models that share a single group # leave this as None and share one context slot mapping. self._layer_group_idx: list[int] | None = None + # Per-KV-group causal, falling back to the scalar dflash_causal. + self._group_causal: dict[int, bool] | bool = self.dflash_causal if hasattr(self.model, "get_draft_kv_cache_layer_names"): + layer_names = self.model.get_draft_kv_cache_layer_names() name_to_gid = { ln: gid for gid, group in enumerate(kv_cache_config.kv_cache_groups) @@ -159,9 +160,15 @@ class DFlashSpeculator(DraftModelSpeculator): } gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)} self._layer_group_idx = [ - gid_to_idx[name_to_gid[name]] - for name in self.model.get_draft_kv_cache_layer_names() + gid_to_idx[name_to_gid[name]] for name in layer_names ] + if hasattr(self.model, "get_draft_attn_causal"): + self._group_causal = { + name_to_gid[name]: layer_causal + for name, layer_causal in zip( + layer_names, self.model.get_draft_attn_causal() + ) + } @torch.inference_mode() def _run_model( @@ -229,7 +236,7 @@ class DFlashSpeculator(DraftModelSpeculator): num_reqs_padded: int, num_tokens_padded: int, num_query_per_req: int | None = None, - causal: bool = False, + causal: bool | Mapping[int, bool] = False, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None @@ -337,7 +344,7 @@ class DFlashSpeculator(DraftModelSpeculator): last_sampled, next_prefill_tokens, self.block_tables.input_block_tables[gid], - self.block_tables.block_sizes[gid], + self.block_tables.kernel_block_sizes[gid], self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, @@ -386,7 +393,7 @@ class DFlashSpeculator(DraftModelSpeculator): num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, num_tokens_padded=num_tokens_padded, - causal=self.dflash_causal, + causal=self._group_causal, ) draft_slot_mappings_by_layer = build_slot_mappings_by_layer( self.block_tables.slot_mappings[:, :num_tokens_padded], diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index c4f98e715b9..37fe693bbbd 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -5,7 +5,10 @@ import torch.nn as nn from vllm.config import ModelConfig, VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.model_executor.model_loader import get_model -from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share +from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _should_share, + get_target_lm_head, +) def get_dflash_causal(draft_model_config: ModelConfig) -> bool: @@ -57,7 +60,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo del draft_inner.embed_tokens draft_inner.embed_tokens = target_embed - target_lm_head = getattr(target_model, "lm_head", None) + target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(dflash_model, "lm_head", None) if target_lm_head is not None and _should_share( dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index acc32dafcd4..08ff30e5fdf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -6,7 +6,10 @@ import torch.nn as nn from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.model_executor.model_loader import get_model -from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share +from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _should_share, + get_target_lm_head, +) def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: @@ -52,7 +55,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo del draft_inner.embed_tokens draft_inner.embed_tokens = target_embed - target_lm_head = getattr(target_model, "lm_head", None) + target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(draft_model, "lm_head", None) if target_lm_head is not None and _should_share( draft_model, "has_own_lm_head", draft_lm_head, target_lm_head diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index 11961ceef4d..bdd588e5786 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -25,6 +25,14 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool: return torch.equal(w, target.weight) +def get_target_lm_head(target_model: nn.Module, target_language_model: nn.Module): + """The target's lm_head — from get_language_model() for + *ForConditionalGeneration targets, else the top-level module.""" + return getattr(target_language_model, "lm_head", None) or getattr( + target_model, "lm_head", None + ) + + def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag @@ -64,7 +72,7 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod del draft_inner.embed_tokens draft_inner.embed_tokens = target_embed - target_lm_head = getattr(target_model, "lm_head", None) + target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(eagle_model, "lm_head", None) if target_lm_head is not None and _should_share( eagle_model, "has_own_lm_head", draft_lm_head, target_lm_head diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 9e36c7cc655..4b7d7fef410 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Any import torch @@ -202,7 +203,7 @@ class DraftModelSpeculator(BaseSpeculator): num_reqs_padded: int, num_tokens_padded: int, num_query_per_req: int = 1, - causal: bool = True, + causal: bool | Mapping[int, bool] = True, ) -> dict[str, Any] | None: # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req. # Clamp keeps the series non-decreasing past num_reqs, which some diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 37ca1665937..e25672e953b 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -12,7 +12,8 @@ class DraftTokensHandler: def __init__(self, device: torch.device | None = None): self.device = device self.copy_stream = torch.cuda.Stream(device) - self.copy_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index be7bae7f17e..7f0ae33c809 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -110,16 +110,6 @@ class RequestState: self.num_computed_tokens_np[req_idx] = num_computed_tokens self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens) - if 0 < num_computed_tokens <= prefill_len: - # For PD disagg or resumed requests: set last_sampled to the last - # computed token so the first decode step gets the right input_id. - # For fresh prefill requests (num_computed_tokens == 0) the tensor - # is not read by combine_sampled_and_draft_tokens so we skip the - # write. Use a slice assignment rather than scalar indexing so the - # write is dispatched through fill_ without a host/device sync. - self.last_sampled_tokens[req_idx : req_idx + 1] = all_token_ids[ - num_computed_tokens - 1 - ] self.draft_tokens[req_idx].zero_() def apply_staged_writes(self) -> None: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3930a07b248..199470aaf87 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -132,6 +132,9 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.linear_attn import ( + BailingLinearAttentionMetadataBuilder, +) from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, @@ -256,7 +259,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self._invalid_req_indices = invalid_req_indices # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensor to avoid it being # deallocated until we finish copying it to the host. @@ -389,7 +393,8 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): self._model_runner_output = model_runner_output # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensors to avoid them being # deallocated until we finish copying it to the host. @@ -697,11 +702,9 @@ class GPUModelRunner( custom_logitsprocs, ), # We currently don't know whether a particular custom logits processor - # uses output token ids so we set this conservatively. - # ThinkingTokenBudgetLogitsProcessor also needs output token ids to - # correctly track think start/end token sequences in async scheduling. - logitsprocs_need_output_token_ids=bool(custom_logitsprocs) - or self.vllm_config.reasoning_config is not None, + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), is_pooling_model=self.is_pooling_model, cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, reasoning_config=self.vllm_config.reasoning_config, @@ -715,7 +718,9 @@ class GPUModelRunner( self.prepare_inputs_event: torch.Event | None = None if self.use_async_scheduling: self.async_output_copy_stream = torch.cuda.Stream() - self.prepare_inputs_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock; + # under TP contention that spin can balloon and make the rank a straggler. + self.prepare_inputs_event = torch.cuda.Event(blocking=True) # self.cudagraph_batch_sizes sorts in ascending order. if ( @@ -857,7 +862,7 @@ class GPUModelRunner( # N-gram GPU path: async D2H buffer/event for per-request valid draft counts. self._num_valid_draft_tokens: torch.Tensor | None = None self._num_valid_draft_tokens_cpu: torch.Tensor | None = None - self._num_valid_draft_tokens_event: torch.Event | None = None + self._num_valid_draft_tokens_event: torch.cuda.Event | None = None self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None if ( self.speculative_config is not None @@ -866,7 +871,7 @@ class GPUModelRunner( self._num_valid_draft_tokens_cpu = torch.empty( self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY ) - self._num_valid_draft_tokens_event = torch.Event() + self._num_valid_draft_tokens_event = torch.cuda.Event() self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream() self._draft_token_req_ids: list[str] | None = None @@ -2201,10 +2206,7 @@ class GPUModelRunner( req_idx = self.input_batch.req_id_to_index[req_id] draft_len = len(draft_token_ids) num_draft_tokens[req_idx] = draft_len - if ( - self.input_batch.num_computed_tokens_cpu[req_idx] - >= self.input_batch.num_prompt_tokens[req_idx] - ): + if num_scheduled_tokens[req_idx] == draft_len + 1: num_decode_draft_tokens[req_idx] = draft_len spec_decode_metadata = self._calc_spec_decode_metadata( num_draft_tokens, cu_num_tokens @@ -2442,9 +2444,16 @@ class GPUModelRunner( extra_attn_metadata_args = {} if use_spec_decode and isinstance( - builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder) + builder, + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + BailingLinearAttentionMetadataBuilder, + ), ): - assert ubid is None, "UBatching not supported with GDN yet" + assert ubid is None, ( + "UBatching not supported with GDN or linear attn yet" + ) extra_attn_metadata_args = dict( num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded], num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[ @@ -4517,17 +4526,23 @@ class GPUModelRunner( self._copy_draft_token_ids_to_cpu(scheduler_output) spec_config = self.speculative_config - propose_drafts_after_bookkeeping = False + draft_after_bookkeeping = False if spec_config is not None: # Decide whether to run the drafter or zero out draft tokens. input_fits_in_drafter = self._input_fits_in_drafter( spec_decode_common_attn_metadata ) - use_gpu_toks = ( + # Whether the drafter runs a GPU model forward (and thus carries + # TP/EP/DP collectives), independent of padded-batch timing. + drafter_runs_model_forward = ( spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states() - ) and not spec_config.disable_padded_drafter_batch + ) + use_gpu_toks = ( + drafter_runs_model_forward + and not spec_config.disable_padded_drafter_batch + ) if use_gpu_toks: # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. @@ -4542,19 +4557,23 @@ class GPUModelRunner( sampled_token_ids = sampler_output.sampled_token_ids if input_fits_in_drafter: propose_draft_token_ids(sampled_token_ids) - elif self.valid_sampled_token_count_event is not None: - assert spec_decode_common_attn_metadata is not None - next_token_ids, valid_sampled_tokens_count = ( - self.drafter.prepare_next_token_ids_padded( - sampled_token_ids, - self.requests, - self.input_batch, - self.discard_request_mask.gpu, + else: + if self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count = ( + self.drafter.prepare_next_token_ids_padded( + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) ) - ) - self._copy_valid_sampled_token_count( - next_token_ids, valid_sampled_tokens_count - ) + self._copy_valid_sampled_token_count( + next_token_ids, valid_sampled_tokens_count + ) + if self.parallel_config.data_parallel_size > 1: + # Prevent hang when DP ranks disagree on input_fits_in_drafter + self.drafter.dummy_run(num_tokens=1) elif ( spec_config.use_ngram_gpu() and not spec_config.disable_padded_drafter_batch @@ -4578,7 +4597,9 @@ class GPUModelRunner( next_token_ids, valid_sampled_tokens_count ) else: - propose_drafts_after_bookkeeping = input_fits_in_drafter + # These drafters consume CPU sampled tokens, so they run + # after bookkeeping. + draft_after_bookkeeping = True if not input_fits_in_drafter: # Zero out draft tokens so the scheduler doesn't schedule @@ -4610,10 +4631,25 @@ class GPUModelRunner( scheduler_output.total_num_scheduled_tokens, ) - if propose_drafts_after_bookkeeping: + if draft_after_bookkeeping: # ngram and other speculative decoding methods use the sampled # tokens on the CPU, so they are run after bookkeeping. - propose_draft_token_ids(valid_sampled_token_ids) + if input_fits_in_drafter: + propose_draft_token_ids(valid_sampled_token_ids) + elif ( + drafter_runs_model_forward + and self.parallel_config.data_parallel_size > 1 + ): + # Prevent hang when DP ranks disagree on input_fits_in_drafter + assert isinstance( + self.drafter, + EagleProposer + | DFlashProposer + | DraftModelProposer + | ExtractHiddenStatesProposer + | Gemma4Proposer, + ) + self.drafter.dummy_run(num_tokens=1) # Finalize KV connector (wait_for_save + clear metadata) after # draft model runs. Deferred from target model forward to allow @@ -7154,7 +7190,12 @@ class GPUModelRunner( layer_cache_dtype_str = ( "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE - else self.cache_config.cache_dtype + else getattr( + kv_cache_spec, + "cache_dtype_str", + None, + ) + or self.cache_config.cache_dtype ) kv_cache_shape = attn_backend.get_kv_cache_shape( kernel_num_blocks, diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 76fa12b4121..7b619998435 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -252,7 +252,7 @@ class UBatchWrapper: results: list[tuple[int, torch.Tensor]] = [] compute_stream = ubatch_metadata[0].context.compute_stream - num_tokens = ubatch_metadata[0].num_tokens + ubatch_metadata[1].num_tokens + num_tokens = sum(m.num_tokens for m in ubatch_metadata) # Ubatches will manually manage the forward context, so we override # it to None here so we can have it restored correctly later @@ -268,7 +268,7 @@ class UBatchWrapper: ) ubatch_threads.append(thread) thread.start() - self.ready_barrier.wait() # Wait for both threads to be ready + self.ready_barrier.wait() # Wait for all ubatch threads to be ready # Capture the cudagraph cudagraph_metadata = CUDAGraphMetaData( @@ -332,7 +332,7 @@ class UBatchWrapper: ) ubatch_threads.append(thread) thread.start() - self.ready_barrier.wait() # Wait for both threads to be ready + self.ready_barrier.wait() # Wait for all ubatch threads to be ready ubatch_metadata[0].context.cpu_wait_event.set() for thread in ubatch_threads: thread.join() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 28efee3dee8..182476e2533 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -55,7 +55,7 @@ from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, PYNVVIDEOCODEC_MAX_RETAINED_DECODERS, - PYNVVIDEOCODEC_VIDEO_BACKEND, + VIDEO_LOADER_REGISTRY, ) from vllm.platforms import current_platform from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper @@ -75,6 +75,10 @@ from vllm.v1.outputs import ( ModelRunnerOutput, ) from vllm.v1.utils import compute_iteration_details, report_usage_stats +from vllm.v1.worker.startup_plan import ( + maybe_apply_startup_plan, + maybe_save_startup_plan, +) from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase from vllm.v1.worker.workspace import init_workspace_manager @@ -374,7 +378,7 @@ class Worker(WorkerBase): "worker requested memory: %sGiB", format_gib(self.requested_memory) ) else: - raise RuntimeError(f"Not support device type: {self.device_config.device}") + raise RuntimeError(f"Unsupported device type: {self.device_config.device}") # Initialize workspace manager num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 @@ -439,6 +443,8 @@ class Worker(WorkerBase): You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. """ + maybe_apply_startup_plan(self) + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: # still need a profile run which compiles the model for # max_num_batched_tokens @@ -584,15 +590,15 @@ class Worker(WorkerBase): ) @staticmethod - def _uses_pynvvideocodec_video_backend(mm_config) -> bool: + def _uses_gpu_video_backend(mm_config) -> bool: video_kwargs = mm_config.media_io_kwargs.get("video", {}) video_loader_backend = ( video_kwargs.get("video_backend") or envs.VLLM_VIDEO_LOADER_BACKEND ) codec_backend = video_kwargs.get("backend") - return ( - video_loader_backend == PYNVVIDEOCODEC_VIDEO_BACKEND - or codec_backend == PYNVVIDEOCODEC_VIDEO_BACKEND + return VIDEO_LOADER_REGISTRY.backend_requires_gpu(video_loader_backend) or ( + codec_backend is not None + and VIDEO_LOADER_REGISTRY.backend_requires_gpu(codec_backend) ) def _reserve_mm_ipc_gpu_memory(self, available_kv_cache_memory_bytes: int) -> int: @@ -623,7 +629,7 @@ class Worker(WorkerBase): ) decoder_reserved_bytes = ( num_api_servers * per_server_decoder_bytes - if self._uses_pynvvideocodec_video_backend(mm_config) + if self._uses_gpu_video_backend(mm_config) else 0 ) reserved_bytes = raw_frame_reserved_bytes + decoder_reserved_bytes @@ -832,7 +838,9 @@ class Worker(WorkerBase): f"{format_gib(self.available_kv_cache_memory_bytes)} GiB." ) - logger.debug(msg) + logger.info(msg) + + maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit) if self.use_v2_model_runner: # V2: Run full execute_model + sample_tokens to JIT compile triton kernels. diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py new file mode 100644 index 00000000000..2c494207070 --- /dev/null +++ b/vllm/v1/worker/startup_plan.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Persist and reuse the memory-profiling result across engine boots. + +On startup, vLLM measures how much GPU memory the KV cache can use and +computes the ``--kv-cache-memory`` value that reproduces that allocation. +For a fixed (model, config, hardware, library) combination the result is +deterministic, yet it is re-measured on every boot. + +When ``VLLM_ENABLE_STARTUP_PLAN=1``, each worker persists that value under +``{VLLM_CACHE_ROOT}/startup_plan/`` (regenerable derived state, alongside +the torch.compile cache), keyed by a fingerprint of everything the value +depends on, and later boots apply it automatically -- skipping the +memory-profiling measurement and the CUDA-graph memory estimation pass -- +if and only if the fingerprint matches and the device has at least as much +free memory as when the plan was recorded. On any mismatch the worker +falls back to full profiling, so a stale plan costs nothing and is never +trusted. +""" + +import hashlib +import json +import os +from typing import TYPE_CHECKING + +import torch + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +PLAN_SCHEMA_VERSION = 1 + + +def compute_plan_fingerprint( + vllm_config: VllmConfig, rank: int, world_size: int +) -> str: + """Hash everything the profiled KV-cache memory value depends on. + + ``VllmConfig.compute_hash()`` covers the vLLM version and the model, + cache, parallel, and compilation configs, but deliberately contains no + device identity (``DeviceConfig.compute_hash`` is empty), so device + name, total memory, compute capability, and the torch/CUDA build are + added here. The vLLM version is also pinned as an explicit factor so + version invalidation holds no matter how ``compute_hash`` evolves. + Rank is included because per-rank memory use differs under TP/PP. + Driver-only changes are not part of the key; the free-memory gate at + apply time bounds the residual risk. + """ + # Imported here (as VllmConfig.compute_hash does) to avoid a cycle with + # the top-level vllm package. + from vllm import __version__ as vllm_version + + capability = current_platform.get_device_capability() + factors = { + "schema": PLAN_SCHEMA_VERSION, + "vllm": vllm_version, + "vllm_config": vllm_config.compute_hash(), + "device_name": current_platform.get_device_name(), + "device_total_memory": current_platform.get_device_total_memory(), + "device_capability": str(capability) if capability else "", + "torch": torch.__version__, + "cuda": torch.version.cuda or "", + "rank": rank, + "world_size": world_size, + } + digest = hashlib.sha256(json.dumps(factors, sort_keys=True).encode()).hexdigest() + return digest[:16] + + +def _plan_path(fingerprint: str) -> str: + """Plans are regenerable derived state, so they live under the standard + vLLM cache root (like the torch.compile cache) and relocate with + ``VLLM_CACHE_ROOT`` instead of needing a location knob of their own.""" + # VLLM_CACHE_ROOT is already user-expanded by envs.py. + return os.path.join( + envs.VLLM_CACHE_ROOT, "startup_plan", f"startup_plan_{fingerprint}.json" + ) + + +def _load_plan(fingerprint: str) -> dict | None: + """Load a plan for this fingerprint; None if absent or unreadable.""" + path = _plan_path(fingerprint) + try: + with open(path) as f: + plan = json.load(f) + except FileNotFoundError: + return None + except (OSError, json.JSONDecodeError) as e: + logger.warning("Ignoring unreadable startup plan %s: %s", path, e) + return None + if ( + plan.get("schema") != PLAN_SCHEMA_VERSION + or plan.get("fingerprint") != fingerprint + ): + return None + return plan + + +def _applicable_kv_cache_memory_bytes( + plan: dict, current_free_memory: int +) -> int | None: + """The apply-time OOM-safety gate. + + The recorded value is only valid if the device has at least as much + free memory now as when the plan was measured (co-tenants, leaked + allocations, or MIG changes all reduce it). Outside that envelope, + return None and let the caller re-profile. + """ + kv_bytes = plan.get("kv_cache_memory_bytes") + baseline = plan.get("free_memory_baseline") + if not isinstance(kv_bytes, int) or not isinstance(baseline, int): + return None + if kv_bytes <= 0: + return None + if current_free_memory < baseline: + logger.info( + "Startup plan not applied: current free memory (%.2f GiB) is " + "below the recorded baseline (%.2f GiB); falling back to full " + "memory profiling.", + current_free_memory / (1 << 30), + baseline / (1 << 30), + ) + return None + return kv_bytes + + +def maybe_apply_startup_plan(worker: "Worker") -> None: + """If enabled and ``--kv-cache-memory`` was not set explicitly, apply a + persisted plan by setting ``worker.cache_config.kv_cache_memory_bytes``. + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``.""" + if ( + not envs.VLLM_ENABLE_STARTUP_PLAN + or worker.cache_config.kv_cache_memory_bytes is not None + ): + return + fingerprint = compute_plan_fingerprint( + worker.vllm_config, worker.rank, worker.parallel_config.world_size + ) + plan = _load_plan(fingerprint) + if plan is None: + return + current_free_memory = worker.init_snapshot.free_memory + kv_bytes = _applicable_kv_cache_memory_bytes(plan, current_free_memory) + if kv_bytes is None: + return + logger.info( + "Applying persisted startup plan (fingerprint %s): " + "kv_cache_memory_bytes=%d (%.2f GiB), recorded free-memory " + "baseline %.2f GiB, current %.2f GiB. Memory profiling will " + "be skipped.", + fingerprint, + kv_bytes, + kv_bytes / (1 << 30), + plan["free_memory_baseline"] / (1 << 30), + current_free_memory / (1 << 30), + ) + worker.cache_config.kv_cache_memory_bytes = kv_bytes + + +def maybe_save_startup_plan(worker: "Worker", kv_cache_memory_bytes: int) -> None: + """Atomically persist this boot's profiling result for future boots. + No-op unless ``VLLM_ENABLE_STARTUP_PLAN=1``; failures are logged, + never raised.""" + if not envs.VLLM_ENABLE_STARTUP_PLAN: + return + fingerprint = compute_plan_fingerprint( + worker.vllm_config, worker.rank, worker.parallel_config.world_size + ) + path = _plan_path(fingerprint) + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + payload = { + "schema": PLAN_SCHEMA_VERSION, + "fingerprint": fingerprint, + "kv_cache_memory_bytes": int(kv_cache_memory_bytes), + "free_memory_baseline": int(worker.init_snapshot.free_memory), + } + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as f: + json.dump(payload, f) + os.replace(tmp, path) + logger.info("Saved startup plan to %s", path) + except OSError as e: + logger.warning("Failed to save startup plan to %s: %s", path, e) diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index 6cdca994da5..77edbba58f9 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/vllm/v1/worker/xpu_model_runner.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from contextlib import contextmanager +from functools import partial import torch @@ -40,14 +41,23 @@ class XPUModelRunnerV2(GPUModelRunnerV2): @contextmanager def _torch_cuda_wrapper(): - # replace cuda APIs with xpu APIs, this should work by default + # Replace cuda APIs with xpu APIs. Each callable gets its own functools.partial + # so it is not the same object as torch.xpu.* (Torch Dynamo _get_handlers() + # asserts on duplicate registration when cuda aliases xpu directly). torch.cuda.Stream = torch.xpu.Stream - torch.cuda.default_stream = torch.xpu.current_stream - torch.cuda.current_stream = torch.xpu.current_stream - torch.cuda.stream = torch.xpu.stream - torch.cuda.set_stream = torch.xpu.set_stream + torch.cuda.default_stream = partial(torch.xpu.current_stream) + torch.cuda.current_stream = partial(torch.xpu.current_stream) + torch.cuda.stream = partial(torch.xpu.stream) + torch.cuda.set_stream = partial(torch.xpu.set_stream) + + # torch.xpu.Event does not accept the ``blocking`` kwarg that + # torch.cuda.Event supports, so drop it here. + def _xpu_event(*args, blocking=None, **kwargs): + return torch.xpu.Event(*args, **kwargs) + + torch.cuda.Event = _xpu_event if supports_xpu_graph(): - torch.cuda.graph = torch.xpu.graph + torch.cuda.graph = partial(torch.xpu.graph) torch.cuda.CUDAGraph = torch.xpu.XPUGraph - torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle + torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle) yield diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index e669365890f..091406a22b1 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -82,7 +82,7 @@ class XPUWorker(Worker): self.local_rank ).total_memory else: - raise RuntimeError(f"Not support device type: {self.device_config.device}") + raise RuntimeError(f"Unsupported device type: {self.device_config.device}") ENV_CCL_ATL_TRANSPORT = os.getenv("CCL_ATL_TRANSPORT", "ofi") ENV_LOCAL_WORLD_SIZE = os.getenv(