diff --git a/.buildkite/ci_config_rocm.yaml b/.buildkite/ci_config_rocm.yaml new file mode 100644 index 00000000000..23f32340071 --- /dev/null +++ b/.buildkite/ci_config_rocm.yaml @@ -0,0 +1,23 @@ +name: vllm_rocm_ci +job_dirs: + - ".buildkite/hardware_tests" +run_all_patterns: + - "docker/Dockerfile.rocm" + - "docker/Dockerfile.rocm_base" + - "docker/ci-rocm.hcl" + - "docker/docker-bake-rocm.hcl" + - ".buildkite/hardware_tests/amd.yaml" + - ".buildkite/scripts/ci-bake-rocm.sh" + - ".buildkite/scripts/hardware_ci/run-amd-test.py" + - ".buildkite/scripts/hardware_ci/run-amd-test.sh" + - "CMakeLists.txt" + - "requirements/common.txt" + - "requirements/rocm.txt" + - "requirements/build/rocm.txt" + - "requirements/test/rocm.txt" + - "setup.py" + - "csrc/" + - "cmake/" +run_all_exclude_patterns: + - "csrc/cpu/" + - "cmake/cpu_extension.cmake" diff --git a/.buildkite/hardware_tests/amd.yaml b/.buildkite/hardware_tests/amd.yaml index 1351eba92f2..c2510f38aab 100644 --- a/.buildkite/hardware_tests/amd.yaml +++ b/.buildkite/hardware_tests/amd.yaml @@ -1,42 +1,73 @@ -group: Hardware - AMD Build +group: Hardware - AMD Build steps: - - label: "AMD: :docker: build image" - key: image-build-amd + # 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 + # differ ci_base is rebuilt and pushed automatically. + - label: "AMD: :docker: ensure ci_base" + key: ensure-ci-base-amd depends_on: [] device: amd_cpu no_plugin: true commands: - - > - docker build - --build-arg max_jobs=16 - --build-arg REMOTE_VLLM=1 - --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942;gfx950' - --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT - --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}" - -f docker/Dockerfile.rocm - --target test - --no-cache - --progress plain . - - | - 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 - <&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 - </dev/null 2>&1; then + timeout "${timeout_secs}s" git fetch "$@" 2>/dev/null + else + git fetch "$@" 2>/dev/null + fi +} + +hash_string_short() { + printf '%s' "$1" | sha256sum | cut -c1-16 +} + +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 +} + +compose_dependency_cache_key() { + local prefix="$1" + local material="$2" + local cleaned_prefix="" + + cleaned_prefix=$(clean_docker_tag "${prefix}" | cut -c1-96) + printf '%s-%s\n' "${cleaned_prefix}" "$(hash_string_short "${material}")" +} + +hash_dockerfile_stages() { + local dockerfile="$1" + local stages="$2" + + awk -v wanted_stages="${stages}" ' + BEGIN { + split(wanted_stages, stage_list, /[[:space:]]+/) + for (idx in stage_list) { + if (stage_list[idx] != "") { + wanted[stage_list[idx]] = 1 + } + } + emit = 1 + } + $1 == "FROM" { + stage = "" + for (idx = 1; idx <= NF; idx++) { + if (tolower($idx) == "as" && idx < NF) { + stage = $(idx + 1) + } + } + emit = (stage in wanted) + } + emit { + print + } + ' "${dockerfile}" +} + +discover_dockerfile_stage_args() { + local dockerfile="$1" + local stages="$2" + + [[ -f "${dockerfile}" ]] || return 0 + + awk -v wanted_stages="${stages}" ' + function add_arg(name) { + if (name != "" && !(name in seen)) { + seen[name] = 1 + args[++arg_count] = name + } + } + BEGIN { + split(wanted_stages, stage_list, /[[:space:]]+/) + for (idx in stage_list) { + if (stage_list[idx] != "") { + wanted[stage_list[idx]] = 1 + } + } + emit = 1 + } + { + line = $0 + if ($1 == "FROM") { + stage = "" + for (idx = 1; idx <= NF; idx++) { + if (tolower($idx) == "as" && idx < NF) { + stage = $(idx + 1) + } + } + emit = (stage in wanted) + } + if (emit) { + lines[++line_count] = line + } + } + END { + for (idx = 1; idx <= line_count; idx++) { + line = lines[idx] + arg_name = line + sub(/^[[:space:]]*ARG[[:space:]]+/, "", arg_name) + if (arg_name != line) { + sub(/[=[:space:]].*/, "", arg_name) + if (arg_name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) { + add_arg(arg_name) + } + } + } + + for (idx = 1; idx <= line_count; idx++) { + line = lines[idx] + for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) { + name = args[arg_idx] + if (line ~ "\\$\\{" name "([}:][^}]*)?\\}" \ + || line ~ "\\$" name "([^A-Za-z0-9_]|$)") { + used[name] = 1 + } + } + } + + for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) { + name = args[arg_idx] + if (used[name]) { + print name + } + } + } + ' "${dockerfile}" +} + +get_content_arg_names() { + local dockerfile="$1" + local stages="$2" + local explicit_args="${3:-}" + + if [[ -n "${explicit_args}" ]]; then + tr ' ' '\n' <<< "${explicit_args}" + else + discover_dockerfile_stage_args "${dockerfile}" "${stages}" + fi | awk 'NF && !seen[$0]++' +} + +compute_ci_base_content_hash() { + local -a content_paths=() + local -a content_args=() + local dockerfile="${CI_BASE_DOCKERFILE:-}" + local stages="${CI_BASE_DOCKERFILE_STAGES:-}" + + read -r -a content_paths <<< "${CI_BASE_CONTENT_FILES}" + mapfile -t content_args < <( + get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}" + ) + + { + printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")" + if [[ -n "${dockerfile}" ]]; then + printf 'dockerfile:%s\n' "${dockerfile}" + printf 'resolved-build-args:\n' + hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" + if [[ -n "${stages}" ]]; then + printf 'dockerfile-stages:%s\n' "${stages}" + if [[ -f "${dockerfile}" ]]; then + hash_dockerfile_stages "${dockerfile}" "${stages}" + else + printf 'missing:%s\n' "${dockerfile}" + fi + fi + fi + } | sha256sum | cut -d' ' -f1 +} + +extract_dockerfile_arg_default() { + local dockerfile="$1" + local arg_name="$2" + + 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 +} + +resolve_dockerfile_arg_value() { + local dockerfile="$1" + local arg_name="$2" + local env_name="${arg_name}" + local value="" + + case "${arg_name}" in + ARG_PYTORCH_ROCM_ARCH) + env_name="PYTORCH_ROCM_ARCH" + ;; + esac + + value="${!env_name:-}" + if [[ -z "${value}" && "${env_name}" != "${arg_name}" ]]; then + value="${!arg_name:-}" + fi + if [[ -z "${value}" && -f "${dockerfile}" ]]; then + value=$(extract_dockerfile_arg_default "${dockerfile}" "${arg_name}") + fi + + printf '%s\n' "${value}" +} + +hash_dockerfile_arg_values() { + local dockerfile="$1" + local arg_name="" + local arg_value="" + local digest="" + shift || true + + for arg_name in "$@"; do + [[ -n "${arg_name}" ]] || continue + arg_value=$(resolve_dockerfile_arg_value "${dockerfile}" "${arg_name}") + printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-}" + if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then + digest=$(resolve_image_digest "${arg_value}") + printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest:-unknown}" + fi + done +} + +is_ci_base_target() { + [[ "${TARGET}" == *"ci-base-rocm"* ]] +} + +is_commit_image_target() { + [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1 + is_ci_base_target && return 1 + return 0 +} + +image_tag_is_commit_scoped() { + [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1 + [[ "${IMAGE_TAG}" == *"${BUILDKITE_COMMIT}"* ]] +} + +should_upload_wheel_artifacts() { + [[ "${UPLOAD_ROCM_WHEEL_ARTIFACTS:-0}" == "1" ]] && return 0 + [[ "${TARGET}" == *"with-wheel"* \ + || "${TARGET}" == *"export-wheel"* \ + || "${TARGET}" == *"artifact"* ]] +} + +get_remote_image_label() { + local image_ref="$1" + local label_key="$2" + + docker buildx imagetools inspect "${image_ref}" --raw 2>/dev/null \ + | python3 -c ' +import json +import subprocess +import sys +import urllib.parse +import urllib.request + +image_ref = sys.argv[1] +label_key = sys.argv[2] + + +def docker_hub_repo(image_name): + image_name = image_name.split("@", 1)[0] + last_component = image_name.rsplit("/", 1)[-1] + if ":" in last_component: + image_name = image_name.rsplit(":", 1)[0] + + parts = image_name.split("/") + if len(parts) > 1 and ( + "." in parts[0] or ":" in parts[0] or parts[0] == "localhost" + ): + registry = parts[0] + if registry not in { + "docker.io", + "index.docker.io", + "registry-1.docker.io", + }: + return None + image_name = "/".join(parts[1:]) + elif len(parts) == 1: + image_name = f"library/{image_name}" + + return image_name + + +try: + data = json.load(sys.stdin) + if data.get("manifests"): + manifest = next( + ( + entry + for entry in data["manifests"] + if entry.get("platform", {}).get("os") != "unknown" + and entry.get("platform", {}).get("architecture") != "unknown" + ), + data["manifests"][0], + ) + digest = manifest["digest"] + result = subprocess.run( + [ + "docker", + "buildx", + "imagetools", + "inspect", + image_ref + "@" + digest, + "--raw", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0 or not result.stdout: + raise RuntimeError("digest inspect failed") + data = json.loads(result.stdout) + + annotations = data.get("annotations", {}) + if label_key in annotations: + print(annotations[label_key]) + raise SystemExit(0) + + config_digest = data.get("config", {}).get("digest") + if not config_digest: + print("") + raise SystemExit(0) + + image_name = docker_hub_repo(image_ref) + if not image_name: + print("") + raise SystemExit(0) + + token_url = ( + "https://auth.docker.io/token?" + + urllib.parse.urlencode( + { + "service": "registry.docker.io", + "scope": f"repository:{image_name}:pull", + } + ) + ) + with urllib.request.urlopen(token_url, timeout=30) as response: + token = json.load(response)["token"] + + request = urllib.request.Request( + f"https://registry-1.docker.io/v2/{image_name}/blobs/{config_digest}", + headers={"Authorization": f"Bearer {token}"}, + ) + with urllib.request.urlopen(request, timeout=30) as response: + config_blob = json.load(response) + + labels = config_blob.get("config", {}).get("Labels", {}) + print(labels.get(label_key, "")) +except Exception: + print("") +' "${image_ref}" "${label_key}" 2>/dev/null || echo "" +} + +get_remote_image_label_with_retry() { + local image_ref="$1" + local label_key="$2" + local attempts="${3:-6}" + local delay_secs="${4:-5}" + local label_value="" + local attempt + + for ((attempt = 1; attempt <= attempts; attempt++)); do + label_value=$(get_remote_image_label "${image_ref}" "${label_key}") + if [[ -n "${label_value}" ]]; then + printf '%s\n' "${label_value}" + return 0 + fi + if [[ ${attempt} -lt ${attempts} ]]; then + sleep "${delay_secs}" + fi + done + + return 0 +} + +remote_image_exists() { + local image_ref="$1" + docker manifest inspect "${image_ref}" >/dev/null 2>&1 +} + +use_existing_builder() { + echo "Using existing builder: ${BUILDER_NAME}" + docker buildx use "${BUILDER_NAME}" + docker buildx inspect --bootstrap +} + +buildx_driver() { + local builder="${1:-}" + + if [[ -n "${builder}" ]]; then + docker buildx inspect "${builder}" 2>/dev/null + else + docker buildx inspect 2>/dev/null + fi | awk -F': *' '$1 == "Driver" { print $2; exit }' +} + +builder_supports_registry_cache() { + local driver="$1" + + [[ -n "${driver}" && "${driver}" != "docker" ]] +} + +create_and_bootstrap_builder() { + local driver="$1" + local endpoint="${2:-}" + + echo "Creating builder '${BUILDER_NAME}' with ${driver} driver" + if [[ -n "${endpoint}" ]]; then + docker buildx create \ + --name "${BUILDER_NAME}" \ + --driver "${driver}" \ + --use \ + "${endpoint}" + else + docker buildx create --name "${BUILDER_NAME}" --driver "${driver}" --use + fi + docker buildx inspect --bootstrap +} + +init_config() { + TARGET="${1:-test-ci}" + BAKE_TARGETS=("${TARGET}") + DEPENDENCY_CACHE_TARGETS=() + CI_HCL_SOURCE="${CI_HCL_SOURCE:-${CI_HCL_FILE:-${DEFAULT_CI_HCL_SOURCE}}}" + VLLM_BAKE_FILE="${VLLM_BAKE_FILE:-docker/docker-bake-rocm.hcl}" + BUILDER_NAME="${BUILDER_NAME:-vllm-builder}" + BUILDKIT_SOCKET="${BUILDKIT_SOCKET:-/run/buildkit/buildkitd.sock}" + PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950}" + CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" + CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" + CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" + export PYTORCH_ROCM_ARCH + + SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX) + CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl" + CI_BASE_LABEL_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/ci-base-label-override.hcl" + CSRC_CACHE_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-csrc-cache-override.hcl" + ROCM_ARG_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-arg-override.hcl" + BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json" +} + +print_header() { + echo "--- :docker: Setting up Docker buildx bake" + echo "Target: ${TARGET}" + echo "CI HCL source: ${CI_HCL_SOURCE}" + echo "vLLM bake file: ${VLLM_BAKE_FILE}" + if is_ci_base_target; then + echo "Build mode: ci_base" + elif is_commit_image_target; then + echo "Build mode: commit image" + else + echo "Build mode: generic" + fi + if [[ "${USE_SCCACHE:-0}" == "1" ]]; then + echo "Compiler cache: sccache enabled" + fi +} + +validate_inputs() { + if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then + echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}" + echo "Make sure you're running from the vLLM repository root" + exit 1 + fi + + if [[ -n "${CI_HCL_SOURCE:-}" ]] && is_url_like "${CI_HCL_SOURCE}"; then + echo "Error: remote CI HCL sources are not supported: ${CI_HCL_SOURCE}" + echo "Use the vLLM-owned docker/ci-rocm.hcl or set CI_HCL_SOURCE to a local file." + exit 1 + fi + + if [[ -n "${CI_HCL_SOURCE:-}" && ! -f "${CI_HCL_SOURCE}" ]]; then + echo "Error: CI HCL file not found at ${CI_HCL_SOURCE}" + echo "Set CI_HCL_SOURCE to a local file if you need an override." + exit 1 + fi +} + +load_ci_hcl() { + echo "--- :page_facing_up: Loading ci.hcl" + cp "${CI_HCL_SOURCE}" "${CI_HCL_PATH}" + echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}" +} + +compute_ci_base_hash_if_needed() { + if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then + return 0 + fi + + CI_BASE_CONTENT_HASH=$(compute_ci_base_content_hash) + export CI_BASE_CONTENT_HASH + echo "ci_base content hash: ${CI_BASE_CONTENT_HASH:0:16}..." +} + +should_push_stable_ci_base_tag() { + if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then + return 0 + fi + if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then + return 1 + fi + + [[ "${NIGHTLY:-0}" == "1" && "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]] +} + +ci_base_tag_with_suffix() { + local base_tag="$1" + local suffix="$2" + + printf '%s-%s\n' "${base_tag}" "$(clean_docker_tag "${suffix}")" +} + +configure_ci_base_image_refs() { + local stable_tag="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" + local content_tag="" + local commit_tag="" + local primary_tag="" + + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then + CI_BASE_IMAGE="${CI_BASE_IMAGE:-${stable_tag}}" + export CI_BASE_IMAGE + return 0 + fi + + content_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${CI_BASE_CONTENT_HASH}") + if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then + commit_tag=$(ci_base_tag_with_suffix "${stable_tag}" "${BUILDKITE_COMMIT}") + CI_BASE_IMAGE_TAG_COMMIT="${commit_tag}" + export CI_BASE_IMAGE_TAG_COMMIT + fi + + if should_push_stable_ci_base_tag; then + primary_tag="${content_tag}" + CI_BASE_IMAGE_TAG_STABLE="${stable_tag}" + else + primary_tag="${commit_tag:-${content_tag}}" + CI_BASE_IMAGE_TAG_STABLE="" + fi + CI_BASE_IMAGE_TAG="${primary_tag}" + if [[ "${primary_tag}" == "${content_tag}" ]]; then + CI_BASE_IMAGE_TAG_CONTENT="" + else + CI_BASE_IMAGE_TAG_CONTENT="${content_tag}" + fi + export CI_BASE_IMAGE_TAG CI_BASE_IMAGE_TAG_CONTENT CI_BASE_IMAGE_TAG_STABLE + + if is_ci_base_target; then + IMAGE_TAG="${primary_tag}" + export IMAGE_TAG + + echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}" + if [[ -n "${CI_BASE_IMAGE_TAG_COMMIT:-}" ]]; then + echo "ci_base commit image tag: ${CI_BASE_IMAGE_TAG_COMMIT}" + fi + echo "ci_base content image tag: ${content_tag}" + if [[ -n "${CI_BASE_IMAGE_TAG_STABLE}" ]]; then + echo "ci_base stable alias will also be pushed: ${CI_BASE_IMAGE_TAG_STABLE}" + else + 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 + return 0 + fi + + if [[ -z "${CI_BASE_IMAGE:-}" || "${CI_BASE_IMAGE}" == "${stable_tag}" ]]; then + CI_BASE_IMAGE="${primary_tag}" + export CI_BASE_IMAGE + echo "Using ci_base image: ${CI_BASE_IMAGE}" + else + echo "Using provided CI_BASE_IMAGE override: ${CI_BASE_IMAGE}" + fi +} + +ci_base_candidate_refs() { + printf '%s\n' \ + "${IMAGE_TAG:-}" \ + "${CI_BASE_IMAGE_TAG:-}" \ + "${CI_BASE_IMAGE_TAG_COMMIT:-}" \ + "${CI_BASE_IMAGE_TAG_CONTENT:-}" \ + "${CI_BASE_IMAGE_TAG_STABLE:-}" \ + | awk 'NF && !seen[$0]++' +} + +find_matching_ci_base_ref() { + local candidate="" + local candidate_hash="" + + while IFS= read -r candidate; do + [[ -n "${candidate}" ]] || continue + remote_image_exists "${candidate}" || continue + candidate_hash=$(get_remote_image_label "${candidate}" "vllm.ci_base.content_hash") + if [[ "${candidate_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + printf '%s\n' "${candidate}" + return 0 + fi + done < <(ci_base_candidate_refs) + + return 1 +} + +refresh_ci_base_tags_from_ref() { + local source_ref="$1" + local tag="" + local tag_hash="" + + while IFS= read -r tag; do + [[ -n "${tag}" ]] || continue + [[ "${tag}" != "${source_ref}" ]] || continue + tag_hash=$(get_remote_image_label "${tag}" "vllm.ci_base.content_hash") + if [[ "${tag_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + echo "ci_base tag is already current: ${tag}" + continue + fi + echo "Updating ci_base tag ${tag} -> ${source_ref}" + docker buildx imagetools create -t "${tag}" "${source_ref}" + done < <(ci_base_candidate_refs) +} + +maybe_skip_existing_image() { + local remote_hash="" + local remote_revision="" + local matching_ref="" + + if [[ -z "${IMAGE_TAG:-}" ]]; then + return 0 + fi + + if [[ "${FORCE_BUILD:-0}" == "1" ]]; then + echo "FORCE_BUILD=1 set; skipping existing-image check" + return 0 + fi + + echo "--- :mag: Checking image tag" + echo "Image tag: ${IMAGE_TAG}" + + if ! remote_image_exists "${IMAGE_TAG}"; then + if is_ci_base_target && [[ -n "${CI_BASE_CONTENT_HASH:-}" ]]; then + matching_ref=$(find_matching_ci_base_ref || true) + if [[ -n "${matching_ref}" ]]; then + echo "Found existing ci_base image with matching content hash: ${matching_ref}" + if ! refresh_ci_base_tags_from_ref "${matching_ref}"; then + echo "ci_base tag refresh failed; rebuilding to push expected tags" + return 0 + fi + echo "Content hashes match -- ci_base is current" + echo "Skipping build" + exit 0 + fi + fi + echo "Image not found, proceeding with build" + return 0 + fi + + IMAGE_EXISTED_BEFORE_BUILD=1 + + if is_ci_base_target; then + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then + echo "ci_base image already exists and no content hash was configured" + echo "Skipping build" + exit 0 + fi + + remote_hash=$(get_remote_image_label "${IMAGE_TAG}" "vllm.ci_base.content_hash") + if [[ -n "${remote_hash}" ]]; then + echo "Remote ci_base content hash: ${remote_hash:0:16}..." + if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + if ! refresh_ci_base_tags_from_ref "${IMAGE_TAG}"; then + echo "ci_base tag refresh failed; rebuilding to push expected tags" + return 0 + fi + echo "Content hashes match -- ci_base is current" + echo "Skipping build" + exit 0 + fi + + echo "Content hashes differ -- ci_base is stale, rebuilding" + return 0 + fi + + echo "Remote ci_base has no content-hash label; rebuilding to add one" + return 0 + fi + + if is_commit_image_target; then + remote_revision=$(get_remote_image_label "${IMAGE_TAG}" "org.opencontainers.image.revision") + if [[ -n "${remote_revision}" && "${remote_revision}" != "${BUILDKITE_COMMIT}" ]]; then + echo "Existing image revision does not match ${BUILDKITE_COMMIT}" + echo " found revision: ${remote_revision}" + echo "Rebuilding image" + return 0 + fi + + if should_upload_wheel_artifacts; then + echo "Commit image already exists: ${IMAGE_TAG}" + echo "Continuing build because this target uploads per-build ROCm artifacts" + return 0 + fi + + echo "Commit image already exists: ${IMAGE_TAG}" + echo "Skipping build" + exit 0 + fi + + echo "Image already exists: ${IMAGE_TAG}" + echo "Skipping build" + exit 0 +} + +setup_builder() { + echo "--- :buildkite: Setting up buildx builder" + + local setup_mode="${ROCM_SETUP_BUILDX_BUILDER:-auto}" + local current_driver="" + local named_driver="" + + if [[ "${setup_mode}" == "0" || "${setup_mode}" == "false" ]]; then + echo "Using current Docker buildx builder" + echo "ROCM_SETUP_BUILDX_BUILDER=${setup_mode}; cache exporters may fail if the driver is docker" + docker buildx inspect --bootstrap + echo "Active builder:" + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls + return 0 + fi + + current_driver=$(buildx_driver || true) + if [[ "${setup_mode}" != "1" ]] && builder_supports_registry_cache "${current_driver}"; then + echo "Using current Docker buildx builder with ${current_driver} driver" + docker buildx inspect --bootstrap + echo "Active builder:" + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls + return 0 + fi + + if [[ "${setup_mode}" != "1" ]]; then + echo "Current buildx driver '${current_driver:-unknown}' cannot export registry caches" + echo "Creating or using a cache-capable builder: ${BUILDER_NAME}" + fi + + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then + named_driver=$(buildx_driver "${BUILDER_NAME}" || true) + if ! builder_supports_registry_cache "${named_driver}"; then + echo "Builder '${BUILDER_NAME}' uses ${named_driver:-unknown} driver; using ${BUILDER_NAME}-cache instead" + BUILDER_NAME="${BUILDER_NAME}-cache" + fi + fi + + if [[ -S "${BUILDKIT_SOCKET}" ]]; then + echo "Found local buildkitd socket at ${BUILDKIT_SOCKET}" + echo "Using remote driver to connect to buildkitd" + + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then + use_existing_builder + else + create_and_bootstrap_builder remote "unix://${BUILDKIT_SOCKET}" + fi + elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then + use_existing_builder + else + echo "No local buildkitd found, using docker-container driver" + create_and_bootstrap_builder docker-container + fi + + echo "Active builder:" + docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls +} + +prepare_git_cache_metadata() { + local cache_branch_name="" + local cache_base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" + local target_repo_slug="" + local target_repo_url="" + local merge_base_ref="" + + if [[ -z "${PARENT_COMMIT:-}" || -z "${VLLM_MERGE_BASE_COMMIT:-}" ]] \ + && git rev-parse --is-shallow-repository 2>/dev/null | grep -q "true"; then + echo "Shallow clone detected - deepening for cache key computation" + git_fetch_for_cache --deepen=1 origin || true + fi + + if [[ -z "${PARENT_COMMIT:-}" ]]; then + PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "") + if [[ -n "${PARENT_COMMIT}" ]]; then + export PARENT_COMMIT + echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}" + else + echo "Could not determine parent commit" + fi + else + echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}" + fi + + if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then + cache_branch_name=$(select_cache_branch_name) + if [[ -z "${cache_branch_name}" && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + cache_branch_name="pr-${BUILDKITE_PULL_REQUEST}" + echo "Using pull request number for ROCm branch cache tag: ${cache_branch_name}" + fi + fi + + if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" && -n "${cache_branch_name}" ]]; then + ROCM_CACHE_BRANCH_TAG=$( + compose_cache_branch_tag "$(get_buildkite_repo_slug)" "${cache_branch_name}" + ) + export ROCM_CACHE_BRANCH_TAG + echo "Computed ROCm branch cache tag: ${ROCM_CACHE_BRANCH_TAG} (from ${cache_branch_name})" + elif [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then + echo "Using provided ROCM_CACHE_BRANCH_TAG: ${ROCM_CACHE_BRANCH_TAG}" + elif [[ -n "${BUILDKITE_BRANCH:-}" ]]; then + echo "Skipping ROCm branch cache tag: no usable branch name found" + echo " BUILDKITE_BRANCH=${BUILDKITE_BRANCH}" + fi + + if [[ -z "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" \ + && -n "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-}" \ + && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + target_repo_slug=$(get_buildkite_target_repo_slug) + ROCM_CACHE_UPSTREAM_BRANCH_TAG=$( + compose_cache_branch_tag "${target_repo_slug}" "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}" + ) + export ROCM_CACHE_UPSTREAM_BRANCH_TAG + echo "Computed ROCm upstream branch cache tag: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" + elif [[ -n "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" ]]; then + echo "Using provided ROCM_CACHE_UPSTREAM_BRANCH_TAG: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" + fi + + if [[ -z "${VLLM_MERGE_BASE_COMMIT:-}" ]]; then + target_repo_url=$(get_buildkite_target_repo_url) + merge_base_ref="refs/remotes/vllm-cache-upstream/${cache_base_branch}" + git_fetch_for_cache --no-tags --depth=200 "${target_repo_url}" \ + "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true + VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "") + if [[ -z "${VLLM_MERGE_BASE_COMMIT}" ]]; then + git_fetch_for_cache --no-tags --deepen=1000 "${target_repo_url}" \ + "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true + VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "") + fi + if [[ -n "${VLLM_MERGE_BASE_COMMIT}" ]]; then + export VLLM_MERGE_BASE_COMMIT + echo "Computed merge base commit for cache fallback: ${VLLM_MERGE_BASE_COMMIT}" + else + echo "Could not determine merge base with ${cache_base_branch}" + fi + else + echo "Using provided VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}" + fi +} + +write_ci_base_label_override() { + local target_name="" + local -a ci_base_targets=() + + BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") + + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then + return 0 + fi + + mapfile -t ci_base_targets < <( + { + printf '%s\n' "ci-base-rocm" + sed -n -E 's/^target "(ci-base-rocm[^"]+)".*/\1/p' "${CI_HCL_PATH}" 2>/dev/null || true + } | awk '!seen[$0]++' + ) + + if [[ ${#ci_base_targets[@]} -eq 0 ]]; then + return 0 + fi + + : > "${CI_BASE_LABEL_OVERRIDE_PATH}" + for target_name in "${ci_base_targets[@]}"; do + cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" < "${ROCM_ARG_OVERRIDE_PATH}" + + BAKE_FILES+=(-f "${ROCM_ARG_OVERRIDE_PATH}") + echo "Appended resolved ROCm Docker ARG override" +} + +write_hcl_string_list_attr() { + local indent="$1" + local attr="$2" + shift 2 + + printf '%s%s = [\n' "${indent}" "${attr}" + write_hcl_string_list_entries "${indent} " "$@" + printf '%s]\n' "${indent}" +} + +validate_cache_export_mode() { + local mode="$1" + local env_name="$2" + + case "${mode}" in + min|max) + ;; + *) + echo "Error: ${env_name} must be one of: min, max" + exit 1 + ;; + esac +} + +write_rocm_cache_override() { + local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" + local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}" + local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}" + local -a content_cache_from=() + local -a csrc_cache_to=() + local -a rocm_cache_to=() + local -a export_wheel_cache_to=() + + if ! uses_rocm_csrc_cache; then + return 0 + fi + + validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE" + validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE" + echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}" + echo "ROCm final image cache export mode: ${rocm_cache_to_mode}" + + if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then + content_cache_from+=("type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}") + csrc_cache_to+=( + "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true" + ) + fi + + # Docker Hub cache exports are best-effort. A cache-only target failure can + # otherwise cancel the sibling image target before its manifest is pushed. + if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then + csrc_cache_to+=( + "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true" + ) + rocm_cache_to+=( + "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true" + ) + fi + + if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then + csrc_cache_to+=( + "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true" + ) + rocm_cache_to+=( + "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true" + ) + fi + + if [[ "${TARGET}" == "test-rocm-ci-with-wheel" ]]; then + export_wheel_cache_to=() + else + export_wheel_cache_to=("${rocm_cache_to[@]}") + fi + + { + cat < "${CSRC_CACHE_OVERRIDE_PATH}" + + BAKE_FILES+=(-f "${CSRC_CACHE_OVERRIDE_PATH}") + echo "Appended ROCm cache override with non-fatal registry exports" +} + +extract_dependency_pins() { + local bake_dir="" + local dockerfile_rocm="" + local var="" + local val="" + + bake_dir=$(dirname "${VLLM_BAKE_FILE}") + dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + if [[ ! -f "${dockerfile_rocm}" ]]; then + return 0 + fi + + for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do + if [[ -n "${!var:-}" ]]; then + echo "Using provided ${var}: ${!var}" + continue + fi + + val=$( + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + "${dockerfile_rocm}" | head -1 + ) + if [[ -n "${val}" ]]; then + export "${var}=${val}" + echo "Extracted ${var}=${val} from Dockerfile.rocm" + fi + done +} + +compute_dependency_cache_keys() { + local bake_dir="" + local dockerfile_rocm="" + local rixl_branch="" + local ucx_branch="" + local rocshmem_branch="" + local deepep_branch="" + local rixl_material="" + local rocshmem_material="" + local deepep_material="" + + bake_dir=$(dirname "${VLLM_BAKE_FILE}") + dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH") + ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") + rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") + deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH") + + if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then + rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl") + RIXL_CACHE_KEY=$( + compose_dependency_cache_key \ + "${rixl_branch}-ucx-${ucx_branch}" \ + "${rixl_material}" + ) + export RIXL_CACHE_KEY + echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}" + fi + + if [[ -n "${rocshmem_branch}" ]]; then + rocshmem_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem") + ROCSHMEM_CACHE_KEY=$( + compose_dependency_cache_key \ + "${rocshmem_branch}" \ + "${rocshmem_material}" + ) + export ROCSHMEM_CACHE_KEY + echo "ROCShmem dependency cache key: ${ROCSHMEM_CACHE_KEY}" + fi + + if [[ -n "${deepep_branch}" && -n "${rocshmem_branch}" ]]; then + deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rocshmem build_deepep") + DEEPEP_CACHE_KEY=$( + compose_dependency_cache_key \ + "${deepep_branch}-rocshmem-${rocshmem_branch}" \ + "${deepep_material}" + ) + export DEEPEP_CACHE_KEY + echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}" + fi +} + +compose_stage_cache_material() { + local dockerfile="$1" + local stages="$2" + local -a content_args=() + + mapfile -t content_args < <(get_content_arg_names "${dockerfile}" "${stages}" "") + { + printf 'dockerfile:%s\n' "${dockerfile}" + printf 'dockerfile-stages:%s\n' "${stages}" + hash_dockerfile_stages "${dockerfile}" "${stages}" + printf 'resolved-build-args:\n' + hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" + } +} + +dependency_cache_ref_exists() { + local cache_ref="$1" + docker buildx imagetools inspect "${cache_ref}" >/dev/null 2>&1 +} + +dependency_cache_ref_for_target() { + local target="$1" + local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" + + case "${target}" in + rixl-rocm-ci) + if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then + printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}" + elif [[ -n "${RIXL_BRANCH:-}" ]]; then + printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}" + fi + ;; + rocshmem-rocm-ci) + if [[ -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then + printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" + elif [[ -n "${ROCSHMEM_BRANCH:-}" ]]; then + printf '%s\n' "${cache_repo}:rocshmem-rocm-${ROCSHMEM_BRANCH}" + fi + ;; + deepep-rocm-ci) + if [[ -n "${DEEPEP_CACHE_KEY:-}" ]]; then + printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_CACHE_KEY}" + elif [[ -n "${DEEPEP_BRANCH:-}" ]]; then + printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH:-}" + fi + ;; + esac +} + +add_dependency_cache_target() { + local target="$1" + + if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "${target}"; then + return 0 + fi + DEPENDENCY_CACHE_TARGETS+=("${target}") +} + +resolve_ci_base_dependency_targets() { + local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}" + local rixl_ref="" + local rocshmem_ref="" + local deepep_ref="" + + [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0 + + case "${mode}" in + always) + echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially" + for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do + if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then + add_dependency_cache_target "${target}" + fi + done + ;; + never) + BAKE_TARGETS=("ci-base-rocm-ci") + DEPENDENCY_CACHE_TARGETS=() + echo "ROCM_DEP_CACHE_EXPORT_MODE=never; building ci_base without dependency cache exports" + return 0 + ;; + missing|"") + ;; + *) + echo "Error: ROCM_DEP_CACHE_EXPORT_MODE must be one of: missing, always, never" + exit 1 + ;; + esac + + if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then + rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci") + if dependency_cache_ref_exists "${rixl_ref}"; then + echo "RIXL dependency cache exists: ${rixl_ref}" + else + echo "RIXL dependency cache missing; will seed: ${rixl_ref}" + add_dependency_cache_target "rixl-rocm-ci" + fi + fi + + if [[ "${mode}" != "always" && -n "${ROCSHMEM_CACHE_KEY:-}" ]]; then + rocshmem_ref=$(dependency_cache_ref_for_target "rocshmem-rocm-ci") + if dependency_cache_ref_exists "${rocshmem_ref}"; then + echo "ROCShmem dependency cache exists: ${rocshmem_ref}" + else + echo "ROCShmem dependency cache missing; will seed: ${rocshmem_ref}" + add_dependency_cache_target "rocshmem-rocm-ci" + fi + fi + + if [[ "${mode}" != "always" && -n "${DEEPEP_CACHE_KEY:-}" ]]; then + deepep_ref=$(dependency_cache_ref_for_target "deepep-rocm-ci") + if dependency_cache_ref_exists "${deepep_ref}"; then + echo "DeepEP dependency cache exists: ${deepep_ref}" + else + echo "DeepEP dependency cache missing; will seed: ${deepep_ref}" + add_dependency_cache_target "deepep-rocm-ci" + fi + fi + + # DeepEP inherits from ROCShmem. If ROCShmem is being seeded, seed DeepEP too + # so the pair stays consistent for future ci_base rebuilds. + if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "rocshmem-rocm-ci" \ + && ! printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "deepep-rocm-ci" \ + && [[ -n "${DEEPEP_BRANCH:-}" ]]; then + echo "ROCShmem cache is missing; also seeding DeepEP cache" + add_dependency_cache_target "deepep-rocm-ci" + fi + + BAKE_TARGETS=("ci-base-rocm-ci") + if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then + echo "All dependency caches exist; building ci_base without dependency cache exports" + else + echo "Resolved dependency cache seed targets: ${DEPENDENCY_CACHE_TARGETS[*]}" + echo "Resolved ci_base bake targets: ${BAKE_TARGETS[*]}" + fi +} + +bake_config_targets() { + printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" "${BAKE_TARGETS[@]}" \ + | awk 'NF && !seen[$0]++' +} + +print_bake_config() { + local -a print_targets=() + + echo "--- :page_facing_up: Resolved bake configuration" + mapfile -t print_targets < <(bake_config_targets) + docker buildx bake "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}" + + if command -v buildkite-agent >/dev/null 2>&1 && [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then + buildkite-agent artifact upload "${BAKE_CONFIG_FILE}" || true + echo "Uploaded ${BAKE_CONFIG_FILE} as Buildkite artifact" + else + echo "Saved bake config to ${BAKE_CONFIG_FILE} (not in Buildkite, skipping upload)" + fi +} + +confirm_remote_image_push() { + local image_ref="$1" + local remote_hash="" + local remote_revision="" + + if ! remote_image_exists "${image_ref}"; then + return 1 + fi + + if is_ci_base_target; then + if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then + return 0 + fi + + remote_hash=$(get_remote_image_label_with_retry "${image_ref}" "vllm.ci_base.content_hash") + if [[ -n "${remote_hash}" && "${remote_hash}" == "${CI_BASE_CONTENT_HASH}" ]]; then + return 0 + fi + + echo "Remote image exists but does not have the expected ci_base content hash." + echo " expected: ${CI_BASE_CONTENT_HASH:0:16}..." + echo " found: ${remote_hash:0:16}..." + return 1 + fi + + if is_commit_image_target; then + remote_revision=$(get_remote_image_label_with_retry "${image_ref}" "org.opencontainers.image.revision") + if [[ -n "${remote_revision}" && "${remote_revision}" == "${BUILDKITE_COMMIT}" ]]; then + return 0 + fi + + if [[ -z "${remote_revision}" \ + && ${IMAGE_EXISTED_BEFORE_BUILD} -eq 0 \ + && image_tag_is_commit_scoped ]]; then + echo "Remote image exists under a commit-scoped tag; accepting push despite missing revision label." + return 0 + fi + + echo "Remote image exists but revision label does not match ${BUILDKITE_COMMIT}." + echo " found revision: ${remote_revision:-}" + return 1 + fi + + return 0 +} + +verify_dependency_cache_ref() { + local cache_ref="$1" + local attempts="${ROCM_DEP_CACHE_VERIFY_ATTEMPTS:-6}" + local delay_secs="${ROCM_DEP_CACHE_VERIFY_DELAY:-5}" + local attempt + + for ((attempt = 1; attempt <= attempts; attempt++)); do + if dependency_cache_ref_exists "${cache_ref}"; then + echo "Dependency cache confirmed: ${cache_ref}" + return 0 + fi + if [[ ${attempt} -lt ${attempts} ]]; then + echo "Dependency cache not visible yet (${attempt}/${attempts}): ${cache_ref}" + sleep "${delay_secs}" + fi + done + + echo "ERROR: dependency cache was not confirmed after upload: ${cache_ref}" + return 1 +} + +seed_dependency_caches_if_needed() { + local target="" + local cache_ref="" + + if [[ "${TARGET}" != "ci-base-rocm-ci-with-deps" ]]; then + return 0 + fi + if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then + return 0 + fi + + echo "--- :docker: Seeding ROCm dependency caches" + echo "Dependency cache uploads are required for this build." + echo "Seeding serially to avoid concurrent Docker Hub cache exporters." + + for target in "${DEPENDENCY_CACHE_TARGETS[@]}"; do + cache_ref=$(dependency_cache_ref_for_target "${target}") + if [[ -z "${cache_ref}" ]]; then + echo "ERROR: could not resolve dependency cache ref for ${target}" + return 1 + fi + + echo "--- :docker: Seeding ${target}" + echo "Expected cache ref: ${cache_ref}" + docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}" + verify_dependency_cache_ref "${cache_ref}" + done +} + +annotate_cache_export_warning() { + local build_rc="$1" + + if ! command -v buildkite-agent >/dev/null 2>&1; then + return 0 + fi + + buildkite-agent annotate \ + --style warning \ + --context "cache-export-warning" \ + "### :warning: Docker cache export failed (non-fatal) + +Image was pushed successfully: \`${IMAGE_TAG}\` + +The BuildKit build returned exit code ${build_rc}, but the expected image +is present in the registry. Treating this as a registry cache export failure +so tests can continue with the pushed image." 2>/dev/null || true +} + +run_bake() { + local build_rc=0 + + echo "--- :docker: Building ${TARGET}" + docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$? + + if [[ ${build_rc} -eq 0 ]]; then + echo "--- :white_check_mark: Build complete" + return 0 + fi + + echo "" + echo "WARNING: docker buildx bake exited with code ${build_rc}" + + if [[ -n "${IMAGE_TAG:-}" ]]; then + echo "Checking if image was pushed successfully..." + if confirm_remote_image_push "${IMAGE_TAG}"; then + echo "" + echo "WARNING: Build reported failure (rc=${build_rc}) but the" + echo " image was pushed successfully: ${IMAGE_TAG}" + echo "" + echo " Treating this as a non-fatal registry cache export failure." + echo " The image is usable, but registry cache may be cold on the next build." + echo "" + annotate_cache_export_warning "${build_rc}" + echo "--- :white_check_mark: Build complete" + return 0 + fi + + echo "" + echo "ERROR: Build failed and image was NOT confirmed: ${IMAGE_TAG}" + echo " This is a real build failure, not a cache export warning." + echo "" + fi + + return "${build_rc}" +} + +upload_wheel_artifacts_if_present() { + local wheel_dir="./wheel-export" + local artifact_dir="artifacts/vllm-rocm-install" + local archive_name="vllm-rocm-install.tar.gz" + local whl="" + local whl_name="" + + if ! should_upload_wheel_artifacts; then + return 0 + fi + + if [[ ! -d "${wheel_dir}" ]] || ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then + echo "No ROCm wheel artifacts found in ${wheel_dir}" + return 0 + fi + + echo "--- :package: Uploading ROCm vLLM install artifact" + mkdir -p "${artifact_dir}" + + tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" . + echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)" + printf '%s\n' "${CI_BASE_IMAGE:-}" > "${artifact_dir}/ci-base-image.txt" + printf '%s\n' "${IMAGE_TAG:-}" > "${artifact_dir}/fallback-image.txt" + + for whl in "${wheel_dir}"/*.whl; do + [[ -f "${whl}" ]] || continue + whl_name=$(basename "${whl}") + cp "${whl}" "${artifact_dir}/${whl_name}" + echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)" + done + + if command -v buildkite-agent >/dev/null 2>&1; then + buildkite-agent artifact upload "${artifact_dir}/*" + echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/" + else + echo "Not in Buildkite, skipping artifact upload" + fi + + rm -rf "${wheel_dir}" +} + +main() { + init_config "$@" + print_header + validate_inputs + load_ci_hcl + compute_ci_base_hash_if_needed + configure_ci_base_image_refs + maybe_skip_existing_image + setup_builder + prepare_git_cache_metadata + write_ci_base_label_override + extract_dependency_pins + write_rocm_build_arg_override + compute_dependency_cache_keys + compute_rocm_csrc_content_hash_if_needed + write_rocm_cache_override + resolve_ci_base_dependency_targets + print_bake_config + if [[ "${BAKE_PRINT_ONLY:-0}" == "1" ]]; then + echo "BAKE_PRINT_ONLY=1 set; skipping build" + return 0 + fi + seed_dependency_caches_if_needed + run_bake + upload_wheel_artifacts_if_present +} + +main "$@" diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 94bbc15fcff..953074c3882 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -52,6 +52,108 @@ cleanup_network() { fi } +prepare_artifact_image() { + if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then + return 1 + fi + if ! command -v buildkite-agent >/dev/null 2>&1; then + echo "buildkite-agent not found; cannot download ROCm wheel artifact" + return 1 + fi + + local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}" + local archive="" + local metadata_file="" + local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}" + local artifact_image="" + local artifact_key="" + local base_digest="" + local wheel_dir="" + local context_dir="" + local workspace_dir="" + + artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) + wheel_dir="${artifact_work_dir}/wheels" + context_dir="${artifact_work_dir}/context" + workspace_dir="${context_dir}/workspace" + mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}" + + echo "--- Downloading ROCm wheel artifact" + if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then + echo "Failed to download ${artifact_glob}" + return 1 + fi + buildkite-agent artifact download \ + "artifacts/vllm-rocm-install/ci-base-image.txt" \ + "${artifact_work_dir}" >/dev/null 2>&1 || true + + archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1) + if [[ -z "${archive}" || ! -f "${archive}" ]]; then + echo "ROCm wheel artifact archive was not found" + return 1 + fi + + metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1) + if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then + base_image=$(tr -d '[:space:]' < "${metadata_file}") + fi + + echo "--- Preparing local ROCm test image" + echo "Base image: ${base_image}" + docker pull "${base_image}" || return 1 + base_digest=$( + docker image inspect \ + --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \ + "${base_image}" 2>/dev/null || printf '%s' "${base_image}" + ) + + artifact_key=$( + { + printf 'base-image:%s\n' "${base_digest}" + sha256sum "${archive}" + } | sha256sum | cut -c1-24 + ) + artifact_image="rocm/vllm-ci-artifact:${artifact_key}" + + if docker image inspect "${artifact_image}" >/dev/null 2>&1; then + echo "Using existing local ROCm artifact image: ${artifact_image}" + image_name="${artifact_image}" + return 0 + fi + + tar -xzf "${archive}" -C "${wheel_dir}" || return 1 + if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then + echo "ROCm wheel artifact did not contain a wheel" + return 1 + fi + if [[ ! -d "${wheel_dir}/tests" ]]; then + echo "ROCm wheel artifact did not contain the test workspace" + return 1 + fi + + cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1 + tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \ + | tar -C "${workspace_dir}" -xf - || return 1 + cat > "${context_dir}/Dockerfile" <<'EOF' +ARG BASE_IMAGE +FROM ${BASE_IMAGE} +COPY wheels/ /tmp/vllm-wheels/ +COPY workspace/ /vllm-workspace/ +RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \ + && rm -rf /tmp/vllm-wheels +WORKDIR /vllm-workspace +EOF + + echo "--- Building local ROCm test image" + docker build \ + --pull=false \ + --build-arg "BASE_IMAGE=${base_image}" \ + -t "${artifact_image}" \ + "${context_dir}" || return 1 + image_name="${artifact_image}" + return 0 +} + is_multi_node() { local cmds="$1" # Primary signal: NUM_NODES environment variable set by the pipeline @@ -243,22 +345,30 @@ report_docker_usage # --- Pull test image --- echo "--- Pulling container" -image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}" +image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}" +artifact_work_dir="" container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)" -docker pull "${image_name}" remove_docker_container() { - # docker run uses --rm, so the container is normally already gone when the - # EXIT trap runs. Cleanup is best-effort and must not affect the test result. - docker rm -f "${container_name}" >/dev/null 2>&1 || true + if docker container inspect "${container_name}" >/dev/null 2>&1; then + docker rm -f "${container_name}" || true + fi + if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then + docker image rm -f "${image_name}" || true + else + # Keep images by default so later jobs on the same AMD node can reuse layers. + echo "Keeping ROCm test image locally: ${image_name}" + fi + if [[ -n "${artifact_work_dir}" ]]; then + rm -rf "${artifact_work_dir}" + fi } +trap remove_docker_container EXIT -on_exit() { - local exit_code=$? - remove_docker_container - exit "$exit_code" -} -trap on_exit EXIT +if ! prepare_artifact_image; then + echo "Using full ROCm CI image: ${image_name}" + docker pull "${image_name}" || exit 1 +fi # --- Prepare commands --- echo "--- Running container" diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 9c13fa79fcb..35513727f16 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -37,7 +37,8 @@ function cpu_tests() { pytest -x -v -s tests/kernels/test_onednn.py 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_moe.py -k test_cpu_fused_moe_basic" + pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits diff --git a/.buildkite/scripts/install-kv-connectors.sh b/.buildkite/scripts/install-kv-connectors.sh new file mode 100755 index 00000000000..34c502e6b9a --- /dev/null +++ b/.buildkite/scripts/install-kv-connectors.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}" + +uv pip install --system -r "${REQUIREMENTS_FILE}" + +NIXL_METADATA=$(python3 - <<'PY' +import importlib.metadata as metadata + +import torch + +cuda_version = torch.version.cuda +if cuda_version is None: + raise SystemExit("torch.version.cuda is not set") + +print(cuda_version.split(".", 1)[0], metadata.version("nixl")) +PY +) +read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}" + +# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant +# matching this CI image so nixl_ep_cpp links against the available libcudart. +uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true +uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}" + +python3 - <<'PY' +import importlib.metadata as metadata + +for package_name in ("nixl", "nixl-cu12", "nixl-cu13"): + try: + version = metadata.version(package_name) + except metadata.PackageNotFoundError: + version = "not installed" + print(f"{package_name}: {version}") +PY diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a04e88b3d7e..a7e26280c90 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1238,14 +1238,11 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - tests/entrypoints/serve commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use + - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - label: Entrypoints Integration (API Server openai - Part 1) # TBD timeout_in_minutes: 180 @@ -1276,11 +1273,13 @@ steps: - tests/entrypoints/openai - tests/entrypoints/test_chat_utils - tests/entrypoints/generate + - tests/tool_use commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/test_chat_utils.py - pytest -v -s entrypoints/generate + - pytest -v -s tool_use - label: Entrypoints Integration (API Server openai - Part 3) # TBD timeout_in_minutes: 180 @@ -1370,7 +1369,7 @@ steps: - vllm/platforms/rocm.py commands: - pytest -v -s entrypoints/openai/tool_parsers - - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate + - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate - label: OpenAI API correctness # TBD timeout_in_minutes: 180 @@ -2747,14 +2746,11 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - tests/entrypoints/serve commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use + - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - label: Entrypoints Integration (API Server openai - Part 1) # TBD timeout_in_minutes: 180 @@ -2785,11 +2781,13 @@ steps: - tests/entrypoints/openai - tests/entrypoints/test_chat_utils - tests/entrypoints/generate + - tests/tool_use commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/test_chat_utils.py - pytest -v -s entrypoints/generate + - pytest -v -s tool_use - label: Entrypoints Integration (API Server openai - Part 3) # TBD timeout_in_minutes: 180 diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index d3e02be2398..c9d5237b67b 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -11,7 +11,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs) key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus @@ -22,7 +22,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) @@ -34,7 +34,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) @@ -46,7 +46,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) @@ -58,7 +58,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) @@ -73,7 +73,7 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) @@ -87,7 +87,7 @@ steps: - vllm/v1/worker/kv_connector_model_runner_mixin.py - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh - label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) @@ -102,5 +102,5 @@ steps: - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - tests/v1/kv_connector/nixl_integration/ commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh \ No newline at end of file + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index ebaec9954a3..548174ed748 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -11,7 +11,7 @@ steps: - tests/entrypoints/ commands: - pytest -v -s entrypoints/openai/tool_parsers - - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate + - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate - label: Entrypoints Integration (LLM) key: entrypoints-integration-llm @@ -61,10 +61,12 @@ steps: - tests/entrypoints/openai - tests/entrypoints/test_chat_utils - tests/entrypoints/generate + - tests/tool_use commands: - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/test_chat_utils.py - pytest -v -s entrypoints/generate + - pytest -v -s tool_use mirror: amd: device: mi325_1 @@ -100,14 +102,11 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ - - tests/entrypoints/rpc - - tests/entrypoints/serve/instrumentator - - tests/tool_use + - tests/entrypoints/serve commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve/instrumentator - - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc - - pytest -v -s tool_use + - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc + - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc mirror: amd: device: mi325_1 @@ -155,6 +154,5 @@ steps: source_file_dependencies: - csrc/ - vllm/entrypoints/openai/ - - vllm/model_executor/models/whisper.py commands: # LMEval - pytest -s entrypoints/openai/correctness/ diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index ddeb692d831..e04016d6dcc 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -86,7 +86,7 @@ steps: - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - export VLLM_WORKER_MULTIPROC_METHOD=spawn # split the test to avoid interference - pytest -v -s -m 'not cpu_test' v1/core @@ -281,6 +281,7 @@ steps: - vllm/model_executor/layers/quantization/quark/ - vllm/multimodal/ - vllm/outputs.py + - vllm/parser/ - vllm/platforms/ - vllm/pooling_params.py - vllm/ray/ diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 2964762b346..617c80b2fec 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -94,11 +94,13 @@ steps: - vllm/v1/worker/gpu_worker.py - tests/distributed/test_pipeline_parallel.py - tests/distributed/test_pp_cudagraph.py + - tests/v1/distributed/test_pp_dp_v2.py commands: - set -x - export VLLM_USE_V2_MODEL_RUNNER=1 - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba" - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray" + - pytest -v -s v1/distributed/test_pp_dp_v2.py - label: Model Runner V2 Spec Decode device: h200_35gb diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index df37022725f..16d69f77345 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -45,19 +45,19 @@ steps: - vllm/entrypoints/serve/ - vllm/v1/engine/ - tests/utils.py - # - tests/entrypoints/rpc/test_collective_rpc.py + # - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py - tests/entrypoints/serve/disagg/test_serving_tokens.py - tests/entrypoints/serve/instrumentator/test_basic.py - tests/entrypoints/serve/instrumentator/test_metrics.py - # - tests/entrypoints/serve/instrumentator/test_sleep.py + # - tests/entrypoints/serve/dev/test_sleep.py commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - # - pytest -v -s entrypoints/rpc/test_collective_rpc.py + # - pytest -v -s entrypoints/serve/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/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - # - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py + # - pytest -v -s entrypoints/serve/dev/test_sleep.py - label: Rust Frontend Core Correctness timeout_in_minutes: 30 diff --git a/.dockerignore b/.dockerignore index 66447272e95..fb010600db9 100644 --- a/.dockerignore +++ b/.dockerignore @@ -33,3 +33,10 @@ share/python-wheels/ *.egg MANIFEST rust/target/ +# Not needed in Docker builds +docs/ +.github/ +.pre-commit-config.yaml +.clang-format +.gitattributes +format.sh diff --git a/.github/workflows/add_label_automerge.yml b/.github/workflows/add_label_automerge.yml index d8bbedef317..28e6c526245 100644 --- a/.github/workflows/add_label_automerge.yml +++ b/.github/workflows/add_label_automerge.yml @@ -10,7 +10,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Add label - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | github.rest.issues.addLabels({ diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 3efa582f670..4eac3d7b789 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -14,7 +14,7 @@ jobs: steps: - name: Label issues based on keywords id: label-step - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Add new labels and keywords here @@ -315,7 +315,7 @@ jobs: - name: CC users for labeled issues if: steps.label-step.outputs.labels_added != '[]' - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Map labels to GitHub users to CC @@ -392,7 +392,7 @@ jobs: - name: Request missing ROCm info from issue author if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug') - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const body = (context.payload.issue.body || '').toLowerCase(); diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index 27100f9f4da..4124583d96d 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -12,7 +12,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Update PR description - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const { owner, repo } = context.repo; @@ -55,7 +55,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Post welcome comment for first-time contributors - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const { owner, repo } = context.repo; diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 1dd31b0e50f..93a5a5ff0ae 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -20,7 +20,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Check PR label and author merge count - uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const { data: pr } = await github.rest.pulls.get({ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 05625e8f667..c11a80683f8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: rev: v21.1.2 hooks: - id: clang-format - exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*' + exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 diff --git a/CMakeLists.txt b/CMakeLists.txt index 86c2214b249..0652a5f066e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -112,6 +112,8 @@ endif() # # spinloop extension (pure CXX; must stay above the non-CUDA device branch so # CPU builds define the target before the early return) +# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading +# failure is handled gracefully on vLLM side for lower Python versions. # set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp") set(SPINLOOP_COMPILE_FLAGS "") @@ -309,14 +311,9 @@ set(VLLM_EXT_SRC "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" "csrc/quantization/activation_kernels.cu" "csrc/cuda_utils_kernels.cu" - "csrc/custom_all_reduce.cu" - "csrc/torch_bindings.cpp" - "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu") + "csrc/torch_bindings.cpp") if(VLLM_GPU_LANG STREQUAL "CUDA") - list(APPEND VLLM_EXT_SRC - "csrc/minimax_reduce_rms_kernel.cu") - SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. @@ -503,12 +500,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS) set(SRCS - "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu" - "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu") + "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu" + "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu") set_gencode_flags_for_srcs( SRCS "${SRCS}" CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}") - list(APPEND VLLM_EXT_SRC "${SRCS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1") message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}") else() @@ -598,7 +595,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") endif() if (VLLM_GPU_LANG STREQUAL "HIP") - # Add QuickReduce kernels + # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration). list(APPEND VLLM_EXT_SRC "csrc/custom_quickreduce.cu" ) @@ -633,6 +630,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/activation_kernels.cu" "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu" + "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" + "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" "csrc/libtorch_stable/quantization/gptq/q_gemm.cu" "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu" "csrc/libtorch_stable/pos_encoding_kernels.cu" @@ -647,7 +646,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/attention/paged_attention_v1.cu" "csrc/libtorch_stable/attention/paged_attention_v2.cu" "csrc/libtorch_stable/cache_kernels.cu" - "csrc/libtorch_stable/cache_kernels_fused.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") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC @@ -659,7 +661,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/permute_cols.cu" "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu" "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu" - "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu") + "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu" + "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_STABLE_EXT_SRC}" diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index c51384e3196..6f836ff5354 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -369,6 +369,18 @@ else() add_compile_definitions(-DVLLM_NUMA_DISABLED) endif() +# check if the pytorch wheel ships libopenblas.so. +set(VLLM_OPENBLAS_LIB "") +if (NOT ENABLE_X86_ISA) + file(GLOB _VLLM_TORCH_OPENBLAS_LIBS + "${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*") + # Note: we don't link openblas directly to _C extension, as it's available through libtorch.so + if (_VLLM_TORCH_OPENBLAS_LIBS) + list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB) + message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}") + endif() +endif() + # # Generate CPU attention dispatch header # @@ -387,6 +399,7 @@ endif() # set(VLLM_EXT_SRC "csrc/cpu/activation.cpp" + "csrc/cpu/sgl-kernels/fla.cpp" "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/layernorm.cpp" @@ -410,6 +423,12 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) ${VLLM_EXT_SRC}) endif() +if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) + set(VLLM_EXT_SRC + "csrc/cpu/shm.cpp" + ${VLLM_EXT_SRC}) +endif() + if(USE_ONEDNN) set(VLLM_EXT_SRC "csrc/cpu/dnnl_kernels.cpp" @@ -418,7 +437,6 @@ endif() if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_SGL - "csrc/cpu/sgl-kernels/fla.cpp" "csrc/cpu/sgl-kernels/conv.cpp" "csrc/cpu/sgl-kernels/gemm.cpp" "csrc/cpu/sgl-kernels/gemm_int8.cpp" @@ -430,6 +448,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/moe_fp8.cpp") set(VLLM_EXT_SRC_AVX512 + "csrc/cpu/sgl-kernels/fla.cpp" "csrc/cpu/shm.cpp" "csrc/cpu/cpu_wna16.cpp" "csrc/cpu/cpu_fused_moe.cpp" @@ -446,6 +465,7 @@ if (ENABLE_X86_ISA) "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") set(VLLM_EXT_SRC_AVX2 + "csrc/cpu/sgl-kernels/fla.cpp" "csrc/cpu/utils.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" @@ -519,6 +539,9 @@ else() USE_SABI 3 WITH_SOABI ) + if (VLLM_OPENBLAS_LIB) + target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS) + endif() endif() message(STATUS "Enabling C extension.") diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index b38917a7b0b..1e4feb0ff9e 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -31,7 +31,7 @@ endif() if(VLLM_FLASH_ATTN_SRC_DIR) FetchContent_Declare( - vllm-flash-attn SOURCE_DIR + vllm-flash-attn SOURCE_DIR ${VLLM_FLASH_ATTN_SRC_DIR} BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn ) @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG bce29425653ec0fbc579d329883030e832d15ada + GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/cmake/hipify.py b/cmake/hipify.py index 8504f9defee..f4932260c9e 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -14,7 +14,18 @@ import argparse import os import shutil -from torch.utils.hipify.hipify_python import hipify +from torch.utils.hipify.hipify_python import get_hip_file_path, hipify + + +def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: + """Match torch.utils.hipify.hipify_python.preprocessor fout_path naming.""" + rel = os.path.relpath(source_abs, output_directory) + return os.path.abspath( + os.path.join( + output_directory, get_hip_file_path(rel, is_pytorch_extension=True) + ) + ) + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -53,7 +64,11 @@ if __name__ == "__main__": hipify_result = hipify( project_directory=args.project_dir, output_directory=args.output_dir, - header_include_dirs=[], + # Hipify resolves quoted includes next to the including file first; vLLM + # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h" + # from quantization/w8a8/fp8/*.cu). Without an include root here, those + # headers are never found and are not hipified or rewritten in dependents. + header_include_dirs=["."], includes=includes, extra_files=extra_files, show_detailed=True, @@ -64,14 +79,20 @@ if __name__ == "__main__": hipified_sources = [] for source in args.sources: s_abs = os.path.abspath(source) - hipified_s_abs = ( - hipify_result[s_abs].hipified_path - if ( - s_abs in hipify_result - and hipify_result[s_abs].hipified_path is not None - ) - else s_abs - ) + if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: + path = hipify_result[s_abs].hipified_path + # PyTorch skips writing when is_pytorch_extension and text unchanged; + # hipified_path then stays *.cu. CMake expects *.hip under output_dir. + if s_abs.endswith(".cu") and path.endswith(".cu"): + dest = _expected_hip_build_path(s_abs, args.output_dir) + if os.path.normpath(path) != os.path.normpath(dest): + os.makedirs(os.path.dirname(dest), exist_ok=True) + shutil.copy2(path, dest) + hipified_s_abs = dest + else: + hipified_s_abs = path + else: + hipified_s_abs = s_abs hipified_sources.append(hipified_s_abs) assert len(hipified_sources) == len(args.sources) diff --git a/cmake/utils.cmake b/cmake/utils.cmake index f10ba93f7c6..dd2034c1c5e 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -81,6 +81,14 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS) set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS}) set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS}) + # Chain hipify targets so they run sequentially. Parallel hipify + # invocations race on shutil.copytree, overwriting .hip files + # produced by another target back to .cu originals. + if (DEFINED _VLLM_LAST_HIPIFY_TARGET) + add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET}) + endif() + set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE) + # Swap out original extension sources with hipified sources. list(APPEND HIP_SRCS ${CXX_SRCS}) set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE) diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 0dc5060fe99..5839d6c2aaf 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -30,7 +30,12 @@ }() namespace { -enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul }; +enum class FusedMOEAct { + SiluAndMul, + SwigluOAIAndMul, + GeluAndMul, + GeluTanhAndMul, +}; FusedMOEAct get_act_type(const std::string& act) { if (act == "silu") { @@ -39,6 +44,8 @@ FusedMOEAct get_act_type(const std::string& act) { return FusedMOEAct::SwigluOAIAndMul; } else if (act == "gelu") { return FusedMOEAct::GeluAndMul; + } else if (act == "gelu_tanh") { + return FusedMOEAct::GeluTanhAndMul; } else { TORCH_CHECK(false, "Invalid act type: " + act); } @@ -143,6 +150,44 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, } } +template +void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, + const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(0.7978845608028654); + vec_op::FP32Vec16 w2_vec(0.5); + vec_op::FP32Vec16 w3_vec(0.044715); + alignas(64) float temp[16]; + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto gate_pow3_vec = gate_vec * gate_vec * gate_vec; + auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec); + + inner_vec.save(temp); + for (int32_t i = 0; i < 16; ++i) { + temp[i] = std::tanh(temp[i]); + } + vec_op::FP32Vec16 tanh_vec(temp); + auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec); + auto gated_output_fp32 = up_vec * gelu_tanh; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + template FORCE_INLINE void apply_gated_act(const FusedMOEAct act, float* __restrict__ input, @@ -160,6 +205,9 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act, case FusedMOEAct::GeluAndMul: gelu_and_mul(input, output, m, n, input_stride, output_stride); return; + case FusedMOEAct::GeluTanhAndMul: + gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride); + return; default: TORCH_CHECK(false, "Unsupported act type."); } diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 87c7a9dd51f..ba65e27a15e 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -89,6 +89,35 @@ struct BF16Vec8 : public Vec { } }; +struct FP16Vec16 : public Vec { + constexpr static int VEC_ELEM_NUM = 16; + ss16x8x2_t reg; + + explicit FP16Vec16(const void* ptr) { + reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr); + reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr); + } + + explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {} + + explicit FP16Vec16(const FP32Vec16&); + + void save(void* ptr) const { + vec_xst(reg.val[0], 0, (signed short*)ptr); + vec_xst(reg.val[1], 16, (signed short*)ptr); + } + + void save(void* ptr, int elem_num) const { + int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM)); + if (num <= 8) { + vec_xst_len(reg.val[0], (signed short*)ptr, num * 2); + } else { + vec_xst(reg.val[0], 0, (signed short*)ptr); + vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2); + } + } +}; + struct BF16Vec16 : public Vec { constexpr static int VEC_ELEM_NUM = 16; @@ -100,6 +129,8 @@ struct BF16Vec16 : public Vec { reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr); } + explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {} + explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { @@ -379,6 +410,8 @@ struct FP32Vec16 : public Vec { reg.val[3] = vec_xl(48, ptr); } + explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {} + explicit FP32Vec16(f32x4x4_t data) : reg(data) {} explicit FP32Vec16(const FP32Vec16& data) { @@ -402,6 +435,7 @@ struct FP32Vec16 : public Vec { reg.val[3] = data.reg.val[1]; } + explicit FP32Vec16(const FP16Vec16& v); explicit FP32Vec16(const BF16Vec16& v) { reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]); reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]); @@ -735,6 +769,40 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) { #endif } +inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { + alignas(16) float temp_fp32[16]; + alignas(16) c10::Half temp_fp16[16]; + + vec_xst(v.reg.val[0], 0, temp_fp32); + vec_xst(v.reg.val[1], 16, temp_fp32); + vec_xst(v.reg.val[2], 32, temp_fp32); + vec_xst(v.reg.val[3], 48, temp_fp32); + + for (int i = 0; i < 16; i++) { + temp_fp16[i] = c10::Half(temp_fp32[i]); + } + + reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16); + reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16); +} + +inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { + alignas(16) c10::Half temp_fp16[16]; + alignas(16) float temp_fp32[16]; + + vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16); + vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16); + + for (int i = 0; i < 16; i++) { + temp_fp32[i] = float(temp_fp16[i]); + } + + reg.val[0] = vec_xl(0, temp_fp32); + reg.val[1] = vec_xl(16, temp_fp32); + reg.val[2] = vec_xl(32, temp_fp32); + reg.val[3] = vec_xl(48, temp_fp32); +} + inline BF16Vec16::BF16Vec16(const FP32Vec16& v) { #ifdef _ARCH_PWR10 __vector signed short ret[4]; @@ -794,6 +862,43 @@ inline void prefetch(const void* addr) { __asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory"); } -}; // namespace vec_op +struct INT8Vec64 { + __vector signed char data[4]; + + INT8Vec64() = default; + + explicit INT8Vec64(const int8_t* ptr) { + data[0] = vec_xl(0, ptr); + data[1] = vec_xl(16, ptr); + data[2] = vec_xl(32, ptr); + data[3] = vec_xl(48, ptr); + } + + explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {} + + void save(int8_t* ptr) const { + vec_xst(data[0], 0, ptr); + vec_xst(data[1], 16, ptr); + vec_xst(data[2], 32, ptr); + vec_xst(data[3], 48, ptr); + } + + void save(int8_t* ptr, int elem_num) const { + if (elem_num <= 0) return; + + int full_vecs = elem_num / 16; + for (int i = 0; i < full_vecs && i < 4; i++) { + vec_xst(data[i], i * 16, ptr); + } + + int remaining = elem_num % 16; + if (remaining > 0 && full_vecs < 4) { + vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining); + } + } + + void nt_save(int8_t* ptr) const { save(ptr); } +}; +} // namespace vec_op #endif diff --git a/csrc/cpu/sgl-kernels/blas_gemm.h b/csrc/cpu/sgl-kernels/blas_gemm.h new file mode 100644 index 00000000000..315eb210b0c --- /dev/null +++ b/csrc/cpu/sgl-kernels/blas_gemm.h @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include + +// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm +// If OpenBLS is available in the PyTorch wheel, we rely on it for fast +// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path. +#if defined(VLLM_HAS_OPENBLAS) +extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k, + float* alpha, const at::BFloat16* a, int* lda, + const at::BFloat16* b, int* ldb, float* beta, float* c, + int* ldc); + +extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k, + float* alpha, const float* a, int* lda, const float* b, + int* ldb, float* beta, float* c, int* ldc); + +inline char blas_transpose(at::native::TransposeType trans) { + switch (trans) { + case at::native::TransposeType::NoTranspose: + return 'n'; + case at::native::TransposeType::Transpose: + return 't'; + case at::native::TransposeType::ConjTranspose: + return 'c'; + } + return 'n'; +} + +inline void blas_gemm(at::native::TransposeType transa, + at::native::TransposeType transb, int64_t m, int64_t n, + int64_t k, float alpha, const at::BFloat16* a, + int64_t lda, const at::BFloat16* b, int64_t ldb, + float beta, float* c, int64_t ldc) { + char transa_ = blas_transpose(transa); + char transb_ = blas_transpose(transb); + int m_ = static_cast(m); + int n_ = static_cast(n); + int k_ = static_cast(k); + int lda_ = static_cast(lda); + int ldb_ = static_cast(ldb); + int ldc_ = static_cast(ldc); + sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta, + c, &ldc_); +} + +inline void blas_gemm(at::native::TransposeType transa, + at::native::TransposeType transb, int64_t m, int64_t n, + int64_t k, float alpha, const float* a, int64_t lda, + const float* b, int64_t ldb, float beta, float* c, + int64_t ldc) { + char transa_ = blas_transpose(transa); + char transb_ = blas_transpose(transb); + int m_ = static_cast(m); + int n_ = static_cast(n); + int k_ = static_cast(k); + int lda_ = static_cast(lda); + int ldb_ = static_cast(ldb); + int ldc_ = static_cast(ldc); + sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta, + c, &ldc_); +} + +inline void blas_gemm(at::native::TransposeType, at::native::TransposeType, + int64_t, int64_t, int64_t, float, const at::Half*, + int64_t, const at::Half*, int64_t, float, float*, + int64_t) { + TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available."); +} +#else +template +inline void blas_gemm(at::native::TransposeType transa, + at::native::TransposeType transb, int64_t m, int64_t n, + int64_t k, float alpha, const scalar_t* a, int64_t lda, + const scalar_t* b, int64_t ldb, float beta, float* c, + int64_t ldc) { + auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT; + gemm(c10::CppTypeToScalarType::value, transa, transb, m, n, k, + at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc); +} +#endif \ No newline at end of file diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index e939e1c5256..bf1b6444bdd 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -301,25 +301,42 @@ void chunk_gated_delta_rule_kernel_impl( // attn = k_beta @ key.transpose(-1, -2) // attn: [B, HV, num_chunk, chunk_size, chunk_size] // transpose and pack for key - pack_vnni( - /* dst */ k_transpose, - /* src */ curr_k_pad, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // k_beta @ key.transpose(-1, -2) - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, - /* add_C */ false, - /* A */ curr_k_beta, - /* B */ k_transpose, - /* C */ curr_attn); + if constexpr (brgemm_supported()) { + pack_vnni( + /* dst */ k_transpose, + /* src */ curr_k_pad, + /* N */ chunk_size, + /* K */ qk_head_size, + /* ld_src */ qk_head_size, + /* ld_dst */ chunk_size); + // k_beta @ key.transpose(-1, -2) + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ chunk_size, + /* K */ qk_head_size, + /* lda */ qk_head_size, + /* ldb */ chunk_size, + /* ldc */ chunk_size, + /* add_C */ false, + /* A */ curr_k_beta, + /* B */ k_transpose, + /* C */ curr_attn); + } else { + blas_gemm( + at::native::TransposeType::Transpose, + at::native::TransposeType::NoTranspose, + chunk_size, + chunk_size, + qk_head_size, + 1.0f, + curr_k_pad, + qk_head_size, + curr_k_beta, + qk_head_size, + 0.0f, + curr_attn, + chunk_size); + } // attn = attn * decay_mask for (int64_t m = 0; m < chunk_size; m++) { at::vec::map2( @@ -413,25 +430,42 @@ void chunk_gated_delta_rule_kernel_impl( // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK] // k_cumdecay: [B, HV, num_chunk, chunk_size, EK] // pack for value - pack_vnni2( - /* dst */ v_pack, - /* src */ curr_v_beta, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // value = attn @ v_beta - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ v_pack, - /* C */ curr_value); + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ v_pack, + /* src */ curr_v_beta, + /* N */ chunk_size, + /* K */ v_head_size, + /* ld_src */ v_head_size, + /* ld_dst */ v_head_size); + // value = attn @ v_beta + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ v_head_size, + /* K */ chunk_size, + /* lda */ chunk_size, + /* ldb */ v_head_size, + /* ldc */ v_head_size, + /* add_C */ false, + /* A */ curr_attn_reduced, + /* B */ v_pack, + /* C */ curr_value); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + v_head_size, + chunk_size, + chunk_size, + 1.0f, + curr_v_beta, + v_head_size, + curr_attn_reduced, + chunk_size, + 0.0f, + curr_value, + v_head_size); + } // k_beta_g = k_beta * g.exp().unsqueeze(-1) for (int64_t j = 0; j < chunk_size; j++) { int64_t i = 0; @@ -445,25 +479,42 @@ void chunk_gated_delta_rule_kernel_impl( } } // pack for k_beta_g - pack_vnni2( - /* dst */ k_beta_g_pack, - /* src */ k_beta_g, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ qk_head_size); - // k_cumdecay = attn @ k_beta_g - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ qk_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ qk_head_size, - /* ldc */ qk_head_size, - /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ k_beta_g_pack, - /* C */ k_cumdecay); + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ k_beta_g_pack, + /* src */ k_beta_g, + /* N */ chunk_size, + /* K */ qk_head_size, + /* ld_src */ qk_head_size, + /* ld_dst */ qk_head_size); + // k_cumdecay = attn @ k_beta_g + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ qk_head_size, + /* K */ chunk_size, + /* lda */ chunk_size, + /* ldb */ qk_head_size, + /* ldc */ qk_head_size, + /* add_C */ false, + /* A */ curr_attn_reduced, + /* B */ k_beta_g_pack, + /* C */ k_cumdecay); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + qk_head_size, + chunk_size, + chunk_size, + 1.0f, + k_beta_g, + qk_head_size, + curr_attn_reduced, + chunk_size, + 0.0f, + k_cumdecay, + qk_head_size); + } for (int i = 0; i < chunk_size; i++) { at::vec::map( [](fVec x) { return x; }, @@ -551,25 +602,42 @@ void chunk_gated_delta_rule_kernel_impl( // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) // k_transpose_i = k_i.transpose(-1, -2) - pack_vnni( - /* dst */ k_transpose_i, - /* src */ k_i, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // attn_i = q_i @ k_transpose_i - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, - /* add_C */ false, - /* A */ q_i, - /* B */ k_transpose_i, - /* C */ attn_i); + if constexpr (brgemm_supported()) { + pack_vnni( + /* dst */ k_transpose_i, + /* src */ k_i, + /* N */ chunk_size, + /* K */ qk_head_size, + /* ld_src */ qk_head_size, + /* ld_dst */ chunk_size); + // attn_i = q_i @ k_transpose_i + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ chunk_size, + /* K */ qk_head_size, + /* lda */ qk_head_size, + /* ldb */ chunk_size, + /* ldc */ chunk_size, + /* add_C */ false, + /* A */ q_i, + /* B */ k_transpose_i, + /* C */ attn_i); + } else { + blas_gemm( + at::native::TransposeType::Transpose, + at::native::TransposeType::NoTranspose, + chunk_size, + chunk_size, + qk_head_size, + 1.0f, + k_i, + qk_head_size, + q_i, + qk_head_size, + 0.0f, + attn_i, + chunk_size); + } // attn_i = attn_i * decay_mask_i for (int64_t m = 0; m < chunk_size; m++) { auto attn_i_m = attn_i + m * chunk_size; @@ -609,28 +677,45 @@ void chunk_gated_delta_rule_kernel_impl( } // pack for curr_last_recurrent_state - pack_vnni2( - /* dst */ curr_last_recurrent_state_pack_reduced, - /* src */ curr_last_recurrent_state_reduced, - /* N */ qk_head_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ curr_last_recurrent_state_pack_reduced, + /* src */ curr_last_recurrent_state_reduced, + /* N */ qk_head_size, + /* K */ v_head_size, + /* ld_src */ v_head_size, + /* ld_dst */ v_head_size); - // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV] - // k_cumdecay_i: [chunk_size, EK] - // curr_last_recurrent_state: [EK, EV] - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ k_cumdecay_i_reduced, - /* B */ curr_last_recurrent_state_pack_reduced, - /* C */ v_prime); + // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV] + // k_cumdecay_i: [chunk_size, EK] + // curr_last_recurrent_state: [EK, EV] + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ v_head_size, + /* K */ qk_head_size, + /* lda */ qk_head_size, + /* ldb */ v_head_size, + /* ldc */ v_head_size, + /* add_C */ false, + /* A */ k_cumdecay_i_reduced, + /* B */ curr_last_recurrent_state_pack_reduced, + /* C */ v_prime); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + v_head_size, + chunk_size, + qk_head_size, + 1.0f, + curr_last_recurrent_state_reduced, + v_head_size, + k_cumdecay_i_reduced, + qk_head_size, + 0.0f, + v_prime, + v_head_size); + } // v_new = v_prime = v_i - v_prime // v_i: [chunk_size, EV] @@ -663,41 +748,75 @@ void chunk_gated_delta_rule_kernel_impl( } // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV] // curr_last_recurrent_state: [EK, EV] - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ qg, - /* B */ curr_last_recurrent_state_pack_reduced, - /* C */ attn_inter); + if constexpr (brgemm_supported()) { + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ v_head_size, + /* K */ qk_head_size, + /* lda */ qk_head_size, + /* ldb */ v_head_size, + /* ldc */ v_head_size, + /* add_C */ false, + /* A */ qg, + /* B */ curr_last_recurrent_state_pack_reduced, + /* C */ attn_inter); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + v_head_size, + chunk_size, + qk_head_size, + 1.0f, + curr_last_recurrent_state_reduced, + v_head_size, + qg, + qk_head_size, + 0.0f, + attn_inter, + v_head_size); + } // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new // pack for v_prime - pack_vnni2( - /* dst */ v_prime_pack_reduced, - /* src */ v_prime_reduced, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV] - // attn_i: [chunk_size, chunk_size] - // v_new: [chunk_size, EV] - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ true, - /* A */ attn_i_reduced, - /* B */ v_prime_pack_reduced, - /* C */ attn_inter); + if constexpr (brgemm_supported()) { + pack_vnni2( + /* dst */ v_prime_pack_reduced, + /* src */ v_prime_reduced, + /* N */ chunk_size, + /* K */ v_head_size, + /* ld_src */ v_head_size, + /* ld_dst */ v_head_size); + // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV] + // attn_i: [chunk_size, chunk_size] + // v_new: [chunk_size, EV] + at::native::cpublas::brgemm( + /* M */ chunk_size, + /* N */ v_head_size, + /* K */ chunk_size, + /* lda */ chunk_size, + /* ldb */ v_head_size, + /* ldc */ v_head_size, + /* add_C */ true, + /* A */ attn_i_reduced, + /* B */ v_prime_pack_reduced, + /* C */ attn_inter); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + v_head_size, + chunk_size, + chunk_size, + 1.0f, + v_prime_reduced, + v_head_size, + attn_i_reduced, + chunk_size, + 1.0f, + attn_inter, + v_head_size); + } // core_attn_out[:, :, i] = attn_inter for (int64_t m = 0; m < chunk_size; m++) { @@ -762,17 +881,34 @@ void chunk_gated_delta_rule_kernel_impl( /* ld_dst */ chunk_size); // kgv = kg.transpose(-1, -2) @ v_new // v_new: [chunk_size, EV] - at::native::cpublas::brgemm( - /* M */ qk_head_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ kg_transpose, - /* B */ v_prime_pack_reduced, - /* C */ kgv); + if constexpr (brgemm_supported()) { + at::native::cpublas::brgemm( + /* M */ qk_head_size, + /* N */ v_head_size, + /* K */ chunk_size, + /* lda */ chunk_size, + /* ldb */ v_head_size, + /* ldc */ v_head_size, + /* add_C */ false, + /* A */ kg_transpose, + /* B */ v_prime_pack_reduced, + /* C */ kgv); + } else { + blas_gemm( + at::native::TransposeType::NoTranspose, + at::native::TransposeType::NoTranspose, + v_head_size, + qk_head_size, + chunk_size, + 1.0f, + v_prime_reduced, + v_head_size, + kg_transpose, + chunk_size, + 0.0f, + kgv, + v_head_size); + } // last_recurrent_state = 1) + 2) for (int64_t m = 0; m < qk_head_size; m++) { at::vec::map2( @@ -921,7 +1057,8 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f; int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH; int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim; - float beta_val = 1 / (1 + std::exp(-b_ptr[ni])); + // See: https://github.com/sgl-project/sglang/pull/26634 + float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni])); fVec beta_vec = fVec(beta_val); int64_t dvi = 0; for (; dvi <= v_head_dim - VecSize; dvi += VecSize) { diff --git a/csrc/cpu/sgl-kernels/gemm.h b/csrc/cpu/sgl-kernels/gemm.h index f3fb37a5f61..494ccfafc4a 100644 --- a/csrc/cpu/sgl-kernels/gemm.h +++ b/csrc/cpu/sgl-kernels/gemm.h @@ -4,9 +4,12 @@ // clang-format off #pragma once -#include - #include "common.h" +#include "blas_gemm.h" + +#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__) +#define CPU_CAPABILITY_AVX512 +#endif // amx-bf16 #define TILE_M 16 @@ -21,31 +24,39 @@ constexpr int block_size_n() { return 2 * TILE_N; } +constexpr bool brgemm_supported() { +#if defined(CPU_CAPABILITY_AVX512) + return true; +#else + return false; +#endif +} + // define threshold using brgemm (intel AMX) template inline bool can_use_brgemm(int M); template <> inline bool can_use_brgemm(int M) { - return M > 4; + return brgemm_supported() && M > 4; } template <> inline bool can_use_brgemm(int M) { - return true; + return brgemm_supported(); } // this requires PyTorch 2.7 or above template <> inline bool can_use_brgemm(int M) { - return M > 4; + return brgemm_supported() && M > 4; } template <> inline bool can_use_brgemm(int M) { - return M > 4; + return brgemm_supported() && M > 4; } template <> inline bool can_use_brgemm(int M) { - return M > 4; + return brgemm_supported() && M > 4; } // work around compiler internal error diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 52b5ff7bedb..77ffeec9fe7 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -11,7 +11,9 @@ #include #include +#if defined(CPU_CAPABILITY_AVX512) #include +#endif namespace { using namespace at::vec; diff --git a/csrc/cpu/shm.cpp b/csrc/cpu/shm.cpp index a7fdd0c9d9d..f1538d27646 100644 --- a/csrc/cpu/shm.cpp +++ b/csrc/cpu/shm.cpp @@ -5,7 +5,7 @@ #include #include -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) #include #endif @@ -38,7 +38,7 @@ struct KernelVecType { }; struct ThreadSHMContext { -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) // memory model is weaker on AArch64, so we use atomic variables for // consumer (load-acquire) and producer (store-release) to make sure // that a stamp cannot be ready before the corresponding data is ready. @@ -75,7 +75,7 @@ struct ThreadSHMContext { TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM); TORCH_CHECK((size_t)this % 64 == 0); TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0); -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) _curr_thread_stamp[0].store(1, std::memory_order_relaxed); _curr_thread_stamp[1].store(1, std::memory_order_relaxed); _ready_thread_stamp[0].store(0, std::memory_order_relaxed); @@ -124,7 +124,7 @@ struct ThreadSHMContext { } char get_curr_stamp(int idx) const { -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) return _curr_thread_stamp[idx].load(std::memory_order_acquire); #else return _curr_thread_stamp[idx]; @@ -132,7 +132,7 @@ struct ThreadSHMContext { } char get_ready_stamp(int idx) const { -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) return _ready_thread_stamp[idx].load(std::memory_order_acquire); #else return _ready_thread_stamp[idx]; @@ -140,7 +140,7 @@ struct ThreadSHMContext { } void next_stamp() { -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) _curr_thread_stamp[local_stamp_buffer_idx].fetch_add( 1, std::memory_order_release); #else @@ -150,7 +150,7 @@ struct ThreadSHMContext { } void commit_ready_stamp() { -#ifdef __aarch64__ +#if defined(__aarch64__) || defined(__powerpc64__) _ready_thread_stamp[local_stamp_buffer_idx].store( _curr_thread_stamp[local_stamp_buffer_idx].load( std::memory_order_relaxed), @@ -186,8 +186,10 @@ struct ThreadSHMContext { break; } ++_spinning_count; -#ifdef __aarch64__ +#if defined(__aarch64__) __asm__ __volatile__("yield"); +#elif defined(__powerpc64__) + __asm__ __volatile__("or 1,1,1"); #else _mm_pause(); #endif // __aarch64__ diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 29f86ba0eb9..7a8188b8c8c 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -378,7 +378,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif // SHM CCL -#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) +#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \ + defined(__powerpc64__) ops.def( "init_shm_manager(str name, int group_size, int rank, int thread_num) -> " "int", @@ -447,6 +448,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "bool is_vnni) -> Tensor"); ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); + // Adapted from sglang: casual_conv1d kernels + ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor"); + ops.impl("causal_conv1d_weight_pack", torch::kCPU, + &causal_conv1d_weight_pack); + ops.def( + "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? " + "conv_states, Tensor? query_start_loc," + "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, " + "int pad_slot_id, bool is_vnni) -> " + "Tensor"); + ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); + ops.def( + "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " + "weight, Tensor? bias, bool silu_activation," + "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " + "bool is_vnni) -> Tensor"); + ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); +#endif + // Adapted from sglang: GDN kernels ops.def( "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, " @@ -470,25 +490,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "-> (Tensor, Tensor)"); ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); - // Adapted from sglang: casual_conv1d kernels - ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor"); - ops.impl("causal_conv1d_weight_pack", torch::kCPU, - &causal_conv1d_weight_pack); - ops.def( - "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? " - "conv_states, Tensor? query_start_loc," - "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, " - "int pad_slot_id, bool is_vnni) -> " - "Tensor"); - ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); - ops.def( - "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " - "weight, Tensor? bias, bool silu_activation," - "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " - "bool is_vnni) -> Tensor"); - ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); -#endif - // CPU attention kernels ops.def( "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " diff --git a/csrc/async_util.cuh b/csrc/libtorch_stable/async_util.cuh similarity index 100% rename from csrc/async_util.cuh rename to csrc/libtorch_stable/async_util.cuh diff --git a/csrc/custom_all_reduce.cu b/csrc/libtorch_stable/custom_all_reduce.cu similarity index 58% rename from csrc/custom_all_reduce.cu rename to csrc/libtorch_stable/custom_all_reduce.cu index a38d6fa24a2..0f7f759949a 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/libtorch_stable/custom_all_reduce.cu @@ -1,7 +1,11 @@ -#include -#include -#include -#include +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include #include "custom_all_reduce.cuh" @@ -11,7 +15,7 @@ using fptr_t = int64_t; static_assert(sizeof(void*) == sizeof(fptr_t)); fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::Tensor& rank_data, int64_t rank, + torch::stable::Tensor& rank_data, int64_t rank, bool fully_connected) { int world_size = fake_ipc_ptrs.size(); if (world_size > 8) @@ -25,9 +29,9 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, for (int i = 0; i < world_size; i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); } - return (fptr_t) new vllm::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(), - rank_data.numel(), rank, world_size, - fully_connected); + return (fptr_t) new vllm::CustomAllreduce( + ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank, + world_size, fully_connected); } /** @@ -46,10 +50,14 @@ fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, * 5. A[None].expand(2, -1, -1, -1): Not OK * 6. A[:, 1:, 1:]: Not OK */ -bool _is_weak_contiguous(torch::Tensor& t) { - return t.is_contiguous() || - (t.storage().nbytes() - t.storage_offset() * t.element_size() == - t.numel() * t.element_size()); +bool _is_weak_contiguous(torch::stable::Tensor& t) { + if (t.is_contiguous()) { + return true; + } + int64_t storage_nbytes = 0; + TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes)); + return storage_nbytes - t.storage_offset() * t.element_size() == + static_cast(t.numel() * t.element_size()); } /** @@ -59,42 +67,45 @@ bool _is_weak_contiguous(torch::Tensor& t) { * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first * copied into _reg_buffer. */ -void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) { +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t _reg_buffer, + int64_t reg_buffer_sz_bytes) { auto fa = reinterpret_cast(_fa); - const at::cuda::OptionalCUDAGuard device_guard(device_of(inp)); - auto stream = c10::cuda::getCurrentCUDAStream().stream(); + const torch::stable::accelerator::DeviceGuard device_guard( + inp.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); - TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type()); - TORCH_CHECK_EQ(inp.numel(), out.numel()); - TORCH_CHECK(_is_weak_contiguous(out)); - TORCH_CHECK(_is_weak_contiguous(inp)); + STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); + STD_TORCH_CHECK((inp.numel()) == (out.numel())); + STD_TORCH_CHECK(_is_weak_contiguous(out)); + STD_TORCH_CHECK(_is_weak_contiguous(inp)); auto input_size = inp.numel() * inp.element_size(); auto reg_buffer = reinterpret_cast(_reg_buffer); if (reg_buffer) { - TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes); - AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size, - cudaMemcpyDeviceToDevice, stream)); + STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); + STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, + cudaMemcpyDeviceToDevice, stream)); } else { - reg_buffer = inp.data_ptr(); + reg_buffer = inp.mutable_data_ptr(); } switch (out.scalar_type()) { - case at::ScalarType::Float: { + case torch::headeronly::ScalarType::Float: { fa->allreduce(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), out.numel()); break; } - case at::ScalarType::Half: { + case torch::headeronly::ScalarType::Half: { fa->allreduce(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.data_ptr()), out.numel()); + reinterpret_cast(out.mutable_data_ptr()), + out.numel()); break; } #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) - case at::ScalarType::BFloat16: { + case torch::headeronly::ScalarType::BFloat16: { fa->allreduce( stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.data_ptr()), out.numel()); + reinterpret_cast(out.mutable_data_ptr()), out.numel()); break; } #endif @@ -112,7 +123,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); } void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs) { auto fa = reinterpret_cast(_fa); - TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); + STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_); void* ipc_ptrs[8]; for (int i = 0; i < fake_ipc_ptrs.size(); i++) { ipc_ptrs[i] = reinterpret_cast(fake_ipc_ptrs[i]); @@ -143,47 +154,49 @@ void register_graph_buffers(fptr_t _fa, fa->register_graph_buffers(bytes, offsets); } -std::tuple allocate_shared_buffer_and_handle( +std::tuple allocate_shared_buffer_and_handle( int64_t size) { - auto device_index = c10::cuda::current_device(); - at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index)); + int device_index; + STD_CUDA_CHECK(cudaGetDevice(&device_index)); + const torch::stable::accelerator::DeviceGuard device_guard(device_index); void* buffer; cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; - auto stream = c10::cuda::getCurrentCUDAStream().stream(); - AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); + const cudaStream_t stream = get_current_cuda_stream(device_index); + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); // Allocate buffer #if defined(USE_ROCM) // data buffers need to be "uncached" for signal on MI200 - AT_CUDA_CHECK( + STD_CUDA_CHECK( hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached)); #else - AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size)); + STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size)); #endif - AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream)); - AT_CUDA_CHECK(cudaStreamSynchronize(stream)); - AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); + STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream)); + STD_CUDA_CHECK(cudaStreamSynchronize(stream)); + STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode)); // Create IPC memhandle for the allocated buffer. // Will use it in open_mem_handle. - auto options = - torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU); - auto handle = - torch::empty({static_cast(sizeof(cudaIpcMemHandle_t))}, options); - AT_CUDA_CHECK( - cudaIpcGetMemHandle((cudaIpcMemHandle_t*)handle.data_ptr(), buffer)); + auto handle = torch::stable::empty( + {static_cast(sizeof(cudaIpcMemHandle_t))}, + torch::headeronly::ScalarType::Byte, std::nullopt, + torch::stable::Device(torch::stable::DeviceType::CPU)); + STD_CUDA_CHECK(cudaIpcGetMemHandle( + (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer)); return std::make_tuple(reinterpret_cast(buffer), handle); } -fptr_t open_mem_handle(torch::Tensor& mem_handle) { +fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) { void* ipc_ptr; - AT_CUDA_CHECK(cudaIpcOpenMemHandle( - (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()), + STD_CUDA_CHECK(cudaIpcOpenMemHandle( + (void**)&ipc_ptr, + *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()), cudaIpcMemLazyEnablePeerAccess)); return reinterpret_cast(ipc_ptr); } void free_shared_buffer(fptr_t buffer) { - AT_CUDA_CHECK(cudaFree(reinterpret_cast(buffer))); + STD_CUDA_CHECK(cudaFree(reinterpret_cast(buffer))); } diff --git a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp b/csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp similarity index 100% rename from csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp rename to csrc/libtorch_stable/cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp diff --git a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp index f6737a73d48..6091cbc5e94 100644 --- a/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp +++ b/csrc/libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp @@ -2,7 +2,7 @@ #include -#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp" +#include "broadcast_load_epilogue_c2x.hpp" /* This file defines custom epilogues for fusing channel scales, token scales, diff --git a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu similarity index 89% rename from csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu rename to csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index e4d432cac97..a5f3f03de00 100644 --- a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -28,7 +28,20 @@ * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token */ +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + #include +#include "cuda_compat.h" +#include "dispatch_utils.h" +#include "type_convert.cuh" + #ifndef USE_ROCM #include #else @@ -37,14 +50,6 @@ #include #include -#include -#include -#include - -#include "cuda_compat.h" -#include "dispatch_utils.h" -#include "type_convert.cuh" - #ifndef FINAL_MASK #ifdef USE_ROCM #define FINAL_MASK 0xffffffffffffffffULL @@ -70,7 +75,7 @@ namespace deepseek_v4_fused_ops { namespace { inline int getSMVersion() { - auto* props = at::cuda::getCurrentDeviceProperties(); + auto* props = get_device_prop(); return props->major * 10 + props->minor; } } // namespace @@ -564,7 +569,7 @@ static void launchFusedDeepseekV4Templated( // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert is // unavailable there. Refuse the launch loudly instead of silently // skipping the work. - TORCH_CHECK( + STD_TORCH_CHECK( sm_version >= 80, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ " "(Ampere or newer); got sm_", @@ -635,7 +640,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( DISPATCH(64) DISPATCH(128) default: - TORCH_CHECK(false, + STD_TORCH_CHECK(false, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: " "unsupported num_heads_q_padded=", num_heads_q_padded, @@ -650,71 +655,80 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── -torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::Tensor const& q_in, // [N, num_heads_q, 512] bf16 - torch::Tensor const& kv, // [N, 512] bf16 (read-only) - torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8 - torch::Tensor const& slot_mapping, // [N] int64 - torch::Tensor const& position_ids, // [N] int64 - torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 - int64_t q_head_padded, // padded Q head count for output +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 + torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) + torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8 + torch::stable::Tensor const& slot_mapping, // [N] int64 + torch::stable::Tensor const& position_ids, // [N] int64 + torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 + int64_t q_head_padded, // padded Q head count for output double eps, int64_t cache_block_size) { - TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(), - "q_in must be contiguous CUDA"); - TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA"); - TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA"); - TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64, - "slot_mapping must be int64 CUDA"); - TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64, - "position_ids must be int64 CUDA"); - TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA"); - TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512, - "q_in shape [N, num_heads_q, 512]"); - TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); - TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match"); - TORCH_CHECK(q_head_padded >= q_in.size(1), - "q_head_padded must be >= q_in.size(1) (num_heads_q)"); - TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8"); - TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, - "cos_sin_cache shape [max_pos, 64]"); - TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32, - "cos_sin_cache must be float32"); + STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), + "q_in must be contiguous CUDA"); + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), + "kv must be contiguous CUDA"); + STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA"); + STD_TORCH_CHECK(slot_mapping.device().is_cuda() && + slot_mapping.scalar_type() == + torch::headeronly::ScalarType::Long, + "slot_mapping must be int64 CUDA"); + STD_TORCH_CHECK(position_ids.device().is_cuda() && + position_ids.scalar_type() == + torch::headeronly::ScalarType::Long, + "position_ids must be int64 CUDA"); + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda(), "cos_sin_cache must be CUDA"); + STD_TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512, + "q_in shape [N, num_heads_q, 512]"); + STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]"); + STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(), + "q_in and kv dtype must match"); + STD_TORCH_CHECK(q_head_padded >= q_in.size(1), + "q_head_padded must be >= q_in.size(1) (num_heads_q)"); + STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte, + "k_cache must be uint8"); + STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64, + "cos_sin_cache shape [max_pos, 64]"); + STD_TORCH_CHECK(cos_sin_cache.scalar_type() == + torch::headeronly::ScalarType::Float, + "cos_sin_cache must be float32"); // With DP padding, slot_mapping can be shorter than q/kv/positions. // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them); // KV quant+insert runs only on the first slot_mapping.size(0) rows. int const num_tokens_full = static_cast(q_in.size(0)); int const num_tokens_insert = static_cast(slot_mapping.size(0)); - TORCH_CHECK(static_cast(kv.size(0)) == num_tokens_full && - static_cast(position_ids.size(0)) == num_tokens_full, - "q/kv/position_ids row counts must match"); - TORCH_CHECK(num_tokens_insert <= num_tokens_full, - "slot_mapping must not exceed q row count"); + STD_TORCH_CHECK(static_cast(kv.size(0)) == num_tokens_full && + static_cast(position_ids.size(0)) == num_tokens_full, + "q/kv/position_ids row counts must match"); + STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full, + "slot_mapping must not exceed q row count"); int const num_heads_q = static_cast(q_in.size(1)); int const num_heads_q_padded = static_cast(q_head_padded); int const cache_block_size_i = static_cast(cache_block_size); int const kv_block_stride = static_cast(k_cache.stride(0)); - at::cuda::OptionalCUDAGuard device_guard(device_of(q_in)); - auto stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + q_in.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index()); // Allocate the padded q output. The kernel writes every element (live // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe. - torch::Tensor q_out = torch::empty( - {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options()); + auto q_out = torch::stable::new_empty( + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type()); - VLLM_DISPATCH_HALF_TYPES( + VLLM_STABLE_DISPATCH_HALF_TYPES( q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] { using qkv_scalar_t = scalar_t; vllm::deepseek_v4_fused_ops:: launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( - reinterpret_cast(q_in.data_ptr()), - reinterpret_cast(q_out.data_ptr()), - reinterpret_cast(kv.data_ptr()), - reinterpret_cast(k_cache.data_ptr()), - reinterpret_cast(slot_mapping.data_ptr()), - reinterpret_cast(position_ids.data_ptr()), - cos_sin_cache.data_ptr(), static_cast(eps), + reinterpret_cast(q_in.const_data_ptr()), + reinterpret_cast(q_out.mutable_data_ptr()), + reinterpret_cast(kv.const_data_ptr()), + reinterpret_cast(k_cache.mutable_data_ptr()), + slot_mapping.const_data_ptr(), + position_ids.const_data_ptr(), + cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, stream); diff --git a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu b/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu index bcf0ae58547..c9b7ee9e4e9 100644 --- a/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu +++ b/csrc/libtorch_stable/fused_qknorm_rope_kernel.cu @@ -20,7 +20,7 @@ #include "torch_utils.h" -#include "../async_util.cuh" +#include "async_util.cuh" #include "../cuda_compat.h" #include "../type_convert.cuh" #include "dispatch_utils.h" diff --git a/csrc/launch_bounds_utils.h b/csrc/libtorch_stable/launch_bounds_utils.h similarity index 100% rename from csrc/launch_bounds_utils.h rename to csrc/libtorch_stable/launch_bounds_utils.h diff --git a/csrc/minimax_reduce_rms_kernel.cu b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu similarity index 87% rename from csrc/minimax_reduce_rms_kernel.cu rename to csrc/libtorch_stable/minimax_reduce_rms_kernel.cu index 6245b02d6e9..d9af0f5efe0 100644 --- a/csrc/minimax_reduce_rms_kernel.cu +++ b/csrc/libtorch_stable/minimax_reduce_rms_kernel.cu @@ -15,16 +15,19 @@ * limitations under the License. */ +#include "torch_utils.h" + +#include +#include +#include +#include +#include +#include + #include #include -#include -#include -#include - #include "cuda_compat.h" -#include "cuda_utils.h" -#include "core/registration.h" #include "minimax_reduce_rms_kernel.h" #include @@ -611,7 +614,7 @@ int get_sm_count() { static int sm_count = 0; if (sm_count == 0) { int device_id; - CUDA_CHECK(cudaGetDevice(&device_id)); + STD_CUDA_CHECK(cudaGetDevice(&device_id)); cudaDeviceProp device_prop; cudaGetDeviceProperties(&device_prop, device_id); sm_count = device_prop.multiProcessorCount; @@ -621,13 +624,13 @@ int get_sm_count() { inline int getSMVersion(bool queryRealSmArch = false) { int device{-1}; - CUDA_CHECK(cudaGetDevice(&device)); + STD_CUDA_CHECK(cudaGetDevice(&device)); int sm_major = 0; int sm_minor = 0; - CUDA_CHECK(cudaDeviceGetAttribute(&sm_major, - cudaDevAttrComputeCapabilityMajor, device)); - CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor, - cudaDevAttrComputeCapabilityMinor, device)); + STD_CUDA_CHECK(cudaDeviceGetAttribute( + &sm_major, cudaDevAttrComputeCapabilityMajor, device)); + STD_CUDA_CHECK(cudaDeviceGetAttribute( + &sm_minor, cudaDevAttrComputeCapabilityMinor, device)); int sm = sm_major * 10 + sm_minor; if (sm == 121 && !queryRealSmArch) { return 120; @@ -639,7 +642,7 @@ template int get_max_active_blocks(KernelFunc kernel, int block_size, int dynamic_smem = 0) { int max_active = 0; - CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &max_active, kernel, block_size, dynamic_smem)); return std::max(max_active, 1); } @@ -678,27 +681,27 @@ void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) { cfg.attrs = attribute; cfg.numAttrs = SM >= 90 ? 2 : 0; - CUDA_CHECK(cudaLaunchKernelEx( + STD_CUDA_CHECK(cudaLaunchKernelEx( &cfg, minimax_reduce_rms_kernel_lamport, params)); } template void minimax_reduce_rms_kernel_launcher_float4( MiniMaxReduceRMSParams const& params) { - TORCH_CHECK(params.size_q % params.hidden_dim == 0); - TORCH_CHECK(params.hidden_dim % kElemsPerAccess == 0); + STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0); + STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess == 0); if (params.stride_q > 0) { - TORCH_CHECK(params.stride_q % kElemsPerAccess == 0); + STD_TORCH_CHECK(params.stride_q % kElemsPerAccess == 0); } - TORCH_CHECK(params.allreduce_in_k != nullptr, - "float4 QK kernel requires K input"); - TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k); - TORCH_CHECK(params.size_k % params.hidden_dim_k == 0); - TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess == 0); - TORCH_CHECK(params.size_q / params.hidden_dim == - params.size_k / params.hidden_dim_k); + STD_TORCH_CHECK(params.allreduce_in_k != nullptr, + "float4 QK kernel requires K input"); + STD_TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k); + STD_TORCH_CHECK(params.size_k % params.hidden_dim_k == 0); + STD_TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess == 0); + STD_TORCH_CHECK(params.size_q / params.hidden_dim == + params.size_k / params.hidden_dim_k); if (params.stride_k > 0) { - TORCH_CHECK(params.stride_k % kElemsPerAccess == 0); + STD_TORCH_CHECK(params.stride_k % kElemsPerAccess == 0); } int token_num = params.size_q / params.hidden_dim; @@ -746,7 +749,7 @@ void minimax_reduce_rms_kernel_launcher_float4( cfg.attrs = attribute; cfg.numAttrs = SM >= 90 ? 2 : 0; - CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params)); + STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params)); } template @@ -759,21 +762,21 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) { (params.hidden_dim * params.nranks == 6144) && (params.hidden_dim_k * params.nranks == 1024); - if (params.dtype == at::ScalarType::Half) { + if (params.dtype == torch::headeronly::ScalarType::Half) { if (use_float4) { minimax_reduce_rms_kernel_launcher_float4( params); } else { minimax_reduce_rms_kernel_launcher(params); } - } else if (params.dtype == at::ScalarType::BFloat16) { + } else if (params.dtype == torch::headeronly::ScalarType::BFloat16) { if (use_float4) { minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144, 1024>(params); } else { minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params); } - } else if (params.dtype == at::ScalarType::Float) { + } else if (params.dtype == torch::headeronly::ScalarType::Float) { if (use_float4) { minimax_reduce_rms_kernel_launcher_float4( params); @@ -781,7 +784,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) { minimax_reduce_rms_kernel_launcher(params); } } else { - TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op"); + STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op"); } } @@ -795,16 +798,18 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { } else if (params.nranks == 16) { dispatch_dtype<16>(params); } else { - TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!"); + STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!"); } } } // namespace tensorrt_llm } // namespace vllm -torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, - torch::Tensor const& norm_weight, - torch::Tensor workspace, int64_t const rank, - int64_t const nranks, double const eps) { +torch::stable::Tensor minimax_allreduce_rms( + torch::stable::Tensor const& input, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, + int64_t const rank, int64_t const nranks, double const eps) { + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); allreduce_params.nranks = static_cast(nranks); @@ -815,12 +820,12 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, allreduce_params.stride_q = allreduce_params.hidden_dim; allreduce_params.workspace = reinterpret_cast(workspace.mutable_data_ptr()); - allreduce_params.allreduce_in = input.data_ptr(); - allreduce_params.rms_gamma = norm_weight.data_ptr(); + allreduce_params.allreduce_in = const_cast(input.const_data_ptr()); + allreduce_params.rms_gamma = const_cast(norm_weight.const_data_ptr()); allreduce_params.rms_eps = static_cast(eps); - allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + allreduce_params.stream = get_current_cuda_stream(input.get_device_index()); - torch::Tensor rms_norm_out = torch::empty_like(input); + torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input); allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr(); vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params); @@ -828,26 +833,33 @@ torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, return rms_norm_out; } -std::tuple minimax_allreduce_rms_qk( - torch::Tensor qkv, torch::Tensor const& norm_weight_q, - torch::Tensor const& norm_weight_k, torch::Tensor workspace, - int64_t const q_size, int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps) { - TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D"); - TORCH_CHECK(qkv.is_contiguous(), - "minimax_allreduce_rms_qk: qkv must be contiguous"); +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps) { + STD_TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D"); + STD_TORCH_CHECK(qkv.is_contiguous(), + "minimax_allreduce_rms_qk: qkv must be contiguous"); int64_t qkv_dim = qkv.size(-1); - TORCH_CHECK(qkv_dim == q_size + 2 * kv_size, - "minimax_allreduce_rms_qk: qkv last dim must equal " - "q_size + 2 * kv_size"); - TORCH_CHECK(rank < nranks, - "minimax_allreduce_rms_qk: rank must be less than nranks"); + STD_TORCH_CHECK(qkv_dim == q_size + 2 * kv_size, + "minimax_allreduce_rms_qk: qkv last dim must equal " + "q_size + 2 * kv_size"); + STD_TORCH_CHECK(rank < nranks, + "minimax_allreduce_rms_qk: rank must be less than nranks"); + + const torch::stable::accelerator::DeviceGuard device_guard( + qkv.get_device_index()); int64_t num_tokens = qkv.size(0); int elem_bytes = qkv.element_size(); - torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options()); - torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options()); + torch::stable::Tensor q_out = + torch::stable::new_empty(qkv, {num_tokens, q_size}, qkv.scalar_type()); + torch::stable::Tensor k_out = + torch::stable::new_empty(qkv, {num_tokens, kv_size}, qkv.scalar_type()); auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); params.nranks = static_cast(nranks); @@ -863,13 +875,14 @@ std::tuple minimax_allreduce_rms_qk( params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k params.workspace = reinterpret_cast(workspace.mutable_data_ptr()); - uint8_t* base = static_cast(qkv.data_ptr()); + uint8_t* base = + const_cast(static_cast(qkv.const_data_ptr())); params.allreduce_in = base; params.allreduce_in_k = base + q_size * elem_bytes; - params.rms_gamma = norm_weight_q.data_ptr(); - params.rms_gamma_k = norm_weight_k.data_ptr(); + params.rms_gamma = const_cast(norm_weight_q.const_data_ptr()); + params.rms_gamma_k = const_cast(norm_weight_k.const_data_ptr()); params.rms_eps = static_cast(eps); - params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + params.stream = get_current_cuda_stream(qkv.get_device_index()); params.rms_norm_out = q_out.mutable_data_ptr(); params.rms_norm_out_k = k_out.mutable_data_ptr(); diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu new file mode 100644 index 00000000000..fda9bc020da --- /dev/null +++ b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Adapted from SGLang: +// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu + +#include +#include +#include "libtorch_stable/torch_utils.h" + +#include "cutlass_mxfp8_grouped_mm_launcher.cuh" + +void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a, + const torch::stable::Tensor& b, + const torch::stable::Tensor& sfa, + const torch::stable::Tensor& sfb, + torch::stable::Tensor& d, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& blockscale_offsets) { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); + STD_TORCH_CHECK(problem_sizes.size(1) == 3, + "problem_sizes must have shape (num_experts, 3)"); + STD_TORCH_CHECK( + problem_sizes.size(0) == expert_offsets.size(0), + "Number of experts in problem_sizes must match expert_offsets"); + STD_TORCH_CHECK( + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, + "problem_sizes must be int32"); + STD_TORCH_CHECK( + expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int, + "expert_offsets must be int32"); + STD_TORCH_CHECK( + blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int, + "blockscale_offsets must be int32"); + STD_TORCH_CHECK(a.dim() == 2, + "a must be a 2D tensor of shape (num_tokens, k)"); + STD_TORCH_CHECK(b.dim() == 3, + "b must be a 3D tensor of shape (num_experts, k, n)"); + STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0, + "k should align 128"); + STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128"); + STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major"); + STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major"); + + const torch::stable::accelerator::DeviceGuard device_guard( + a.get_device_index()); + auto stream = get_current_cuda_stream(a.get_device_index()); + if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< + cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, + blockscale_offsets, stream); + } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) { + expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< + cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, + blockscale_offsets, stream); + } else { + STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); + } +#else + STD_TORCH_CHECK(false, + "No implemented cutlass_mxfp8_grouped_mm for " + "current device"); +#endif +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm)); +} diff --git a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh similarity index 100% rename from csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh rename to csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh diff --git a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh similarity index 54% rename from csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh rename to csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh index 2c46e1fa725..82d6543b288 100644 --- a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh +++ b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh @@ -4,9 +4,9 @@ // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh #pragma once -#include -#include -#include + +#include +#include #include #include @@ -15,18 +15,22 @@ #include "cute/tensor.hpp" #include "cutlass_mxfp8_grouped_mm_functor.cuh" #include "cutlass_mxfp8_grouped_mm_traits.cuh" +#include "libtorch_stable/torch_utils.h" namespace expert_specialization { template void cutlass_mxfp8_grouped_mm_pre_compute( - torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs, - torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a, - torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa, - torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b, - const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d, - const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets, - const torch::Tensor& blockscale_offsets, cudaStream_t stream) { + torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs, + torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs, + torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a, + torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d, + torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb, + const torch::stable::Tensor& a, const torch::stable::Tensor& b, + const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb, + const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) { using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor; using ElementA = typename OffsetFunctor::ElementA; using ElementB = typename OffsetFunctor::ElementB; @@ -42,10 +46,10 @@ void cutlass_mxfp8_grouped_mm_pre_compute( using StrideB = typename StrideFunctor::StrideB; using StrideD = typename StrideFunctor::StrideD; - int num_experts = (int)expert_offsets.size(0); - TORCH_CHECK(num_experts <= 1024, - "Number of experts cannot exceed 1024, the maximum number of " - "threads per block."); + int num_experts = static_cast(expert_offsets.size(0)); + STD_TORCH_CHECK(num_experts <= 1024, + "Number of experts cannot exceed 1024, the maximum number of " + "threads per block."); OffsetFunctor offset_functor( reinterpret_cast(expert_offsets.data_ptr()), @@ -72,13 +76,18 @@ void cutlass_mxfp8_grouped_mm_pre_compute( } template -void cutlass_mxfp8_grouped_mm( - const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs, - const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs, - const torch::Tensor& d_ptrs, const torch::Tensor& stride_a, - const torch::Tensor& stride_b, const torch::Tensor& stride_d, - const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb, - const torch::Tensor& problem_sizes, cudaStream_t stream) { +void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs, + const torch::stable::Tensor& b_ptrs, + const torch::stable::Tensor& sfa_ptrs, + const torch::stable::Tensor& sfb_ptrs, + const torch::stable::Tensor& d_ptrs, + const torch::stable::Tensor& stride_a, + const torch::stable::Tensor& stride_b, + const torch::stable::Tensor& stride_d, + const torch::stable::Tensor& layout_sfa, + const torch::stable::Tensor& layout_sfb, + const torch::stable::Tensor& problem_sizes, + cudaStream_t stream) { using Gemm = typename GemmTraits::Gemm; using ElementA = typename Gemm::ElementA; using ElementB = typename Gemm::ElementB; @@ -93,13 +102,12 @@ void cutlass_mxfp8_grouped_mm( typename GemmTraits::ProblemShape::UnderlyingProblemShape; cutlass::KernelHardwareInfo hw_info; - hw_info.device_id = c10::cuda::current_device(); - hw_info.sm_count = - at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + hw_info.device_id = d_ptrs.get_device_index(); + hw_info.sm_count = get_device_prop()->multiProcessorCount; hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster; hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster; - int num_experts = (int)problem_sizes.size(0); + int num_experts = static_cast(problem_sizes.size(0)); UnderlyingProblemShape* underlying_problem_shape = reinterpret_cast(problem_sizes.data_ptr()); @@ -127,44 +135,55 @@ void cutlass_mxfp8_grouped_mm( Gemm gemm; auto can_implement_status = gemm.can_implement(arguments); - TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, - "Failed to implement GEMM"); + STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess, + "Failed to implement GEMM"); - torch::TensorOptions options_uint8 = - torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device()); size_t workspace_size = gemm.get_workspace_size(arguments); - torch::Tensor workspace = torch::empty(workspace_size, options_uint8); + torch::stable::Tensor workspace = torch::stable::empty( + {static_cast(workspace_size)}, + torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device()); auto status = gemm.initialize(arguments, workspace.data_ptr(), stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM"); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, + "Failed to initialize GEMM"); status = gemm.run(stream, nullptr, true); // Enable PDL - TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); + STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); } template void cutlass_mxfp8_grouped_mm_dispatch_out_dtype( - const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa, - const torch::Tensor& sfb, torch::Tensor& d, - const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets, - const torch::Tensor& blockscale_offsets, cudaStream_t stream) { - int num_experts = (int)problem_sizes.size(0); - torch::TensorOptions options_int64 = - torch::TensorOptions().dtype(torch::kInt64).device(a.device()); - torch::TensorOptions options_int32 = - torch::TensorOptions().dtype(torch::kInt32).device(a.device()); + const torch::stable::Tensor& a, const torch::stable::Tensor& b, + const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb, + torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) { + int num_experts = static_cast(problem_sizes.size(0)); + auto device = a.device(); - torch::Tensor a_ptrs = torch::empty(num_experts, options_int64); - torch::Tensor b_ptrs = torch::empty(num_experts, options_int64); - torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64); - torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64); - torch::Tensor d_ptrs = torch::empty(num_experts, options_int64); + torch::stable::Tensor a_ptrs = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor b_ptrs = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor sfa_ptrs = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor sfb_ptrs = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor d_ptrs = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); - torch::Tensor stride_a = torch::empty(num_experts, options_int64); - torch::Tensor stride_b = torch::empty(num_experts, options_int64); - torch::Tensor stride_d = torch::empty(num_experts, options_int64); - torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32); - torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32); + torch::stable::Tensor stride_a = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor stride_b = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor stride_d = torch::stable::empty( + num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device); + torch::stable::Tensor layout_sfa = + torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int, + std::nullopt, device); + torch::stable::Tensor layout_sfb = + torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int, + std::nullopt, device); using GemmTraits = CutlassMxfp8GroupedMmGemmTraits; cutlass_mxfp8_grouped_mm_pre_compute( @@ -176,4 +195,4 @@ void cutlass_mxfp8_grouped_mm_dispatch_out_dtype( layout_sfa, layout_sfb, problem_sizes, stream); } -} // namespace expert_specialization \ No newline at end of file +} // namespace expert_specialization diff --git a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh similarity index 100% rename from csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh rename to csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh diff --git a/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu new file mode 100644 index 00000000000..e075721c2a3 --- /dev/null +++ b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Adapted from SGLang: +// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu + +#include +#include +#include "libtorch_stable/torch_utils.h" + +#include "mxfp8_experts_quant.cuh" + +void mxfp8_experts_quant(const torch::stable::Tensor& input, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& blockscale_offsets, + torch::stable::Tensor& quant_output, + torch::stable::Tensor& scale_factor) { +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor"); + STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128"); + STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major"); + STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); + STD_TORCH_CHECK( + problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, + "problem_sizes must be int32"); + STD_TORCH_CHECK( + expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int, + "expert_offsets must be int32"); + STD_TORCH_CHECK( + blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int, + "blockscale_offsets must be int32"); + + auto groups = problem_sizes.size(0); + STD_TORCH_CHECK( + expert_offsets.dim() == 1 && expert_offsets.size(0) == groups, + "expert_offsets must be 1D and have size equal to the number of groups"); + STD_TORCH_CHECK( + blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups, + "blockscale_offsets must be 1D and have size equal to the number of " + "groups"); + + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>( + input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, + scale_factor); + } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) { + expert_specialization::launch_mxfp8_experts_quant<__half>( + input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, + scale_factor); + } else { + STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); + } +#else + STD_TORCH_CHECK(false, + "No implemented mxfp8_experts_quant for " + "current device"); +#endif +} + +// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM +// is applied only under COMPILE_LANGUAGE:CUDA. +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant)); +} diff --git a/csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh similarity index 95% rename from csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh rename to csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh index 9a85852080f..a57e00e76c3 100644 --- a/csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh +++ b/csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh @@ -4,16 +4,19 @@ // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh #pragma once -#include -#include #include #include #include -#include + +#include +#include +#include +#include #include #include "cute/tensor.hpp" +#include "libtorch_stable/torch_utils.h" namespace expert_specialization { @@ -356,12 +359,12 @@ __global__ void mxfp8_experts_quant_kernel( } template -void launch_mxfp8_experts_quant(const torch::Tensor& input, - const torch::Tensor& problem_sizes, - const torch::Tensor& expert_offsets, - const torch::Tensor& blockscale_offsets, - torch::Tensor& quant_output, - torch::Tensor& scale_factor) { +void launch_mxfp8_experts_quant(const torch::stable::Tensor& input, + const torch::stable::Tensor& problem_sizes, + const torch::stable::Tensor& expert_offsets, + const torch::stable::Tensor& blockscale_offsets, + torch::stable::Tensor& quant_output, + torch::stable::Tensor& scale_factor) { ThrLayout thr_layout{}; ValLayout val_layout{}; SfR2SThrLayout r2s_thr_layout{}; @@ -386,19 +389,18 @@ void launch_mxfp8_experts_quant(const torch::Tensor& input, CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4) int max_active_blocks_per_sm = -1; - AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( &max_active_blocks_per_sm, mxfp8_experts_quant_kernel, THREAD_BLOCK_SIZE, 0)); - dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount * - max_active_blocks_per_sm, + dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm, 1, 1); dim3 block(THREAD_BLOCK_SIZE, 1, 1); - int num_experts = (int)problem_sizes.size(0); - auto stream = at::cuda::getCurrentCUDAStream(); + int num_experts = static_cast(problem_sizes.size(0)); + auto stream = get_current_cuda_stream(input.get_device_index()); mxfp8_experts_quant_kernel <<>>( diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 2b64cab597a..dd27a6968d0 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -3,10 +3,6 @@ #include #include -#ifndef USE_ROCM -torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, - torch::stable::Tensor const& perm); - void per_token_group_quant_fp8(const torch::stable::Tensor& input, torch::stable::Tensor& output_q, torch::stable::Tensor& output_s, @@ -28,6 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input, int64_t group_size, double eps, double int8_min, double int8_max); +#ifndef USE_ROCM +torch::stable::Tensor permute_cols(torch::stable::Tensor const& A, + torch::stable::Tensor const& perm); + bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability); bool cutlass_group_gemm_supported(int64_t cuda_device_capability); @@ -231,6 +231,27 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, + torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, + double eps, int64_t cache_block_size); + +#ifndef USE_ROCM +torch::stable::Tensor minimax_allreduce_rms( + torch::stable::Tensor const& input, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace, + int64_t const rank, int64_t const nranks, double const eps); +std::tuple +minimax_allreduce_rms_qk(torch::stable::Tensor qkv, + torch::stable::Tensor const& norm_weight_q, + torch::stable::Tensor const& norm_weight_k, + torch::stable::Tensor workspace, int64_t const q_size, + int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif + // Sampler kernels (shared CUDA/ROCm) void apply_repetition_penalties_( torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, @@ -273,6 +294,26 @@ void selective_scan_fwd( const std::optional& cu_chunk_seqlen, const std::optional& last_chunk_indices); +using fptr_t = int64_t; +fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, + torch::stable::Tensor& rank_data, int64_t rank, + bool fully_connected); +void all_reduce(fptr_t _fa, torch::stable::Tensor& inp, + torch::stable::Tensor& out, fptr_t reg_buffer, + int64_t reg_buffer_sz_bytes); +void dispose(fptr_t _fa); +int64_t meta_size(); +void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); +std::tuple, std::vector> +get_graph_buffer_ipc_meta(fptr_t _fa); +void register_graph_buffers(fptr_t _fa, + const std::vector>& handles, + const std::vector>& offsets); +std::tuple allocate_shared_buffer_and_handle( + int64_t size); +int64_t open_mem_handle(torch::stable::Tensor& mem_handle); +void free_shared_buffer(int64_t buffer); + // Activation kernels (shared CUDA/ROCm) void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void silu_and_mul_clamp(torch::stable::Tensor& out, diff --git a/csrc/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh similarity index 100% rename from csrc/persistent_topk.cuh rename to csrc/libtorch_stable/persistent_topk.cuh diff --git a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index d0cafa0c6df..6238a27191b 100644 --- a/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -25,7 +25,7 @@ #include #include "cuda_utils.h" -#include "launch_bounds_utils.h" +#include "libtorch_stable/launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index ae4f4829d93..062f6018653 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -34,7 +34,7 @@ static_assert(CVT_FP4_ELTS_PER_THREAD == 16, "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)"); -#include "launch_bounds_utils.h" +#include "libtorch_stable/launch_bounds_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu index 8f2b4405cf6..92b139b7e40 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu @@ -26,7 +26,7 @@ #include "cuda_utils.h" #include "nvfp4_utils.cuh" -#include "launch_bounds_utils.h" +#include "libtorch_stable/launch_bounds_utils.h" namespace vllm { diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu index ccefe294022..f7c965dbc1b 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu @@ -26,7 +26,7 @@ #include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" -#include "launch_bounds_utils.h" +#include "libtorch_stable/launch_bounds_utils.h" // Define before including nvfp4_utils.cuh so the header // can use this macro during compilation. diff --git a/csrc/quantization/gguf/dequantize.cuh b/csrc/libtorch_stable/quantization/gguf/dequantize.cuh similarity index 100% rename from csrc/quantization/gguf/dequantize.cuh rename to csrc/libtorch_stable/quantization/gguf/dequantize.cuh diff --git a/csrc/quantization/gguf/ggml-common.h b/csrc/libtorch_stable/quantization/gguf/ggml-common.h similarity index 100% rename from csrc/quantization/gguf/ggml-common.h rename to csrc/libtorch_stable/quantization/gguf/ggml-common.h diff --git a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu b/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu index 0fdfcafab8c..2a56d7a18f4 100644 --- a/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu +++ b/csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu @@ -7,14 +7,11 @@ #include -// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not -// moved to libtorch_stable) to avoid unnecessary reformatting that would break -// git rename detection and pollute blame history. -#include "../../../quantization/gguf/ggml-common.h" -#include "../../../quantization/gguf/vecdotq.cuh" -#include "../../../quantization/gguf/dequantize.cuh" -#include "../../../quantization/gguf/mmvq.cuh" -#include "../../../quantization/gguf/mmq.cuh" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "dequantize.cuh" +#include "mmvq.cuh" +#include "mmq.cuh" #include "moe.cuh" #include "moe_vec.cuh" diff --git a/csrc/quantization/gguf/mmq.cuh b/csrc/libtorch_stable/quantization/gguf/mmq.cuh similarity index 100% rename from csrc/quantization/gguf/mmq.cuh rename to csrc/libtorch_stable/quantization/gguf/mmq.cuh diff --git a/csrc/quantization/gguf/mmvq.cuh b/csrc/libtorch_stable/quantization/gguf/mmvq.cuh similarity index 100% rename from csrc/quantization/gguf/mmvq.cuh rename to csrc/libtorch_stable/quantization/gguf/mmvq.cuh diff --git a/csrc/quantization/gguf/vecdotq.cuh b/csrc/libtorch_stable/quantization/gguf/vecdotq.cuh similarity index 100% rename from csrc/quantization/gguf/vecdotq.cuh rename to csrc/libtorch_stable/quantization/gguf/vecdotq.cuh diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index d388b475b66..dea11c0e717 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -7,7 +7,11 @@ #include -#include +#ifdef USE_ROCM + #include +#else + #include +#endif #include "libtorch_stable/quantization/vectorization.cuh" #include "libtorch_stable/quantization/vectorization_utils.cuh" @@ -15,12 +19,23 @@ #include "libtorch_stable/torch_utils.h" __device__ __forceinline__ float GroupReduceMax(float val) { +#ifdef USE_ROCM + // 16-thread logical groups may pack up to four per 64-lane wavefront; use a + // 64-bit mask and explicit width so shuffles stay within each group. + const int lane_in_wave = threadIdx.x % warpSize; + const unsigned long long mask = 0xFFFFull << ((lane_in_wave / 16) * 16); + val = fmaxf(val, __shfl_xor_sync(mask, val, 8, 16)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 16)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 16)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 16)); +#else unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff; val = fmaxf(val, __shfl_xor_sync(mask, val, 8)); val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); +#endif return val; } @@ -103,10 +118,18 @@ __device__ __forceinline__ float LoadRegisterGroupAndComputeAbsmax( } __device__ __forceinline__ float GroupReduceMax8(float val) { +#ifdef USE_ROCM + const int lane_in_wave = threadIdx.x % warpSize; + const unsigned long long mask = 0xFFull << (lane_in_wave & ~7); + val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 8)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 8)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 8)); +#else unsigned mask = 0xffu << (threadIdx.x & 24u); val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); +#endif return val; } @@ -684,15 +707,12 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, VLLM_STABLE_DISPATCH_HALF_TYPES( input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] { - if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) { - LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3); - } else if (dst_type == torch::headeronly::ScalarType::Char) { + if (dst_type == torch::headeronly::ScalarType::Char) { LAUNCH_REG_KERNEL(scalar_t, int8_t); } else { - STD_TORCH_CHECK( - false, - "per_token_group_quant_8bit_packed only supports FP8/INT8 " - "outputs."); + VLLM_STABLE_DISPATCH_FP8_TYPES( + dst_type, "per_token_group_quant_8bit_packed_fp8", + ([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); })); } })); diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 15af18118f3..7656ba8cf8f 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -7,7 +7,7 @@ #include "torch_utils.h" #ifndef USE_ROCM - #include "../persistent_topk.cuh" + #include "persistent_topk.cuh" #endif namespace { diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 13d75445009..e9a62a8666c 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -7,11 +7,6 @@ // Note: We register under namespace "_C" so ops are accessible as // torch.ops._C. for compatibility with existing code. STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { -#ifndef USE_ROCM - ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); -#endif - -#ifndef USE_ROCM // Compute per-token-group FP8 quantized tensor and scaling factor. // The dummy arguments are here so we can correctly fuse with RMSNorm. ops.def( @@ -32,6 +27,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "output_s, int group_size, float eps, float int8_min, float int8_max) -> " "()"); +#ifndef USE_ROCM + ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor"); +#endif + +#ifndef USE_ROCM // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column // quantization, as well as bias ops.def( @@ -337,6 +337,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + ops.def( + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" + "Tensor q_in, Tensor kv, Tensor! k_cache, " + "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " + "int q_head_padded, float eps, int cache_block_size) -> Tensor"); + +#ifndef USE_ROCM + ops.def( + "minimax_allreduce_rms(" + "Tensor input, Tensor norm_weight, Tensor workspace, " + "int rank, int nranks, float eps) -> Tensor"); + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, " + "Tensor workspace, int q_size, int kv_size, int rank, int nranks, " + "float eps) -> (Tensor, Tensor)"); +#endif + // Apply repetition penalties to logits in-place. ops.def( "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " @@ -508,11 +526,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { -#ifndef USE_ROCM - ops.impl("permute_cols", TORCH_BOX(&permute_cols)); -#endif - -#ifndef USE_ROCM // Per-token group quantization ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8)); ops.impl("per_token_group_fp8_quant_packed", @@ -520,6 +533,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("per_token_group_quant_int8", TORCH_BOX(&per_token_group_quant_int8)); +#ifndef USE_ROCM + ops.impl("permute_cols", TORCH_BOX(&permute_cols)); +#endif + +#ifndef USE_ROCM // CUTLASS scaled_mm ops ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm)); ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp)); @@ -571,6 +589,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Positional encoding kernels (shared CUDA/ROCm) ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); +#ifndef USE_ROCM + ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms)); + ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk)); +#endif // Sampler kernels (shared CUDA/ROCm) ops.impl("apply_repetition_penalties_", @@ -725,6 +749,45 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); } +STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) { + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.def("dispose(int fa) -> ()"); + custom_ar.def("meta_size() -> int"); + custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()"); + custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])"); + custom_ar.def( + "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()"); + custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)"); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int"); + custom_ar.def("free_shared_buffer(int ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) { + custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar)); + custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) { + custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) { + custom_ar.impl("dispose", TORCH_BOX(&dispose)); + custom_ar.impl("meta_size", TORCH_BOX(&meta_size)); + custom_ar.impl("register_buffer", TORCH_BOX(®ister_buffer)); + custom_ar.impl("get_graph_buffer_ipc_meta", + TORCH_BOX(&get_graph_buffer_ipc_meta)); + custom_ar.impl("register_graph_buffers", TORCH_BOX(®ister_graph_buffers)); + custom_ar.impl("allocate_shared_buffer_and_handle", + TORCH_BOX(&allocate_shared_buffer_and_handle)); + custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer)); +} + STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); } diff --git a/csrc/libtorch_stable/torch_utils.h b/csrc/libtorch_stable/torch_utils.h index 1adbb4d4986..f02346739e0 100644 --- a/csrc/libtorch_stable/torch_utils.h +++ b/csrc/libtorch_stable/torch_utils.h @@ -6,11 +6,7 @@ #include #include -#ifndef USE_ROCM - #include -#else - #include -#endif +#include #include #include diff --git a/csrc/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h index e8c2d012247..c3d2dd5c599 100644 --- a/csrc/minimax_reduce_rms_kernel.h +++ b/csrc/minimax_reduce_rms_kernel.h @@ -19,7 +19,7 @@ #include #include -#include +#include namespace vllm { namespace tensorrt_llm { @@ -51,7 +51,7 @@ static constexpr int kElemsPerAccess = ElemsPerAccess::value; struct MiniMaxReduceRMSParams { int nranks{}; int rank{}; - at::ScalarType dtype{at::ScalarType::Undefined}; + torch::headeronly::ScalarType dtype{torch::headeronly::ScalarType::Undefined}; int size_q{}; int hidden_dim{}; int size_k{}; diff --git a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu b/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu deleted file mode 100644 index f507f9299b0..00000000000 --- a/csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu +++ /dev/null @@ -1,60 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu - -#include - -#include "cutlass_mxfp8_grouped_mm_launcher.cuh" - -void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b, - const torch::Tensor& sfa, - const torch::Tensor& sfb, torch::Tensor& d, - const torch::Tensor& problem_sizes, - const torch::Tensor& expert_offsets, - const torch::Tensor& blockscale_offsets) { -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); - TORCH_CHECK(problem_sizes.size(1) == 3, - "problem_sizes must have shape (num_experts, 3)"); - TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0), - "Number of experts in problem_sizes must match expert_offsets"); - TORCH_CHECK(problem_sizes.dtype() == torch::kInt32, - "problem_sizes must be int32"); - TORCH_CHECK(expert_offsets.dtype() == torch::kInt32, - "expert_offsets must be int32"); - TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32, - "blockscale_offsets must be int32"); - TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)"); - TORCH_CHECK(b.dim() == 3, - "b must be a 3D tensor of shape (num_experts, k, n)"); - TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0, - "k should align 128"); - TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128"); - TORCH_CHECK(a.strides()[1] == 1, "a must be row major"); - TORCH_CHECK(b.strides()[1] == 1, "b must be column major"); - - auto stream = at::cuda::getCurrentCUDAStream(); - if (d.dtype() == torch::kBFloat16) { - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< - cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, - blockscale_offsets, stream); - } else if (d.dtype() == torch::kFloat16) { - expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype< - cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets, - blockscale_offsets, stream); - } else { - TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); - } -#else - TORCH_CHECK(false, - "No implemented cutlass_mxfp8_grouped_mm for " - "current device"); -#endif -} - -#include "core/registration.h" - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm); -} \ No newline at end of file diff --git a/csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu b/csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu deleted file mode 100644 index 2a93ab94d5c..00000000000 --- a/csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu +++ /dev/null @@ -1,60 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// Adapted from SGLang: -// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu - -#include - -#include "mxfp8_experts_quant.cuh" - -void mxfp8_experts_quant(const torch::Tensor& input, - const torch::Tensor& problem_sizes, - const torch::Tensor& expert_offsets, - const torch::Tensor& blockscale_offsets, - torch::Tensor& quant_output, - torch::Tensor& scale_factor) { -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - TORCH_CHECK(input.dim() == 2, "input must be 2D tensor"); - TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128"); - TORCH_CHECK(input.strides()[1] == 1, "input must be row major"); - TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor"); - TORCH_CHECK(problem_sizes.dtype() == torch::kInt32, - "problem_sizes must be int32"); - TORCH_CHECK(expert_offsets.dtype() == torch::kInt32, - "expert_offsets must be int32"); - TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32, - "blockscale_offsets must be int32"); - - auto groups = problem_sizes.size(0); - TORCH_CHECK( - expert_offsets.dim() == 1 && expert_offsets.size(0) == groups, - "expert_offsets must be 1D and have size equal to the number of groups"); - TORCH_CHECK( - blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups, - "blockscale_offsets must be 1D and have size equal to the number of " - "groups"); - - auto stream = at::cuda::getCurrentCUDAStream(); - if (input.dtype() == torch::kBFloat16) { - expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>( - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, - scale_factor); - } else if (input.dtype() == torch::kFloat16) { - expert_specialization::launch_mxfp8_experts_quant<__half>( - input, problem_sizes, expert_offsets, blockscale_offsets, quant_output, - scale_factor); - } else { - TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16"); - } -#else - TORCH_CHECK(false, - "No implemented mxfp8_experts_quant for " - "current device"); -#endif -} - -#include "core/registration.h" - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("mxfp8_experts_quant", mxfp8_experts_quant); -} \ No newline at end of file diff --git a/csrc/ops.h b/csrc/ops.h index f458f79d6f4..ed2fca26b0d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -40,12 +40,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight, void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, torch::Tensor& weight, double epsilon); -torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache, - torch::Tensor const& slot_mapping, torch::Tensor const& position_ids, - torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps, - int64_t cache_block_size); - void silu_and_mul_per_block_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scales, int64_t group_size, @@ -107,24 +101,6 @@ torch::Tensor dynamic_4bit_int_moe_cpu( int64_t activation_kind); using fptr_t = int64_t; -fptr_t init_custom_ar(const std::vector& fake_ipc_ptrs, - torch::Tensor& rank_data, int64_t rank, - bool fully_connected); -void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, - fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); -void dispose(fptr_t _fa); -int64_t meta_size(); -void register_buffer(fptr_t _fa, const std::vector& fake_ipc_ptrs); -std::tuple, std::vector> -get_graph_buffer_ipc_meta(fptr_t _fa); -void register_graph_buffers(fptr_t _fa, - const std::vector>& handles, - const std::vector>& offsets); -std::tuple allocate_shared_buffer_and_handle( - int64_t size); -int64_t open_mem_handle(torch::Tensor& mem_handle); -void free_shared_buffer(int64_t buffer); - #ifdef USE_ROCM fptr_t init_custom_qr(int64_t rank, int64_t world_size, std::optional qr_max_size = std::nullopt); @@ -135,15 +111,3 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, int64_t quant_level, bool cast_bf2half = false); int64_t qr_max_size(); #endif - -#ifndef USE_ROCM -torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, - torch::Tensor const& norm_weight, - torch::Tensor workspace, int64_t const rank, - int64_t const nranks, double const eps); -std::tuple minimax_allreduce_rms_qk( - torch::Tensor qkv, torch::Tensor const& norm_weight_q, - torch::Tensor const& norm_weight_k, torch::Tensor workspace, - int64_t const q_size, int64_t const kv_size, int64_t const rank, - int64_t const nranks, double const eps); -#endif diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 01869474e0f..c078222bca0 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -55,14 +55,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one - // kernel launch. - ops.def( - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" - "Tensor q_in, Tensor kv, Tensor! k_cache, " - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " - "int q_head_padded, float eps, int cache_block_size) -> Tensor"); - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA, - &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert); + // kernel launch. Registered in _C_stable_libtorch. // Quantization ops #ifndef USE_ROCM @@ -163,34 +156,27 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // conditionally compiled so impl registration is in source file #endif - -#ifndef USE_ROCM - ops.def( - "minimax_allreduce_rms(" - "Tensor input," - "Tensor norm_weight," - "Tensor workspace," - "int rank," - "int nranks," - "float eps) -> Tensor"); - ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms); - ops.def( - "minimax_allreduce_rms_qk(" - "Tensor qkv," - "Tensor norm_weight_q," - "Tensor norm_weight_k," - "Tensor workspace," - "int q_size," - "int kv_size," - "int rank," - "int nranks," - "float eps) -> (Tensor, Tensor)"); - ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk); - - // conditionally compiled so impl in source file -#endif } +#ifdef USE_ROCM +TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { + // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). + custom_ar.def( + "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " + "cast_bf2half) -> ()"); + custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); + + custom_ar.def("init_custom_qr", &init_custom_qr); + custom_ar.def("qr_destroy", &qr_destroy); + custom_ar.def("qr_get_handle", &qr_get_handle); + + custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); + custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); + + custom_ar.def("qr_max_size", &qr_max_size); +} +#endif + TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { // Cuda utils @@ -205,48 +191,4 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { &get_max_shared_memory_per_block_device_attribute); } -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { - // Custom all-reduce kernels - custom_ar.def( - "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " - "int rank, bool fully_connected) -> int"); - custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar); - custom_ar.def( - "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " - "int reg_buffer_sz_bytes) -> ()"); - custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce); - - custom_ar.def("dispose", &dispose); - custom_ar.def("meta_size", &meta_size); - - custom_ar.def("register_buffer", ®ister_buffer); - custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta); - custom_ar.def("register_graph_buffers", ®ister_graph_buffers); - - custom_ar.def("allocate_shared_buffer_and_handle", - &allocate_shared_buffer_and_handle); - custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle); - custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle); - - custom_ar.def("free_shared_buffer", &free_shared_buffer); -#ifdef USE_ROCM - // Quick Reduce all-reduce kernels - custom_ar.def( - "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " - "cast_bf2half) -> ()"); - custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); - - custom_ar.def("init_custom_qr", &init_custom_qr); - custom_ar.def("qr_destroy", &qr_destroy); - - custom_ar.def("qr_get_handle", &qr_get_handle); - - custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); - custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); - - // Max input size in bytes - custom_ar.def("qr_max_size", &qr_max_size); -#endif -} - REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/type_convert.cuh b/csrc/type_convert.cuh index 9d939bb828f..8093c4bc871 100644 --- a/csrc/type_convert.cuh +++ b/csrc/type_convert.cuh @@ -50,7 +50,7 @@ struct _typeConvert { #if defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) // CUDA < 12.0 runs into issues with packed type conversion template <> -struct _typeConvert { +struct _typeConvert { static constexpr bool exists = true; using hip_type = __half; using packed_hip_type = __half2; @@ -73,7 +73,7 @@ struct _typeConvert { // CUDA_ARCH < 800 does not have BF16 support // ROCm 7.0+ supports bfloat16 template <> -struct _typeConvert { +struct _typeConvert { static constexpr bool exists = true; using hip_type = __nv_bfloat16; using packed_hip_type = __nv_bfloat162; diff --git a/docker/Dockerfile b/docker/Dockerfile index 06cdc0b667f..9b4227cdf65 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -757,10 +757,10 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.11.post2 +ARG FLASHINFER_VERSION=0.6.12 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ - --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') + --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') # ============================================================ # OPENAI API SERVER DEPENDENCIES diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 0d5a9cc5f83..4fbfe832ac3 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2. # build flashinfer for torch nightly from source around 10 mins -# release version: v0.6.11.post2 +# release version: v0.6.12 # todo(elainewy): cache flashinfer build result for faster build ENV CCACHE_DIR=/root/.cache/ccache RUN --mount=type=cache,target=/root/.cache/ccache \ --mount=type=cache,target=/root/.cache/uv \ echo "git clone flashinfer..." \ - && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \ + && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \ && cd flashinfer \ && git submodule update --init --recursive \ && echo "finish git clone flashinfer..." \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 61d73cd1527..1e39306e39f 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -2,6 +2,7 @@ ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. @@ -16,7 +17,8 @@ ARG NIC_BACKEND=all ARG AINIC_VERSION=1.117.3-hydra ARG UBUNTU_CODENAME=jammy -# Sccache configuration (only used in release pipeline) +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. ARG USE_SCCACHE ARG SCCACHE_DOWNLOAD_URL ARG SCCACHE_ENDPOINT @@ -29,12 +31,16 @@ FROM ${BASE_IMAGE} AS base ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} -# Install some basic utilities +# Install build dependencies and utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ apt-transport-https ca-certificates wget curl \ - libnuma-dev -RUN python3 -m pip install --upgrade pip + libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE RUN if [ "$USE_SCCACHE" != "1" ]; then \ @@ -55,6 +61,12 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} # Install sccache if USE_SCCACHE is enabled (for release builds) ARG USE_SCCACHE @@ -86,6 +98,7 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \ ARG USE_SCCACHE ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_ENDPOINT=${USE_SCCACHE:+${SCCACHE_ENDPOINT}} ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} @@ -114,8 +127,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain or protoc. Runs in parallel with the main wheel -# build for faster end-to-end builds. +# don't need the rust toolchain or protoc. FROM fetch_vllm AS rust-build ARG COMMON_WORKDIR @@ -144,24 +156,74 @@ ENV RUSTUP_MAX_RETRIES=10 # layer for later COPY --from=rust-build. RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \ && test -x /tmp/vllm-rs # ----------------------- -# vLLM build stages +# vLLM native build stages +# +# csrc-build intentionally copies only files that affect ROCm native extension +# compilation. That keeps unrelated CI/test/docs edits from invalidating the +# expensive HIP/C++ build layer. +FROM base AS csrc-build +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY requirements/rocm.txt requirements/rocm.txt +COPY requirements/common.txt requirements/common.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/rocm.txt + +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do +# not invalidate the expensive native build layer. +COPY setup.py CMakeLists.txt ./ +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/envs.py +COPY vllm/__init__.py vllm/__init__.py + +ENV VLLM_TARGET_DEVICE=rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" + +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + export CCACHE_BASEDIR="$PWD" \ + && echo "=== ccache stats before ROCm native build ===" \ + && (ccache --show-stats || true) \ + && (ccache --zero-stats || true) \ + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null \ + && echo "=== ccache stats after ROCm native build ===" \ + && (ccache --show-stats || true) + +# Build the full vLLM ROCm wheel by reusing the native extension wheel from +# csrc-build. This stage still rebuilds for Python/package changes, but skips +# the expensive HIP/C++ compile when native inputs are unchanged. FROM fetch_vllm AS build_vllm ARG COMMON_WORKDIR +ENV VLLM_TARGET_DEVICE=rocm + +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels # Drop the pre-built rust frontend binary into the source tree. setup.py # detects it and ships it as-is, skipping the local cargo build. COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -# Build vLLM (setup.py auto-detects sccache in PATH) -RUN cd vllm \ - && python3 -m pip install -r requirements/rocm.txt \ - && python3 setup.py clean --all \ - && python3 setup.py bdist_wheel --dist-dir=dist +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system -r requirements/rocm.txt \ + && export VLLM_USE_PRECOMPILED=1 \ + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ + && export VLLM_DOCKER_BUILD_CONTEXT=1 \ + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null FROM scratch AS export_vllm ARG COMMON_WORKDIR COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / @@ -171,6 +233,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ 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/vllm/v1 /vllm_v1 # RIXL/UCX build stages @@ -201,14 +264,17 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ ibverbs-providers \ && rm -rf /var/lib/apt/lists/* -RUN uv pip install --system meson auditwheel patchelf tomlkit +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit -RUN cd /usr/local/src && \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ git clone ${UCX_REPO} && \ cd ucx && \ git checkout ${UCX_BRANCH} && \ ./autogen.sh && \ mkdir build && cd build && \ + CC="ccache gcc" CXX="ccache g++" \ ../configure \ --prefix=/usr/local/ucx \ --enable-shared \ @@ -220,20 +286,22 @@ RUN cd /usr/local/src && \ --with-verbs \ --with-dm \ --enable-mt && \ - make -j && \ + make -j$(nproc) && \ make install ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} -RUN git clone ${RIXL_REPO} /opt/rixl && \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ cd /opt/rixl && \ git checkout ${RIXL_BRANCH} && \ + CC="ccache gcc" CXX="ccache g++" \ meson setup build --prefix=${RIXL_HOME} \ -Ducx_path=${UCX_HOME} \ -Drocm_path=${ROCM_PATH} && \ cd build && \ - ninja && \ + ninja -j$(nproc) && \ ninja install # Generate RIXL wheel @@ -250,30 +318,44 @@ RUN cd /opt/rixl && \ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins -# DeepEP build stage -FROM base AS build_deep +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem ARG ROCSHMEM_BRANCH="f0acb0c6" ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -ARG DEEPEP_BRANCH="a9ea9774" -ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" -ARG DEEPEP_NIC="cx7" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +ENV ROCM_PATH=/opt/rocm ENV ROCSHMEM_DIR=/opt/rocshmem -RUN git clone ${ROCSHMEM_REPO} \ +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ && git checkout ${ROCSHMEM_BRANCH} \ && mkdir -p projects/rocshmem/build \ && cd projects/rocshmem/build \ - && INSTALL_PREFIX=${ROCSHMEM_DIR} \ - ../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF -# Build DeepEP wheel. -# DeepEP looks for rocshmem at ROCSHMEM_DIR. -RUN git clone ${DEEPEP_REPO} \ +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep +ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" +ARG DEEPEP_NIC="cx7" + +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +RUN --mount=type=cache,target=/root/.cache/ccache \ + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ + && git clone ${DEEPEP_REPO} \ && cd DeepEP \ && git checkout ${DEEPEP_BRANCH} \ - && python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install # MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do # not force users to rebuild the long-lived Dockerfile.rocm_base image. @@ -372,8 +454,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) # This ensures setuptools_scm sees clean repo state for version detection RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ cd vllm \ - && pip install setuptools_scm regex \ + && uv pip install --system setuptools_scm regex \ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt @@ -409,18 +492,20 @@ RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt # Install dependencies using custom wheels from /install -RUN cd vllm \ +RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ && echo "Building vLLM with custom wheels from /install" \ - && python3 -m pip install --find-links /install -r requirements/rocm.txt \ - && python3 setup.py clean --all + && uv pip install --system --find-links /install -r requirements/rocm.txt # Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt -# (setup.py auto-detects sccache in PATH) +# (setup.py auto-detects ccache/sccache in PATH) RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ + && export CCACHE_BASEDIR="$PWD" \ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ - && python3 setup.py bdist_wheel --dist-dir=dist + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist FROM scratch AS export_vllm_wheel_release ARG COMMON_WORKDIR @@ -431,112 +516,118 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ 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/vllm/v1 /vllm_v1 # ----------------------- -# Test vLLM image -FROM mori_base AS test +# CI base image (Tier 1) - stable, rarely changing CI dependencies. +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers +# in the vLLM artifacts for the current commit. +FROM mori_base AS ci_base +ARG COMMON_WORKDIR -RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* - -# Install vLLM using uv (inherited from base stage) -# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && uv pip install --system -r requirements/rocm.txt \ - && uv pip install --system -r requirements/test/rocm.txt \ - && pip uninstall -y vllm \ - && uv pip install --system *.whl - -# Persist the built wheel in the image so python_only_compile_rocm.sh can -# reinstall it after removing compilers. The bind-mounted /install contents -# above are not available once that RUN step completes. -COPY --from=export_vllm /*.whl /opt/vllm-wheels/ - -# Update rdma-core to support latest rocshmem +# Update rdma-core to support latest rocshmem. ARG DEEPEP_NIC RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ cd /tmp/rdma-core && \ mkdir -p build && cd build && \ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ - ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ fi -# Install RIXL wheel +# Install RIXL + DeepEP wheels. RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ - uv pip install --system /rixl_install/*.whl + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl -# Install DeepEP wheel -RUN --mount=type=bind,from=build_deep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /deep_install/*.whl -COPY --from=build_deep /opt/rocshmem /opt/rocshmem +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem -# RIXL/MoRIIO runtime dependencies (RDMA userspace libraries) -RUN apt-get update -q -y && apt-get install -q -y \ +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ librdmacm1 \ libibverbs1 \ ibverbs-providers \ ibverbs-utils \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ && rm -rf /var/lib/apt/lists/* -WORKDIR /vllm-workspace -ARG COMMON_WORKDIR -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace - -# install development dependencies (for testing) -RUN cd /vllm-workspace \ - && python3 -m pip install -e tests/vllm_test_utils \ - && python3 -m pip install pytest-shard - -# enable fast downloads from hf (for testing) -ENV HF_XET_HIGH_PERFORMANCE=1 - -# increase timeout for hf downloads (for testing) -ENV HF_HUB_DOWNLOAD_TIMEOUT 60 - -# install audio decode package `torchcodec` from source (required due to -# ROCm and torch version mismatch) for tests with datasets package +# Install torchcodec from source for ROCm/torch ABI compatibility. COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh -RUN bash /tmp/install_torchcodec.sh \ +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ && rm /tmp/install_torchcodec.sh \ - && apt-get clean \ - && rm -rf /var/lib/apt/lists/* + && apt-get clean && rm -rf /var/lib/apt/lists/* -# Copy in the v1 package (for python-only install test group) -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 +# Pre-install shared ROCm runtime dependencies. +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + && rm -rf /tmp/ci-base-requirements -# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel +# Enable fast and less brittle model downloads in tests. +ENV HF_XET_HIGH_PERFORMANCE=1 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60 + +# Pre-install vLLM test dependencies. +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/rocm-test-reqs.txt + +# Rebuild fastsafetensors from source so its C++ extension is compiled with +# USE_ROCM and can detect libamdhip64.so at runtime. +RUN --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ + && test -n "${FASTSAFETENSORS_REQ}" \ + && python3 -m pip install --force-reinstall --no-deps \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + && rm /tmp/rocm-test-reqs.txt + +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # See: https://github.com/pytorch/pytorch/issues/169857 ENV MIOPEN_DEBUG_CONV_DIRECT=0 ENV MIOPEN_DEBUG_CONV_GEMM=0 -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. # See: https://github.com/ROCm/rocm-libraries/issues/6266 ENV HSA_ENABLE_IPC_MODE_LEGACY=1 -# Source code is used in the `python_only_compile.sh` test -# We hide it inside `src/` so that this source code -# will not be imported by other tests -RUN mkdir src && mv vllm src/vllm +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -# This is a workaround to ensure pytest exits with the correct status code in CI tests. -RUN printf '%s\n' \ - 'import os' \ - '' \ - '_exit_code = 1' \ - '' \ - 'def pytest_sessionfinish(session, exitstatus):' \ - ' global _exit_code' \ - ' _exit_code = int(exitstatus)' \ - '' \ - 'def pytest_unconfigure(config):' \ - ' import sys' \ - ' sys.stdout.flush()' \ - ' sys.stderr.flush()' \ - ' os._exit(_exit_code)' \ - > /vllm-workspace/conftest.py +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + +# Hide source under src/ so it won't shadow the installed package in tests. +RUN mkdir src && mv vllm src/vllm # ----------------------- # Final vLLM image @@ -553,6 +644,7 @@ RUN rm -f /usr/bin/sccache || true \ # This prevents S3 bucket config from leaking into production images ENV SCCACHE_BUCKET= ENV SCCACHE_REGION= +ENV SCCACHE_ENDPOINT= ENV SCCACHE_S3_NO_CREDENTIALS= ENV SCCACHE_IDLE_TIMEOUT= diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 195067b51a2..208ce863f6b 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.13" +ARG AITER_BRANCH="v0.1.13.post1" 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/ci-rocm.hcl b/docker/ci-rocm.hcl new file mode 100644 index 00000000000..138adcffcad --- /dev/null +++ b/docker/ci-rocm.hcl @@ -0,0 +1,376 @@ +# ci-rocm.hcl - CI-specific configuration for vLLM ROCm Docker builds +# +# This file lives in the vLLM repo at docker/ci-rocm.hcl so ROCm Docker +# build mechanics can evolve with Dockerfile.rocm and docker-bake-rocm.hcl. +# Used with: docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci +# +# Registry cache: Docker Hub (rocm/vllm-ci-cache) is used exclusively. +# AMD build agents already have Docker Hub credentials (they push the test +# image to rocm/vllm-ci), so no additional credential setup is required. +# ROCm CI uses Docker Hub for BuildKit layer cache by default. A separate +# compiler cache can be enabled with USE_SCCACHE=1 when AMD provides a shared +# S3-compatible cache endpoint. + +# CI metadata + +variable "BUILDKITE_COMMIT" { + default = "" +} + +variable "BUILDKITE_BUILD_NUMBER" { + default = "" +} + +variable "BUILDKITE_BUILD_ID" { + default = "" +} + +variable "PARENT_COMMIT" { + default = "" +} + +# Merge-base of HEAD with main - provides a more stable cache fallback than +# parent commit for long-lived PRs. Mirrors the VLLM_MERGE_BASE_COMMIT +# pattern used in the shared ci.hcl file. Auto-computed by ci-bake-rocm.sh +# when unset. +variable "VLLM_MERGE_BASE_COMMIT" { + default = "" +} + +# Bridge to vLLM's COMMIT variable for OCI labels +variable "COMMIT" { + default = BUILDKITE_COMMIT +} + +# Image tags (set by CI) + +variable "IMAGE_TAG" { + default = "" +} + +variable "IMAGE_TAG_LATEST" { + default = "" +} + +# ROCm-specific GPU architecture targets + +variable "PYTORCH_ROCM_ARCH" { + default = "gfx90a;gfx942;gfx950" +} + +# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of +# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in +# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test +# stage and is irrelevant when building --target ci_base itself. +variable "CI_BASE_IMAGE" { + default = "rocm/vllm-dev:ci_base" +} + +# Leave CI_MAX_JOBS empty so the Dockerfile falls back to $(nproc) and uses +# the full builder parallelism. Operators can still override this per build. +variable "CI_MAX_JOBS" { + default = "" +} + +# Upstream dependency commit pins -- extracted from Dockerfile.rocm by +# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache +# functions produce no entries when the variable is empty. +variable "RIXL_BRANCH" { + default = "" +} + +variable "UCX_BRANCH" { + default = "" +} + +variable "ROCSHMEM_BRANCH" { + default = "" +} + +variable "DEEPEP_BRANCH" { + default = "" +} + +variable "RIXL_CACHE_KEY" { + default = "" +} + +variable "ROCSHMEM_CACHE_KEY" { + default = "" +} + +variable "DEEPEP_CACHE_KEY" { + default = "" +} + +# Docker Hub registry cache for AMD builds. +# +# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache. +# Final-image cache exports use mode=min to reduce the volume of data pushed. +# Source-scoped csrc cache exports default to mode=max so fresh workers can +# recover more of the native build graph when ROCm extension inputs change. +# NOTE: mode=min still includes all layers referenced by the final image +# manifest, including inherited base layers (~7.25GB ROCm runtime). +# Docker Hub auto-creates the repo on first push. +# +# Final-image cache stays commit-scoped. Branch-to-branch reuse for the test +# image comes from importing the parent and merge-base commit cache refs. +# +# The source-scoped native cache is exported both per-commit and per-branch so +# ROCm extension rebuilds are shareable within the same commit reruns and across +# consecutive commits on the same branch without depending on a single global +# latest tag. + +variable "DOCKERHUB_CACHE_REPO" { + default = "rocm/vllm-ci-cache" +} + +variable "DOCKERHUB_CACHE_TO" { + default = "" +} + +variable "ROCM_CACHE_BRANCH_TAG" { + default = "" +} + +variable "ROCM_CACHE_UPSTREAM_BRANCH_TAG" { + default = "" +} + +variable "ROCM_CSRC_CACHE_TO_MODE" { + default = "max" +} + +variable "ROCM_FINAL_CACHE_TO_MODE" { + default = "min" +} + +# Functions + +function "get_cache_from_rocm" { + params = [] + result = compact([ + # Exact commit hit - fastest cache on re-runs of the same commit + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT}" : "", + # Parent commit - useful cache for incremental changes + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${PARENT_COMMIT}" : "", + # Merge-base with main - stable fallback for long-lived or rebased PRs; + # maps to a real main-branch commit whose cache layers are likely warm + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${VLLM_MERGE_BASE_COMMIT}" : "", + # Import the source-scoped native build cache as well so builds whose + # Python/package layers changed can still reuse compiled ROCm objects. + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "", + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "", + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", + # Branch-scoped full image cache - fallback when parent-commit cache is evicted + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", + ]) +} + +function "get_cache_to_rocm" { + params = [] + result = compact([ + # Commit-scoped cache for exact re-runs. + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "", + # Branch-scoped cache so later commits on the same branch can reuse the full + # image layers when the parent-commit cache is evicted. Unlike the old + # rocm-latest tag (which caused duplicate exporter 400s), this is per-branch. + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "", + ]) +} + +function "get_cache_from_rocm_csrc" { + params = [] + result = compact([ + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "", + PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "", + VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", + ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", + ]) +} + +function "get_cache_to_rocm_csrc" { + params = [] + result = compact([ + # Export the exact-commit native cache for same-commit reruns. + BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "", + # Export the branch-scoped native cache so later commits on the same branch + # can reuse compiled ROCm objects even when the exact parent cache is absent. + ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "", + ]) +} + +# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP). +# These stages are pinned to specific upstream commit hashes, so cache keys use +# those hashes rather than the Buildkite commit. This means the cache persists +# across all vLLM commits as long as the upstream dependency pins don't change. + +function "get_cache_from_rocm_deps" { + params = [] + result = compact([ + RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), + ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""), + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""), + ]) +} + +function "get_cache_to_rocm_rixl" { + params = [] + result = compact([ + RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""), + ]) +} + +function "get_cache_to_rocm_rocshmem" { + params = [] + result = compact([ + ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY},mode=min" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH},mode=min" : ""), + ]) +} + +function "get_cache_to_rocm_deepep" { + params = [] + result = compact([ + DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH},mode=min" : ""), + ]) +} + +# CI targets + +target "_ci-rocm" { + annotations = [ + "manifest:vllm.buildkite.build_number=${BUILDKITE_BUILD_NUMBER}", + "manifest:vllm.buildkite.build_id=${BUILDKITE_BUILD_ID}", + ] + args = { + ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH + CI_BASE_IMAGE = CI_BASE_IMAGE + max_jobs = CI_MAX_JOBS + } +} + +target "test-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm", "_labels"] + target = "test" + cache-from = get_cache_from_rocm() + cache-to = get_cache_to_rocm() + tags = compact([ + IMAGE_TAG, + IMAGE_TAG_LATEST, + ]) + output = ["type=registry"] +} + +# Cache-only target for the source-scoped ROCm native build stage. +# This persists the csrc-build stage in the registry cache even though the +# final test image only consumes it indirectly while packaging the wheel. +target "csrc-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "csrc-build" + cache-from = get_cache_from_rocm_csrc() + cache-to = get_cache_to_rocm_csrc() + output = ["type=cacheonly"] +} + +# Keep wheel export on the same CI graph as the test image build so the +# shared build_vllm/export_vllm stages resolve identically within one bake +# invocation. Without this, export-wheel-rocm uses the plain local target +# args while test-rocm-ci uses CI-only args, which can lead to separate +# cache lineages and inconsistent export_vllm results. +target "export-wheel-rocm" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "export_vllm" + cache-from = get_cache_from_rocm() + cache-to = get_cache_to_rocm() + output = ["type=local,dest=./wheel-export"] +} + +# Artifact-only vLLM build. GPU test jobs consume this artifact on top of +# ci_base, avoiding a per-commit multi-GB image push/pull. +group "test-rocm-ci-with-artifacts" { + targets = ["csrc-rocm-ci", "export-wheel-rocm"] +} + +# Full test image + wheel export. Kept for fallback/debugging when a pushed +# per-commit image is useful. +group "test-rocm-ci-with-wheel" { + targets = ["csrc-rocm-ci", "test-rocm-ci", "export-wheel-rocm"] +} + +# Image tags for the ci_base build. ci-bake-rocm.sh rewrites CI_BASE_IMAGE_TAG +# to the primary tag for this build. Non-nightly builds use a commit-scoped tag +# and also publish a content tag for reuse. NIGHTLY=1 builds on the stable branch +# can additionally set CI_BASE_IMAGE_TAG_STABLE to refresh rocm/vllm-dev:ci_base. +variable "CI_BASE_IMAGE_TAG" { + default = "rocm/vllm-dev:ci_base" +} + +variable "CI_BASE_IMAGE_TAG_CONTENT" { + default = "" +} + +variable "CI_BASE_IMAGE_TAG_STABLE" { + default = "" +} + +# Cache-only targets for upstream dependency stages. These persist each stage +# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds +# (e.g., requirements change), these stages are cache hits if their upstream +# pins haven't changed -- saving ~35min of compilation. +target "rixl-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "build_rixl" + cache-from = get_cache_from_rocm_deps() + cache-to = get_cache_to_rocm_rixl() + output = ["type=cacheonly"] +} + +target "rocshmem-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "build_rocshmem" + cache-from = get_cache_from_rocm_deps() + cache-to = get_cache_to_rocm_rocshmem() + output = ["type=cacheonly"] +} + +target "deepep-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm"] + target = "build_deepep" + cache-from = get_cache_from_rocm_deps() + cache-to = get_cache_to_rocm_deepep() + output = ["type=cacheonly"] +} + +# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.) +# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting +# files drifts from the remote image label. Per-PR builds then pull the result +# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit. +# Uses inline cache metadata on the ci_base image itself instead of exporting a +# separate registry cache artifact. +target "ci-base-rocm-ci" { + inherits = ["_common-rocm", "_ci-rocm", "_labels"] + target = "ci_base" + cache-from = concat( + compact([ + CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "", + CI_BASE_IMAGE_TAG_CONTENT != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT}" : "", + CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "", + ]), + # Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages + # are cache hits even when ci_base itself needs rebuilding. + get_cache_from_rocm_deps(), + ) + cache-to = ["type=inline"] + tags = compact([CI_BASE_IMAGE_TAG, CI_BASE_IMAGE_TAG_CONTENT, CI_BASE_IMAGE_TAG_STABLE]) + output = ["type=registry"] +} + +# Group for ci_base builds -- exports dependency stage caches alongside the +# ci_base image so future rebuilds can reuse them independently. +group "ci-base-rocm-ci-with-deps" { + targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] +} diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl new file mode 100644 index 00000000000..6b51781834b --- /dev/null +++ b/docker/docker-bake-rocm.hcl @@ -0,0 +1,143 @@ +# docker-bake-rocm.hcl - vLLM ROCm Docker build configuration +# +# This file lives in the vLLM repo at docker/docker-bake-rocm.hcl +# Equivalent of docker-bake.hcl for ROCm builds. +# +# Usage: +# docker buildx bake -f docker/docker-bake-rocm.hcl # Build test (default) +# docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm # Build final image +# docker buildx bake -f docker/docker-bake-rocm.hcl --print # Show resolved config +# +# CI usage (with the vLLM-owned CI overlay): +# docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci + +variable "MAX_JOBS" { + # Empty string lets the Dockerfile fall back to $(nproc) via + # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all + # available cores on whatever machine the build runs on. + # Override with --set '*.args.max_jobs=8' for local builds on small machines. + default = "" +} + +variable "PYTORCH_ROCM_ARCH" { + default = "gfx90a;gfx942;gfx950" +} + +variable "COMMIT" { + default = "" +} + +# Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and +# embedded as a label so future builds can compare without rebuilding. +variable "CI_BASE_CONTENT_HASH" { + default = "" +} + +# REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/) +# REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source) +variable "REMOTE_VLLM" { + default = "0" +} + +variable "VLLM_BRANCH" { + default = "main" +} + +# CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds. +# Defaults to the local "ci_base" stage for standalone/local builds. +# CI overrides this to "rocm/vllm-dev:ci_base" via environment variable. +variable "CI_BASE_IMAGE" { + default = "rocm/vllm-dev:ci_base" +} + +# Upstream dependency commit pins. Plain local bake builds use the Dockerfile +# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env +# overrides) and writes a small HCL override before invoking CI targets. +variable "RIXL_BRANCH" { + default = "" +} + +variable "UCX_BRANCH" { + default = "" +} + +variable "ROCSHMEM_BRANCH" { + default = "" +} + +variable "DEEPEP_BRANCH" { + default = "" +} + +group "default" { + targets = ["test-rocm"] +} + +target "_common-rocm" { + dockerfile = "docker/Dockerfile.rocm" + context = "." + args = { + max_jobs = MAX_JOBS + ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH + REMOTE_VLLM = REMOTE_VLLM + VLLM_BRANCH = VLLM_BRANCH + CI_BASE_IMAGE = CI_BASE_IMAGE + } +} + +target "_labels" { + labels = { + "org.opencontainers.image.source" = "https://github.com/vllm-project/vllm" + "org.opencontainers.image.vendor" = "vLLM" + "org.opencontainers.image.title" = "vLLM ROCm" + "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)" + "org.opencontainers.image.licenses" = "Apache-2.0" + "org.opencontainers.image.revision" = COMMIT + } + annotations = [ + "manifest:org.opencontainers.image.revision=${COMMIT}", + ] +} + +target "test-rocm" { + inherits = ["_common-rocm", "_labels"] + target = "test" + tags = ["rocm/vllm:test"] + output = ["type=docker"] +} + +# CI base image target - builds only the ci_base stage (RIXL, DeepEP, +# torchcodec, requirements, etc.). Used by the weekly scheduled build and +# the auto-rebuild trigger when requirements change in a PR. +target "ci-base-rocm" { + inherits = ["_common-rocm", "_labels"] + target = "ci_base" + labels = { + "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH + } + tags = ["rocm/vllm-dev:ci_base"] + output = ["type=docker"] +} + +# Wheel export target - extracts the built vLLM wheel + test workspace +# to local disk. Used by CI to upload the wheel as a Buildkite artifact +# so test jobs can assemble images locally from ci_base + wheel instead +# of pulling the full large image from Docker Hub. +# +# Usage: +# docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm +# # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc. +# +# After a full bake build, BuildKit cache makes this nearly instant. +target "export-wheel-rocm" { + inherits = ["_common-rocm"] + target = "export_vllm" + output = ["type=local,dest=./wheel-export"] +} + +target "final-rocm" { + inherits = ["_common-rocm", "_labels"] + target = "final" + tags = ["rocm/vllm:latest"] + output = ["type=docker"] +} diff --git a/docker/versions.json b/docker/versions.json index ee23b5baf04..15f77648a9c 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.11.post2" + "default": "0.6.12" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 1b598aea38c..6d0b2a01aca 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -246,6 +246,12 @@ Every image listed in "image_files" is added to the request in the listed order The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string. +By default, image references are sent to the serving endpoint as provided, with local image paths converted to `file://` URLs. + +If the benchmark client should load local and HTTP(S) images before sending requests, pass `--custom-ensure-client-side-data` to encode them as base64 data URLs on the client side. + +Existing `data:image/...` URLs are already self-contained and are kept unchanged. + ```bash # need a model with vision capability here vllm serve Qwen/Qwen2-VL-7B-Instruct @@ -253,13 +259,13 @@ vllm serve Qwen/Qwen2-VL-7B-Instruct ```bash # run benchmarking script -vllm bench serve--save-result --save-detailed \ +vllm bench serve --save-result --save-detailed \ --backend openai-chat \ --model Qwen/Qwen2-VL-7B-Instruct \ --endpoint /v1/chat/completions \ --dataset-name custom_image \ --dataset-path \ - --allowed-local-media-path /path/to/image/folder + --custom-ensure-client-side-data ``` Note that we need to use the `openai-chat` backend and `/v1/chat/completions` endpoint for multimodal inputs. diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index ce46445a983..c9bd0e5bdd9 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -35,8 +35,7 @@ Traces can be visualized using . !!! tip To stop the profiler - it flushes out all the profile trace files to the directory. This takes time, for example for about 100 requests worth of data for a llama 70b, it takes about 10 minutes to flush out on a H100. - Set the env variable VLLM_RPC_TIMEOUT to a big number before you start the server. Say something like 30 minutes. - `export VLLM_RPC_TIMEOUT=1800000` + The engine client waits for this flush to complete without timing out, so simply allow the stop call to run to completion. ### Example commands and usage diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 14781f6a5a3..1fb5c2ba651 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -17,6 +17,7 @@ The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, m * [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]: orchestrates capture, replay, greedy packing, and data-parallel execution for encoder CUDA Graphs. * [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph]: a runtime-checkable protocol that models implement to opt-in to encoder CUDA Graphs. +* [EncoderItemSpec][vllm.v1.worker.encoder_cudagraph_defs.EncoderItemSpec]: describes a single encoder input item (image or video) with its input size and output token count. * [BudgetGraphMetadata][vllm.v1.worker.encoder_cudagraph.BudgetGraphMetadata]: holds the captured CUDA Graph and its associated I/O buffers for a single token budget level. ### Budget-based graph capture @@ -30,8 +31,7 @@ class BudgetGraphMetadata: max_batch_size: int max_frames_per_batch: int graph: torch.cuda.CUDAGraph - input_buffer: torch.Tensor # e.g. pixel_values - metadata_buffers: dict[str, torch.Tensor] # e.g. embeddings, seq metadata + input_buffers: dict[str, torch.Tensor] # e.g. pixel_values, embeddings, seq metadata output_buffer: torch.Tensor # encoder hidden states ``` @@ -43,8 +43,8 @@ When a batch of images arrives, the manager sorts images by output token count ( For each graph replay: -1. Zero the pre-allocated `input_buffer`, then copy input tensors (e.g., `pixel_values`) into it. -2. Zero `metadata_buffers`, then slice-copy precomputed values (e.g., rotary embeddings, sequence metadata). +1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs. +2. Zero the pre-allocated `input_buffers`, then slice-copy the replay values into them. 3. Replay the CUDA Graph. 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays). @@ -65,19 +65,15 @@ Following (ViT full CUDA graph Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods: -* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, input key, buffer keys, output hidden size). +* `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video). * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets. -* `get_encoder_cudagraph_num_items(mm_kwargs)` — returns the number of items (e.g. images) in the batch. -* `get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)` — returns per-item output token counts, used for greedy packing. -* `get_encoder_cudagraph_per_item_input_sizes(mm_kwargs)` — returns per-item input sizes (e.g. patch counts), used for DP load balancing. +* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size and output token count. Replaces the former three separate methods (`get_num_items`, `get_per_item_output_tokens`, `get_per_item_input_sizes`). * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding. -* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. -* `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay. -* `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay). -* `encoder_eager_forward(...)` — fallback eager forward when no graph fits. -* `get_input_modality(...)` - return the modality of the inputs. -* `get_max_frames_per_video()` - return model-specific max frames per video. -* `postprocess_encoder_output(...)` - post process encoder output, directly call scatter_output_slices by default +* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph. +* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch)` — computes buffer values from actual batch inputs. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match `buffer_keys` in the config. +* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor])` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `pixel_values` tensor is included in `inputs` alongside metadata buffers. +* `encoder_eager_forward(mm_kwargs)` — fallback eager forward when no graph fits. +* `postprocess_encoder_output(...)` — post-process encoder output, delegates to `scatter_output_slices` by default. !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. @@ -103,7 +99,7 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs: * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level. * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference. * `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. -* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value according to its `processing_info`). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode). +* `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode). ## Usage guide diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index c8437704447..9fa1763108c 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -100,14 +100,44 @@ For further details on renderer APIs, please refer to [this page](renderer.md). - `/version` - Version information - `/load` - Server load metrics -## Sleep Mode APIs +## Server in development mode + +When using the flag VLLM_SERVER_DEV_MODE=1, you enable development endpoints. + +**SECURITY WARNING: These endpoints should NOT be used in production!** + +### Cache Management APIs + +- `/reset_prefix_cache` - Reset prefix cache (can disrupt service) +- `/reset_mm_cache` - Reset multimodal cache (can disrupt service) +- `/reset_encoder_cache` - Reset encoder cache (can disrupt service) + +### Weight Transfer APIs (RL Training) + +For further details on Weight Transfer, please refer to [this page](../../training/weight_transfer/README.md). + +- `/pause` - Pause generation (causes denial of service) +- `/resume` - Resume generation +- `/is_paused` - Check if generation is paused +- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF +- `/update_weights` - Update model weights (can alter model behavior) +- `/get_world_size` - Get distributed world size + +### Collective RPC + +- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous) + +### Server info + +- `/server_info` - Get detailed server configuration + +### Sleep Mode APIs For further details on sleep mode, please refer to [this page](../../features/sleep_mode.md). - `/sleep` - Put engine to sleep (causes denial of service) - `/wake_up` - Wake engine from sleep - `/is_sleeping` - Check if engine is sleeping -- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous) ## Chat Template diff --git a/docs/training/weight_transfer/base.md b/docs/training/weight_transfer/base.md index 6c768c87fd9..ace228b0091 100644 --- a/docs/training/weight_transfer/base.md +++ b/docs/training/weight_transfer/base.md @@ -156,5 +156,6 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory engine = WeightTransferEngineFactory.create_engine( config=weight_transfer_config, parallel_config=parallel_config, + model=model, ) ``` diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 99a45c9d3ca..b0e16d11c75 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -9,8 +9,8 @@ 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 # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.11.post2 -flashinfer-cubin==0.6.11.post2 +flashinfer-python==0.6.12 +flashinfer-cubin==0.6.12 apache-tvm-ffi==0.1.9 tilelang==0.1.9 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 3da176cd1c6..7639b9cc13a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5622,6 +5622,7 @@ dependencies = [ "expect-test", "futures", "half", + "indexmap 2.13.0", "itertools 0.14.0", "llm-multimodal", "minijinja", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index e742b68b2ad..9ca38d0ae79 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -43,6 +43,7 @@ half = { version = "2.7.1", features = ["bytemuck"] } hex = "0.4.3" hf-hub = { version = "0.5.0", features = ["tokio"] } http-body = "1.0.1" +indexmap = "2.13.0" itertools = "0.14.0" libc = "0.2.177" llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" } @@ -69,7 +70,7 @@ rustc-hash = "1.1.0" serde = { version = "1.0.228", features = ["derive"] } serde-json-fmt = "0.1.0" serde_default = "0.2.0" -serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] } +serde_json = { version = "1.0.145", features = ["preserve_order"] } serde_repr = "0.1.20" serde_tuple = "1.1.3" serde_with = "3.18.0" diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 1548c6f5926..0523b9defe9 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -10,6 +10,7 @@ asynk-strim-attr.workspace = true easy-ext.workspace = true futures.workspace = true half.workspace = true +indexmap.workspace = true itertools.workspace = true llm-multimodal.workspace = true minijinja.workspace = true diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 0669af8daff..5b6f66cf417 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -189,6 +189,7 @@ impl ChatLlm { cache_salt: request.cache_salt, add_special_tokens: request.add_special_tokens, data_parallel_rank: request.data_parallel_rank, + lora_request: request.lora_request, }; let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed(); diff --git a/rust/src/chat/src/parser/mod.rs b/rust/src/chat/src/parser/mod.rs index 52e83b3e047..244a87cc7a7 100644 --- a/rust/src/chat/src/parser/mod.rs +++ b/rust/src/chat/src/parser/mod.rs @@ -6,10 +6,10 @@ use std::convert::Infallible; use std::fmt; use std::str::FromStr; -use serde_with::DeserializeFromStr; +use serde_with::{DeserializeFromStr, SerializeDisplay}; /// Specify which reasoning or tool-call parser implementation to use. -#[derive(Debug, Clone, PartialEq, Eq, Default, DeserializeFromStr)] +#[derive(Debug, Clone, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)] pub enum ParserSelection { /// Use model-based auto-detection. #[default] diff --git a/rust/src/chat/src/renderer/hf/format.rs b/rust/src/chat/src/renderer/hf/format.rs index a9b35d0f41c..2c990fb37ba 100644 --- a/rust/src/chat/src/renderer/hf/format.rs +++ b/rust/src/chat/src/renderer/hf/format.rs @@ -5,7 +5,7 @@ use std::str::FromStr; use minijinja::machinery::ast::{Expr, ForLoop, Set, Stmt}; use minijinja::machinery::{WhitespaceConfig, parse}; use minijinja::syntax::SyntaxConfig; -use serde_with::DeserializeFromStr; +use serde_with::{DeserializeFromStr, SerializeDisplay}; /// Chat template content format. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -18,7 +18,7 @@ pub enum ChatTemplateContentFormat { } /// Configurable chat-template content format selection. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)] pub enum ChatTemplateContentFormatOption { /// Detect the format from the template source. #[default] diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 851df5068f4..47c10c0219e 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use serde::Serialize; -use serde_json::Value; +use serde_json::Value as JsonValue; use thiserror_ext::AsReport as _; use tracing::{info, trace, warn}; use vllm_text::Prompt; @@ -13,6 +13,7 @@ use self::format::{ ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption, }; use self::template::{CompiledChatTemplate, TemplateContext}; +use self::value::{TemplateValue, to_template_value}; use super::{ChatRenderer, RenderedPrompt}; use crate::error::Result; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; @@ -24,6 +25,7 @@ mod error; mod format; mod template; mod tojson; +mod value; pub use template::{load_chat_template, resolve_chat_template}; @@ -38,7 +40,7 @@ pub struct MultimodalRenderInfo { /// state. pub struct HfChatRenderer { default_template: Option, - default_template_kwargs: HashMap, + default_template_kwargs: HashMap, content_format: ContentFormatOption, special_tokens: Option, multimodal: Option, @@ -48,7 +50,7 @@ impl HfChatRenderer { /// Create a renderer from the given template string. pub fn new( template: Option, - default_template_kwargs: HashMap, + default_template_kwargs: HashMap, content_format: ContentFormatOption, ) -> Result { Ok(Self { @@ -245,7 +247,7 @@ struct TemplateToolCall { #[derive(Debug, Serialize)] struct TemplateToolFunction { name: String, - arguments: Value, + arguments: TemplateValue, } #[derive(Debug, Serialize)] @@ -259,7 +261,7 @@ pub(super) struct TemplateTool { struct TemplateToolDefinition { name: String, description: Option, - parameters: Value, + parameters: TemplateValue, strict: Option, } @@ -345,13 +347,14 @@ fn to_template_tool_calls( let mut tool_calls = Vec::new(); for tool_call in content.tool_calls() { - let arguments = serde_json::from_str::(&tool_call.arguments).map_err(|error| { + let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| { Error::ChatTemplate(format!( "assistant tool call `{}` has invalid JSON arguments: {}", tool_call.id, error.as_report() )) })?; + let arguments = to_template_value(arguments); tool_calls.push(TemplateToolCall { id: tool_call.id.clone(), @@ -434,7 +437,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec { function: TemplateToolDefinition { name: tool.name.clone(), description: tool.description.clone(), - parameters: tool.parameters.clone(), + parameters: to_template_value(tool.parameters.clone()), strict: tool.strict, }, }) @@ -909,6 +912,29 @@ mod tests { assert_eq!(rendered, "get_weather|Paris|call_1|Sunny"); } + #[test] + fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() { + let request = sample_request(vec![ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(crate::AssistantToolCall { + id: "call_1".to_string(), + name: "add".to_string(), + arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(), + }), + ])]); + + let rendered = render( + Some( + "{%- set arguments = messages[0].tool_calls[0].function.arguments -%} +{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%} +|{{ arguments['items'] }}", + ), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands"); + } + #[test] fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() { let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); diff --git a/rust/src/chat/src/renderer/hf/tojson.rs b/rust/src/chat/src/renderer/hf/tojson.rs index f04c954a79f..cd53108c579 100644 --- a/rust/src/chat/src/renderer/hf/tojson.rs +++ b/rust/src/chat/src/renderer/hf/tojson.rs @@ -208,11 +208,27 @@ mod tests { } #[test] - fn tojson_preserves_arbitrary_precision_number_spelling() { + fn tojson_uses_standard_serde_json_number_spelling() { let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap(); let rendered = render("{{ payload|tojson }}", payload); - assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}"); + // TODO: we cannot preserve the original number precision by enabling `serde_json`'s + // `arbitrary_precision` feature, otherwise the following test + // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail. + // See issue: https://github.com/mitsuhiko/minijinja/issues/641 + assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}"); + } + + #[test] + fn serialized_json_numbers_do_not_leak_serde_private_representation() { + let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap(); + let rendered = render("{{ payload }}", payload); + + // TODO: we cannot preserve the original number precision by enabling `serde_json`'s + // `arbitrary_precision` feature, otherwise this will fail. + // See issue: https://github.com/mitsuhiko/minijinja/issues/641 + assert!(!rendered.contains("$serde_json::private::Number")); + assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#); } #[test] diff --git a/rust/src/chat/src/renderer/hf/value.rs b/rust/src/chat/src/renderer/hf/value.rs new file mode 100644 index 00000000000..65064705e01 --- /dev/null +++ b/rust/src/chat/src/renderer/hf/value.rs @@ -0,0 +1,77 @@ +use std::sync::Arc; + +use indexmap::IndexMap; +use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr}; +use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State}; +use serde::Serialize; +use serde_json::Value as JsonValue; + +/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used +/// as a value in the chat template. +#[derive(Debug, Serialize)] +#[serde(transparent)] +pub(super) struct TemplateValue(minijinja::Value); + +pub(super) fn to_template_value(value: JsonValue) -> TemplateValue { + TemplateValue(match value { + JsonValue::Array(values) => values + .into_iter() + .map(to_template_value) + .map(|value| value.0) + .collect::(), + JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap( + values + .into_iter() + .map(|(key, value)| (key, to_template_value(value).0)) + .collect(), + )), + // For primitive values, directly convert them to `minijinja::Value` using `from_serialize`. + value => minijinja::Value::from_serialize(value), + }) +} + +/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can +/// always handle dict methods through the unknown-method callback. +/// +/// Use `IndexMap` to preserve the original key order when iterating. +/// +/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates +/// commonly call `dict.items()`, which would fail if the map had an `items` field. +/// See issue: https://github.com/mitsuhiko/minijinja/issues/903 +#[derive(Debug)] +struct TemplateMap(IndexMap); + +impl Object for TemplateMap { + fn repr(self: &Arc) -> ObjectRepr { + ObjectRepr::Map + } + + fn get_value(self: &Arc, key: &minijinja::Value) -> Option { + self.0.get(key.as_str()?).cloned() + } + + fn get_value_by_str(self: &Arc, key: &str) -> Option { + self.0.get(key).cloned() + } + + fn enumerate(self: &Arc) -> Enumerator { + self.mapped_rev_enumerator(|this| { + Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str()))) + }) + } + + fn enumerator_len(self: &Arc) -> Option { + Some(self.0.len()) + } + + fn call_method( + self: &Arc, + _state: &State<'_, '_>, + _method: &str, + _args: &[minijinja::Value], + ) -> std::result::Result { + // Always return `UnknownMethod` for method calls, + // so that pycompat can handle dict methods through the unknown-method callback. + Err(TemplateError::from(TemplateErrorKind::UnknownMethod)) + } +} diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index f4bd565bafd..cb22f95de0d 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -1,10 +1,10 @@ use std::fmt; use std::str::FromStr; -use serde_with::DeserializeFromStr; +use serde_with::{DeserializeFromStr, SerializeDisplay}; /// Specify which chat renderer implementation to use. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, DeserializeFromStr, SerializeDisplay)] pub enum RendererSelection { /// Use model-based auto-detection. #[default] diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index c1cb83b8dc3..842c941a6c0 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use llm_multimodal::ImageDetail; use serde::{Deserialize, Serialize}; use serde_json::Value; +use vllm_engine_core_client::protocol::lora::LoraRequest; pub use vllm_text::SamplingParams; use vllm_text::TextDecodeOptions; pub use vllm_tool_parser::Tool as ChatTool; @@ -426,6 +427,9 @@ pub struct ChatRequest { /// Override data parallel rank. #[serde(default)] pub data_parallel_rank: Option, + /// LoRA adapter selected for this request. + #[serde(default)] + pub lora_request: Option, } impl ChatRequest { @@ -445,6 +449,7 @@ impl ChatRequest { cache_salt: None, add_special_tokens: false, data_parallel_rank: None, + lora_request: None, } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 74491cd0243..ab2ca06cb37 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -20,6 +20,7 @@ use vllm_chat::{ use vllm_text::{DecodedTextEvent, Finished, Prompt}; /// One model/parser configuration used to run the fixed roundtrip fixtures. +#[derive(Clone)] struct RoundtripCase { /// Hugging Face model id resolved through the production backend loader. model_id: &'static str, @@ -31,11 +32,45 @@ struct RoundtripCase { tool_call_parser: ParserSelection, /// Reasoning parser selection used by the output processor. reasoning_parser: ParserSelection, + /// How this model's chat template handles thinking mode. + thinking_behavior: ThinkingBehavior, /// JSON formatting expected after this model's template has materialized /// tool-call arguments. json_fmt: JsonFmt, } +#[derive(Clone, Copy)] +enum ThinkingBehavior { + /// The chat template accepts explicit thinking on/off kwargs, and uses + /// `default` when the request does not specify either kwarg. + Toggleable { default: bool }, + /// The chat template always behaves as `value` for this fixture. + Always { value: bool }, +} + +impl ThinkingBehavior { + fn default(self) -> bool { + match self { + Self::Toggleable { default } => default, + Self::Always { value } => value, + } + } + + fn fixtures(self) -> Vec> { + match self { + Self::Toggleable { .. } => vec![ + Some(true), // explicitly enable thinking + Some(false), // explicitly disable thinking + None, // use default template behavior + ], + Self::Always { value } => vec![ + Some(value), // explicitly request the supported thinking behavior + None, // use default template behavior + ], + } + } +} + impl RoundtripCase { /// Qwen3 XML tool-call format with `qwen3` reasoning tags. fn qwen3() -> Self { @@ -44,6 +79,7 @@ impl RoundtripCase { assistant_stop_suffix: "<|im_end|>\n", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), } } @@ -55,6 +91,7 @@ impl RoundtripCase { assistant_stop_suffix: "<|im_end|>\n", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), } } @@ -66,6 +103,7 @@ impl RoundtripCase { assistant_stop_suffix: "[e~[\n", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Always { value: true }, json_fmt: compact_json_fmt(), } } @@ -77,6 +115,7 @@ impl RoundtripCase { 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(), } } @@ -88,6 +127,7 @@ impl RoundtripCase { assistant_stop_suffix: "", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: compact_json_fmt(), } } @@ -100,6 +140,7 @@ impl RoundtripCase { assistant_stop_suffix: "<|im_end|>", tool_call_parser: ParserSelection::Auto, reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, json_fmt: spaced_json_fmt(), } } @@ -135,35 +176,44 @@ roundtrip_tests! { /// Run the fixed reasoning+content fixture for one model/parser case. async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> { + for thinking in case.thinking_behavior.fixtures() { + run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?; + } + Ok(()) +} + +async fn run_roundtrip_reasoning_and_content_inner( + case: RoundtripCase, + thinking: Option, +) -> Result<()> { let backends = load_roundtrip_backends(&case).await?; let request = roundtrip_request( "roundtrip-reasoning-content", vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")], Vec::new(), + thinking, ); let expected_reasoning = "Need compute 2 + 2 directly."; let expected_text = "The answer is 4."; + let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default()); - let result = run_roundtrip( - &case, - &backends, - &request, - AssistantMessage { - content: vec![ - AssistantContentBlock::Reasoning { - text: expected_reasoning.to_string(), - }, - AssistantContentBlock::Text { - text: expected_text.to_string(), - }, - ], - }, - ) - .await?; + let assistant = { + let mut content = Vec::new(); + if effective_thinking { + content.push(AssistantContentBlock::Reasoning { + text: expected_reasoning.to_string(), + }); + } + content.push(AssistantContentBlock::Text { + text: expected_text.to_string(), + }); + AssistantMessage { content } + }; + let result = run_roundtrip(&case, &backends, &request, assistant).await?; assert_eq!( result.parsed_message.reasoning().as_deref().map(str::trim), - Some(expected_reasoning) + effective_thinking.then_some(expected_reasoning) ); assert_eq!(result.parsed_message.text().trim(), expected_text); assert_eq!(result.parsed_message.tool_calls().count(), 0); @@ -183,9 +233,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> { "roundtrip-reasoning-tools", vec![ChatMessage::text( ChatRole::User, - "Check Shanghai weather and add 1.00 plus 2.", + "Check Shanghai weather and add 1.0 plus 2.", )], test_tools(), + Some(true), // always enable thinking in this fixture ); let expected_reasoning = "Need call the weather and add tools."; let expected_text = "I will call the tools."; @@ -210,9 +261,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> { AssistantContentBlock::ToolCall(AssistantToolCall { id: "functions.add:1".to_string(), name: "add".to_string(), - // Intentionally use a non-lexical order of keys and a different number - // formatting style to verify text-level fidelity of the roundtrip. - arguments: r#"{"y":1.00,"x":2}"#.to_string(), + // Intentionally use a non-lexical order of keys to verify text-level + // fidelity of the roundtrip where JSON formatting remains stable. The + // `items` key also exercises templates that call `arguments.items()`. + arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(), }), ], }, @@ -240,7 +292,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> { assert_eq!(tool_calls[1].name, "add"); assert_eq!( tool_calls[1].arguments, - expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?, + expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?, ); assert_eq!( @@ -487,6 +539,7 @@ fn roundtrip_request( request_id: impl Into, messages: Vec, tools: Vec, + thinking: Option, ) -> ChatRequest { let mut request = ChatRequest { request_id: request_id.into(), @@ -500,10 +553,12 @@ fn roundtrip_request( ..ChatRequest::for_test() }; - // Enable thinking for some models so that rendering and parsing the reasoning block is - // exercised in the roundtrip. - for key in ["thinking", "enable_thinking"] { - request.chat_options.template_kwargs.insert(key.to_string(), true.into()); + // Explicitly enable or disable thinking so that rendering and parsing the reasoning block is + // exercised or skipped in the roundtrip. If unspecified, use the default template behavior. + if let Some(thinking) = thinking { + for key in ["thinking", "enable_thinking"] { + request.chat_options.template_kwargs.insert(key.to_string(), thinking.into()); + } } request @@ -531,9 +586,13 @@ fn test_tools() -> Vec { "type": "object", "properties": { "y": { "type": "number" }, - "x": { "type": "number" } + "x": { "type": "number" }, + "items": { + "type": "array", + "items": { "type": "string" } + } }, - "required": ["y", "x"] + "required": ["y", "x", "items"] }), strict: None, }, diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 70ac8440453..ee7848fe0be 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -165,6 +165,15 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub enable_log_requests: bool, + /// If specified, API server will add X-Request-Id header to responses. + #[arg( + long, + default_missing_value = "true", + num_args = 0..=1 + )] + #[serde(default)] + pub enable_request_id_headers: bool, + /// Disable periodic logging of engine statistics (throughput, queue depth, /// cache usage). #[arg(long)] @@ -238,6 +247,7 @@ impl SharedRuntimeArgs { default_chat_template_kwargs: self.default_chat_template_kwargs, chat_template_content_format: self.chat_template_content_format, enable_log_requests: self.enable_log_requests, + enable_request_id_headers: self.enable_request_id_headers, disable_log_stats: self.disable_log_stats, grpc_port: self.grpc_port, shutdown_timeout, @@ -278,6 +288,7 @@ impl SharedRuntimeArgs { default_chat_template_kwargs: self.default_chat_template_kwargs, chat_template_content_format: self.chat_template_content_format, enable_log_requests: self.enable_log_requests, + enable_request_id_headers: self.enable_request_id_headers, disable_log_stats: self.disable_log_stats, grpc_port: self.grpc_port, shutdown_timeout, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 0762468456e..ea867e4673a 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -43,6 +43,7 @@ fn serve_args_forward_python_flags_with_separator() { default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, served_model_name: [], }, @@ -86,6 +87,17 @@ fn serve_args_auto_forward_python_flags_without_separator() { ); } +#[test] +fn serve_args_auto_forward_enable_lora_to_python() { + let cli = + Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--enable-lora"]).unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + assert_eq!(args.managed_engine.python_args, vec!["--enable-lora"]); +} + #[test] fn serve_args_auto_forward_python_multi_char_alias_without_separator() { let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap(); @@ -116,6 +128,46 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() { assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32); } +#[test] +fn serve_passes_enable_request_id_headers_into_config() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-request-id-headers", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string()); + assert!(config.enable_request_id_headers); +} + +#[test] +fn frontend_args_json_passes_enable_request_id_headers_into_config() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#, + ]) + .unwrap(); + + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args"); + }; + let config = args.into_config(); + assert!(config.enable_request_id_headers); +} + #[test] fn serve_args_reject_unknown_renderer_value() { let error = Cli::try_parse_from([ @@ -218,6 +270,7 @@ fn frontend_args_accept_json() { default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, served_model_name: [], }, @@ -616,6 +669,7 @@ fn serve_args_accept_handshake_aliases() { default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, served_model_name: [], }, @@ -733,6 +787,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, grpc_port: None, shutdown_timeout: 0ns, @@ -795,6 +850,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, grpc_port: None, shutdown_timeout: 0ns, @@ -872,6 +928,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present default_chat_template_kwargs: None, chat_template_content_format: Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, grpc_port: None, shutdown_timeout: 0ns, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index eeaa0832888..8bd972ae17a 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -326,15 +326,6 @@ pub struct EngineUnsupportedArgs { #[arg(long)] pub mm_processor_cache_type: Option, - /// If True, enable handling of LoRA adapters. - #[arg( - long, - visible_alias = "no-enable-lora", - default_missing_value = "true", - num_args = 0..=1 - )] - pub enable_lora: Option, - /// Dictionary mapping specific modalities to LoRA model paths. #[arg(long)] pub default_mm_loras: Option, @@ -620,15 +611,6 @@ pub struct ServerUnsupportedArgs { #[arg(long)] pub middleware: Option, - /// If specified, API server will add X-Request-Id header to responses. - #[arg( - long, - visible_alias = "no-enable-request-id-headers", - default_missing_value = "true", - num_args = 0..=1 - )] - pub enable_request_id_headers: Option, - /// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint. #[arg( long, diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 94d5ab1c628..2a8c3c74188 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use futures::future::{join_all, try_join_all}; +use serde::Serialize; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info, trace}; @@ -10,6 +11,7 @@ use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop use crate::coordinator::CoordinatorHandle; use crate::error::{Error, Result}; use crate::protocol::handshake::EngineCoreReadyResponse; +use crate::protocol::lora::LoraRequest; use crate::protocol::utility::EngineCoreUtilityRequest; use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype}; use crate::transport::{self, ConnectedEngine}; @@ -22,7 +24,7 @@ pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput}; /// How the frontend acquires its request/response transport with Python /// `EngineCoreProc`s. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub enum TransportMode { /// The Rust process owns the startup handshake and allocates or binds the /// frontend transport addresses itself before replying to engine @@ -659,6 +661,24 @@ impl EngineCoreClient { Ok(results.into_iter().all(|ok| ok)) } + /// Load or refresh one LoRA adapter on every connected engine. + pub async fn add_lora(&self, lora_request: &LoraRequest) -> Result { + Ok(self + .call_utility::("add_lora", (lora_request,)) + .await? + .into_iter() + .all(|loaded| loaded)) + } + + /// Remove one LoRA adapter from every connected engine. + pub async fn remove_lora(&self, lora_id: u64) -> Result { + Ok(self + .call_utility::("remove_lora", (lora_id,)) + .await? + .into_iter() + .all(|removed| removed)) + } + /// Put the engine to sleep. pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> { self.call_utility::<(), _>("sleep", (level, mode)).await?; diff --git a/rust/src/engine-core-client/src/protocol/lora.rs b/rust/src/engine-core-client/src/protocol/lora.rs new file mode 100644 index 00000000000..27b74cf991c --- /dev/null +++ b/rust/src/engine-core-client/src/protocol/lora.rs @@ -0,0 +1,42 @@ +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::protocol::OpaqueValue; + +/// Request for a LoRA adapter. +/// +/// Mirrors Python `vllm.lora.request.LoRARequest`, which is a msgspec +/// `array_like=True` struct. Keep the field order aligned with Python. +#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)] +pub struct LoraRequest { + pub lora_name: String, + pub lora_int_id: u64, + pub lora_path: String, + #[serde(default)] + pub base_model_name: Option, + #[serde(default)] + pub tensorizer_config_dict: Option, + #[serde(default)] + pub load_inplace: bool, + #[serde(default)] + pub is_3d_lora_weight: bool, +} + +impl LoraRequest { + pub fn new( + lora_name: String, + lora_int_id: u64, + lora_path: String, + load_inplace: bool, + is_3d_lora_weight: bool, + ) -> Self { + Self { + lora_name, + lora_int_id, + lora_path, + base_model_name: None, + tensorizer_config_dict: None, + load_inplace, + is_3d_lora_weight, + } + } +} diff --git a/rust/src/engine-core-client/src/protocol/mod.rs b/rust/src/engine-core-client/src/protocol/mod.rs index 4a00d9d31c5..e87bc334fd0 100644 --- a/rust/src/engine-core-client/src/protocol/mod.rs +++ b/rust/src/engine-core-client/src/protocol/mod.rs @@ -48,6 +48,7 @@ mod classified_outputs; pub mod dtype; pub mod handshake; pub mod logprobs; +pub mod lora; pub mod multimodal; pub mod stats; pub mod tensor; @@ -349,7 +350,7 @@ pub struct EngineCoreRequest { pub pooling_params: Option, pub arrival_time: f64, #[serde(default)] - pub lora_request: Option, + pub lora_request: Option, #[serde(default)] pub cache_salt: Option, #[serde(default)] diff --git a/rust/src/llm/src/request.rs b/rust/src/llm/src/request.rs index b17035b0512..af5d257774b 100644 --- a/rust/src/llm/src/request.rs +++ b/rust/src/llm/src/request.rs @@ -2,8 +2,9 @@ use std::collections::BTreeMap; use std::time::{SystemTime, UNIX_EPOCH}; use uuid::Uuid; +use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; -use vllm_engine_core_client::protocol::{EngineCoreRequest, EngineCoreSamplingParams, OpaqueValue}; +use vllm_engine_core_client::protocol::{EngineCoreRequest, EngineCoreSamplingParams}; use crate::error::{Error, Result}; @@ -34,7 +35,7 @@ pub struct GenerateRequest { pub priority: i32, pub data_parallel_rank: Option, pub reasoning_ended: Option, - pub lora_request: Option, + pub lora_request: Option, } #[derive(Debug)] diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 6ef2e1a883e..50d6fc1be40 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -68,6 +68,7 @@ async fn main() -> Result<()> { default_chat_template_kwargs: None, chat_template_content_format: ChatTemplateContentFormatOption::Auto, enable_log_requests: false, + enable_request_id_headers: false, disable_log_stats: false, grpc_port: None, shutdown_timeout: Duration::ZERO, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 522133427f4..f1599d18793 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -2,12 +2,13 @@ use std::collections::HashMap; use std::time::Duration; use anyhow::Result; +use serde::Serialize; use serde_json::Value; use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection}; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; /// How the HTTP server obtains its listening socket. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub enum HttpListenerMode { /// Bind a fresh TCP listener on the given host/port. BindTcp { host: String, port: u16 }, @@ -20,7 +21,7 @@ pub enum HttpListenerMode { /// Which coordinator implementation should be active when one is present for a /// frontend client. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub enum CoordinatorMode { /// Do not run a coordinator at all. None, @@ -32,7 +33,7 @@ pub enum CoordinatorMode { } /// Normalized runtime configuration for the minimal OpenAI-compatible server. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct Config { /// Frontend-to-engine transport setup. pub transport_mode: TransportMode, @@ -61,6 +62,8 @@ pub struct Config { pub chat_template_content_format: ChatTemplateContentFormatOption, /// Log a summary line for each completed request. pub enable_log_requests: bool, + /// When `true`, set `X-Request-Id` on every HTTP response. + pub enable_request_id_headers: bool, /// When `true`, suppress periodic stats logging (throughput, queue depth, /// cache usage). pub disable_log_stats: bool, diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index ed21dd3d339..0246064b48d 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -91,6 +91,7 @@ pub fn to_text_request( cache_salt: kv.map(|k| &k.cache_salt).filter(|s| !s.is_empty()).cloned(), add_special_tokens: true, data_parallel_rank: None, + lora_request: None, }) } diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 2b684287ba2..8d779da132f 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -4,15 +4,17 @@ mod config; mod error; mod grpc; mod listener; +mod lora; mod middleware; mod routes; +mod server_info; mod state; mod utils; use std::sync::{Arc, OnceLock}; use anyhow::{Context as _, Result}; -use axum::serve::ListenerExt as _; +use axum::{Router, serve::ListenerExt as _}; pub use config::{Config, CoordinatorMode, HttpListenerMode}; use tokio::net::TcpListener; use tokio::time::{Instant, sleep_until}; @@ -29,6 +31,7 @@ use vllm_text::TextLlm; use crate::listener::Listener; use crate::routes::build_router; +use crate::server_info::ServerInfoSnapshot; use crate::state::AppState; /// Build the shared application state for one configured model and one engine @@ -85,7 +88,10 @@ async fn build_state(config: &Config) -> Result> { }; Ok(Arc::new( - AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests), + AppState::new(served_model_names, chat) + .with_log_requests(config.enable_log_requests) + .with_request_id_headers(config.enable_request_id_headers) + .with_server_info(ServerInfoSnapshot::from_config(config)), )) } @@ -95,6 +101,21 @@ async fn build_state(config: &Config) -> Result> { /// The server owns one `vllm-chat` facade, which in turn owns the lower /// `vllm-text` and `vllm-llm` layers, and shuts them down before returning. pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> { + serve_with_router_extension(config, shutdown, |router| router).await +} + +/// Run the OpenAI-compatible HTTP server with an opt-in router extension. +/// +/// The extension receives the finalized vLLM router and can merge additional +/// routes before the server starts accepting requests. +pub async fn serve_with_router_extension( + config: Config, + shutdown: CancellationToken, + extend_router: F, +) -> Result<()> +where + F: FnOnce(Router) -> Router, +{ config.validate().context("invalid OpenAI frontend configuration")?; // Also check shutdown during the (potentially long) startup handshake. @@ -107,7 +128,7 @@ pub async fn serve(config: Config, shutdown: CancellationToken) -> Result<()> { .context("failed to bind listener for OpenAI server")?; let bind_address = listener.local_addr()?; let model = state.primary_model_name().to_owned(); - let app = build_router(state.clone()); + let app = extend_router(build_router(state.clone())); // Optionally bind the gRPC Generate server on a separate port. Bind // synchronously here so bind errors (port in use, permission denied, ...) diff --git a/rust/src/server/src/lora.rs b/rust/src/server/src/lora.rs new file mode 100644 index 00000000000..d58a61df862 --- /dev/null +++ b/rust/src/server/src/lora.rs @@ -0,0 +1,168 @@ +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +use tokio::sync::{Mutex, RwLock}; +use vllm_engine_core_client::EngineCoreClient; +use vllm_engine_core_client::protocol::lora::LoraRequest; + +/// Snapshot of the currently served model names plus the requested LoRA, if +/// the model name resolves to a dynamic adapter. +#[derive(Debug, Clone)] +pub(crate) struct LoraModelResolution { + pub model_names: Vec, + pub lora_request: Option, +} + +/// Runtime registry for dynamically loaded LoRA adapters. +pub(crate) struct LoraManager { + /// Dynamically loaded LoRA adapters keyed by public model name. + requests: RwLock>, + /// Monotonic adapter id allocator. LoRA ids are one-indexed. + id_counter: AtomicU64, + /// Serialize dynamic LoRA registry updates around engine utility calls. + update_lock: Mutex<()>, +} + +#[derive(Debug)] +pub(crate) enum LoadLoraError { + AlreadyLoaded { lora_name: String }, + BaseModelName { lora_name: String }, + Engine(vllm_engine_core_client::Error), + NotLoaded { lora_name: String }, +} + +#[derive(Debug)] +pub(crate) enum UnloadLoraError { + NotFound { + lora_name: String, + }, + IntIdMismatch { + lora_name: String, + expected: u64, + actual: u64, + }, + Engine(vllm_engine_core_client::Error), + NotRemoved { + lora_name: String, + lora_int_id: u64, + }, +} + +impl LoraManager { + pub fn new() -> Self { + Self { + requests: RwLock::new(BTreeMap::new()), + id_counter: AtomicU64::new(0), + update_lock: Mutex::new(()), + } + } + + /// Return base served model names plus dynamically loaded LoRA adapter + /// names. + pub async fn served_model_names(&self, base_model_names: &[String]) -> Vec { + let mut names = base_model_names.to_vec(); + names.extend(self.requests.read().await.keys().cloned()); + names + } + + /// Resolve the requested model against one consistent LoRA registry + /// snapshot. + pub async fn resolve_model( + &self, + base_model_names: &[String], + model_name: Option<&str>, + ) -> LoraModelResolution { + let requests = self.requests.read().await; + let mut model_names = base_model_names.to_vec(); + model_names.extend(requests.keys().cloned()); + let lora_request = model_name.and_then(|name| requests.get(name).cloned()); + + LoraModelResolution { + model_names, + lora_request, + } + } + + /// Load one dynamic LoRA adapter and register it as a public model name. + pub async fn load_lora( + &self, + engine_core_client: &EngineCoreClient, + base_model_names: &[String], + lora_name: String, + lora_path: String, + load_inplace: bool, + is_3d_lora_weight: bool, + ) -> Result { + let _guard = self.update_lock.lock().await; + if base_model_names.iter().any(|name| name == &lora_name) { + return Err(LoadLoraError::BaseModelName { lora_name }); + } + if !load_inplace && self.requests.read().await.contains_key(&lora_name) { + return Err(LoadLoraError::AlreadyLoaded { lora_name }); + } + + let lora_int_id = self + .requests + .read() + .await + .get(&lora_name) + .map(|request| request.lora_int_id) + .unwrap_or_else(|| self.id_counter.fetch_add(1, Ordering::Relaxed) + 1); + let lora_request = LoraRequest::new( + lora_name.clone(), + lora_int_id, + lora_path, + load_inplace, + is_3d_lora_weight, + ); + + let loaded = engine_core_client + .add_lora(&lora_request) + .await + .map_err(LoadLoraError::Engine)?; + if !loaded { + return Err(LoadLoraError::NotLoaded { lora_name }); + } + self.requests.write().await.insert(lora_name, lora_request.clone()); + Ok(lora_request) + } + + /// Remove one dynamic LoRA adapter from the engine and public model + /// registry. + pub async fn unload_lora( + &self, + engine_core_client: &EngineCoreClient, + lora_name: &str, + requested_lora_int_id: Option, + ) -> Result { + let _guard = self.update_lock.lock().await; + let lora_request = self.requests.read().await.get(lora_name).cloned().ok_or_else(|| { + UnloadLoraError::NotFound { + lora_name: lora_name.to_string(), + } + })?; + + if let Some(actual) = requested_lora_int_id + && actual != lora_request.lora_int_id + { + return Err(UnloadLoraError::IntIdMismatch { + lora_name: lora_name.to_string(), + expected: lora_request.lora_int_id, + actual, + }); + } + + let removed = engine_core_client + .remove_lora(lora_request.lora_int_id) + .await + .map_err(UnloadLoraError::Engine)?; + if !removed { + return Err(UnloadLoraError::NotRemoved { + lora_name: lora_request.lora_name, + lora_int_id: lora_request.lora_int_id, + }); + } + + Ok(self.requests.write().await.remove(lora_name).unwrap_or(lora_request)) + } +} diff --git a/rust/src/server/src/middleware/mod.rs b/rust/src/server/src/middleware/mod.rs index acb3dd1fdb7..1f9647c4efa 100644 --- a/rust/src/server/src/middleware/mod.rs +++ b/rust/src/server/src/middleware/mod.rs @@ -1,5 +1,7 @@ mod load; mod metrics; +mod request_id; pub use load::track_server_load; pub use metrics::track_http_metrics; +pub use request_id::set_request_id_header; diff --git a/rust/src/server/src/middleware/request_id.rs b/rust/src/server/src/middleware/request_id.rs new file mode 100644 index 00000000000..f96b483a165 --- /dev/null +++ b/rust/src/server/src/middleware/request_id.rs @@ -0,0 +1,24 @@ +use axum::extract::Request; +use axum::http::HeaderValue; +use axum::http::header::HeaderName; +use axum::middleware::Next; +use axum::response::Response; +use uuid::Uuid; + +const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id"); + +/// Echo the request's `X-Request-Id` on the response, or generate a fresh +/// `uuid4` hex if the request did not provide one. +/// +/// Original Python: +/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`. +pub async fn set_request_id_header(req: Request, next: Next) -> Response { + let incoming = req.headers().get(&X_REQUEST_ID).cloned(); + let mut response = next.run(req).await; + let value = incoming.unwrap_or_else(|| { + HeaderValue::from_str(&Uuid::new_v4().simple().to_string()) + .expect("uuid hex is valid header value") + }); + response.headers_mut().insert(X_REQUEST_ID, value); + response +} diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index b9549e7144f..a0473c783a0 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -3,8 +3,10 @@ mod collective_rpc; mod health; mod inference; mod load; +mod lora; mod metrics; pub(crate) mod openai; +mod server_info; mod sleep; mod version; @@ -25,12 +27,40 @@ fn server_dev_mode_enabled() -> bool { .is_some_and(|value| value != 0) } -/// Build the minimal OpenAI-compatible router for one configured model. -pub fn build_router(state: Arc) -> Router { - build_router_with_dev_mode(state, server_dev_mode_enabled()) +fn runtime_lora_updating_enabled() -> bool { + std::env::var("VLLM_ALLOW_RUNTIME_LORA_UPDATING") + .ok() + .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true")) } +/// Build the minimal OpenAI-compatible router for one configured model. +pub fn build_router(state: Arc) -> Router { + build_router_with_options( + state, + server_dev_mode_enabled(), + runtime_lora_updating_enabled(), + ) +} + +#[cfg(test)] fn build_router_with_dev_mode(state: Arc, dev_mode_enabled: bool) -> Router { + build_router_with_dev_mode_and_lora(state, dev_mode_enabled, false) +} + +#[cfg(test)] +fn build_router_with_dev_mode_and_lora( + state: Arc, + dev_mode_enabled: bool, + runtime_lora_updating_enabled: bool, +) -> Router { + build_router_with_options(state, dev_mode_enabled, runtime_lora_updating_enabled) +} + +fn build_router_with_options( + state: Arc, + dev_mode_enabled: bool, + runtime_lora_updating_enabled: bool, +) -> Router { let mut router = Router::new() // Health & monitoring .route("/health", get(health::health)) @@ -44,6 +74,12 @@ fn build_router_with_dev_mode(state: Arc, dev_mode_enabled: bool) -> R // vLLM specific inference endpoints .route("/inference/v1/generate", post(inference::generate)); + if runtime_lora_updating_enabled { + router = router + .route("/v1/load_lora_adapter", post(lora::load_lora_adapter)) + .route("/v1/unload_lora_adapter", post(lora::unload_lora_adapter)); + } + if dev_mode_enabled { // Development-only router = router @@ -54,13 +90,21 @@ fn build_router_with_dev_mode(state: Arc, dev_mode_enabled: bool) -> R .route("/sleep", post(sleep::sleep)) .route("/wake_up", post(sleep::wake_up)) .route("/is_sleeping", get(sleep::is_sleeping)) + .route("/server_info", get(server_info::server_info)) } - router + let enable_request_id_headers = state.enable_request_id_headers; + let mut router = router .with_state(state.clone()) .layer(from_fn_with_state(state, middleware::track_server_load)) .layer(from_fn(middleware::track_http_metrics)) - .layer(TraceLayer::new_for_http()) + .layer(TraceLayer::new_for_http()); + + if enable_request_id_headers { + router = router.layer(from_fn(middleware::set_request_id_header)); + } + + router } #[cfg(test)] diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index b256b7721f9..f15f757c09a 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -42,8 +42,8 @@ pub async fn generate( ValidatedJson(body): ValidatedJson, ) -> Response { let request_context = resolve_request_context(&headers, body.request_id.as_deref()); - let prepared = match prepare_generate_request(body, state.served_model_names(), request_context) - { + let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await; + let prepared = match prepare_generate_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, Err(error) => return error.into_response(), }; diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 844a606bcfd..f87ff403a7b 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -3,6 +3,7 @@ use vllm_text::{Prompt, TextDecodeOptions, TextRequest}; use super::types::GenerateRequest; use super::validate; use crate::error::ApiError; +use crate::lora::LoraModelResolution; use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params}; /// Lowered generate request plus the response request ID. @@ -21,10 +22,10 @@ pub struct PreparedRequest { /// text-generation format. pub fn prepare_generate_request( request: GenerateRequest, - served_model_names: &[String], + lora_resolution: &LoraModelResolution, ctx: ResolvedRequestContext, ) -> Result { - validate::validate_request_compat(&request, served_model_names)?; + validate::validate_request_compat(&request, &lora_resolution.model_names)?; let stream = request.stream; let include_usage = request @@ -57,6 +58,7 @@ pub fn prepare_generate_request( cache_salt: request.cache_salt, add_special_tokens: false, data_parallel_rank: ctx.data_parallel_rank, + lora_request: lora_resolution.lora_request.clone(), }; Ok(PreparedRequest { @@ -76,9 +78,17 @@ mod tests { use vllm_text::Prompt; use super::prepare_generate_request; + use crate::lora::LoraModelResolution; use crate::routes::inference::generate::types::GenerateRequest; use crate::utils::ResolvedRequestContext; + fn served(names: &[&str]) -> LoraModelResolution { + LoraModelResolution { + model_names: names.iter().map(|s| s.to_string()).collect(), + lora_request: None, + } + } + #[test] fn prepare_generate_request_maps_token_prompt_and_sampling_params() { let request: GenerateRequest = serde_json::from_value(json!({ @@ -100,7 +110,7 @@ mod tests { let prepared = prepare_generate_request( request, - &["Qwen/Qwen1.5-0.5B-Chat".to_string()], + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), ResolvedRequestContext::default(), ) .expect("prepare"); @@ -143,7 +153,7 @@ mod tests { let prepared = prepare_generate_request( request, - &["Qwen/Qwen1.5-0.5B-Chat".to_string()], + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), ResolvedRequestContext::default(), ) .expect("prepare"); diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs new file mode 100644 index 00000000000..99f1c0fe320 --- /dev/null +++ b/rust/src/server/src/routes/lora.rs @@ -0,0 +1,296 @@ +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; + +use axum::extract::State; +use serde::Deserialize; +use thiserror_ext::AsReport; +use validator::Validate; + +use crate::error::ApiError; +use crate::lora::{LoadLoraError, UnloadLoraError}; +use crate::routes::openai::utils::types::Normalizable; +use crate::routes::openai::utils::validated_json::ValidatedJson; +use crate::state::AppState; + +const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES"; + +#[derive(Debug, Deserialize, Validate)] +pub(crate) struct LoadLoraAdapterRequest { + lora_name: String, + lora_path: String, + #[serde(default)] + load_inplace: bool, + #[serde(default)] + is_3d_lora_weight: bool, +} + +impl Normalizable for LoadLoraAdapterRequest {} + +#[derive(Debug, Deserialize, Validate)] +pub(crate) struct UnloadLoraAdapterRequest { + lora_name: String, + #[serde(default)] + lora_int_id: Option, +} + +impl Normalizable for UnloadLoraAdapterRequest {} + +fn runtime_lora_allowed_path_prefixes() -> Option> { + let prefixes = std::env::var_os(RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV)?; + let prefixes: Vec<_> = std::env::split_paths(&prefixes) + .filter(|path| !path.as_os_str().is_empty()) + .collect(); + (!prefixes.is_empty()).then_some(prefixes) +} + +fn looks_like_local_lora_path(lora_path: &str) -> bool { + let path = Path::new(lora_path); + path.is_absolute() + || lora_path.starts_with('~') + || lora_path.starts_with('.') + || path.components().any(|component| matches!(component, Component::ParentDir)) +} + +fn validate_lora_path_access( + lora_path: &str, + allowed_prefixes: Option<&[PathBuf]>, +) -> Result, ApiError> { + let path = Path::new(lora_path); + if !looks_like_local_lora_path(lora_path) && !path.exists() { + return Ok(None); + } + + let Some(allowed_prefixes) = allowed_prefixes else { + return Err(ApiError::invalid_request( + format!( + "Local LoRA adapter paths require {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} to be configured." + ), + Some("lora_path"), + )); + }; + + if !path.is_absolute() { + return Err(ApiError::invalid_request( + format!( + "Local LoRA adapter paths must be absolute and under one of the prefixes configured by {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV}." + ), + Some("lora_path"), + )); + } + + let canonical_path = path.canonicalize().map_err(|_| { + ApiError::invalid_request( + "Local LoRA adapter path must exist and be accessible.".to_string(), + Some("lora_path"), + ) + })?; + let canonical_prefixes = allowed_prefixes + .iter() + .map(|prefix| { + prefix.canonicalize().map_err(|_| { + ApiError::server_error(format!( + "configured {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} path prefix must exist and be accessible" + )) + }) + }) + .collect::, _>>()?; + + if !canonical_prefixes.iter().any(|prefix| canonical_path.starts_with(prefix)) { + return Err(ApiError::invalid_request( + "Local LoRA adapter path is outside the configured allowed prefixes.".to_string(), + Some("lora_path"), + )); + } + + Ok(Some(canonical_path.to_string_lossy().into_owned())) +} + +/// Dynamically load one LoRA adapter and expose it as an OpenAI model id. +pub async fn load_lora_adapter( + State(state): State>, + ValidatedJson(request): ValidatedJson, +) -> Result { + if request.lora_name.is_empty() || request.lora_path.is_empty() { + return Err(ApiError::invalid_request( + "Both 'lora_name' and 'lora_path' must be provided.".to_string(), + None, + )); + } + let allowed_prefixes = runtime_lora_allowed_path_prefixes(); + let lora_path = validate_lora_path_access(&request.lora_path, allowed_prefixes.as_deref())? + .unwrap_or(request.lora_path); + + let lora_name = request.lora_name; + state + .load_lora( + lora_name.clone(), + lora_path, + request.load_inplace, + request.is_3d_lora_weight, + ) + .await + .map_err(|error| match error { + LoadLoraError::AlreadyLoaded { lora_name } => ApiError::invalid_request( + format!( + "The lora adapter '{lora_name}' has already been loaded. If you want to load the adapter in place, set 'load_inplace' to true." + ), + Some("lora_name"), + ), + LoadLoraError::BaseModelName { lora_name } => ApiError::invalid_request( + format!("The lora adapter name '{lora_name}' conflicts with a served base model."), + Some("lora_name"), + ), + LoadLoraError::Engine(error) => ApiError::server_error(format!( + "failed to load LoRA adapter '{lora_name}': {}", + error.to_report_string() + )), + LoadLoraError::NotLoaded { lora_name } => ApiError::server_error(format!( + "failed to load LoRA adapter '{lora_name}': engine rejected the adapter" + )), + })?; + + Ok(format!( + "Success: LoRA adapter '{lora_name}' added successfully." + )) +} + +/// Remove one LoRA adapter from the engine and frontend registry. +pub async fn unload_lora_adapter( + State(state): State>, + ValidatedJson(request): ValidatedJson, +) -> Result { + if request.lora_name.is_empty() { + return Err(ApiError::invalid_request( + "'lora_name' needs to be provided to unload a LoRA adapter.".to_string(), + Some("lora_name"), + )); + } + + let lora_request = state + .unload_lora(&request.lora_name, request.lora_int_id) + .await + .map_err(|error| match error { + UnloadLoraError::NotFound { lora_name } => ApiError::model_not_found(lora_name), + UnloadLoraError::IntIdMismatch { + lora_name, + expected, + actual, + } => ApiError::invalid_request( + format!( + "The requested lora_int_id {actual} does not match loaded adapter '{lora_name}' with id {expected}." + ), + Some("lora_int_id"), + ), + UnloadLoraError::Engine(error) => ApiError::server_error(format!( + "failed to unload LoRA adapter '{}': {}", + request.lora_name, + error.to_report_string() + )), + UnloadLoraError::NotRemoved { + lora_name, + lora_int_id, + } => ApiError::server_error(format!( + "failed to unload LoRA adapter '{lora_name}' with id {lora_int_id}" + )), + })?; + + Ok(format!( + "Success: LoRA adapter '{}' removed successfully.", + lora_request.lora_name + )) +} + +#[cfg(test)] +mod tests { + use std::fs; + use std::path::PathBuf; + use std::time::{SystemTime, UNIX_EPOCH}; + + use super::validate_lora_path_access; + + fn temp_lora_dir(test_name: &str) -> PathBuf { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock should be after unix epoch") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "vllm-lora-{test_name}-{}-{suffix}", + std::process::id() + )); + fs::create_dir_all(&path).expect("create temp lora dir"); + path + } + + #[test] + fn lora_path_allows_hf_repo_ids_without_prefixes() { + assert_eq!( + validate_lora_path_access("org/adapter-a", None).expect("hf repo id should be allowed"), + None + ); + } + + #[test] + fn lora_path_rejects_local_paths_without_prefixes() { + assert!(validate_lora_path_access("/tmp/adapter-a", None).is_err()); + assert!(validate_lora_path_access("./adapter-a", None).is_err()); + assert!(validate_lora_path_access("~/adapter-a", None).is_err()); + assert!(validate_lora_path_access("subdir/../../../etc/sensitive", None).is_err()); + } + + #[test] + fn lora_path_rejects_existing_bare_relative_paths_without_prefixes() { + let root = + PathBuf::from("target").join(format!("vllm-lora-relative-{}", std::process::id())); + let adapter = root.join("adapter-a"); + fs::create_dir_all(&adapter).expect("create relative adapter dir"); + + assert!( + validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), None).is_err() + ); + + fs::remove_dir_all(root).ok(); + } + + #[test] + fn lora_path_allows_absolute_paths_under_configured_prefixes() { + let root = temp_lora_dir("allowed-prefix"); + let allowed = root.join("allowed"); + let adapter = allowed.join("adapter-a"); + fs::create_dir_all(&adapter).expect("create adapter dir"); + + let prefixes = [allowed]; + let resolved = + validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .expect("path under configured prefix should be allowed"); + assert_eq!( + resolved.as_deref(), + Some( + adapter + .canonicalize() + .expect("canonical adapter") + .to_str() + .expect("utf-8 temp path") + ) + ); + + fs::remove_dir_all(root).ok(); + } + + #[test] + fn lora_path_rejects_parent_escape_from_configured_prefixes() { + let root = temp_lora_dir("parent-escape"); + let allowed = root.join("allowed"); + let private_adapter = root.join("private").join("adapter-a"); + fs::create_dir_all(&allowed).expect("create allowed dir"); + fs::create_dir_all(&private_adapter).expect("create private adapter dir"); + + let escaped = allowed.join("../private/adapter-a"); + let prefixes = [allowed]; + assert!( + validate_lora_path_access(escaped.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .is_err() + ); + + fs::remove_dir_all(root).ok(); + } +} diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index c0894bb70c9..543a7e806c4 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -49,8 +49,9 @@ pub async fn chat_completions( ) -> Response { let stream = body.stream; let request_context = resolve_request_context(&headers, body.request_id.as_deref()); + let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; - let prepared = match prepare_chat_request(body, state.served_model_names(), request_context) { + let prepared = match prepare_chat_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, Err(error) => return error.into_response(), }; 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 a7884b11e52..2701bef809c 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -8,6 +8,7 @@ use vllm_chat::{ use super::types::ChatCompletionRequest; use super::validate; use crate::error::{ApiError, bail_invalid_request}; +use crate::lora::LoraModelResolution; use crate::routes::openai::utils::structured_outputs::convert_from_response_format; use crate::routes::openai::utils::types::{ ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, @@ -41,16 +42,21 @@ pub struct PreparedRequest { /// Validate and lower one OpenAI chat completion request into the internal chat /// format. /// -/// `served_model_names` must be non-empty; the first entry is used as the -/// `model` field in responses. +/// `lora_resolution.model_names` must be non-empty; the first entry is used as +/// the base `model` field in responses when no LoRA adapter is selected. pub(crate) fn prepare_chat_request( request: ChatCompletionRequest, - served_model_names: &[String], + lora_resolution: &LoraModelResolution, ctx: ResolvedRequestContext, ) -> Result { - validate::validate_request_compat(&request, served_model_names)?; + validate::validate_request_compat(&request, &lora_resolution.model_names)?; let request_id = format!("chatcmpl-{}", ctx.request_id); + let response_model = lora_resolution + .lora_request + .as_ref() + .map(|request| request.lora_name.clone()) + .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default()); let echo = request .echo .then(|| extract_last_assistant_content(&request.messages)) @@ -131,11 +137,12 @@ pub(crate) fn prepare_chat_request( cache_salt: request.cache_salt, add_special_tokens: request.add_special_tokens, data_parallel_rank: ctx.data_parallel_rank, + lora_request: lora_resolution.lora_request.clone(), }; Ok(PreparedRequest { request_id, - response_model: served_model_names.first().cloned().unwrap_or_default(), + response_model, include_usage, requested_logprobs, include_prompt_logprobs, @@ -352,6 +359,7 @@ mod tests { use vllm_text::output::TextDecodeOptions; use super::prepare_chat_request; + use crate::lora::LoraModelResolution; use crate::routes::openai::chat_completions::types::{ AssistantRole, ChatCompletionMessage, ChatCompletionRequest, }; @@ -365,8 +373,11 @@ mod tests { resolve_request_context(headers, request_id) } - fn served(names: &[&str]) -> Vec { - names.iter().map(|s| s.to_string()).collect() + fn served(names: &[&str]) -> LoraModelResolution { + LoraModelResolution { + model_names: names.iter().map(|s| s.to_string()).collect(), + lora_request: None, + } } fn base_request() -> ChatCompletionRequest { diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 33813e67687..9eda8b9d2a5 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -44,12 +44,12 @@ pub async fn completions( let stream = body.stream; let logprobs = body.logprobs; let request_context = resolve_request_context(&headers, body.request_id.as_deref()); + let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; - let prepared = - match prepare_completion_request(body, state.served_model_names(), request_context) { - Ok(prepared) => prepared, - Err(error) => return error.into_response(), - }; + let prepared = match prepare_completion_request(body, &lora_resolution, request_context) { + Ok(prepared) => prepared, + Err(error) => return error.into_response(), + }; let request_span = tracing::info_span!( "completions", request_id = %prepared.request_id, diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 066c4c046f4..2d4ff089397 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -2,6 +2,7 @@ use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest}; use super::types::CompletionRequest; use crate::error::ApiError; +use crate::lora::LoraModelResolution; use crate::routes::openai::completions::validate; use crate::routes::openai::utils::structured_outputs::convert_from_response_format_value; use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer_params}; @@ -30,16 +31,21 @@ pub struct PreparedRequest { /// Validate and lower one OpenAI completions request into the internal /// text-generation format. /// -/// `served_model_names` must be non-empty; the first entry is used as the -/// `model` field in responses. +/// `lora_resolution.model_names` must be non-empty; the first entry is used as +/// the base `model` field in responses when no LoRA adapter is selected. pub(crate) fn prepare_completion_request( request: CompletionRequest, - served_model_names: &[String], + lora_resolution: &LoraModelResolution, ctx: ResolvedRequestContext, ) -> Result { - validate::validate_request_compat(&request, served_model_names)?; + validate::validate_request_compat(&request, &lora_resolution.model_names)?; let request_id = format!("cmpl-{}", ctx.request_id); + let response_model = lora_resolution + .lora_request + .as_ref() + .map(|request| request.lora_name.clone()) + .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default()); let logprobs = match request.logprobs { Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| { @@ -104,11 +110,12 @@ pub(crate) fn prepare_completion_request( cache_salt: request.cache_salt, add_special_tokens: request.add_special_tokens, data_parallel_rank: ctx.data_parallel_rank, + lora_request: lora_resolution.lora_request.clone(), }; Ok(PreparedRequest { request_id, - response_model: served_model_names.first().cloned().unwrap_or_default(), + response_model, include_usage, text_request, echo, @@ -124,6 +131,7 @@ mod tests { use vllm_text::Prompt; use super::prepare_completion_request; + use crate::lora::LoraModelResolution; use crate::routes::openai::completions::types::CompletionRequest; use crate::utils::{ResolvedRequestContext, resolve_request_context}; @@ -131,8 +139,11 @@ mod tests { resolve_request_context(headers, request_id) } - fn served(names: &[&str]) -> Vec { - names.iter().map(|s| s.to_string()).collect() + fn served(names: &[&str]) -> LoraModelResolution { + LoraModelResolution { + model_names: names.iter().map(|s| s.to_string()).collect(), + lora_request: None, + } } fn base_request_json() -> serde_json::Value { diff --git a/rust/src/server/src/routes/openai/models.rs b/rust/src/server/src/routes/openai/models.rs index 42e3098fc9e..42efd259e1b 100644 --- a/rust/src/server/src/routes/openai/models.rs +++ b/rust/src/server/src/routes/openai/models.rs @@ -8,13 +8,13 @@ use crate::state::AppState; /// Return all configured served model names in OpenAI `list models` format. pub async fn list_models(State(state): State>) -> Json { + let model_names = state.served_model_names_with_loras().await; Json(ListModelsResponse { object: "list".to_string(), - data: state - .served_model_names() - .iter() + data: model_names + .into_iter() .map(|name| ModelObject { - id: name.clone(), + id: name, object: "model".to_string(), created: 0, owned_by: "vllm-frontend-rs".to_string(), diff --git a/rust/src/server/src/routes/server_info.rs b/rust/src/server/src/routes/server_info.rs new file mode 100644 index 00000000000..aefb17a25fa --- /dev/null +++ b/rust/src/server/src/routes/server_info.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use axum::Json; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use serde::Deserialize; + +use crate::server_info::ServerInfoConfigFormat; +use crate::state::AppState; + +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(rename_all = "lowercase")] +enum ConfigFormat { + Text, + Json, +} + +impl From for ServerInfoConfigFormat { + fn from(value: ConfigFormat) -> Self { + match value { + ConfigFormat::Text => Self::Text, + ConfigFormat::Json => Self::Json, + } + } +} + +fn default_config_format() -> ConfigFormat { + ConfigFormat::Text +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ServerInfoParams { + #[serde(default = "default_config_format")] + config_format: ConfigFormat, +} + +/// Get server configuration and environment metadata. +pub async fn server_info( + State(state): State>, + Query(params): Query, +) -> Response { + match state.server_info_response(params.config_format.into()) { + Some(response) => Json(response).into_response(), + None => StatusCode::NOT_FOUND.into_response(), + } +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index b1a4f0705fd..a3e437e0480 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -43,7 +43,8 @@ use vllm_text::{Prompt, TextBackend}; use zeromq::prelude::{SocketRecv, SocketSend}; use zeromq::{DealerSocket, PushSocket, ZmqMessage}; -use super::{build_router, build_router_with_dev_mode}; +use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora}; +use crate::lora::LoraModelResolution; use crate::routes::openai::chat_completions::convert::prepare_chat_request; use crate::state::AppState; @@ -141,6 +142,13 @@ fn default_stream_output_specs() -> Vec<(Vec, Option Vec<&str> { text.lines().filter_map(|line| line.strip_prefix("data: ")).collect() } @@ -734,16 +742,37 @@ async fn test_chat_with_engine_outputs( } async fn test_app() -> axum::Router { + test_app_with_dev_mode(false).await +} + +async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router { let (chat, _engine_task) = test_models_with_engine_outputs_and_backend( b"engine-openai", default_stream_output_specs(), Arc::new(FakeChatBackend::new()), ) .await; - build_router(Arc::new(AppState::new( - vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], - chat, - ))) + build_router_with_dev_mode( + Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + )), + dev_mode_enabled, + ) +} + +async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) { + let (chat, engine_task) = test_models_with_engine_outputs_and_backend( + b"engine-openai-request-id", + default_stream_output_specs(), + Arc::new(FakeChatBackend::new()), + ) + .await; + let app = build_router(Arc::new( + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat) + .with_request_id_headers(true), + )); + (app, engine_task) } async fn test_health_app_with_engine_script( @@ -808,12 +837,13 @@ where let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); ( - build_router_with_dev_mode( + build_router_with_dev_mode_and_lora( Arc::new(AppState::new( vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat, )), true, + true, ), engine_task, ) @@ -947,6 +977,18 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) { (status, body) } +async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response { + let mut builder = Request::builder().method("GET").uri("/health"); + if let Some(request_id) = request_id { + builder = builder.header("X-Request-Id", request_id); + } + + app.clone() + .call(builder.body(Body::empty()).expect("build request")) + .await + .expect("call app") +} + fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option { rendered.lines().find_map(|line| { let rest = line.strip_prefix(metric)?; @@ -994,6 +1036,43 @@ async fn list_models_returns_configured_model() { assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn request_id_header_is_absent_by_default() { + let app = test_app().await; + let response = health_response(&app, None).await; + + assert_eq!(response.status(), StatusCode::OK); + assert!(!response.headers().contains_key("x-request-id")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn request_id_header_generates_uuid_hex_when_enabled() { + let (app, _engine_task) = test_app_with_request_id_headers().await; + let response = health_response(&app, None).await; + + assert_eq!(response.status(), StatusCode::OK); + let request_id = response + .headers() + .get("x-request-id") + .expect("x-request-id header") + .to_str() + .expect("header is ascii"); + assert_eq!(request_id.len(), 32); + assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit())); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn request_id_header_echoes_incoming_header_when_enabled() { + let (app, _engine_task) = test_app_with_request_id_headers().await; + let response = health_response(&app, Some("req-123")).await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn version_returns_engine_vllm_version() { @@ -1015,6 +1094,401 @@ async fn version_returns_engine_vllm_version() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn server_info_endpoint_is_dev_mode_only() { + let mut app = test_app().await; + let response = app + .call( + Request::builder() + .uri("/server_info") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn load_lora_adapter_registers_model_and_forwards_lora_request() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + assert_eq!(array[2], Value::from("add_lora")); + + let args = array[3].as_array().expect("utility args"); + let lora = args[0].as_array().expect("lora request tuple"); + assert_eq!(lora[0], Value::from("adapter-a")); + assert_eq!(lora[1], Value::from(1)); + assert_eq!(lora[2], Value::from("org/adapter-a")); + + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; + + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode engine request"); + assert_adapter_a_lora_request(&request); + + send_outputs( + push, + engine_outputs_for_request(&request.request_id, default_stream_output_specs()), + ) + .await; + + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode engine request"); + assert_adapter_a_lora_request(&request); + + send_outputs( + push, + engine_outputs_for_request(&request.request_id, default_stream_output_specs()), + ) + .await; + + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode engine request"); + assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..])); + assert_adapter_a_lora_request(&request); + + send_outputs( + push, + engine_outputs_for_request(&request.request_id, default_stream_output_specs()), + ) + .await; + + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + assert_eq!(array[2], Value::from("remove_lora")); + + let args = array[3].as_array().expect("utility args"); + assert_eq!(args[0], Value::from(1)); + + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; + }) + }) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/load_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_path": "org/adapter-a" + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request")) + .await + .expect("call app"); + let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["data"][1]["id"], "adapter-a"); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "adapter-a", + "prompt": "hello", + "max_tokens": 2 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "adapter-a", + "stream": false, + "messages": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "adapter-a", + "token_ids": [11, 22], + "stream": false, + "sampling_params": { + "max_tokens": 2 + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/unload_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_int_id": 1 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request")) + .await + .expect("call app"); + let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["data"].as_array().expect("model data").len(), 1); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "adapter-a", + "prompt": "hello", + "max_tokens": 2 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn server_info_endpoint_returns_not_found_without_snapshot() { + let mut app = test_app_with_dev_mode(true).await; + let response = app + .call( + Request::builder() + .uri("/server_info") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unload_lora_adapter_rejects_mismatched_lora_int_id() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + assert_eq!(array[2], Value::from("add_lora")); + + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; + }) + }) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/load_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_path": "org/adapter-a" + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/unload_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_int_id": 99 + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request")) + .await + .expect("call app"); + let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["data"][1]["id"], "adapter-a"); + + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn load_lora_adapter_rejects_engine_false_result() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + assert_eq!(array[2], Value::from("add_lora")); + + send_outputs(push, utility_outputs(call_id, utility_result_value(false))).await; + }) + }) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/load_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_path": "org/adapter-a" + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).expect("build request")) + .await + .expect("call app"); + let body = to_bytes(models.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["data"].as_array().expect("model data").len(), 1); + + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn load_lora_adapter_rejects_base_model_name_collision() { + let (mut app, engine_task) = + test_admin_app_with_engine_script(|_, _| boxed_test_future(async move {})).await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/v1/load_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "Qwen/Qwen1.5-0.5B-Chat", + "lora_path": "org/adapter-a" + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + + drop(app); + engine_task.finish().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn http_metrics_record_list_models_requests() { @@ -2937,7 +3411,10 @@ async fn prepared_openai_request_streams_text_events() { "messages": [{"role": "user", "content": "hello"}] })) .expect("decode request"), - &["Qwen/Qwen1.5-0.5B-Chat".to_string()], + &LoraModelResolution { + model_names: vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + lora_request: None, + }, crate::utils::ResolvedRequestContext::default(), ) .expect("prepare request"); diff --git a/rust/src/server/src/server_info.rs b/rust/src/server/src/server_info.rs new file mode 100644 index 00000000000..cca1b0f3795 --- /dev/null +++ b/rust/src/server/src/server_info.rs @@ -0,0 +1,144 @@ +use std::collections::BTreeMap; + +use serde_json::{Value, json}; + +use crate::config::Config; + +const SENSITIVE_VLLM_ENV_PATTERNS: &[&str] = + &["KEY", "SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "AUTH"]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ServerInfoConfigFormat { + Text, + Json, +} + +/// Snapshot returned by `/server_info`. +#[derive(Debug, Clone)] +pub(crate) struct ServerInfoSnapshot { + vllm_config_text: String, + vllm_config_json: Value, + vllm_env: BTreeMap, + system_env: BTreeMap, +} + +impl ServerInfoSnapshot { + /// Capture the runtime configuration fields available to the Rust frontend. + pub(crate) fn from_config(config: &Config) -> Self { + let vllm_config_json = + serde_json::to_value(config).expect("server info value must serialize"); + + Self { + vllm_config_text: render_config_text(&vllm_config_json), + vllm_config_json, + vllm_env: collect_vllm_env(), + system_env: collect_system_env(), + } + } + + pub(crate) fn response(&self, config_format: ServerInfoConfigFormat) -> Value { + let vllm_config = match config_format { + ServerInfoConfigFormat::Text => Value::String(self.vllm_config_text.clone()), + ServerInfoConfigFormat::Json => self.vllm_config_json.clone(), + }; + + json!({ + "vllm_config": vllm_config, + "vllm_env": self.vllm_env.clone(), + "system_env": self.system_env.clone(), + }) + } +} + +fn render_config_text(config: &Value) -> String { + match config { + Value::Object(fields) => fields + .iter() + .map(|(key, value)| format!("{key}={}", render_config_text_value(value))) + .collect::>() + .join("\n"), + _ => render_config_text_value(config), + } +} + +fn render_config_text_value(value: &Value) -> String { + match value { + Value::Null => "None".to_string(), + Value::String(value) => value.clone(), + _ => value.to_string(), + } +} + +fn collect_vllm_env() -> BTreeMap { + std::env::vars().filter(|(key, _)| is_public_vllm_env_key(key)).collect() +} + +fn is_public_vllm_env_key(key: &str) -> bool { + let key = key.to_ascii_uppercase(); + key.starts_with("VLLM_") + && !SENSITIVE_VLLM_ENV_PATTERNS.iter().any(|pattern| key.contains(pattern)) +} + +fn collect_system_env() -> BTreeMap { + BTreeMap::from([ + ("arch".to_string(), std::env::consts::ARCH.to_string()), + ("family".to_string(), std::env::consts::FAMILY.to_string()), + ("os".to_string(), std::env::consts::OS.to_string()), + ]) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use serde_json::{Value, json}; + + use super::{is_public_vllm_env_key, render_config_text}; + + #[test] + fn render_config_text_formats_config_snapshot() { + let rendered = render_config_text(&json!({ + "model": "test-model", + "served_model_name": ["served-model"], + "chat_template": null, + "enable_log_requests": true, + })); + let lines = rendered.lines().collect::>(); + + assert_eq!( + lines, + BTreeSet::from([ + "chat_template=None", + "enable_log_requests=true", + "model=test-model", + "served_model_name=[\"served-model\"]", + ]) + ); + assert_eq!( + render_config_text(&Value::String("inline".to_string())), + "inline" + ); + assert_eq!(render_config_text(&Value::Null), "None"); + } + + #[test] + fn server_info_env_filter_excludes_sensitive_vllm_keys() { + for key in [ + "VLLM_API_KEY", + "VLLM_AUTH_TOKEN", + "VLLM_SECRET", + "VLLM_PASSWORD", + "VLLM_CREDENTIAL_FILE", + "vllm_token", + ] { + assert!(!is_public_vllm_env_key(key), "{key}"); + } + } + + #[test] + fn server_info_env_filter_includes_public_vllm_keys() { + assert!(is_public_vllm_env_key("VLLM_LOGGING_LEVEL")); + assert!(is_public_vllm_env_key("VLLM_USE_MODELSCOPE")); + assert!(!is_public_vllm_env_key("OTHER_ENV")); + } +} diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 04d37f1a5d4..c73ca04c5d6 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -1,10 +1,16 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use serde_json::Value; use tokio::time::{Duration, Instant, sleep_until}; use tracing::warn; use vllm_chat::ChatLlm; use vllm_engine_core_client::EngineCoreClient; +use vllm_engine_core_client::protocol::lora::LoraRequest; + +use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError}; + +use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot}; const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100); @@ -17,8 +23,14 @@ pub struct AppState { pub chat: ChatLlm, /// Whether to log a summary line for each completed request. pub enable_log_requests: bool, + /// Whether to set X-Request-Id on every HTTP response. + pub enable_request_id_headers: bool, + /// Runtime server information returned by `/server_info`, when available. + server_info: Option, /// Number of in-flight inference requests currently owned by this frontend. server_load: AtomicU64, + /// Dynamic LoRA adapter registry. + lora_manager: LoraManager, } impl AppState { @@ -39,7 +51,10 @@ impl AppState { served_model_names, chat, enable_log_requests: false, + enable_request_id_headers: false, + server_info: None, server_load: AtomicU64::new(0), + lora_manager: LoraManager::new(), } } @@ -49,6 +64,26 @@ impl AppState { self } + /// Enable X-Request-Id response headers. + pub fn with_request_id_headers(mut self, enabled: bool) -> Self { + self.enable_request_id_headers = enabled; + self + } + + /// Attach the runtime server information snapshot used by `/server_info`. + pub(crate) fn with_server_info(mut self, server_info: ServerInfoSnapshot) -> Self { + self.server_info = Some(server_info); + self + } + + /// Build a `/server_info` response payload. + pub(crate) fn server_info_response( + &self, + config_format: ServerInfoConfigFormat, + ) -> Option { + self.server_info.as_ref().map(|server_info| server_info.response(config_format)) + } + /// The primary model name echoed back in API responses (the first served /// name). pub fn primary_model_name(&self) -> &str { @@ -60,6 +95,49 @@ impl AppState { &self.served_model_names } + /// Return base served model names plus dynamically loaded LoRA adapter + /// names. + pub async fn served_model_names_with_loras(&self) -> Vec { + self.lora_manager.served_model_names(&self.served_model_names).await + } + + /// Resolve the requested model against one dynamic LoRA registry snapshot. + pub async fn resolve_model_with_loras(&self, model_name: Option<&str>) -> LoraModelResolution { + self.lora_manager.resolve_model(&self.served_model_names, model_name).await + } + + /// Load one dynamic LoRA adapter and register it as a public model name. + pub async fn load_lora( + &self, + lora_name: String, + lora_path: String, + load_inplace: bool, + is_3d_lora_weight: bool, + ) -> Result { + self.lora_manager + .load_lora( + self.engine_core_client(), + &self.served_model_names, + lora_name, + lora_path, + load_inplace, + is_3d_lora_weight, + ) + .await + } + + /// Remove one dynamic LoRA adapter from the engine and public model + /// registry. + pub async fn unload_lora( + &self, + lora_name: &str, + lora_int_id: Option, + ) -> Result { + self.lora_manager + .unload_lora(self.engine_core_client(), lora_name, lora_int_id) + .await + } + /// Return a reference to the underlying engine core client for utility /// calls. pub(crate) fn engine_core_client(&self) -> &EngineCoreClient { diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 54ffe1ff4b8..d661c99606b 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -40,11 +40,11 @@ pub fn lower_text_request( cache_salt: request.cache_salt.clone(), priority: request.priority, data_parallel_rank: request.data_parallel_rank, + lora_request: request.lora_request.clone(), // Fields below are currently placeholders. arrival_time: None, trace_headers: None, reasoning_ended: None, - lora_request: None, }; Ok(PreparedTextRequest { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 9e2464f14af..1ca8f8a924a 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -4,6 +4,7 @@ use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; use serde_json::Value; use vllm_engine_core_client::protocol::StructuredOutputsParams; +use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; use crate::error::{Error, Result}; @@ -166,6 +167,9 @@ pub struct TextRequest { /// Override data parallel rank. #[serde(default)] pub data_parallel_rank: Option, + /// LoRA adapter selected for this request. + #[serde(default)] + pub lora_request: Option, } impl TextRequest { @@ -182,6 +186,7 @@ impl TextRequest { cache_salt: None, add_special_tokens: false, data_parallel_rank: None, + lora_request: None, } } diff --git a/rust/src/tool-parser/src/deepseek_dsml/mod.rs b/rust/src/tool-parser/src/deepseek_dsml/mod.rs index eba36e33db6..c332037f451 100644 --- a/rust/src/tool-parser/src/deepseek_dsml/mod.rs +++ b/rust/src/tool-parser/src/deepseek_dsml/mod.rs @@ -104,7 +104,7 @@ impl DeepSeekDsmlToolParser { self.tool_parameters.convert_param_with_schema( &name, ¶m.name, - ¶m.value, + param.value, ) }; arguments.insert(param.name, value); diff --git a/rust/src/tool-parser/src/json/internlm2.rs b/rust/src/tool-parser/src/json/internlm2.rs index 5c3b024a936..8284a4d0e1d 100644 --- a/rust/src/tool-parser/src/json/internlm2.rs +++ b/rust/src/tool-parser/src/json/internlm2.rs @@ -34,26 +34,23 @@ const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig { /// This Rust port intentionally diverges from /// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways: /// -/// - **Parallel tool calls are supported.** Python silently drops every -/// `<|action_start|>` block after the first (`current_tool_id > 0` returns -/// an empty delta); this parser emits every well-formed block with -/// incrementing `tool_index`. Models that legitimately emit multiple action -/// blocks therefore produce more tool calls under Rust than under Python. -/// - **End-marker bytes inside JSON string values are preserved.** Python -/// does `action.split("<|action_end|>")[0]` which truncates regardless of -/// JSON context; this parser scans matched braces and quotes so a literal -/// `<|action_end|>` inside an arguments string is forwarded intact. +/// - **Parallel tool calls are supported.** Python silently drops every `<|action_start|>` block +/// after the first (`current_tool_id > 0` returns an empty delta); this parser emits every +/// well-formed block with incrementing `tool_index`. Models that legitimately emit multiple +/// action blocks therefore produce more tool calls under Rust than under Python. +/// - **End-marker bytes inside JSON string values are preserved.** Python does +/// `action.split("<|action_end|>")[0]` which truncates regardless of JSON context; this parser +/// scans matched braces and quotes so a literal `<|action_end|>` inside an arguments string is +/// forwarded intact. /// - **Only whitespace is allowed before the `{`.** Python's non-streaming -/// `action[action.find("{"):]` drops any bytes before the first `{`, but -/// its streaming path has no equivalent and the model format always emits -/// `<|plugin|>{...`; this parser allows only whitespace there, matching the -/// other JSON parsers in this crate. -/// - **Truncated tool calls error rather than silently dropping.** Python's -/// streaming wrapper swallows mid-stream errors with `except Exception: -/// return None` (logging a traceback) while its non-streaming path raises -/// `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool -/// call` error from `finish()`, matching the other JSON parsers and Python's -/// non-streaming behavior. +/// `action[action.find("{"):]` drops any bytes before the first `{`, but its streaming path has +/// no equivalent and the model format always emits `<|plugin|>{...`; this parser allows only +/// whitespace there, matching the other JSON parsers in this crate. +/// - **Truncated tool calls error rather than silently dropping.** Python's streaming wrapper +/// swallows mid-stream errors with `except Exception: return None` (logging a traceback) while +/// its non-streaming path raises `JSONDecodeError`; this parser returns an `incomplete InternLM2 +/// tool call` error from `finish()`, matching the other JSON parsers and Python's non-streaming +/// behavior. /// /// # Known unaddressed divergences (TODO) /// @@ -63,21 +60,18 @@ const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig { /// Qwen as well. If a real-world InternLM2 deployment hits one of these, /// prioritize the corresponding fix. /// -/// - **Arguments value type.** The shared core requires the arguments value -/// to be a JSON object (`take_json_object` rejects anything not starting -/// with `{`). Python's `json.dumps(action_dict.get("parameters", ...))` -/// accepts `null`, arrays, strings, and numbers and round-trips them -/// verbatim. Models that legitimately emit `"parameters":null` will hard- +/// - **Arguments value type.** The shared core requires the arguments value to be a JSON object +/// (`take_json_object` rejects anything not starting with `{`). Python's +/// `json.dumps(action_dict.get("parameters", ...))` accepts `null`, arrays, strings, and numbers +/// and round-trips them verbatim. Models that legitimately emit `"parameters":null` will hard- /// fail under Rust. -/// - **Unknown arguments key.** Python falls back to `{}` via -/// `action_dict.get("parameters", action_dict.get("arguments", {}))` when -/// neither key is present; the Rust header parser raises -/// `parsing failed: invalid InternLM2` for any unrecognized key. A model -/// that emits a typo (e.g. `"params"`) breaks the whole response. -/// - **Field order independence.** The header parser requires the JSON keys -/// to appear in the order `name` then arguments key. Python's -/// `json.loads` + `dict.get` is order-independent, so a model emitting -/// `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust. +/// - **Unknown arguments key.** Python falls back to `{}` via `action_dict.get("parameters", +/// action_dict.get("arguments", {}))` when neither key is present; the Rust header parser raises +/// `parsing failed: invalid InternLM2` for any unrecognized key. A model that emits a typo (e.g. +/// `"params"`) breaks the whole response. +/// - **Field order independence.** The header parser requires the JSON keys to appear in the order +/// `name` then arguments key. Python's `json.loads` + `dict.get` is order-independent, so a model +/// emitting `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust. pub struct Internlm2ToolParser { inner: JsonToolCallParser, } diff --git a/rust/src/tool-parser/src/parameters.rs b/rust/src/tool-parser/src/parameters.rs index 98ded7532fb..f857c147cb6 100644 --- a/rust/src/tool-parser/src/parameters.rs +++ b/rust/src/tool-parser/src/parameters.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use serde_json::{Number, Value}; +use serde_json::{Map, Number, Value}; use crate::Tool; @@ -21,6 +21,29 @@ pub(super) struct ToolSchema { params: BTreeMap, } +/// Parameter input for schema-aware conversion. +/// +/// It can be either a raw text string, or a structured input with named child elements. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum ParamInput { + Text(String), + #[allow(dead_code)] + Elements(Vec), +} + +impl From for ParamInput { + fn from(value: String) -> Self { + Self::Text(value) + } +} + +/// One named structured parameter child. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct ParamElement { + pub name: String, + pub value: ParamInput, +} + /// Normalized JSON parameter type used for raw string coercion. #[derive(Debug, Clone, PartialEq, Eq)] pub(super) enum JsonParamType { @@ -28,8 +51,13 @@ pub(super) enum JsonParamType { Integer, Number, Boolean, - Object, - Array, + Object { + properties: BTreeMap, + additional_properties: Option>, + }, + Array { + items: Option>, + }, Null, OneOf(Vec), } @@ -45,33 +73,39 @@ impl ToolSchemas { Self { tools } } - /// Convert raw string parameter values for one named tool. + /// Convert parameter values for one named tool. /// /// Unknown tool names use an empty schema, so all parameters fall back to - /// strings. - pub(super) fn convert_params_with_schema( + /// strings or object-like JSON for structured inputs. + pub(super) fn convert_params_with_schema

( &self, function_name: &str, - params: Vec<(String, String)>, - ) -> serde_json::Map { + params: Vec<(String, P)>, + ) -> Map + where + P: Into, + { let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty()); - let mut converted = serde_json::Map::with_capacity(params.len()); + let mut converted = Map::with_capacity(params.len()); for (name, value) in params { - let value = tool_schema.convert(&name, &value); + let value = tool_schema.convert(&name, value.into()); converted.insert(name, value); } converted } - /// Convert one raw string parameter value for one named tool. - pub(super) fn convert_param_with_schema( + /// Convert one parameter value for one named tool. + pub(super) fn convert_param_with_schema

( &self, function_name: &str, name: &str, - value: &str, - ) -> Value { + value: P, + ) -> Value + where + P: Into, + { let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty()); - tool_schema.convert(name, value) + tool_schema.convert(name, value.into()) } } @@ -101,21 +135,13 @@ impl ToolSchema { Self { params } } - /// Convert one raw parameter value using its normalized schema type. + /// Convert one parameter value using its normalized schema type. /// /// If the parameter name is unknown, or we don't have a schema for it, or /// the value fails to convert, this falls back to returning the raw - /// string as a JSON string value. - fn convert(&self, name: &str, value: &str) -> Value { - if value.eq_ignore_ascii_case("null") { - return Value::Null; - } - - let Some(param_type) = self.params.get(name) else { - return Value::String(value.to_string()); - }; - - convert_value(param_type, value).unwrap_or_else(|| Value::String(value.to_string())) + /// string as a JSON string value, or object-like JSON for structured input. + fn convert(&self, name: &str, input: ParamInput) -> Value { + convert_with_optional_schema(self.params.get(name), &input) } } @@ -125,7 +151,7 @@ impl JsonParamType { let schema = schema.as_object()?; if let Some(type_value) = schema.get("type") { - return Self::from_type_value(type_value); + return Self::from_type_value(type_value, schema); } if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) { @@ -134,32 +160,34 @@ impl JsonParamType { .map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::>()) .filter(|types| !types.is_empty()) .map(Self::one_of) - .unwrap_or(Self::Object); + .unwrap_or_else(|| Self::object_from_schema(Some(schema))); return Some(param_type); } + // Typically, these types are already handled by checking the "type" field, but + // we can also infer them from their characteristic fields if "type" is missing. if schema.contains_key("enum") { return Some(Self::String); } if schema.contains_key("items") { - return Some(Self::Array); + return Some(Self::array_from_schema(Some(schema))); } - if schema.contains_key("properties") { - return Some(Self::Object); + if schema.contains_key("properties") || schema.contains_key("additionalProperties") { + return Some(Self::object_from_schema(Some(schema))); } None } /// Normalize a JSON schema `type` value. - fn from_type_value(type_value: &Value) -> Option { + fn from_type_value(type_value: &Value, schema: &Map) -> Option { match type_value { - Value::String(kind) => Self::from_type_name(kind), + Value::String(kind) => Self::from_type_name(kind, Some(schema)), Value::Array(kinds) => { let types = kinds .iter() .filter_map(Value::as_str) - .filter_map(Self::from_type_name) + .filter_map(|kind| Self::from_type_name(kind, Some(schema))) .collect::>(); if types.is_empty() { None @@ -172,15 +200,15 @@ impl JsonParamType { } /// Normalize one JSON schema type name. - fn from_type_name(kind: &str) -> Option { + fn from_type_name(kind: &str, schema: Option<&Map>) -> Option { let kind = kind.trim().to_ascii_lowercase(); match kind.as_str() { "string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String), "integer" | "int" => Some(Self::Integer), "number" | "float" | "double" => Some(Self::Number), "boolean" | "bool" | "binary" => Some(Self::Boolean), - "object" | "dict" | "map" => Some(Self::Object), - "array" | "arr" | "list" | "sequence" => Some(Self::Array), + "object" | "dict" | "map" => Some(Self::object_from_schema(schema)), + "array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)), "null" => Some(Self::Null), _ if kind.starts_with("int") || kind.starts_with("uint") @@ -191,12 +219,52 @@ impl JsonParamType { Some(Self::Integer) } _ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number), - _ if kind.starts_with("dict") => Some(Self::Object), - _ if kind.starts_with("list") => Some(Self::Array), + _ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)), + _ if kind.starts_with("list") => Some(Self::array_from_schema(schema)), _ => None, } } + /// Normalize object schema fields. + fn object_from_schema(schema: Option<&Map>) -> Self { + let properties = schema + .and_then(|schema| schema.get("properties")) + .and_then(Value::as_object) + .map(|properties| { + properties + .iter() + .filter_map(|(name, schema)| { + Self::from_schema(schema).map(|param_type| (name.clone(), param_type)) + }) + .collect() + }) + .unwrap_or_default(); + + let additional_properties = + schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| { + if schema.is_object() { + Self::from_schema(schema).map(Box::new) + } else { + None + } + }); + + Self::Object { + properties, + additional_properties, + } + } + + /// Normalize array schema fields. + fn array_from_schema(schema: Option<&Map>) -> Self { + let items = schema + .and_then(|schema| schema.get("items")) + .and_then(Self::from_schema) + .map(Box::new); + + Self::Array { items } + } + /// Collapse a candidate type list into one normalized type. fn one_of(mut types: Vec) -> Self { if types.len() == 1 { @@ -207,23 +275,126 @@ impl JsonParamType { } } -/// Convert one raw string value to a normalized JSON type. -fn convert_value(param_type: &JsonParamType, value: &str) -> Option { - match param_type { - JsonParamType::String => Some(Value::String(value.to_string())), - JsonParamType::Integer => value.parse::().ok().map(Number::from).map(Value::Number), - JsonParamType::Number => convert_number(value), - JsonParamType::Boolean => convert_boolean(value), - JsonParamType::Object | JsonParamType::Array => serde_json::from_str(value).ok(), - JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null), - JsonParamType::OneOf(types) => { - types.iter().find_map(|param_type| convert_value(param_type, value)) +/// Convert one parameter input to a normalized JSON value. +fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value { + // For literal `null`, always convert to JSON null value. + if let ParamInput::Text(value) = input + && value.eq_ignore_ascii_case("null") + { + return Value::Null; + } + + // If we have a schema, try to convert the value using it. + if let Some(param_type) = param_type + && let Some(value) = try_convert_value(param_type, input) + { + return value; + } + // We don't have a schema, or conversion failed, use fallback logic. + match input { + ParamInput::Text(value) => Value::String(value.clone()), + ParamInput::Elements(elements) => { + // Convert structured input to object without a schema. + Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None)) } } } +/// Convert one parameter input to a normalized JSON type. +fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option { + match input { + ParamInput::Text(value) => try_convert_text_value(param_type, value), + ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements), + } +} + +/// Convert one raw string value to a normalized JSON type. +fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option { + match param_type { + JsonParamType::String => Some(Value::String(value.to_string())), + JsonParamType::Integer => value.parse::().ok().map(Number::from).map(Value::Number), + JsonParamType::Number => try_convert_number(value), + JsonParamType::Boolean => try_convert_boolean(value), + JsonParamType::Object { .. } if value.is_empty() => Some(Value::Object(Map::new())), + JsonParamType::Array { .. } if value.is_empty() => Some(Value::Array(Vec::new())), + JsonParamType::Object { .. } | JsonParamType::Array { .. } => { + // For composite types with string input, simply interpret the string as JSON. + serde_json::from_str(value).ok() + } + JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null), + JsonParamType::OneOf(types) => { + types.iter().find_map(|param_type| try_convert_text_value(param_type, value)) + } + } +} + +/// Convert one structured parameter input to a normalized JSON type. +fn try_convert_elements_value( + param_type: &JsonParamType, + elements: &[ParamElement], +) -> Option { + match param_type { + JsonParamType::Object { + properties, + additional_properties, + } => Some(Value::Object(convert_elements_to_object( + elements, + properties, + additional_properties.as_deref(), + ))), + JsonParamType::Array { items } => Some(Value::Array( + // Collect all child elements into an array, regardless of their names. + elements + .iter() + .map(|element| convert_with_optional_schema(items.as_deref(), &element.value)) + .collect(), + )), + JsonParamType::OneOf(types) => types + .iter() + .find_map(|param_type| try_convert_elements_value(param_type, elements)), + + // Primitive types can't be converted from structured input. + JsonParamType::String + | JsonParamType::Integer + | JsonParamType::Number + | JsonParamType::Boolean + | JsonParamType::Null => None, + } +} + +/// Convert structured elements to an object, using field schemas when present. +fn convert_elements_to_object( + elements: &[ParamElement], + properties: &BTreeMap, + additional_properties: Option<&JsonParamType>, +) -> Map { + let mut object = Map::with_capacity(elements.len()); + for element in elements { + let param_type = properties.get(&element.name).or(additional_properties); + let value = convert_with_optional_schema(param_type, &element.value); + insert_object_value(&mut object, element.name.clone(), value); + } + object +} + +/// Insert an object field while preserving duplicate keys as arrays. +fn insert_object_value(object: &mut Map, key: String, value: Value) { + if let Some(existing) = object.get_mut(&key) { + match existing { + // Collect values under the same key into an array. + Value::Array(values) => values.push(value), + existing => { + let first = std::mem::replace(existing, Value::Null); + *existing = Value::Array(vec![first, value]); + } + } + } else { + object.insert(key, value); + } +} + /// Convert one raw string value to a JSON number. -fn convert_number(value: &str) -> Option { +fn try_convert_number(value: &str) -> Option { serde_json::from_str::(value) .or_else(|_| value.parse::().map(Number::from)) .or_else(|_| value.parse::().ok().and_then(Number::from_f64).ok_or(())) @@ -232,7 +403,7 @@ fn convert_number(value: &str) -> Option { } /// Convert one raw string value to a boolean. -fn convert_boolean(value: &str) -> Option { +fn try_convert_boolean(value: &str) -> Option { match value.trim().to_ascii_lowercase().as_str() { "true" | "1" => Some(Value::Bool(true)), "false" | "0" => Some(Value::Bool(false)), @@ -242,9 +413,9 @@ fn convert_boolean(value: &str) -> Option { #[cfg(test)] mod tests { - use serde_json::json; + use serde_json::{Value, json}; - use super::{ToolSchema, ToolSchemas}; + use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas}; use crate::Tool; fn test_tool(name: &str, parameters: serde_json::Value) -> Tool { @@ -260,8 +431,8 @@ mod tests { fn invalid_schema_converts_everything_as_string() { let params = ToolSchema::from_schema(&json!({ "type": "object" })); - assert_eq!(params.convert("count", "42"), json!("42")); - assert_eq!(params.convert("count", "null"), json!(null)); + assert_eq!(params.convert("count", text("42")), json!("42")); + assert_eq!(params.convert("count", text("null")), json!(null)); } #[test] @@ -275,9 +446,9 @@ mod tests { } })); - assert_eq!(params.convert("unknown_schema", "42"), json!("42")); - assert_eq!(params.convert("unknown_type", "42"), json!("42")); - assert_eq!(params.convert("known", "42"), json!(42)); + assert_eq!(params.convert("unknown_schema", text("42")), json!("42")); + assert_eq!(params.convert("unknown_type", text("42")), json!("42")); + assert_eq!(params.convert("known", text("42")), json!(42)); } #[test] @@ -298,16 +469,25 @@ mod tests { } })); - assert_eq!(params.convert("text", "42"), json!("42")); - assert_eq!(params.convert("count", "42"), json!(42)); - assert_eq!(params.convert("size", "5.0"), json!(5.0)); - assert_eq!(params.convert("ratio", "2.5"), json!(2.5)); - assert_eq!(params.convert("enabled", "1"), json!(true)); - assert_eq!(params.convert("payload", r#"{"k":1}"#), json!({ "k": 1 })); - assert_eq!(params.convert("mapping", r#"{"k":1}"#), json!({ "k": 1 })); - assert_eq!(params.convert("items", "[1,2]"), json!([1, 2])); - assert_eq!(params.convert("names", r#"["a","b"]"#), json!(["a", "b"])); - assert_eq!(params.convert("nothing", "null"), json!(null)); + assert_eq!(params.convert("text", text("42")), json!("42")); + assert_eq!(params.convert("count", text("42")), json!(42)); + assert_eq!(params.convert("size", text("5.0")), json!(5.0)); + assert_eq!(params.convert("ratio", text("2.5")), json!(2.5)); + assert_eq!(params.convert("enabled", text("1")), json!(true)); + assert_eq!( + params.convert("payload", text(r#"{"k":1}"#)), + json!({ "k": 1 }) + ); + assert_eq!( + params.convert("mapping", text(r#"{"k":1}"#)), + json!({ "k": 1 }) + ); + assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2])); + assert_eq!( + params.convert("names", text(r#"["a","b"]"#)), + json!(["a", "b"]) + ); + assert_eq!(params.convert("nothing", text("null")), json!(null)); } #[test] @@ -321,19 +501,40 @@ mod tests { assert_eq!(converted_number_text(¶ms, "5"), "5"); assert_eq!(converted_number_text(¶ms, "5.0"), "5.0"); - assert_eq!(converted_number_text(¶ms, "5.00"), "5.00"); - assert_eq!(converted_number_text(¶ms, "1e0"), "1e+0"); assert_eq!(converted_number_text(¶ms, "5."), "5.0"); assert_eq!(converted_number_text(¶ms, "+1"), "1"); assert_eq!(converted_number_text(¶ms, "+1.0"), "1.0"); - assert_eq!( - converted_number_text(¶ms, "9223372036854775807.5"), - "9223372036854775807.5" - ); + + // TODO: we cannot preserve the original number precision by enabling `serde_json`'s + // `arbitrary_precision` feature, otherwise the test + // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail. + // See issue: https://github.com/mitsuhiko/minijinja/issues/641 + + // assert_eq!(converted_number_text(¶ms, "5.00"), "5.00"); + // assert_eq!(converted_number_text(¶ms, "1e0"), "1e+0"); + // assert_eq!( + // converted_number_text(¶ms, "9223372036854775807.5"), + // "9223372036854775807.5" + // ); } fn converted_number_text(params: &ToolSchema, value: &str) -> String { - serde_json::to_string(¶ms.convert("value", value)).unwrap() + serde_json::to_string(¶ms.convert("value", text(value))).unwrap() + } + + fn text(value: &str) -> ParamInput { + ParamInput::Text(value.to_string()) + } + + fn elem(name: &str, value: ParamInput) -> ParamElement { + ParamElement { + name: name.to_string(), + value, + } + } + + fn elements(elements: Vec) -> ParamInput { + ParamInput::Elements(elements) } #[test] @@ -350,12 +551,12 @@ mod tests { } })); - assert_eq!(params.convert("s", "x"), json!("x")); - assert_eq!(params.convert("i", "7"), json!(7)); - assert_eq!(params.convert("n", "7.5"), json!(7.5)); - assert_eq!(params.convert("b", "true"), json!(true)); - assert_eq!(params.convert("a", "[1]"), json!([1])); - assert_eq!(params.convert("o", r#"{"x":1}"#), json!({ "x": 1 })); + assert_eq!(params.convert("s", text("x")), json!("x")); + assert_eq!(params.convert("i", text("7")), json!(7)); + assert_eq!(params.convert("n", text("7.5")), json!(7.5)); + assert_eq!(params.convert("b", text("true")), json!(true)); + assert_eq!(params.convert("a", text("[1]")), json!([1])); + assert_eq!(params.convert("o", text(r#"{"x":1}"#)), json!({ "x": 1 })); } #[test] @@ -373,8 +574,8 @@ mod tests { } })); - assert_eq!(integer_first.convert("value", "42"), json!(42)); - assert_eq!(string_first.convert("value", "42"), json!("42")); + assert_eq!(integer_first.convert("value", text("42")), json!(42)); + assert_eq!(string_first.convert("value", text("42")), json!("42")); } #[test] @@ -396,9 +597,9 @@ mod tests { } })); - assert_eq!(params.convert("choice", "42"), json!(42)); + assert_eq!(params.convert("choice", text("42")), json!(42)); assert_eq!( - params.convert("fallback_object", r#"{"x":1}"#), + params.convert("fallback_object", text(r#"{"x":1}"#)), json!({ "x": 1 }) ); } @@ -414,9 +615,12 @@ mod tests { } })); - assert_eq!(params.convert("choice", "a"), json!("a")); - assert_eq!(params.convert("items", "[1,2]"), json!([1, 2])); - assert_eq!(params.convert("payload", r#"{"x":1}"#), json!({ "x": 1 })); + assert_eq!(params.convert("choice", text("a")), json!("a")); + assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2])); + assert_eq!( + params.convert("payload", text(r#"{"x":1}"#)), + json!({ "x": 1 }) + ); } #[test] @@ -518,4 +722,162 @@ mod tests { assert_eq!(converted.get("topn"), Some(&json!("5"))); assert_eq!(converted.get("nullish"), Some(&json!(null))); } + + #[test] + fn converts_structured_inputs_with_recursive_schema() { + let schemas = ToolSchemas::from_tools(&[test_tool( + "create_order", + json!({ + "type": "object", + "properties": { + "user_id": { "type": "integer" }, + "urgent": { "type": "boolean" }, + "note": { "type": "string" }, + "nil": { "type": "string" }, + "shipping": { + "type": "object", + "properties": { + "city": { "type": "string" }, + "zip": { "type": "integer" } + } + }, + "items": { + "type": "array", + "items": { + "type": "object", + "properties": { + "sku": { "type": "string" }, + "qty": { "type": "integer" } + } + } + }, + "metadata": { + "type": "object", + "additionalProperties": { "type": "integer" } + }, + "duplicate_demo": { + "type": "object", + "properties": { + "tag": { "type": "string" } + } + }, + "schema_mismatch_array": { + "type": "array", + "items": { "type": "integer" } + }, + "closed_object": { + "type": "object", + "additionalProperties": false + }, + "open_object": { + "type": "object", + "additionalProperties": true + }, + "payload_text": { "type": "object" }, + "items_text": { "type": "array" } + } + }), + )]); + + let converted = schemas.convert_params_with_schema( + "create_order", + vec![ + ("user_id".to_string(), text("42")), + ("urgent".to_string(), text("true")), + ("note".to_string(), text("Please leave at front desk.")), + ("nil".to_string(), text("NULL")), + ( + "shipping".to_string(), + elements(vec![ + elem("city", text("Singapore")), + elem("zip", text("018956")), + ]), + ), + ( + "items".to_string(), + elements(vec![ + elem( + "item1", + elements(vec![elem("sku", text("book-001")), elem("qty", text("2"))]), + ), + elem( + "item2", + elements(vec![elem("sku", text("pen-007")), elem("qty", text("5"))]), + ), + ]), + ), + ( + "metadata".to_string(), + elements(vec![elem("score", text("42")), elem("rank", text("7"))]), + ), + ( + "duplicate_demo".to_string(), + elements(vec![elem("tag", text("a")), elem("tag", text("b"))]), + ), + ( + "closed_object".to_string(), + elements(vec![elem("unknown", text("x"))]), + ), + ( + "open_object".to_string(), + elements(vec![elem("unknown", text("y"))]), + ), + ("payload_text".to_string(), text(r#"{"x":1}"#)), + ("items_text".to_string(), text("[1,2]")), + ( + "unknown_struct".to_string(), + elements(vec![ + elem("a", text("1")), + elem("a", text("2")), + elem("nil", text("null")), + ]), + ), + ], + ); + + assert_eq!( + Value::Object(converted), + json!({ + "user_id": 42, + "urgent": true, + "note": "Please leave at front desk.", + "nil": null, + "shipping": { + "city": "Singapore", + "zip": 18956 + }, + "items": [ + { + "sku": "book-001", + "qty": 2 + }, + { + "sku": "pen-007", + "qty": 5 + } + ], + "metadata": { + "score": 42, + "rank": 7 + }, + "duplicate_demo": { + "tag": ["a", "b"] + }, + "closed_object": { + "unknown": "x" + }, + "open_object": { + "unknown": "y" + }, + "payload_text": { + "x": 1 + }, + "items_text": [1, 2], + "unknown_struct": { + "a": ["1", "2"], + "nil": null + } + }) + ); + } } diff --git a/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py new file mode 100644 index 00000000000..149ea6625a9 --- /dev/null +++ b/tests/benchmarks/test_custom_dataset_chat_template_kwargs.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse +import json +from pathlib import Path + +import pytest + +from vllm.benchmarks.datasets import get_samples + + +class _RecordingTokenizer: + """Minimal tokenizer stub that records the kwargs forwarded to + apply_chat_template, so we can assert chat_template_kwargs propagation + without loading a real model/template.""" + + def __init__(self) -> None: + self.captured_kwargs: dict | None = None + self.chat_template = "dummy-template" + + def apply_chat_template( + self, + conversation, + add_generation_prompt: bool = True, + tokenize: bool = False, + **kwargs, + ) -> str: + self.captured_kwargs = kwargs + return conversation[0]["content"] + + def __call__(self, text: str): + return argparse.Namespace(input_ids=list(range(len(text.split())))) + + +def _args(dataset_path: str, chat_template_kwargs) -> argparse.Namespace: + return argparse.Namespace( + dataset_name="custom", + dataset_path=dataset_path, + disable_shuffle=True, + num_prompts=1, + custom_output_len=32, + skip_chat_template=False, + chat_template_kwargs=chat_template_kwargs, + no_oversample=False, + seed=0, + request_id_prefix="", + ) + + +def _write_one(path: Path) -> None: + path.write_text(json.dumps({"prompt": "hello world"}) + "\n") + + +@pytest.mark.benchmark +def test_chat_template_kwargs_forwarded(tmp_path: Path) -> None: + """--chat-template-kwargs must reach the client-side apply_chat_template.""" + jsonl = tmp_path / "data.jsonl" + _write_one(jsonl) + + tok = _RecordingTokenizer() + get_samples(_args(str(jsonl), {"thinking": True}), tok) + + assert tok.captured_kwargs == {"thinking": True} + + +@pytest.mark.benchmark +def test_chat_template_kwargs_default_is_noop(tmp_path: Path) -> None: + """When not provided, no extra kwargs are passed (existing behavior).""" + jsonl = tmp_path / "data.jsonl" + _write_one(jsonl) + + tok = _RecordingTokenizer() + get_samples(_args(str(jsonl), None), tok) + + assert tok.captured_kwargs == {} diff --git a/tests/benchmarks/test_custom_image_dataset.py b/tests/benchmarks/test_custom_image_dataset.py index 336bac93d0b..f2a48abe604 100644 --- a/tests/benchmarks/test_custom_image_dataset.py +++ b/tests/benchmarks/test_custom_image_dataset.py @@ -2,11 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json from argparse import Namespace +from io import BytesIO from pathlib import Path from typing import Any +import pybase64 as base64 import pytest +from PIL import Image +import vllm.benchmarks.datasets.datasets as datasets_module from vllm.benchmarks.datasets import CustomImageDataset, get_samples from vllm.benchmarks.lib.endpoint_request_func import ( RequestFuncInput, @@ -33,6 +37,22 @@ def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: f.write(json.dumps(row) + "\n") +def _write_png(path: Path, color: tuple[int, int, int] = (255, 0, 0)) -> None: + Image.new("RGB", (1, 1), color=color).save(path) + + +def _decode_data_url(data_url: str) -> tuple[str, bytes]: + prefix, image_base64 = data_url.split(",", 1) + return prefix, base64.b64decode(image_base64) + + +def _assert_png_data_url(data_url: str) -> None: + prefix, image_bytes = _decode_data_url(data_url) + assert prefix == "data:image/png;base64" + with Image.open(BytesIO(image_bytes)) as image: + image.verify() + + def _args_for_custom_image(dataset_path: Path) -> Namespace: return Namespace( dataset_name="custom_image", @@ -42,6 +62,7 @@ def _args_for_custom_image(dataset_path: Path) -> Namespace: num_prompts=2, custom_output_len=32, enable_multimodal_chat=False, + custom_ensure_client_side_data=False, request_id_prefix="req-", no_oversample=False, ) @@ -230,6 +251,125 @@ def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat( assert _get_chat_messages(request_input) == sample.prompt +@pytest.mark.benchmark +def test_custom_image_dataset_encodes_image_media_when_requested( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + image_a = tmp_path / "chart_a.png" + image_b = tmp_path / "chart b.png" + _write_png(image_a, color=(255, 0, 0)) + _write_png(image_b, color=(0, 255, 0)) + data_url = "data:image/png;base64,Zm9v" + remote_url = "https://example.com/chart.png" + original_fetch_image = datasets_module.fetch_image + + def fake_fetch_image(image_url: str) -> Image.Image: + if image_url == remote_url: + return Image.new("RGB", (1, 1), color=(0, 0, 255)) + return original_fetch_image(image_url) + + monkeypatch.setattr(datasets_module, "fetch_image", fake_fetch_image) + + jsonl = tmp_path / "images.jsonl" + _write_jsonl( + jsonl, + [ + { + "prompt": "Compare the charts.", + "image_files": [ + str(image_a), + image_b.as_uri(), + remote_url, + data_url, + ], + } + ], + ) + + dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True) + samples = dataset.sample( + tokenizer=_Tokenizer(), + num_requests=1, + output_len=32, + ensure_client_side_data=True, + ) + + assert len(samples) == 1 + assert isinstance(samples[0].multi_modal_data, list) + image_urls = [part["image_url"]["url"] for part in samples[0].multi_modal_data] + + _assert_png_data_url(image_urls[0]) + _assert_png_data_url(image_urls[1]) + _assert_png_data_url(image_urls[2]) + assert image_urls[3] == data_url + + +@pytest.mark.benchmark +def test_custom_image_dataset_encodes_interleaved_image_media( + tmp_path: Path, +) -> None: + image_a = tmp_path / "chart_a.png" + image_b = tmp_path / "chart_b.png" + _write_png(image_a, color=(255, 0, 0)) + _write_png(image_b, color=(0, 255, 0)) + jsonl = tmp_path / "images.jsonl" + _write_jsonl( + jsonl, + [ + { + "content": [ + {"type": "text", "text": "Compare "}, + {"type": "image", "image": str(image_a)}, + { + "type": "image_url", + "image_url": { + "url": image_b.as_uri(), + "detail": "low", + }, + }, + ], + } + ], + ) + + dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True) + samples = dataset.sample( + tokenizer=_Tokenizer(), + num_requests=1, + output_len=32, + ensure_client_side_data=True, + ) + + sample = samples[0] + assert isinstance(sample.prompt, list) + _assert_png_data_url(sample.prompt[1]["image_url"]["url"]) + _assert_png_data_url(sample.prompt[2]["image_url"]["url"]) + assert sample.prompt[2]["image_url"]["detail"] == "low" + + +@pytest.mark.benchmark +def test_custom_image_dataset_rejects_invalid_image_media( + tmp_path: Path, +) -> None: + invalid_image = tmp_path / "not_an_image.png" + invalid_image.write_text("not an image") + jsonl = tmp_path / "images.jsonl" + _write_jsonl( + jsonl, + [{"prompt": "Describe the image.", "image_files": [str(invalid_image)]}], + ) + + dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True) + with pytest.raises(ValueError, match="Invalid image URL"): + dataset.sample( + tokenizer=_Tokenizer(), + num_requests=1, + output_len=32, + ensure_client_side_data=True, + ) + + @pytest.mark.benchmark def test_custom_image_dataset_rejects_invalid_content_part( tmp_path: Path, diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index e45e5cf425f..b8c18fa6cdc 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -55,6 +55,15 @@ def test_dynamic_shapes_compilation( evaluate_guards, ): """Test that all dynamic shapes types compile successfully""" + if shapes_type == DynamicShapesType.UNBACKED and not is_torch_equal_or_newer( + "2.11.0" + ): + # NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only + # landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On + # older torch there's no way to express it, so unbacked shapes go + # data-dependent and compilation blows up -- nothing to test. + pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11") + if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED: pytest.skip("unbacked dynamic shapes do not add guards") diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 29d7f94c510..20dd4f36393 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -284,7 +284,7 @@ def test_multiproc_executor_pipeline_parallel(): assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)" # Verify max_concurrent_batches for pipeline parallel - assert executor.max_concurrent_batches == 2, ( + assert vllm_config.max_concurrent_batches == 2, ( "Max concurrent batches should equal PP size" ) diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py index 5daec22df6f..398ee30c068 100644 --- a/tests/distributed/test_ray_v2_executor.py +++ b/tests/distributed/test_ray_v2_executor.py @@ -83,7 +83,7 @@ def assert_executor(executor, tp_size, pp_size): assert executor._get_output_rank() == expected_output_rank if pp_size > 1: - assert executor.max_concurrent_batches == pp_size + assert executor.vllm_config.max_concurrent_batches == pp_size executor.check_health() assert not executor.is_failed diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 467e3934a05..2df0d9e71c3 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -291,7 +291,9 @@ def test_nccl_receive_sparse_weights_without_init_raises(): config = WeightTransferConfig(backend="nccl") parallel_config = create_mock_parallel_config() - engine = NCCLWeightTransferEngine(config, parallel_config) + engine = NCCLWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) update_info = NCCLWeightTransferUpdateInfo( names=["w"], @@ -526,7 +528,9 @@ def inference_receive_sparse_tensor( parallel_config.data_parallel_rank = 0 parallel_config.data_parallel_index = 0 - engine = NCCLWeightTransferEngine(config, parallel_config) + engine = NCCLWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) engine.init_transfer_engine( NCCLWeightTransferInitInfo( master_address=master_address, diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index eb9798980f0..ad9fed1d355 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -635,3 +635,143 @@ class TestThinkingBlockConversion: # Redacted thinking is ignored, normal thinking still becomes reasoning. assert asst.get("reasoning") == "Thinking..." assert asst.get("content") == "Hi!" + + +class TestInlineSystemMessageInMessagesArray: + """Verify that ``role: system`` messages embedded inside the ``messages`` + array are accepted and merged with the top-level ``system`` prompt. + + This handles clients that place system messages inside the messages array + instead of the Anthropic-standard top-level ``system`` field. + """ + + def test_inline_system_merged_with_top_level_system(self): + """Full integration: inline system + top-level system + user message.""" + request = _make_request( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "\n.....\n\n\n", + }, + { + "type": "text", + "text": "help?", + "cache_control": {"type": "ephemeral"}, + }, + ], + }, + { + "role": "system", + "content": ".....", + }, + ], + system=[ + { + "type": "text", + "text": "x-anthropic-billing-header: " + "cc_version=2.1.160.bca; cc_entrypoint=cli; cch=d1d48;", + }, + { + "type": "text", + "text": "You are Claude Code, Anthropic's official CLI for Claude.", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "text", + "text": "....", + "cache_control": {"type": "ephemeral"}, + }, + ], + tools=[], + ) + + result = _convert(request) + + # First message should be the merged system prompt. + assert result.messages[0]["role"] == "system" + # Billing header stripped, inline system appended. + assert ( + result.messages[0]["content"] + == "You are Claude Code, Anthropic's official CLI for Claude." + "...." + "....." + ) + + # Second message should be the user message, content preserved. + assert result.messages[1]["role"] == "user" + user_content = result.messages[1]["content"] + assert len(user_content) == 2 + assert user_content[0] == { + "type": "text", + "text": "\n.....\n\n\n", + } + assert user_content[1] == { + "type": "text", + "text": "help?", + } + + def test_inline_system_string_only(self): + """Only an inline system string, no top-level system.""" + request = _make_request( + [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Be concise."}, + ] + ) + result = _convert(request) + + assert result.messages[0]["role"] == "system" + assert result.messages[0]["content"] == "Be concise." + assert result.messages[1]["role"] == "user" + + def test_inline_system_list_content(self): + """Inline system with list content blocks.""" + request = _make_request( + [ + {"role": "user", "content": "Hi"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "Part one. "}, + {"type": "text", "text": "Part two."}, + ], + }, + ] + ) + result = _convert(request) + + assert result.messages[0]["role"] == "system" + assert result.messages[0]["content"] == "Part one. Part two." + + def test_multiple_inline_system_messages(self): + """Multiple inline system messages should all be merged.""" + request = _make_request( + [ + {"role": "system", "content": "First system."}, + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Second system."}, + ] + ) + result = _convert(request) + + assert result.messages[0]["role"] == "system" + assert result.messages[0]["content"] == "First system.Second system." + assert result.messages[1]["role"] == "user" + + def test_inline_system_with_top_level_string(self): + """Top-level system is a string, inline system is also present.""" + request = _make_request( + [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Inline hint."}, + ], + system="Top-level prompt.", + ) + result = _convert(request) + + assert result.messages[0]["role"] == "system" + assert result.messages[0]["content"] == "Top-level prompt.Inline hint." + assert result.messages[1]["role"] == "user" diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 582e0792156..e099c282f42 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -6,9 +6,13 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.chat_completion.protocol import ( + BatchChatCompletionRequest, + ChatCompletionRequest, +) from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath @@ -444,3 +448,45 @@ def test_json_schema_response_format_missing_schema(): messages=[{"role": "user", "content": "hello"}], response_format={"type": "json_schema"}, ) + + +@pytest.mark.parametrize("format_value", [None, {}]) +def test_structural_tag_response_format_invalid(format_value): + """Malformed structural tags should be rejected during request validation.""" + with pytest.raises( + ValidationError, + match="Invalid response_format structural_tag", + ): + ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hello"}], + response_format={"type": "structural_tag", "format": format_value}, + ) + + +@pytest.mark.parametrize("format_value", [None, {}]) +def test_batch_structural_tag_response_format_invalid(format_value): + """Batch chat should reject malformed structural tags at request parsing.""" + with pytest.raises( + ValidationError, + match="Invalid response_format structural_tag", + ): + BatchChatCompletionRequest( + model=MODEL_NAME, + messages=[[{"role": "user", "content": "hello"}]], + response_format={"type": "structural_tag", "format": format_value}, + ) + + +@pytest.mark.parametrize("structural_tag", ["not json", ""]) +def test_structured_outputs_structural_tag_invalid(structural_tag): + """Malformed direct structured_outputs structural tags should be rejected.""" + with pytest.raises( + ValidationError, + match="Invalid structured_outputs structural_tag", + ): + ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "hello"}], + structured_outputs={"structural_tag": structural_tag}, + ) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 11793ac8f49..7c0a46a4e63 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1935,8 +1935,10 @@ async def test_streaming_n_gt1_independent_tool_parsers(): finished=True, ) - # Collect tool-call deltas per choice from the SSE stream. + # Collect tool-call deltas and finish_reasons per choice from the SSE + # stream. tc_deltas_by_choice: dict[int, list[dict]] = {i: [] for i in range(num_choices)} + finish_reasons_by_choice: dict[int, list[str]] = {i: [] for i in range(num_choices)} async for chunk_str in serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -1959,6 +1961,8 @@ async def test_streaming_n_gt1_independent_tool_parsers(): if delta.get("tool_calls"): for tc in delta["tool_calls"]: tc_deltas_by_choice[idx].append(tc) + if choice.get("finish_reason") is not None: + finish_reasons_by_choice[idx].append(choice["finish_reason"]) # Both choices must independently produce the correct tool call. for choice_idx in range(num_choices): @@ -1984,141 +1988,11 @@ async def test_streaming_n_gt1_independent_tool_parsers(): f"Choice {choice_idx}: expected {{'city': 'Tokyo'}}, got {parsed_args}" ) - -class TestCreateRemainingArgsDelta: - """Tests for _create_remaining_args_delta helper function. - - This helper is used when streaming tool calls to preserve id/type/name - fields in the finish chunk, which would otherwise be lost. - """ - - def test_preserves_id_type_name(self): - """Test that id, type, and name are preserved from original delta.""" - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat - from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, + reasons = finish_reasons_by_choice[choice_idx] + assert len(reasons) == 1, ( + f"Choice {choice_idx}: expected exactly 1 finish_reason, got {reasons}" ) - - original_delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=0, - id="call_abc123", - type="function", - function=DeltaFunctionCall( - name="get_weather", - arguments='{"location": "Paris"}', - ), - ) - ] + assert reasons[0] == "tool_calls", ( + f"Choice {choice_idx}: expected finish_reason='tool_calls', " + f"got '{reasons[0]}'" ) - - result = OpenAIServingChat._create_remaining_args_delta( - original_delta, '", "unit": "celsius"}', 0 - ) - - assert len(result.tool_calls) == 1 - tc = result.tool_calls[0] - assert tc.index == 0 - assert tc.id == "call_abc123" - assert tc.type == "function" - assert tc.function.name == "get_weather" - assert tc.function.arguments == '", "unit": "celsius"}' - - def test_matches_by_index(self): - """Test that the correct tool call is matched by index.""" - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat - from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ) - - original_delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=0, - id="call_first", - type="function", - function=DeltaFunctionCall(name="func_a", arguments="{}"), - ), - DeltaToolCall( - index=1, - id="call_second", - type="function", - function=DeltaFunctionCall(name="func_b", arguments="{}"), - ), - ] - ) - - result = OpenAIServingChat._create_remaining_args_delta( - original_delta, '{"extra": true}', 1 - ) - - assert len(result.tool_calls) == 1 - tc = result.tool_calls[0] - assert tc.index == 1 - assert tc.id == "call_second" - assert tc.function.name == "func_b" - - def test_no_matching_tool_call(self): - """Test graceful handling when no matching tool call is found.""" - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat - from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ) - - original_delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=0, - id="call_zero", - type="function", - function=DeltaFunctionCall(name="func", arguments="{}"), - ) - ] - ) - - result = OpenAIServingChat._create_remaining_args_delta( - original_delta, '{"arg": 1}', 5 - ) - - assert len(result.tool_calls) == 1 - tc = result.tool_calls[0] - assert tc.index == 5 - assert tc.id is None - assert tc.type is None - assert tc.function.name is None - assert tc.function.arguments == '{"arg": 1}' - - def test_function_is_none(self): - """Test handling when original tool call has no function.""" - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat - from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall - - original_delta = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=0, - id="call_nofunc", - type="function", - function=None, - ) - ] - ) - - result = OpenAIServingChat._create_remaining_args_delta( - original_delta, '{"data": "value"}', 0 - ) - - assert len(result.tool_calls) == 1 - tc = result.tool_calls[0] - assert tc.index == 0 - assert tc.id == "call_nofunc" - assert tc.type == "function" - assert tc.function.name is None - assert tc.function.arguments == '{"data": "value"}' diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index c95e47fa1b1..71a70a4d0eb 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -6,6 +6,7 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +from pydantic import ValidationError from vllm.config.multimodal import MultiModalConfig from vllm.entrypoints.openai.completion.protocol import CompletionRequest @@ -302,6 +303,36 @@ def test_json_schema_response_format_missing_schema(): ) +@pytest.mark.parametrize("format_value", [None, {}]) +def test_structural_tag_response_format_invalid(format_value): + """Malformed structural tags should be rejected during request validation.""" + with pytest.raises( + ValidationError, + match="Invalid response_format structural_tag", + ): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + response_format={"type": "structural_tag", "format": format_value}, + ) + + +@pytest.mark.parametrize("structural_tag", ["not json", ""]) +def test_structured_outputs_structural_tag_invalid(structural_tag): + """Malformed direct structured_outputs structural tags should be rejected.""" + with pytest.raises( + ValidationError, + match="Invalid structured_outputs structural_tag", + ): + CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + structured_outputs={"structural_tag": structural_tag}, + ) + + def test_negative_prompt_token_ids_nested(): """Negative token IDs in prompt (nested list) should raise validation error.""" with pytest.raises(Exception, match="greater than or equal to 0"): diff --git a/tests/entrypoints/openai/test_responses_parser_unified.py b/tests/entrypoints/openai/test_responses_parser_unified.py new file mode 100644 index 00000000000..ecc857e1aac --- /dev/null +++ b/tests/entrypoints/openai/test_responses_parser_unified.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ResponsesParser with the unified Parser interface. + +These tests verify that ResponsesParser correctly delegates to the unified +Parser (via extract_response_outputs) instead of calling separate +ReasoningParser / ToolParser instances directly. +""" + +from collections.abc import Sequence +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.parser.responses_parser import ( + ResponsesParser, + get_responses_parser_for_simple_context, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.outputs import CompletionOutput +from vllm.parser.abstract_parser import DelegatingParser + +pytestmark = pytest.mark.skip_global_cleanup + + +# --------------------------------------------------------------------------- +# Test parser stubs +# --------------------------------------------------------------------------- + + +class _NoOpParser(DelegatingParser): + """Parser that extracts no reasoning and no tool calls.""" + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + return False + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + return input_ids + + def extract_reasoning(self, model_output, request): + return None, model_output + + def extract_reasoning_streaming(self, *args, **kwargs): + return None + + def extract_tool_calls(self, model_output, request): + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + def extract_tool_calls_streaming(self, *args, **kwargs): + return None + + def parse_delta(self, *args, **kwargs) -> DeltaMessage | None: + return None + + +class _ReasoningOnlyParser(DelegatingParser): + """Parser that extracts reasoning but no tool calls.""" + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + return False + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + return input_ids + + def extract_reasoning(self, model_output, request): + if "" in model_output and "" in model_output: + start = model_output.index("") + len("") + end = model_output.index("") + reasoning = model_output[start:end] + content = model_output[end + len("") :] + return reasoning, content.strip() or None + return None, model_output + + def extract_reasoning_streaming(self, *args, **kwargs): + return None + + def extract_tool_calls(self, model_output, request): + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + def extract_tool_calls_streaming(self, *args, **kwargs): + return None + + def parse_delta(self, *args, **kwargs) -> DeltaMessage | None: + return None + + +class _StubToolParser: + """Minimal tool parser stub that always returns a hardcoded tool call.""" + + supports_required_and_named = False + + def __init__(self, tokenizer=None, tools=None): + pass + + def extract_tool_calls(self, model_output, request): + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=[ + ToolCall( + id="call_123", + type="function", + function=FunctionCall( + name="get_weather", + arguments='{"location": "Paris"}', + ), + ) + ], + content=None, + ) + + def extract_tool_calls_streaming(self, *args, **kwargs): + return None + + def adjust_request(self, request): + return request + + +class _ToolCallingParser(DelegatingParser): + """Parser that extracts a hardcoded tool call from any input.""" + + def __init__(self, tokenizer, *args, **kwargs): + super().__init__(tokenizer) + self._tool_parser = _StubToolParser() + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + return False + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + return input_ids + + def extract_reasoning(self, model_output, request): + return None, model_output + + def extract_reasoning_streaming(self, *args, **kwargs): + return None + + def extract_tool_calls_streaming(self, *args, **kwargs): + return None + + def parse_delta(self, *args, **kwargs) -> DeltaMessage | None: + return None + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_request(**overrides) -> ResponsesRequest: + defaults = {"model": "test-model", "input": "test"} + defaults.update(overrides) + return ResponsesRequest.model_validate(defaults) + + +def _make_output( + text: str = "Hello, world!", + token_ids: Sequence[int] = (1, 2, 3), + finish_reason: str = "stop", +) -> CompletionOutput: + return CompletionOutput( + index=0, + text=text, + token_ids=list(token_ids), + cumulative_logprob=None, + logprobs=None, + finish_reason=finish_reason, + ) + + +def _make_parser(parser_cls, **overrides): + defaults = dict( + tokenizer=MagicMock(), + parser_cls=parser_cls, + response_messages=[], + request=_make_request(), + chat_template=None, + chat_template_content_format="auto", + ) + defaults.update(overrides) + return ResponsesParser(**defaults) + + +# --------------------------------------------------------------------------- +# Tests: basic text passthrough +# --------------------------------------------------------------------------- + + +def test_process_text_with_parser(): + """Parser with no reasoning/tools returns a single message item.""" + parser = _make_parser(_NoOpParser) + parser.process(_make_output(text="Hello!")) + + assert len(parser.response_messages) == 1 + msg = parser.response_messages[0] + assert msg.type == "message" + assert msg.content[0].text == "Hello!" + + +def test_process_text_without_parser(): + """parser_cls=None falls back to plain text wrapping.""" + parser = _make_parser(None) + parser.process(_make_output(text="Hello!")) + + assert len(parser.response_messages) == 1 + msg = parser.response_messages[0] + assert msg.type == "message" + assert msg.content[0].text == "Hello!" + + +# --------------------------------------------------------------------------- +# Tests: empty / whitespace output +# --------------------------------------------------------------------------- + + +def test_process_empty_text_without_parser(): + """Empty text with no parser produces no output items.""" + parser = _make_parser(None) + parser.process(_make_output(text="")) + + assert len(parser.response_messages) == 0 + + +def test_process_empty_text_with_parser(): + """Empty text with parser produces no output items.""" + parser = _make_parser(_NoOpParser) + parser.process(_make_output(text="")) + + assert len(parser.response_messages) == 0 + + +# --------------------------------------------------------------------------- +# Tests: reasoning extraction +# --------------------------------------------------------------------------- + + +def test_process_extracts_reasoning(): + """Parser that finds reasoning produces both reasoning and message items.""" + parser = _make_parser(_ReasoningOnlyParser) + parser.process(_make_output(text="Let me checkThe answer is 42")) + + types = [m.type for m in parser.response_messages] + assert "reasoning" in types + assert "message" in types + + reasoning_item = next(m for m in parser.response_messages if m.type == "reasoning") + assert reasoning_item.content[0].text == "Let me check" + + message_item = next(m for m in parser.response_messages if m.type == "message") + assert message_item.content[0].text == "The answer is 42" + + +def test_process_reasoning_only_no_content(): + """When reasoning consumes all text, only a reasoning item is produced.""" + parser = _make_parser(_ReasoningOnlyParser) + parser.process(_make_output(text="Just thinking")) + + types = [m.type for m in parser.response_messages] + assert "reasoning" in types + assert "message" not in types + + +# --------------------------------------------------------------------------- +# Tests: tool call extraction +# --------------------------------------------------------------------------- + + +def test_process_extracts_tool_calls(): + """Parser that finds tool calls produces function_call items.""" + request = _make_request( + tool_choice="auto", + tools=[ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + } + ], + ) + parser = _make_parser(_ToolCallingParser, request=request, enable_auto_tools=True) + parser.process(_make_output(text="calling tool")) + + types = [m.type for m in parser.response_messages] + assert "function_call" in types + + tool_item = next(m for m in parser.response_messages if m.type == "function_call") + assert tool_item.name == "get_weather" + assert tool_item.arguments == '{"location": "Paris"}' + assert tool_item.status == "completed" + + +# --------------------------------------------------------------------------- +# Tests: finish_reason tracking +# --------------------------------------------------------------------------- + + +def test_finish_reason_tracked(): + """finish_reason from CompletionOutput is stored on the parser.""" + parser = _make_parser(_NoOpParser) + assert parser.finish_reason is None + + parser.process(_make_output(finish_reason="stop")) + assert parser.finish_reason == "stop" + + parser.process(_make_output(finish_reason="length")) + assert parser.finish_reason == "length" + + +# --------------------------------------------------------------------------- +# Tests: multi-turn accumulation +# --------------------------------------------------------------------------- + + +def test_multi_turn_accumulation(): + """Multiple process() calls accumulate response_messages.""" + parser = _make_parser(_NoOpParser) + + parser.process(_make_output(text="First turn")) + parser.process(_make_output(text="Second turn")) + + assert len(parser.response_messages) == 2 + texts = [m.content[0].text for m in parser.response_messages] + assert texts == ["First turn", "Second turn"] + + +def test_num_init_messages_offset(): + """Initial messages are preserved and offset works correctly.""" + init_messages = [MagicMock(type="message")] + parser = _make_parser(_NoOpParser, response_messages=init_messages) + + assert parser.num_init_messages == 1 + + parser.process(_make_output(text="New output")) + + assert len(parser.response_messages) == 2 + items = parser.make_response_output_items_from_parsable_context() + assert len(items) == 1 + assert items[0].type == "message" + + +# --------------------------------------------------------------------------- +# Tests: factory function +# --------------------------------------------------------------------------- + + +def test_factory_function_creates_parser(): + """get_responses_parser_for_simple_context returns a working parser.""" + rp = get_responses_parser_for_simple_context( + tokenizer=MagicMock(), + parser_cls=_NoOpParser, + response_messages=[], + request=_make_request(), + chat_template=None, + chat_template_content_format="auto", + ) + assert isinstance(rp, ResponsesParser) + + rp.process(_make_output(text="Works!")) + assert len(rp.response_messages) == 1 + + +def test_factory_function_none_parser(): + """Factory function works with parser_cls=None.""" + rp = get_responses_parser_for_simple_context( + tokenizer=MagicMock(), + parser_cls=None, + response_messages=[], + request=_make_request(), + chat_template=None, + chat_template_content_format="auto", + ) + assert isinstance(rp, ResponsesParser) + assert rp.parser_instance is None diff --git a/tests/entrypoints/openai/test_tool_choice_content_none.py b/tests/entrypoints/openai/test_tool_choice_content_none.py index c1da5918697..75a5c578cca 100644 --- a/tests/entrypoints/openai/test_tool_choice_content_none.py +++ b/tests/entrypoints/openai/test_tool_choice_content_none.py @@ -4,7 +4,6 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.abstract_parser import DelegatingParser @@ -32,11 +31,8 @@ class _DummyDelegatingParser(DelegatingParser): ): return None - def extract_tool_calls(self, model_output: str, request): - return None - -def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_content(): +def test_chat_completion_named_tool_choice_with_none_content(): request = ChatCompletionRequest.model_validate( { "model": "test-model", @@ -53,17 +49,15 @@ def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_conten "tool_choice": {"type": "function", "function": {"name": "get_weather"}}, } ) + parser = _DummyDelegatingParser(tokenizer=None) - tool_calls, content = OpenAIServing._parse_tool_calls_from_content( - request=request, - tokenizer=None, - enable_auto_tools=True, - tool_parser_cls=None, + tool_calls, content = parser._extract_tool_calls( content=None, + request=request, + enable_auto_tools=True, ) assert content is None - assert tool_calls is not None assert tool_calls == [] diff --git a/tests/entrypoints/rpc/__init__.py b/tests/entrypoints/serve/dev/__init__.py similarity index 100% rename from tests/entrypoints/rpc/__init__.py rename to tests/entrypoints/serve/dev/__init__.py diff --git a/vllm/entrypoints/serve/cache/__init__.py b/tests/entrypoints/serve/dev/rpc/__init__.py similarity index 100% rename from vllm/entrypoints/serve/cache/__init__.py rename to tests/entrypoints/serve/dev/rpc/__init__.py diff --git a/tests/entrypoints/rpc/test_collective_rpc.py b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py similarity index 96% rename from tests/entrypoints/rpc/test_collective_rpc.py rename to tests/entrypoints/serve/dev/rpc/test_collective_rpc.py index 56d93a42731..eb9aa7663c9 100644 --- a/tests/entrypoints/rpc/test_collective_rpc.py +++ b/tests/entrypoints/serve/dev/rpc/test_collective_rpc.py @@ -37,7 +37,7 @@ def server(): "--max-num-seqs", "128", "--worker-extension-cls", - "tests.entrypoints.rpc.test_collective_rpc.TestWorkerExtension", + "tests.entrypoints.serve.dev.rpc.test_collective_rpc.TestWorkerExtension", ] with RemoteOpenAIServer( MODEL_NAME, diff --git a/tests/entrypoints/serve/instrumentator/test_sleep.py b/tests/entrypoints/serve/dev/test_sleep.py similarity index 100% rename from tests/entrypoints/serve/instrumentator/test_sleep.py rename to tests/entrypoints/serve/dev/test_sleep.py diff --git a/tests/entrypoints/test_chat_utils.py b/tests/entrypoints/test_chat_utils.py index afda75d4fc1..7738f4c3b04 100644 --- a/tests/entrypoints/test_chat_utils.py +++ b/tests/entrypoints/test_chat_utils.py @@ -13,6 +13,8 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ( + ConversationMessage, + _postprocess_messages, parse_chat_messages, parse_chat_messages_async, ) @@ -2714,3 +2716,29 @@ async def test_parse_chat_messages_video_vision_chunk_with_uuid_async( assert conversation == expected_conversation _assert_mm_data_is_vision_chunk_input(mm_data, 1) _assert_mm_uuids(mm_uuids, 1, expected_uuids=[video_uuid], modality="vision_chunk") + + +def test_postprocess_messages_null_arguments_string(): + """arguments="null" must not reach the chat template as Python None. + + json.loads("null") returns None, which causes Jinja2 templates that call + tc.arguments.items() to raise 'None' has no attribute 'items'. + The function should coerce it to {} instead. + """ + messages: list[ConversationMessage] = [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_current_time", "arguments": "null"}, + } + ], + } + ] + _postprocess_messages(messages) + tool_calls = messages[0]["tool_calls"] + assert tool_calls is not None + assert tool_calls[0]["function"]["arguments"] == {} diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index 6af1bfe1e7a..c3939502551 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -258,10 +258,13 @@ def varlen_with_paged_kv( # KV cache for CPU attention cache_dtype = torch.uint8 if is_fp8 else dtype - packed_key_cache = torch.empty( - num_blocks, num_kv_heads, block_size, head_size, dtype=cache_dtype + packed_key_value_cache = torch.empty( + num_blocks, num_kv_heads, block_size, head_size * 2, dtype=cache_dtype ) - packed_value_cache = torch.empty_like(packed_key_cache) + packed_key_value_cache = packed_key_value_cache.view( + (num_blocks, num_kv_heads, block_size * 2, -1) + ) + packed_key_cache, packed_value_cache = packed_key_value_cache.chunk(2, dim=2) cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( dim=0, dtype=torch.int32 diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py new file mode 100644 index 00000000000..4b800b192b2 --- /dev/null +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import pytest +import torch +import torch.nn.functional as F + +import vllm._custom_ops as ops +from vllm.platforms import 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) + +set_random_seed(12345) + +NUM_HEADS = [ + (2, 4), + (4, 4), +] +HEAD_DIMS = [ + (32, 32), + (64, 32), +] +CHUNK_SIZE = 64 +PREFILL_SEQ_LENS = [ + [1], + [1, 2, 3], + [CHUNK_SIZE - 1], + [CHUNK_SIZE], + [CHUNK_SIZE + 1], + [CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1], + [2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1], + [4 * CHUNK_SIZE + 17], +] +DECODE_BATCH_SIZES = [1, 3, 5] + + +@functools.lru_cache(maxsize=128, typed=False) +def tensor_cache( + elem_num: int, + dtype: torch.dtype, +) -> torch.Tensor: + tensor = torch.rand(elem_num, dtype=dtype) + return tensor + + +def ref_l2norm( + x: torch.Tensor, + dim: int = -1, + eps: float = 1e-5, +) -> torch.Tensor: + inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) + return x * inv_norm + + +def ref_gdn_gating( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0) + g = -torch.exp(A_log.float()) * softplus_x + beta = torch.sigmoid(b.float()).to(dtype=b.dtype) + return g, beta + + +def ref_gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + g, beta = ref_gdn_gating(A_log, a, b, dt_bias) + out = torch.empty_like(value) + final_state = torch.empty_like(initial_state) + + for seq_idx in range(cu_seqlens.numel() - 1): + begin = int(cu_seqlens[seq_idx].item()) + end = int(cu_seqlens[seq_idx + 1].item()) + q_seq = query[:, begin:end] + k_seq = key[:, begin:end] + v_seq = value[:, begin:end] + g_seq = g[begin:end].unsqueeze(0) + beta_seq = beta[begin:end].unsqueeze(0) + initial_dtype = q_seq.dtype + + if use_qk_l2norm_in_kernel: + q_seq = ref_l2norm(q_seq, dim=-1) + k_seq = ref_l2norm(k_seq, dim=-1) + + if q_seq.shape[2] != v_seq.shape[2]: + repeat_factor = v_seq.shape[2] // q_seq.shape[2] + q_seq = q_seq.repeat_interleave(repeat_factor, dim=2) + k_seq = k_seq.repeat_interleave(repeat_factor, dim=2) + + q_seq, k_seq, v_seq, beta_seq, g_seq = [ + x.transpose(1, 2).contiguous().to(torch.float32) + for x in (q_seq, k_seq, v_seq, beta_seq, g_seq) + ] + + batch_size, num_heads, seq_len, head_dim = q_seq.shape + v_head_dim = v_seq.shape[-1] + q_seq = q_seq * (1 / (head_dim**0.5)) + out_seq = torch.empty( + batch_size, + num_heads, + seq_len, + v_head_dim, + dtype=v_seq.dtype, + ) + state = initial_state[seq_idx : seq_idx + 1].to(v_seq) + + for token_idx in range(seq_len): + q_t = q_seq[:, :, token_idx] + k_t = k_seq[:, :, token_idx] + v_t = v_seq[:, :, token_idx] + g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1) + beta_t = beta_seq[:, :, token_idx].unsqueeze(-1) + + state = state * g_t + kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1) + delta = (v_t - kv_mem) * beta_t + state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2) + out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1) + + out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype) + final_state[seq_idx] = state.squeeze(0) + + return out, final_state + + +def gdn_inputs( + num_tokens: int, + num_heads: tuple[int, int], + head_dims: tuple[int, int], +) -> tuple[torch.Tensor, ...]: + num_qk_heads, num_v_heads = num_heads + head_dim, v_head_dim = head_dims + q_shape = (1, num_tokens, num_qk_heads, head_dim) + q_numel = num_tokens * num_qk_heads * head_dim + q = tensor_cache(q_numel, torch.bfloat16).view(q_shape) + k = tensor_cache(q_numel, torch.bfloat16).view(q_shape) + + v_shape = (1, num_tokens, num_v_heads, v_head_dim) + v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view( + v_shape + ) + + gate_shape = (num_tokens, num_v_heads) + gate_numel = num_tokens * num_v_heads + a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) + b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) + A_log = tensor_cache(num_v_heads, torch.float32) + dt_bias = tensor_cache(num_v_heads, torch.bfloat16) + return q, k, v, a, b, A_log, dt_bias + + +@pytest.mark.parametrize("num_tokens", [1, 9]) +@pytest.mark.parametrize("num_v_heads", [4, 8]) +@torch.inference_mode() +def test_fused_gdn_gating_cpu( + num_tokens: int, + num_v_heads: int, +) -> None: + gate_shape = (num_tokens, num_v_heads) + gate_numel = num_tokens * num_v_heads + a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) + b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape) + A_log = tensor_cache(num_v_heads, torch.float32) + dt_bias = tensor_cache(num_v_heads, torch.bfloat16) + + g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias) + g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias) + + torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3 + ) + + +# decode path +@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_dims", HEAD_DIMS) +@torch.inference_mode() +def test_fused_sigmoid_gating_delta_rule_update_cpu( + batch_size: int, + num_heads: tuple[int, int], + head_dims: tuple[int, int], +) -> None: + q, k, v, a, b, A_log, dt_bias = gdn_inputs( + num_tokens=batch_size, + num_heads=num_heads, + head_dims=head_dims, + ) + _, num_v_heads = num_heads + head_dim, v_head_dim = head_dims + state_indices = torch.arange(batch_size, dtype=torch.int32) + cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32) + state_shape = (batch_size, num_v_heads, head_dim, v_head_dim) + state = tensor_cache( + batch_size * num_v_heads * head_dim * v_head_dim, torch.float32 + ).view(state_shape) + state_ref = state[state_indices].transpose(-1, -2).contiguous() + + out_ref, final_state_ref = ref_gated_delta_rule( + query=q, + key=k, + value=v, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + initial_state=state_ref, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + out_ref = out_ref.transpose(0, 1).contiguous() + + state_out = state.clone() + out = ops.fused_sigmoid_gating_delta_rule_update_cpu( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + initial_state_source=state_out, + initial_state_indices=state_indices, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + state_out[state_indices].transpose(-1, -2), + final_state_ref, + atol=1e-2, + rtol=1e-2, + ) + + +# prefill path +@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_dims", HEAD_DIMS) +@torch.inference_mode() +def test_chunk_gated_delta_rule_cpu( + seq_lens: list[int], + num_heads: tuple[int, int], + head_dims: tuple[int, int], +) -> None: + total_tokens = sum(seq_lens) + q, k, v, a, b, A_log, dt_bias = gdn_inputs( + num_tokens=total_tokens, + num_heads=num_heads, + head_dims=head_dims, + ) + _, num_v_heads = num_heads + head_dim, v_head_dim = head_dims + cu_seqlens = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32 + ) + initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim) + initial_state = tensor_cache( + len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32 + ).view(initial_state_shape) + initial_state_ref = initial_state.transpose(-1, -2).contiguous() + + out_ref, final_state_ref = ref_gated_delta_rule( + query=q, + key=k, + value=v, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state_ref, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + + g, beta = ref_gdn_gating(A_log, a, b, dt_bias) + out, final_state = ops.chunk_gated_delta_rule_cpu( + query=q, + key=k, + value=v, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + final_state.transpose(-1, -2), + final_state_ref, + atol=1e-2, + rtol=1e-2, + ) diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 73859175cd1..ca25b8c2e9f 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -20,7 +20,12 @@ EXPERT_NUM = [ HIDDEN_DIM = [128, 2880] INTERMEDIATE_DIM = [128, 2880] BATCH_SIZE = [1, 64, 256] -ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI, MoEActivation.GELU] +ACT = [ + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.GELU, + MoEActivation.GELU_TANH, +] USE_BIAS = [True, False] ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"] DTYPE = [torch.bfloat16] diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 32336f37ac6..1380281bb2e 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import ( fp8_w8a8_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( + CutlassExpertsFp4, CutlassExpertsFp8, run_cutlass_moe_fp8, ) @@ -52,6 +53,12 @@ MNK_FACTORS = [ vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) +def test_cutlass_moe_supports_gelu_tanh_activation_metadata(): + assert CutlassExpertsFp8._supports_activation(MoEActivation.GELU_TANH) + assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH) + assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH_NO_MUL) + + @dataclasses.dataclass class MOETensors: a: torch.Tensor diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index dbb23b9dfad..d294820340d 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -6,6 +6,7 @@ import pytest import torch from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils +from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max from vllm.platforms import current_platform @@ -26,7 +27,9 @@ from vllm.platforms import current_platform @pytest.mark.parametrize("tma_aligned", [False, True]) @pytest.mark.parametrize("scale_ue8m0", [False, True]) @pytest.mark.parametrize("group_size", [64, 128]) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm." +) def test_per_token_group_quant_fp8( shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int ): @@ -47,7 +50,7 @@ def test_per_token_group_quant_fp8( ) # triton ref - with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( x, group_size, @@ -87,7 +90,8 @@ def test_per_token_group_quant_fp8( ) @pytest.mark.parametrize("poisoned_scales", [False, True]) @pytest.mark.skipif( - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" + not current_platform.is_cuda_alike(), + reason="DeepGEMM not available on this platform", ) def test_per_token_group_quant_fp8_packed( num_tokens, hidden_dim, group_size, poisoned_scales @@ -109,8 +113,8 @@ def test_per_token_group_quant_fp8_packed( if poisoned_scales: # Call the kernel with poisoned scale buffer to # ensure padded indices are correctly zeroed. - fp8_dtype = torch.float8_e4m3fn - finfo = torch.finfo(fp8_dtype) + fp8_dtype = current_platform.fp8_dtype() + fp8_min, fp8_max = get_fp8_min_max() out_q = torch.empty_like(x, dtype=fp8_dtype) out_s_packed = torch.empty_strided( (mn, k_num_packed), @@ -125,8 +129,8 @@ def test_per_token_group_quant_fp8_packed( out_s_packed, group_size, 1e-10, - finfo.min, - finfo.max, + fp8_min, + fp8_max, ) else: out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( @@ -136,7 +140,7 @@ def test_per_token_group_quant_fp8_packed( ) # Triton reference (row-major float32 scales, UE8M0) - with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( x, group_size, @@ -167,7 +171,8 @@ def test_per_token_group_quant_fp8_packed( @pytest.mark.skipif( - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" + not current_platform.is_cuda_alike(), + reason="DeepGEMM not available on this platform", ) def test_per_token_group_quant_fp8_packed_all_zero(): """All-zero input must produce well-defined UE8M0 scale bytes via the eps @@ -226,7 +231,8 @@ def test_per_token_group_quant_fp8_packed_all_zero(): @pytest.mark.skipif( - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" + not current_platform.is_cuda_alike(), + reason="DeepGEMM not available on this platform", ) def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): """Inputs whose absmax/max_8bit produces a non-power-of-2 force the @@ -254,7 +260,7 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): use_ue8m0=True, ) - with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( x, group_size, @@ -296,7 +302,8 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): ], ) @pytest.mark.skipif( - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" + not current_platform.is_cuda_alike(), + reason="DeepGEMM not available on this platform", ) def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( num_tokens, hidden_dim @@ -315,8 +322,8 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( k_num_packed = (groups_per_row + 3) // 4 tma_aligned_mn = ((mn + 3) // 4) * 4 - fp8_dtype = torch.float8_e4m3fn - finfo = torch.finfo(fp8_dtype) + fp8_dtype = current_platform.fp8_dtype() + fp8_min, fp8_max = get_fp8_min_max() # Allocate output_q with the padded mn extent and pre-fill with 0xFF # so the kernel cannot rely on a clean buffer. out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype) @@ -330,11 +337,11 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( ) torch.ops._C.per_token_group_fp8_quant_packed( - x, out_q, out_s_packed, group_size, 1e-10, finfo.min, finfo.max + x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max ) # Live rows must match the Triton reference. - with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True) assert torch.equal(out_q[:mn], ref_q), "Live region mismatch" @@ -369,7 +376,7 @@ def test_per_token_group_quant_int8(shape, group_size: int): ) # triton ref - with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): ref_q, ref_s = int8_utils.per_token_group_quant_int8( x, group_size, diff --git a/tests/lora/test_minicpmv_tp.py b/tests/lora/test_minicpmv_tp.py index 0090f9c569b..c552c4a3488 100644 --- a/tests/lora/test_minicpmv_tp.py +++ b/tests/lora/test_minicpmv_tp.py @@ -1,10 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from importlib.metadata import version - import pytest -from packaging.version import Version import vllm from vllm.assets.image import ImageAsset @@ -13,14 +10,6 @@ from vllm.platforms import current_platform from ..utils import multi_gpu_test -pytestmark = pytest.mark.skipif( - Version("5.0") <= Version(version("transformers")), - reason=( - "MiniCPMV custom processor uses tokenizer.im_start_id which is not " - "available on TokenizersBackend in transformers v5.0+" - ), -) - MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5" PROMPT_TEMPLATE = ( diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index 59d53dd606a..f4de9bc9038 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -41,8 +41,10 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch): assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) -def test_rocm_unquantized_gemm_gfx1x_n_gt_4_falls_back(monkeypatch): - x = torch.randn(5, 64, dtype=torch.float16) +def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch): + # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must + # fall back to torch.nn.functional.linear. + x = torch.randn(6, 64, dtype=torch.float16) weight = torch.randn(128, 64, dtype=torch.float16) monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False) diff --git a/tests/model_executor/model_loader/tensorizer_loader/conftest.py b/tests/model_executor/model_loader/tensorizer_loader/conftest.py index 6c85a139919..051890e89a1 100644 --- a/tests/model_executor/model_loader/tensorizer_loader/conftest.py +++ b/tests/model_executor/model_loader/tensorizer_loader/conftest.py @@ -87,10 +87,6 @@ class DummyExecutor(UniProcExecutor): self.collective_rpc("init_worker", args=([kwargs],)) self.collective_rpc("init_device") - @property - def max_concurrent_batches(self) -> int: - return 2 - def shutdown(self): if hasattr(self, "thread_pool"): self.thread_pool.shutdown(wait=False) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index 0f6dccd8477..0a290a00a83 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -28,6 +28,22 @@ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors from vllm.platforms import current_platform +def _fp8_reload_unsupported() -> bool: + """Whether the FP8 reload/online-quantize tests should be skipped. + + ``supports_fp8()`` returns True on MI250 (gfx90a) because the general + quantization paths upcast FP8 weights, but gfx90a has no native FP8 and + cannot run these reload models, so treat it as unsupported here. + """ + if not current_platform.supports_fp8(): + return True + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx90a + + return on_gfx90a() + return False + + class _AliasedBufferLayer(torch.nn.Module): def __init__(self): super().__init__() @@ -284,7 +300,7 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): if current_platform.device_count() < tp_size: pytest.skip(reason="Not enough CUDA devices") - if "FP8" in base_model and not current_platform.supports_fp8(): + if "FP8" in base_model and _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") with vllm_runner( @@ -308,7 +324,7 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): def test_kv_scale_reload(vllm_runner): """Test reloading a checkpoint that contains k_scale/v_scale weights.""" - if not current_platform.supports_fp8(): + if _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV" @@ -378,7 +394,7 @@ def test_online_quantize_reload( if current_platform.device_count() < tp_size: pytest.skip(reason="Not enough GPU devices") - if quantization == "fp8" and not current_platform.supports_fp8(): + if quantization == "fp8" and _fp8_reload_unsupported(): pytest.skip(reason="Requires FP8 support") with vllm_runner( diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 6160280993d..9ac0d4ab446 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -785,8 +785,6 @@ VLM_TEST_SETTINGS = { get_stop_token_ids=lambda tok: [tok.eos_id, tok.eot_id], hf_output_post_proc=model_utils.minicpmv_trunc_hf_output, patch_hf_runner=model_utils.minicpmv_25_patch_hf_runner, - # FIXME: https://huggingface.co/openbmb/MiniCPM-V-2_6/discussions/55 - marks=[pytest.mark.skip("HF import fails")], ), "minicpmo_26": VLMTestInfo( models=["openbmb/MiniCPM-o-2_6"], @@ -800,8 +798,6 @@ VLM_TEST_SETTINGS = { ), hf_output_post_proc=model_utils.minicpmv_trunc_hf_output, patch_hf_runner=model_utils.minicpmo_26_patch_hf_runner, - # FIXME: https://huggingface.co/openbmb/MiniCPM-o-2_6/discussions/49 - marks=[pytest.mark.skip("HF import fails")], ), "minicpmv_26": VLMTestInfo( models=["openbmb/MiniCPM-V-2_6"], diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 7c024052a43..b82bcec9dca 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -10,8 +10,12 @@ import pytest from vllm.assets.base import get_vllm_public_assets from vllm.multimodal.video import ( VIDEO_LOADER_REGISTRY, + DynamicVideoBackend, + Molmo2VideoBackend, VideoLoader, + get_video_loader_backend_for_processor, ) +from vllm.transformers_utils.processor import get_video_processor_cls_name_from_config from .utils import create_long_gop_video, create_video_from_image @@ -54,6 +58,50 @@ def test_video_loader_type_doesnt_exist(): VIDEO_LOADER_REGISTRY.load("non_existing_video_loader") +# ============================================================================ +# Video Processor → Video Loader Tests (via model repo) +# ============================================================================ + + +@pytest.mark.parametrize( + "model_repo, expected_loader_cls", + [ + pytest.param( + "allenai/Molmo2-4B", + Molmo2VideoBackend, + id="molmo2", + ), + pytest.param( + "zai-org/GLM-4.1V-9B-Thinking", + DynamicVideoBackend, + id="glm4v", + ), + ], +) +def test_video_processor_from_model_repo( + model_repo: str, + expected_loader_cls: type, +): + """Test that a model repo resolves to the correct video loader backend. + + The test downloads the preprocessor config from HuggingFace Hub, + extracts the ``video_processor_type`` field, and verifies it maps + to the expected backend and loader class. + """ + video_processor = get_video_processor_cls_name_from_config(model_repo) + assert video_processor is not None, ( + f"Model repo {model_repo!r} did not contain a video_processor_type " + f"in its preprocessor config" + ) + + backend = get_video_loader_backend_for_processor(video_processor) + loader = VIDEO_LOADER_REGISTRY.load(backend) + assert isinstance(loader, expected_loader_cls), ( + f"{model_repo!r}: backend={backend!r} loaded " + f"{type(loader)}, expected {expected_loader_cls}" + ) + + def test_video_backend_handles_broken_frames(monkeypatch: pytest.MonkeyPatch): """ Regression test for handling videos with broken frames. diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py new file mode 100644 index 00000000000..ba8bc1427f2 --- /dev/null +++ b/tests/parser/test_parse.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser + + +class ThinkReasoningParser(BaseThinkingReasoningParser): + @property + def start_token(self) -> str: + return "" + + @property + def end_token(self) -> str: + return "" + + +MODEL_OUTPUT = ( + "let me think about this" + '\n{"name": "get_weather", ' + '"arguments": {"city": "Dallas"}}\n' +) + +PLAIN_TEXT = "The weather in Dallas is sunny and 75°F." + +TOOL_CALL_ONLY = ( + '\n{"name": "get_weather", ' + '"arguments": {"city": "Dallas"}}\n' +) + +TOOL_ARGUMENTS = '{"city": "Dallas"}' + + +@pytest.fixture(scope="module") +def tokenizer(): + from vllm.tokenizers import get_tokenizer + + return get_tokenizer("Qwen/Qwen3-32B") + + +def make_request(**overrides): + base = { + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + } + base.update(overrides) + return ChatCompletionRequest.model_validate(base) + + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {}}, + }, + } +] + + +def make_parser(tokenizer, reasoning=False, tool=False): + class TestParser(DelegatingParser): + reasoning_parser_cls = ThinkReasoningParser if reasoning else None + tool_parser_cls = Hermes2ProToolParser if tool else None + + return TestParser(tokenizer) + + +@pytest.mark.parametrize( + "reasoning,tool", + [(False, False), (False, True)], + ids=["neither", "tool-only"], +) +def test_parse_plain_text_no_reasoning_parser(tokenizer, reasoning, tool): + parser = make_parser(tokenizer, reasoning=reasoning, tool=tool) + request = make_request() + r, content, tool_calls = parser.parse(PLAIN_TEXT, request) + + assert r is None + assert content == PLAIN_TEXT + assert tool_calls is not None + assert len(tool_calls) == 0 + + +@pytest.mark.parametrize( + "reasoning,tool", + [(True, False), (True, True)], + ids=["reasoning-only", "both"], +) +def test_parse_plain_text_with_reasoning_parser(tokenizer, reasoning, tool): + parser = make_parser(tokenizer, reasoning=reasoning, tool=tool) + request = make_request() + r, content, tool_calls = parser.parse(PLAIN_TEXT, request) + + assert r == PLAIN_TEXT + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 0 + + +def test_parse_both_parsers(tokenizer): + parser = make_parser(tokenizer, reasoning=True, tool=True) + request = make_request(tools=TOOLS) + reasoning, content, tool_calls = parser.parse( + MODEL_OUTPUT, request, enable_auto_tools=True + ) + + assert reasoning is not None + assert "let me think about this" in reasoning + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "get_weather" + assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"} + assert not content or content.strip() == "" + + +def test_parse_reasoning_only(tokenizer): + parser = make_parser(tokenizer, reasoning=True, tool=False) + request = make_request() + reasoning, content, tool_calls = parser.parse(MODEL_OUTPUT, request) + + assert reasoning is not None + assert "let me think about this" in reasoning + assert content is not None + assert "" in content + assert "get_weather" in content + assert tool_calls is not None + assert len(tool_calls) == 0 + + +def test_parse_tool_only(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + request = make_request(tools=TOOLS) + reasoning, content, tool_calls = parser.parse( + MODEL_OUTPUT, request, enable_auto_tools=True + ) + + assert reasoning is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "get_weather" + assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"} + + +def test_parse_named_tool_choice(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + request = make_request( + tools=TOOLS, + tool_choice={ + "type": "function", + "function": {"name": "get_weather"}, + }, + ) + reasoning, content, tool_calls = parser.parse( + TOOL_ARGUMENTS, request, enable_auto_tools=True + ) + + assert reasoning is None + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "get_weather" + assert tool_calls[0].arguments == TOOL_ARGUMENTS + + +def test_parse_named_tool_choice_with_reasoning(tokenizer): + parser = make_parser(tokenizer, reasoning=True, tool=True) + model_output = f"thinking{TOOL_ARGUMENTS}" + request = make_request( + tools=TOOLS, + tool_choice={ + "type": "function", + "function": {"name": "get_weather"}, + }, + ) + reasoning, content, tool_calls = parser.parse( + model_output, request, enable_auto_tools=True + ) + + assert reasoning is not None + assert "thinking" in reasoning + assert content is None + assert len(tool_calls) == 1 + assert tool_calls[0].name == "get_weather" + assert tool_calls[0].arguments == TOOL_ARGUMENTS + + +def test_parse_required_tool_choice(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + functions_json = json.dumps( + [ + {"name": "get_weather", "parameters": {"city": "Dallas"}}, + {"name": "get_time", "parameters": {"timezone": "UTC"}}, + ] + ) + request = make_request(tools=TOOLS, tool_choice="required") + reasoning, content, tool_calls = parser.parse( + functions_json, request, enable_auto_tools=True + ) + + assert reasoning is None + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 2 + assert tool_calls[0].name == "get_weather" + assert json.loads(tool_calls[0].arguments) == {"city": "Dallas"} + assert tool_calls[1].name == "get_time" + assert json.loads(tool_calls[1].arguments) == {"timezone": "UTC"} + + +def test_parse_named_tool_choice_content_none(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + request = make_request( + tools=TOOLS, + tool_choice={ + "type": "function", + "function": {"name": "get_weather"}, + }, + ) + reasoning, content, tool_calls = parser.parse("", request, enable_auto_tools=True) + assert reasoning is None + assert content is None + assert tool_calls is not None + + +def test_parse_required_tool_choice_content_none(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + request = make_request(tools=TOOLS, tool_choice="required") + reasoning, content, tool_calls = parser.parse("", request, enable_auto_tools=True) + assert reasoning is None + assert content is None + assert tool_calls is not None + assert len(tool_calls) == 0 + + +def test_parse_auto_tools_no_parser(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=False) + request = make_request() + reasoning, content, tool_calls = parser.parse( + TOOL_CALL_ONLY, request, enable_auto_tools=True + ) + + assert reasoning is None + assert content == TOOL_CALL_ONLY + assert tool_calls is not None + assert len(tool_calls) == 0 + + +def test_parse_auto_tools_no_calls_returns_none(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + request = make_request(tools=TOOLS) + reasoning, content, tool_calls = parser.parse( + PLAIN_TEXT, request, enable_auto_tools=True + ) + + assert reasoning is None + assert content == PLAIN_TEXT + assert tool_calls is None diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index c4409117ad9..2ba2392f8e9 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -7,7 +7,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.parser.abstract_parser import _WrappedParser +from vllm.parser.abstract_parser import DelegatingParser from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser @@ -45,9 +45,11 @@ def request_obj(): def make_parser(tokenizer, reasoning=False, tool=False): - _WrappedParser.reasoning_parser_cls = ThinkReasoningParser if reasoning else None - _WrappedParser.tool_parser_cls = Hermes2ProToolParser if tool else None - return _WrappedParser(tokenizer) + class TestParser(DelegatingParser): + reasoning_parser_cls = ThinkReasoningParser if reasoning else None + tool_parser_cls = Hermes2ProToolParser if tool else None + + return TestParser(tokenizer) def stream_text(parser, tokenizer, text, request, prompt_token_ids=None): @@ -56,7 +58,11 @@ def stream_text(parser, tokenizer, text, request, prompt_token_ids=None): for tid in token_ids: delta_text = tokenizer.decode([tid]) result = parser.parse_delta( - delta_text, [tid], request, prompt_token_ids=prompt_token_ids + delta_text, + [tid], + request, + prompt_token_ids=prompt_token_ids, + finished=False, ) prompt_token_ids = None results.append(result) @@ -144,7 +150,11 @@ def stream_chunks(parser, tokenizer, chunks, request_obj): for chunk in chunks: delta_text = tokenizer.decode(chunk) result = parser.parse_delta( - delta_text, chunk, request_obj, prompt_token_ids=prompt_token_ids + delta_text, + chunk, + request_obj, + prompt_token_ids=prompt_token_ids, + finished=False, ) prompt_token_ids = None results.append(result) @@ -235,3 +245,86 @@ def test_parse_delta_reasoning_only_thinking_disabled(tokenizer, request_obj): assert "Hello" in content assert "assist" in content assert len(tool_calls) == 0 + + +def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, request_obj): + """When finished=True but the final parse_delta produces no + tool-call delta, unstreamed args are not flushed.""" + parser = make_parser(tokenizer, reasoning=False, tool=True) + + results = stream_text( + parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] + ) + _, _, tool_calls = collect_fields(results) + assert len(tool_calls) > 0 + + streamed = parser._tool_parser.streamed_args_for_tool[0] + assert len(streamed) > 5 + parser._tool_parser.streamed_args_for_tool[0] = streamed[:-5] + + # Prevent normal extraction from catching the gap — without a + # tool-call delta to merge into, the flush is skipped. + parser._tool_parser.extract_tool_calls_streaming = lambda *a, **kw: None + + flush_result = parser.parse_delta("", [], request_obj, finished=True) + assert flush_result is None or flush_result.tool_calls is None + + +def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, request_obj): + """When all args have been streamed, finished=True must not + produce extra or duplicate arguments.""" + parser = make_parser(tokenizer, reasoning=False, tool=True) + results = stream_text( + parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[] + ) + _, _, tool_calls = collect_fields(results) + + assert len(tool_calls) > 0 + assert tool_calls[0].function.name == "get_weather" + tool_args = "".join( + tc.function.arguments for tc in tool_calls if tc.function.arguments + ) + assert json.loads(tool_args) == {"city": "Dallas"} + + flush_result = parser.parse_delta("", [], request_obj, finished=True) + assert flush_result is None or flush_result.tool_calls is None + + +def test_parse_delta_finished_appends_remaining_args(tokenizer, request_obj): + """When finished=True and the tool parser has unstreamed args, + parse_delta appends the remaining arguments to the tool-call delta.""" + parser = make_parser(tokenizer, reasoning=False, tool=True) + token_ids = tokenizer.encode(MODEL_OUTPUT, add_special_tokens=False) + + remainder = ',"unit":"celsius"}' + prompt_ids: list[int] | None = [] + results: list[DeltaMessage | None] = [] + for i, tid in enumerate(token_ids): + prev = results[-1] if results else None + prev_had_args = ( + prev + and prev.tool_calls + and any(tc.function and tc.function.arguments for tc in prev.tool_calls) + ) + + if prev_had_args: + parser._tool_parser.get_remaining_unstreamed_args = lambda: remainder + + result = parser.parse_delta( + tokenizer.decode([tid]), + [tid], + request_obj, + prompt_token_ids=prompt_ids, + finished=prev_had_args, + ) + prompt_ids = None + results.append(result) + + if prev_had_args: + break + + _, _, tool_calls = collect_fields(results) + tool_args = "".join( + tc.function.arguments for tc in tool_calls if tc.function.arguments + ) + assert tool_args.endswith(remainder) diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py new file mode 100644 index 00000000000..c4b0ab5a846 --- /dev/null +++ b/tests/quantization/test_moe_wna16.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method +from vllm.platforms import current_platform + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") +def test_moe_wna16_apply_passes_layer_activation(monkeypatch): + captured_kwargs = {} + + def fake_fused_experts(*args, **kwargs): + captured_kwargs.update(kwargs) + return torch.empty(1, 2) + + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.fused_experts", + fake_fused_experts, + ) + + method = object.__new__(MoeWNA16Method) + method.moe = SimpleNamespace(disable_inplace=False) + method.moe_quant_config = object() + layer = SimpleNamespace( + w13_qweight=torch.empty(1, 2), + w2_qweight=torch.empty(1, 2), + activation=MoEActivation.GELU_TANH, + apply_router_weight_on_input=False, + global_num_experts=1, + expert_map=None, + ) + + output = method.apply( + layer, + x=torch.empty(1, 2), + topk_weights=torch.empty(1, 1), + topk_ids=torch.empty(1, 1, dtype=torch.int32), + shared_experts=None, + shared_experts_input=None, + ) + + assert output.shape == (1, 2) + assert captured_kwargs["activation"] is MoEActivation.GELU_TANH diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index c2eb6556394..0545457eb7a 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -7,6 +7,9 @@ from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.renderers.hf import ( + _consolidate_system_messages, + _convert_developer_to_system, + _detect_developer_role_support, _get_hf_base_chat_template_params, _try_extract_ast, resolve_chat_template, @@ -592,3 +595,322 @@ def test_get_gen_prompt( f"The generated prompt does not match the expected output for " f"model {model} and template {template}" ) + + +class TestConvertDeveloperToSystem: + def test_converts_role(self): + conversation = [ + {"role": "developer", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + result = _convert_developer_to_system(conversation) + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful." + assert result[1]["role"] == "user" + + def test_removes_tools_key(self): + conversation = [ + { + "role": "developer", + "content": "Instructions", + "tools": [{"type": "function"}], + }, + ] + result = _convert_developer_to_system(conversation) + assert "tools" not in result[0] + + def test_no_developer_messages_unchanged(self): + conversation = [ + {"role": "system", "content": "System prompt"}, + {"role": "user", "content": "Hello"}, + ] + result = _convert_developer_to_system(conversation) + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + def test_does_not_mutate_original(self): + original = { + "role": "developer", + "content": "Instructions", + "tools": [{"type": "function"}], + } + conversation = [original] + _convert_developer_to_system(conversation) + assert original["role"] == "developer" + assert "tools" in original + + +# --- Developer role detection and conversion tests --- + +CHATML_TEMPLATE = ( + "{% for message in messages %}" + "{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}" + "{% if (loop.last and add_generation_prompt) or not loop.last %}" + "{{ '<|im_end|>' + '\\n'}}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}" + "{{ '<|im_start|>assistant\\n' }}" + "{% endif %}" +) + +TEMPLATE_WITH_DEVELOPER = ( + "{% for message in messages %}" + "{% if message['role'] == 'developer' %}" + "{{'<|im_start|>developer\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'system' %}" + "{{'<|im_start|>system\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'user' %}" + "{{'<|im_start|>user\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'assistant' %}" + "{{'<|im_start|>assistant\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|im_start|>assistant\\n' }}" + "{% endif %}" +) + +STRICT_ROLE_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "{{'<|im_start|>system\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'user' %}" + "{{'<|im_start|>user\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'assistant' %}" + "{{'<|im_start|>assistant\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% else %}" + "{{ raise_exception('Unexpected message role: ' + message['role']) }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|im_start|>assistant\\n' }}" + "{% endif %}" +) + + +class TestDetectDeveloperRoleSupport: + def test_absent_in_chatml(self): + assert _detect_developer_role_support(CHATML_TEMPLATE) is False + + def test_present_double_quotes(self): + assert _detect_developer_role_support(TEMPLATE_WITH_DEVELOPER) is True + + def test_present_single_quotes(self): + template = TEMPLATE_WITH_DEVELOPER.replace('"developer"', "'developer'") + assert _detect_developer_role_support(template) is True + + def test_absent_in_strict_template(self): + assert _detect_developer_role_support(STRICT_ROLE_TEMPLATE) is False + + +class TestSafeApplyChatTemplateDeveloperRole: + @pytest.fixture + def model_config(self): + return ModelConfig( + "facebook/opt-125m", + tokenizer="facebook/opt-125m", + tokenizer_mode="auto", + trust_remote_code=False, + dtype="float16", + ) + + @pytest.fixture + def tokenizer(self): + return get_tokenizer("facebook/opt-125m") + + def test_developer_converted_to_system_for_chatml(self, model_config, tokenizer): + conversation = [ + {"role": "developer", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=CHATML_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>system" in result + assert "You are a helpful assistant." in result + assert "<|im_start|>developer" not in result + + def test_developer_preserved_when_template_supports_it( + self, model_config, tokenizer + ): + conversation = [ + {"role": "developer", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=TEMPLATE_WITH_DEVELOPER, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>developer" in result + assert "You are a helpful assistant." in result + + def test_developer_does_not_crash_strict_template(self, model_config, tokenizer): + conversation = [ + {"role": "developer", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=STRICT_ROLE_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>system" in result + assert "You are a helpful assistant." in result + + def test_no_developer_messages_no_overhead(self, model_config, tokenizer): + conversation = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=CHATML_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>system" in result + assert "You are helpful." in result + + def test_developer_at_non_first_position_consolidated( + self, model_config, tokenizer + ): + conversation = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + {"role": "developer", "content": "Be concise."}, + {"role": "user", "content": "What is 2+2?"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=SYSTEM_FIRST_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>system" in result + assert "You are helpful." in result + assert "Be concise." in result + assert "What is 2+2?" in result + + def test_developer_only_no_prior_system(self, model_config, tokenizer): + conversation = [ + {"role": "user", "content": "Hello"}, + {"role": "developer", "content": "Be concise."}, + {"role": "user", "content": "What is 2+2?"}, + ] + result = safe_apply_chat_template( + model_config, + tokenizer, + conversation, + chat_template=SYSTEM_FIRST_TEMPLATE, + tokenize=False, + add_generation_prompt=True, + ) + assert "<|im_start|>system" in result + assert "Be concise." in result + + +SYSTEM_FIRST_TEMPLATE = ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}" + "{% if not loop.first %}" + "{{ raise_exception('System message must be at the beginning.') }}" + "{% endif %}" + "{{'<|im_start|>system\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'user' %}" + "{{'<|im_start|>user\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% elif message['role'] == 'assistant' %}" + "{{'<|im_start|>assistant\\n' + message['content'] + '<|im_end|>\\n'}}" + "{% else %}" + "{{ raise_exception('Unexpected message role: ' + message['role']) }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "{{ '<|im_start|>assistant\\n' }}" + "{% endif %}" +) + + +class TestConsolidateSystemMessages: + def test_no_system_messages_unchanged(self): + conversation = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + ] + result = _consolidate_system_messages(conversation) + assert result == conversation + + def test_single_system_at_start_unchanged(self): + conversation = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + result = _consolidate_system_messages(conversation) + assert result == conversation + + def test_system_at_non_first_position_moved(self): + conversation = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "You are helpful."}, + ] + result = _consolidate_system_messages(conversation) + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful." + assert result[1]["role"] == "user" + assert result[1]["content"] == "Hello" + + def test_multiple_system_messages_merged(self): + conversation = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Be concise."}, + ] + result = _consolidate_system_messages(conversation) + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful.\n\nBe concise." + assert result[1]["role"] == "user" + + def test_list_content_handled(self): + conversation = [ + {"role": "user", "content": "Hello"}, + { + "role": "system", + "content": [ + {"type": "text", "text": "Rule 1."}, + {"type": "text", "text": "Rule 2."}, + ], + }, + ] + result = _consolidate_system_messages(conversation) + assert result[0]["role"] == "system" + assert result[0]["content"] == "Rule 1.\nRule 2." + assert result[1]["role"] == "user" + + def test_does_not_mutate_original(self): + conversation = [ + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "You are helpful."}, + ] + original_len = len(conversation) + _consolidate_system_messages(conversation) + assert len(conversation) == original_len + assert conversation[0]["role"] == "user" + assert conversation[1]["role"] == "system" diff --git a/tests/test_config.py b/tests/test_config.py index c0bd4b14ff8..b78570e54fb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,6 +90,26 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): ), True, ), + ( + SimpleNamespace( + model="meta-llama/Llama-3.2-1B", + architectures=["LlamaForCausalLM"], + runner_type="generate", + is_moe=False, + is_quantized=False, + ), + True, + ), + ( + SimpleNamespace( + model="mistralai/Mistral-7B-v0.1", + architectures=["MistralForCausalLM"], + runner_type="generate", + is_moe=False, + is_quantized=False, + ), + True, + ), ( SimpleNamespace( model="facebook/opt-125m", diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/tool_parsers/test_mistral_tool_parser.py index f6a5c6bfb26..c9582159abb 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/tool_parsers/test_mistral_tool_parser.py @@ -1382,7 +1382,20 @@ def test_adjust_request_non_mistral_tokenizer( [ {"regex": r"\d+"}, {"choice": ["a", "b"]}, - {"structural_tag": '{"key": "value"}'}, + { + "structural_tag": json.dumps( + { + "structures": [ + { + "begin": "", + "schema": {"type": "object"}, + "end": "", + } + ], + "triggers": [""], + } + ) + }, {"grammar": "start: 'hello'"}, ], ids=["regex", "choice", "structural_tag", "grammar"], @@ -1404,7 +1417,18 @@ def test_adjust_request_unsupported_response_format( ) -> None: request = _make_request( response_format=StructuralTagResponseFormat( - type="structural_tag", format={"some": "config"} + type="structural_tag", + format={ + "type": "triggered_tags", + "tags": [ + { + "begin": "", + "content": {"type": "any_text"}, + "end": "", + } + ], + "triggers": [""], + }, ), ) result = mistral_tool_parser.adjust_request(request) diff --git a/tests/utils_/test_import_utils.py b/tests/utils_/test_import_utils.py index d1f822037ac..464f209f0f2 100644 --- a/tests/utils_/test_import_utils.py +++ b/tests/utils_/test_import_utils.py @@ -83,42 +83,6 @@ class TestHasModule: ): assert _has_module("fake_native_ext") is False - def test_returns_false_on_os_error_during_import(self): - """Some shared-library failures surface as ``OSError``.""" - fake_spec = MagicMock() - - with ( - patch( - "vllm.utils.import_utils.importlib.util.find_spec", - return_value=fake_spec, - ), - patch( - "vllm.utils.import_utils.importlib.import_module", - side_effect=OSError("cannot load library"), - ), - ): - assert _has_module("fake_native_ext_os") is False - - def test_returns_false_on_unexpected_error_during_import(self): - """A broken extension may raise a non-import error (e.g. ``RuntimeError``). - - Such modules are not usable, so ``_has_module`` should still return - ``False`` rather than letting the exception propagate. - """ - fake_spec = MagicMock() - - with ( - patch( - "vllm.utils.import_utils.importlib.util.find_spec", - return_value=fake_spec, - ), - patch( - "vllm.utils.import_utils.importlib.import_module", - side_effect=RuntimeError("CUDA driver version is insufficient"), - ), - ): - assert _has_module("fake_broken_ext") is False - def test_returns_false_when_find_spec_raises(self): """``find_spec`` itself can raise for dotted names whose parent package fails to import. This should be treated as the module being unavailable. diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index 283b74624bb..e225666f844 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -43,6 +43,7 @@ def test_basic_allocate_and_reuse(): assert cache.check_and_update_cache(req, 0) assert "r1" in cache.cached["imgA"] + assert cache.request_cached_ids["r1"] == {0} assert cache.num_free_slots == 6 # Free twice to bring refcount to 0. @@ -50,6 +51,7 @@ def test_basic_allocate_and_reuse(): cache.free_encoder_input(req, 0) assert not cache.cached["imgA"] + assert "r1" not in cache.request_cached_ids assert "imgA" in cache.freeable assert cache.num_freeable_slots == 10 assert cache.num_free_slots == 6 @@ -63,10 +65,12 @@ def test_freeing_decreases_refcount_and_moves_to_freeable(): manager.allocate(req, 0) assert len(manager.cached["img3"]) == 1 + assert manager.request_cached_ids["req2"] == {0} manager.free_encoder_input(req, 0) assert not manager.cached["img3"] + assert "req2" not in manager.request_cached_ids assert "img3" in manager.freeable assert manager.num_freeable_slots == 10 @@ -83,11 +87,13 @@ def test_free_request_frees_all_inputs(): assert len(manager.cached["a"]) == 1 assert len(manager.cached["b"]) == 1 + assert manager.request_cached_ids["req3"] == {0, 1} manager.free(req) assert not manager.cached["a"] assert not manager.cached["b"] + assert "req3" not in manager.request_cached_ids assert "a" in manager.freeable assert "b" in manager.freeable assert manager.num_freeable_slots == 10 @@ -108,6 +114,7 @@ def test_eviction_when_cache_is_full(): # 'x' should have been evicted. assert "x" not in manager.cached + assert "req1" not in manager.request_cached_ids assert "x" in manager.get_freed_mm_hashes() @@ -137,6 +144,7 @@ def test_has_cache_restores_from_freeable(): # Should restore from freeable. assert manager.check_and_update_cache(req, 0) assert len(manager.cached["imgZ"]) == 1 + assert manager.request_cached_ids["reqY"] == {0} assert "imgZ" not in manager.freeable assert manager.num_freeable_slots == 6 @@ -205,6 +213,7 @@ def test_encoder_cache_with_is_embed_mask(): assert manager.num_free_slots == 92 assert "img1" in manager.cached + assert manager.request_cached_ids["r1"] == {0} old_size = 100 new_size = request.mm_features[0].mm_position.get_num_embeds() @@ -276,6 +285,7 @@ def test_reset_clears_all_state(): manager.reset() assert len(manager.cached) == 0 + assert len(manager.request_cached_ids) == 0 assert len(manager.freeable) == 0 assert len(manager.freed) == 0 assert manager.num_free_slots == 20 @@ -298,6 +308,26 @@ def test_reset_allows_fresh_allocations(): assert manager.num_free_slots == 2 assert "img2" in manager.cached assert "img1" not in manager.cached + assert manager.request_cached_ids["req2"] == {0} + assert "req1" not in manager.request_cached_ids + + +def test_free_request_with_duplicate_mm_hashes(): + """Freeing a request whose two inputs share the same mm_hash must fully + clean up request_cached_ids. After the first free_encoder_input call, + cached[mm_hash] becomes empty; the second call must still remove the + remaining input_id from request_cached_ids.""" + manager = EncoderCacheManager(cache_size=20) + + req = MockRequest("r1", ["imgA", "imgA"], [4, 4]) + + manager.allocate(req, 0) + # input 1 has the same hash, so it's already cached. + assert manager.check_and_update_cache(req, 1) + assert manager.request_cached_ids["r1"] == {0, 1} + + manager.free(req) + assert "r1" not in manager.request_cached_ids def test_encoder_decoder_cache_manager_reset(): diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 94e9f6f4c10..68ad7bc42ef 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1447,7 +1447,10 @@ def test_allocate_with_lookahead(): # Test case 1: Requires additional lookahead tokens kv_cache_manager = KVCacheManager( - kv_cache_config=config, max_model_len=100, hash_block_size=block_size + kv_cache_config=config, + max_model_len=100, + scheduler_block_size=block_size, + hash_block_size=block_size, ) blocks = kv_cache_manager.allocate_slots( request, @@ -1458,7 +1461,10 @@ def test_allocate_with_lookahead(): # Test case 2: With precomputed blocks kv_cache_manager = KVCacheManager( - kv_cache_config=config, max_model_len=100, hash_block_size=block_size + kv_cache_config=config, + max_model_len=100, + scheduler_block_size=block_size, + hash_block_size=block_size, ) # required_blocks = ceil((3 + 2) /4) = 2 blocks = kv_cache_manager.allocate_slots( @@ -1471,7 +1477,10 @@ def test_allocate_with_lookahead(): # Test case 3: With precomputed blocks # required_blocks = ceil((3 + 4) / 4) = 2 kv_cache_manager = KVCacheManager( - kv_cache_config=config, max_model_len=100, hash_block_size=block_size + kv_cache_config=config, + max_model_len=100, + scheduler_block_size=block_size, + hash_block_size=block_size, ) blocks = kv_cache_manager.allocate_slots( request, diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 546412b1d2f..91c5f37b417 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -4,6 +4,7 @@ import copy from collections.abc import Callable +from math import lcm import pytest import torch @@ -92,6 +93,18 @@ def make_request( ) +def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheManager: + """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the + config (LCM of group block sizes) unless explicitly provided. This mirrors + ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by + these tests, so callers don't have to pass it at every site.""" + kwargs.setdefault( + "scheduler_block_size", + lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)), + ) + return KVCacheManager(kv_cache_config, **kwargs) + + def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: return KVCacheConfig( num_blocks=num_blocks, @@ -208,7 +221,7 @@ def make_kv_cache_config_three_types( @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_prefill(hash_fn): block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -331,7 +344,7 @@ def test_prefill(hash_fn): def test_prefill_hybrid_model(): block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config_hybrid_model(block_size, 21, 2), max_model_len=8192, enable_caching=True, @@ -500,7 +513,7 @@ def test_prefill_hybrid_model(): def test_prefill_hybrid_model_eagle(): block_size = 16 kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config, max_model_len=8192, enable_caching=True, @@ -837,7 +850,7 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): num_blocks = 10 * num_groups kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config, max_model_len=8192, enable_caching=True, @@ -912,7 +925,7 @@ def test_prefill_hybrid_model_combinations_eagle( num_blocks = 10 * num_groups kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config, max_model_len=8192, enable_caching=True, @@ -984,7 +997,7 @@ def test_prefill_hybrid_model_mamba_align(): kv_cache_config = _make_hybrid_kv_cache_config( block_size, num_blocks, ["full", "mamba_align"] ) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config, max_model_len=8192, enable_caching=True, @@ -1017,7 +1030,7 @@ def test_prefill_plp(): 3. Schedule plp request; no hit should occur; validate blocks """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1125,7 +1138,7 @@ def test_prefill_plp(): def test_decode(): block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1188,7 +1201,7 @@ def test_decode(): def test_evict(): block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1247,7 +1260,7 @@ def test_hash_block_correct_reuse(): its hash metadata should be correctly reset. """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(16, 2), max_model_len=8192, enable_caching=True, @@ -1288,7 +1301,7 @@ def test_computed_blocks_not_evicted(): for a request if there are any other free blocks. """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 3), max_model_len=8192, enable_caching=True, @@ -1347,7 +1360,7 @@ def test_basic_prefix_caching_disabled(): This tests that the prefix caching is disabled. """ block_size = 4 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 5), max_model_len=8192, enable_caching=False, @@ -1531,7 +1544,7 @@ def test_mm_prefix_caching(): """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1639,7 +1652,7 @@ def test_cache_key_salting(): is separated cache as expected. """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1721,7 +1734,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks(): the computed blocks should not be touched. """ block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1794,7 +1807,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks(): def test_reset_prefix_cache(): block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1835,7 +1848,7 @@ def test_reset_prefix_cache(): def test_prefix_cache_stats_disabled(): """Test that prefix_cache_stats is None when log_stats is False.""" block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, 11), max_model_len=8192, enable_caching=True, @@ -1915,7 +1928,7 @@ def test_kv_cache_events(blocks_to_cache: int): # Should see a single block stored event with a blocks_to_cache number of # block hashes # take_events should reset the kv_event_queue - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks), max_model_len=8192, enable_caching=True, @@ -2043,7 +2056,7 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int): num_blocks = blocks_to_cache + 1 # Create KVCacheManager with events enabled - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks), max_model_len=8192, enable_caching=True, @@ -2101,7 +2114,7 @@ def test_block_stored_event_group_idx(group_id: int): block_size = 4 num_tokens = block_size * 2 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config_three_types(block_size, num_blocks=5), max_model_len=8192, enable_caching=True, @@ -2161,7 +2174,7 @@ def test_block_stored_event_group_idx_multiple_groups(): block_size = 4 num_tokens = block_size * 2 - manager = KVCacheManager( + manager = make_kv_cache_manager( KVCacheConfig( num_blocks=5, kv_cache_tensors=[], @@ -2238,7 +2251,7 @@ def test_block_stored_event_group_idx_multiple_groups(): def test_block_stored_event_group_idx_out_of_bounds(monkeypatch): """Out-of-range group_idx events are returned without metadata annotation.""" block_size = 4 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks=5), max_model_len=8192, enable_caching=True, @@ -2328,7 +2341,7 @@ def test_eagle_enabled_removes_last_block(): """Verify Eagle does NOT remove blocks when request length is divisible by block size.""" block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks=10), max_model_len=8192, enable_caching=True, @@ -2361,7 +2374,7 @@ def test_eagle_enabled_removes_last_block(): def test_eagle_with_partial_blocks(): """Test Eagle behavior with requests containing partial blocks.""" block_size = 16 - manager = KVCacheManager( + manager = make_kv_cache_manager( make_kv_cache_config(block_size, num_blocks=10), max_model_len=8192, enable_caching=True, @@ -2397,7 +2410,7 @@ def test_eagle_with_sliding_window(): dtype=torch.float32, sliding_window=block_size, ) - manager = KVCacheManager( + manager = make_kv_cache_manager( KVCacheConfig( num_blocks=10, kv_cache_tensors=[], @@ -2453,6 +2466,201 @@ def test_eagle_with_sliding_window(): assert num_tokens == 0 +def test_eagle_swa_alignment_caches_extra_block(): + """Regression: SWA + EAGLE with `sliding_window <= alignment_tokens`. + + When the cache-hit alignment (lcm of per-group block sizes) is larger than + the SWA window, the SWA mask only kept the last block of each aligned + segment. EAGLE/MTP lookup needs ``tail + 1`` contiguous cached blocks and + that +1 block lives at the next segment's first position, which was left + uncached. The fix caches that extra block when ``use_eagle=True``. + """ + block_size = 8 + # Full group uses 4 * block_size, so lcm/alignment is 4 * block_size. + # SWA group has sliding_window = block_size (i.e., tail = 1 block). + # Without the fix, the second cached block needed for the EAGLE 2-block + # match never exists -> EAGLE cache hit fails entirely. + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4 * block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["swa_mtp"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=block_size, + ), + is_eagle_group=True, + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + use_eagle=True, + ) + + # Prime the cache with a long prompt (16 swa blocks = 4 aligned segments). + token_ids = [i for i in range(16) for _ in range(block_size)] + req0 = make_request("0", token_ids, block_size, sha256) + computed_blocks, _ = manager.get_computed_blocks(req0) + blocks = manager.allocate_slots( + req0, + len(token_ids), + len(computed_blocks.blocks[0]) * block_size, + computed_blocks, + ) + assert blocks is not None + manager.free(req0) + + # Second request with identical prompt should find an EAGLE cache hit. + # Without the fix, ``num_computed_tokens`` is 0; with the fix, it lands at + # an alignment boundary (multiple of 32 tokens, minus the EAGLE drop). + req1 = make_request("1", token_ids, block_size, sha256) + _, num_computed_tokens = manager.get_computed_blocks(req1) + assert num_computed_tokens > 0, ( + "EAGLE + SWA with sliding_window <= alignment failed to find any " + "cache hit; the +1 block past each segment boundary must be cached." + ) + # Each aligned segment contributes 4 * block_size = 32 tokens; EAGLE drops + # the last block (block_size tokens) from the hit. + assert num_computed_tokens % (4 * block_size) == 0 + + +def test_eagle_swa_boundary_caches_post_boundary_block(): + """EAGLE + SWA must cache the first block after an alignment boundary. + + A 40-token computed prefix with 8-token SWA blocks and 32-token hybrid + alignment needs SWA blocks 3 and 4 cached to reuse a 32-token prefix: + block 3 is the segment tail, and block 4 is the EAGLE lookahead block + that gets dropped after lookup. + """ + block_size = 8 + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4 * block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["swa_mtp"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=block_size, + ), + is_eagle_group=True, + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + use_eagle=True, + ) + + token_ids = [i for i in range(5) for _ in range(block_size)] + req0 = make_request("0", token_ids, block_size, sha256) + computed_blocks, _ = manager.get_computed_blocks(req0) + blocks = manager.allocate_slots( + req0, + len(token_ids), + len(computed_blocks.blocks[0]) * block_size, + computed_blocks, + ) + assert blocks is not None + + pool = manager.block_pool + assert pool.get_cached_block(req0.block_hashes[3], kv_cache_group_ids=[1]) + assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1]) + manager.free(req0) + + req1 = make_request("1", token_ids + [999], block_size, sha256) + _, num_computed_tokens = manager.get_computed_blocks(req1) + assert num_computed_tokens == 4 * block_size + + +def test_eagle_grouped_swa_siblings_use_same_cache_mask(): + """Grouped SWA siblings must cache the EAGLE lookahead block together.""" + block_size = 8 + swa_spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=block_size, + ) + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4 * block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec(["swa_main"], swa_spec), + KVCacheGroupSpec(["swa_mtp"], swa_spec, is_eagle_group=True), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + use_eagle=True, + ) + + token_ids = [i for i in range(9) for _ in range(block_size)] + req0 = make_request("0", token_ids, block_size, sha256) + computed_blocks, _ = manager.get_computed_blocks(req0) + blocks = manager.allocate_slots( + req0, + len(token_ids), + len(computed_blocks.blocks[0]) * block_size, + computed_blocks, + ) + assert blocks is not None + + pool = manager.block_pool + assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1, 2]) + assert pool.get_cached_block(req0.block_hashes[8], kv_cache_group_ids=[1, 2]) + manager.free(req0) + + req1 = make_request("1", token_ids + [999], block_size, sha256) + _, num_computed_tokens = manager.get_computed_blocks(req1) + assert num_computed_tokens == 8 * block_size + + def test_different_block_size(): block_size = 16 # full attention and sliding window attention layers have the same page size: @@ -2482,7 +2690,7 @@ def test_different_block_size(): ), ], ) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config=kv_cache_config, max_model_len=8192, enable_caching=True, @@ -2565,7 +2773,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only(): ), ], ) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config=kv_cache_config, max_model_len=8192, enable_caching=True, @@ -2601,7 +2809,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only(): def test_hybrid_cache_blocks_clamped_to_lcm(): - """HybridKVCacheCoordinator.cache_blocks() clamps to lcm_block_size. + """HybridKVCacheCoordinator.cache_blocks() clamps to scheduler_block_size. Chunks past the last lcm-aligned boundary can never participate in a cache hit (find_longest_cache_hit always returns lcm-aligned hits), so caching them only pollutes the prefix-cache hash map and keeps blocks @@ -2633,7 +2841,7 @@ def test_hybrid_cache_blocks_clamped_to_lcm(): ), ], ) - manager = KVCacheManager( + manager = make_kv_cache_manager( kv_cache_config=kv_cache_config, max_model_len=8192, enable_caching=True, @@ -2781,7 +2989,7 @@ def test_can_fit_full_sequence_swa_cap_admits_long_prompt(): ], ) - manager = KVCacheManager( + manager = make_kv_cache_manager( config, max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, @@ -2837,7 +3045,7 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized(): ], ) - manager = KVCacheManager( + manager = make_kv_cache_manager( config, max_model_len=max_model_len, max_num_batched_tokens=max_num_batched_tokens, diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index db33b4c6df1..7fa331747c4 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -4349,3 +4349,128 @@ def test_eagle3_mm_encoder_cache_with_shift(): f"shifted_end={scheduled_end_with_shift}) overlapping MM at " f"{start_pos}. The fix must schedule encoder inputs." ) + + +@pytest.mark.parametrize("use_kv_connector", [False, True]) +def test_ec_connector_ensure_cache_available_defers_request(use_kv_connector): + """Test that ensure_cache_available() returning False defers the request. + + When the EC connector signals a prefetch is in progress (returns False), + the scheduler should: + 1. Not schedule the request (no KV cache or encoder cache allocated) + 2. Still schedule other requests behind the deferred one + 3. Schedule the deferred request on the next step when ensure_cache_available + returns True and has_cache_item returns True + """ + scheduler = create_scheduler( + model="llava-hf/llava-1.5-7b-hf", + enable_prefix_caching=True, + use_kv_connector=use_kv_connector, + use_ec_connector=True, + ec_role="ec_consumer", + ) + + NUM_TOKENS = 200 + NUM_ENCODER_TOKENS = 100 + + request_deferred = create_requests( + num_requests=1, + num_tokens=NUM_TOKENS, + mm_positions=[[PlaceholderRange(offset=0, length=NUM_ENCODER_TOKENS)]], + req_ids=["deferred"], + )[0] + + request_behind = create_requests( + num_requests=1, + num_tokens=20, + req_ids=["behind"], + )[0] + + # --- Step 1: ensure_cache_available returns False → request deferred --- + scheduler.ec_connector.ensure_cache_available = Mock(return_value=False) + + scheduler.add_request(request_deferred) + scheduler.add_request(request_behind) + output = scheduler.schedule() + + # ensure_cache_available must have been called with (request, num_computed_tokens=0) + # for a brand-new request that has no cached tokens yet. + scheduler.ec_connector.ensure_cache_available.assert_called_once_with( + request_deferred, 0 + ) + # Deferred request must NOT be scheduled + assert request_deferred.request_id not in output.num_scheduled_tokens + _assert_right_encoder_cache_allocated(scheduler, expected_total_allocated=0) + # No KV blocks allocated for the deferred request + for mgr in scheduler.kv_cache_manager.coordinator.single_type_managers: + assert request_deferred.request_id not in mgr.req_to_blocks + + # The text-only request behind the deferred one MUST still be scheduled + assert request_behind.request_id in output.num_scheduled_tokens + assert output.num_scheduled_tokens[request_behind.request_id] == 20 + + # --- Step 2: prefetch done, cache exists → request scheduled --- + # has_cache_item is called inside _try_schedule_encoder_inputs (not during + # deferral), so it is only relevant here in step 2. + scheduler.ec_connector.ensure_cache_available = Mock(return_value=True) + scheduler.ec_connector.has_cache_item = Mock(return_value=True) + + output = scheduler.schedule() + + # Now the deferred request should be scheduled + assert request_deferred.request_id in output.num_scheduled_tokens + assert output.num_scheduled_tokens[request_deferred.request_id] == NUM_TOKENS + _assert_right_encoder_cache_allocated(scheduler, requests=[request_deferred]) + # EC connector metadata should carry the deferred request's MM data + _assert_right_ec_connector_metadata( + output, mm_features_list=request_deferred.mm_features + ) + # No local encoder compute — all loaded externally + _assert_right_encoder_inputs(output, expected_total_reqs=0) + + +def test_ec_connector_pending_prefetch_only_checks_future_mm_features(): + """Test that future mm feature filtering only yields features beyond + the computed token frontier. + + Features already within num_computed_tokens (past/boundary) must be + filtered out; only features that extend beyond the frontier (future) should + be yielded so that connector implementations know which items to prefetch. + + Filter cases: + "past": end = 0 + 16 = 16 < 32 → filtered OUT + "boundary": end = 16 + 16 = 32 == 32 → filtered OUT (condition is >, not >=) + "future": end = 48 + 32 = 80 > 32 → yielded + """ + BLOCK_SIZE = 16 + NUM_COMPUTED_TOKENS = BLOCK_SIZE * 2 # 32 + NUM_TOKENS = BLOCK_SIZE * 8 # 128 + + HASH_PAST = "hash_past" + HASH_BOUNDARY = "hash_boundary" + HASH_FUTURE = "hash_future" + + request = create_requests( + num_requests=1, + num_tokens=NUM_TOKENS, + mm_hashes_list=[[HASH_PAST, HASH_BOUNDARY, HASH_FUTURE]], + mm_positions=[ + [ + PlaceholderRange(offset=0, length=BLOCK_SIZE), # end=16 (past) + PlaceholderRange(offset=16, length=BLOCK_SIZE), # end=32 (boundary) + PlaceholderRange(offset=48, length=BLOCK_SIZE * 2), # end=80 (future) + ] + ], + block_size=BLOCK_SIZE, + )[0] + + future_hashes = [ + f.identifier + for f in request.mm_features + if f.mm_position.offset + f.mm_position.length > NUM_COMPUTED_TOKENS + ] + + assert future_hashes == [HASH_FUTURE], ( + f"Expected only {HASH_FUTURE!r} from future mm feature filtering, " + f"got {future_hashes!r}. Past/boundary features must be filtered out." + ) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f59830dcd74..0e3e8879359 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -28,6 +28,7 @@ def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=T block_pool=block_pool, enable_caching=enable_caching, kv_cache_group_id=0, + scheduler_block_size=sliding_window_spec.block_size, max_admission_blocks_per_request=10**9, ) @@ -40,6 +41,7 @@ def get_chunked_local_attention_manager( block_pool=block_pool, enable_caching=enable_caching, kv_cache_group_id=0, + scheduler_block_size=chunked_local_attention_spec.block_size, max_admission_blocks_per_request=10**9, ) @@ -84,7 +86,7 @@ def test_chunked_local_attention_possible_cached_prefix(): kv_cache_group_ids=[0], block_pool=block_pool, kv_cache_spec=chunked_local_attention_spec, - use_eagle=False, + drop_eagle_block=False, alignment_tokens=block_size, )[0] assert len(computed_blocks) == expect_length @@ -155,7 +157,7 @@ def test_sliding_window_possible_cached_prefix(): kv_cache_group_ids=[0], block_pool=block_pool, kv_cache_spec=sliding_window_spec, - use_eagle=False, + drop_eagle_block=False, alignment_tokens=block_size, )[0] assert len(computed_blocks) == expect_length @@ -458,6 +460,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): block_pool=block_pool, enable_caching=False, kv_cache_group_id=0, + scheduler_block_size=spec.block_size, max_admission_blocks_per_request=cap, ) diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 2d9834d2e3a..7213a669c53 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -24,6 +24,7 @@ from vllm.utils.hashing import sha256 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.async_scheduler import AsyncScheduler from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -160,6 +161,7 @@ def create_scheduler( ], ) cache_config.num_gpu_blocks = num_blocks + register_all_kvcache_specs(vllm_config) scheduler_cls = AsyncScheduler if async_scheduling else Scheduler return scheduler_cls( vllm_config=vllm_config, diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 9c315f293da..61134a4f5a2 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -109,6 +109,7 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager: mgr.max_batch_size = 16 mgr.use_dp = False mgr.budget_graphs = {} + mgr.graph_pool = None mgr.graph_hits = 0 mgr.graph_misses = 0 mgr.log_stats_interval = 100 @@ -180,6 +181,10 @@ class TestFindBudgetGraph: # Budget selection still works correctly after sorting assert mgr._find_smallest_fitting_budget_given_tokens(3000) == 4096 + def test_num_graphs_to_capture_tracks_budgets(self): + mgr = _make_manager_with_budgets([8192, 2048, 4096]) + assert mgr.get_num_graphs_to_capture() == 3 + # --------------------------------------------------------------------------- # get_cumulative_stats @@ -409,6 +414,7 @@ def _make_manager_for_gpu( ) mgr.use_dp = False mgr.budget_graphs = {} + mgr.graph_pool = None mgr.graph_hits = 0 mgr.graph_misses = 0 mgr.log_stats_interval = 100 @@ -467,7 +473,8 @@ class TestEncoderCudaGraphCaptureReplay: self.mgr = _make_manager_for_gpu( self.model, _BUDGETS, _MAX_BATCH, self.device, self.dtype ) - self.mgr.capture() + self.graph_pool = current_platform.graph_pool_handle() + self.mgr.capture(graph_pool=self.graph_pool) # --- capture --- @@ -475,6 +482,14 @@ class TestEncoderCudaGraphCaptureReplay: assert len(self.mgr.budget_graphs) == len(_BUDGETS) assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS) + def test_capture_uses_supplied_graph_pool(self): + assert self.mgr.graph_pool is self.graph_pool + + def test_clear_releases_graphs_and_pool(self): + self.mgr.clear() + assert self.mgr.budget_graphs == {} + assert self.mgr.graph_pool is None + # --- output shape --- def test_execute_returns_one_tensor_per_image(self): @@ -742,7 +757,8 @@ class TestEncoderCudaGraphVideoReplay: self.dtype, max_frames_per_batch=_VIDEO_MAX_FRAMES, ) - self.mgr.capture() + self.graph_pool = current_platform.graph_pool_handle() + self.mgr.capture(graph_pool=self.graph_pool) # --- capture --- diff --git a/tests/v1/distributed/test_pp_dp_v2.py b/tests/v1/distributed/test_pp_dp_v2.py new file mode 100644 index 00000000000..35331549976 --- /dev/null +++ b/tests/v1/distributed/test_pp_dp_v2.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""V2 ModelRunner + pipeline parallel + data parallel integration tests. + +Covers the interaction between the V2 model runner's PP sampled-token +broadcast and the DP per-step all-reduce across a few concurrency +regimes. Requires 4 GPUs (DP=2, PP=2, TP=1) on CUDA. +""" + +import asyncio +import contextlib +import os +from contextlib import ExitStack + +import pytest + +from vllm import SamplingParams +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.platforms import current_platform +from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine.async_llm import AsyncLLM + +PP_DP_MODEL = "ibm-research/PowerMoE-3b" # smallest cached MoE that supports PP +PROMPT = "This is a test of data parallel and pipeline parallel together" + + +def _gpu_skip_reason() -> str | None: + if not current_platform.is_cuda(): + return "requires CUDA" + n = current_platform.device_count() + if n < 4: + return f"requires 4 GPUs, got {n}" + return None + + +_GPU_SKIP = _gpu_skip_reason() + +pytestmark = [ + pytest.mark.skipif( + os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0") != "1", + reason="VLLM_USE_V2_MODEL_RUNNER=1 required", + ), + pytest.mark.skipif(_GPU_SKIP is not None, reason=_GPU_SKIP or ""), +] + + +def _engine_args(async_scheduling: bool) -> AsyncEngineArgs: + return AsyncEngineArgs( + model=PP_DP_MODEL, + pipeline_parallel_size=2, + data_parallel_size=2, + data_parallel_backend="mp", + tensor_parallel_size=1, + max_model_len=4096, + max_num_batched_tokens=2048, + max_num_seqs=256, + async_scheduling=async_scheduling, + enable_prefix_caching=False, + enforce_eager=False, + enable_expert_parallel=False, + ) + + +async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int: + """Run one streaming completion and return the number of tokens it yielded.""" + sampling_params = SamplingParams( + max_tokens=max_tokens, + ignore_eos=True, + output_kind=RequestOutputKind.DELTA, + temperature=0.0, + ) + request_id = f"req-{id(prompt):x}-{max_tokens}" + total = 0 + async for out in engine.generate( + request_id=request_id, prompt=prompt, sampling_params=sampling_params + ): + total += len(out.outputs[0].token_ids) + return total + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_scheduling", [True, False]) +async def test_pp_dp_v2_low_concurrency(async_scheduling: bool): + """A single in-flight request at a time, repeated, to exercise the + PP slot ring under empty batches between decodes.""" + with ExitStack() as after: + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) + after.callback(engine.shutdown) + + for _ in range(4): + n = await _generate(engine, PROMPT, max_tokens=16) + assert n == 16 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_scheduling", [True, False]) +async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool): + """64 concurrent requests, staggered, to exercise the steady-state + DP all-reduce + PP slot-ring path.""" + with ExitStack() as after: + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling)) + after.callback(engine.shutdown) + + async def _one(i: int) -> int: + await asyncio.sleep(0.01 * i) # stagger so DP load-balances + return await _generate(engine, f"{PROMPT} {i}", max_tokens=64) + + results = await asyncio.gather(*[_one(i) for i in range(64)]) + assert all(n == 64 for n in results), results + + +@pytest.mark.asyncio +async def test_pp_dp_v2_abort_mid_decode(): + """Cancel half the in-flight requests mid-stream and confirm the + engine survives the abort storm.""" + + with ExitStack() as after: + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) + after.callback(engine.shutdown) + + async def _maybe_cancel(i: int): + sampling_params = SamplingParams( + max_tokens=64, + ignore_eos=True, + output_kind=RequestOutputKind.DELTA, + temperature=0.0, + ) + request_id = f"abort-req-{i}" + count = 0 + cancel_at = 4 if i % 2 == 0 else 64 + async for out in engine.generate( + request_id=request_id, + prompt=f"{PROMPT} {i}", + sampling_params=sampling_params, + ): + count += len(out.outputs[0].token_ids) + if count >= cancel_at: + break + return count, i + + results = await asyncio.gather(*[_maybe_cancel(i) for i in range(32)]) + for count, i in results: + if i % 2 == 0: + assert count >= 4 + else: + assert count == 64 + + # Engine must still serve after the abort storm. + final = await _generate(engine, "post-abort warmup", max_tokens=8) + assert final == 8 + + +@pytest.mark.asyncio +async def test_pp_dp_v2_pause_resume(): + """Pause an engine with a request in flight, then resume and confirm + new requests still work.""" + + with ExitStack() as after: + engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True)) + after.callback(engine.shutdown) + + # Start a long-running generation, let some decoding happen, then + # pause (abort mode) and confirm the in-flight task terminates. + inflight = asyncio.create_task(_generate(engine, PROMPT, max_tokens=128)) + await asyncio.sleep(0.5) + + assert not await engine.is_paused() + await engine.pause_generation(mode="abort") + assert await engine.is_paused() + + with contextlib.suppress(Exception): + await inflight + + await engine.resume_generation() + assert not await engine.is_paused() + + n = await _generate(engine, PROMPT, max_tokens=8) + assert n == 8 diff --git a/tests/v1/engine/test_abort_final_step.py b/tests/v1/engine/test_abort_final_step.py index 8f1e8029955..d8d5b73d45b 100644 --- a/tests/v1/engine/test_abort_final_step.py +++ b/tests/v1/engine/test_abort_final_step.py @@ -184,14 +184,31 @@ async def test_abort_during_final_step(async_scheduling: bool): original_execute_model = Worker.execute_model def execute_model_with_wait(self, scheduler_output): - # Signal that execute_model has been called by deleting ready_file - if ready_file.exists(): - ready_file.unlink() + # V2's `gpu_worker.compile_or_warm_up_model` calls + # `warmup_kernels(...)` during engine init, which itself calls + # `Worker.execute_model` three times (prefill / decode / cleanup) + # to JIT compile triton kernels. None of those carry the test's + # request id, so we only stall when our actual request is being + # processed. + scheduled = scheduler_output.num_scheduled_tokens or {} + finished = scheduler_output.finished_req_ids or set() - # Wait for the block file to be deleted (triggered from test after abort) - # This runs in the worker process (after fork), so we poll the filesystem - while block_file.exists(): - time.sleep(0.01) + def is_target_request(req_ids): + return any( + rid == request_id or rid.startswith(f"{request_id}-") + for rid in req_ids + ) + + if is_target_request(scheduled) or is_target_request(finished): + # Signal that execute_model has been called by deleting ready_file + if ready_file.exists(): + ready_file.unlink() + + # Wait for the block file to be deleted (triggered from test after + # abort). This runs in the worker process (after fork), so we poll + # the filesystem. + while block_file.exists(): + time.sleep(0.01) return original_execute_model(self, scheduler_output) # Patch execute_model to inject the wait diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index ae674919ae9..aa2a70559dd 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -5,6 +5,7 @@ import copy import time import uuid from concurrent.futures import Future, ThreadPoolExecutor +from unittest.mock import PropertyMock, patch import pytest from transformers import AutoTokenizer @@ -293,10 +294,6 @@ def test_engine_core_concurrent_batches(): # Use the thread pool instead of creating a new thread return self.thread_pool.submit(_execute) - @property - def max_concurrent_batches(self) -> int: - return 2 - def shutdown(self): if hasattr(self, "thread_pool"): self.thread_pool.shutdown(wait=False) @@ -314,7 +311,17 @@ def test_engine_core_concurrent_batches(): async_scheduling=False, ) vllm_config = engine_args.create_engine_config() - with set_default_torch_num_threads(1): + # Force two concurrent batches to exercise the batch queue independently + # of async scheduling (which is disabled above). + with ( + set_default_torch_num_threads(1), + patch.object( + VllmConfig, + "max_concurrent_batches", + new_callable=PropertyMock, + return_value=2, + ), + ): engine_core = EngineCore( vllm_config=vllm_config, log_stats=False, executor_class=DummyExecutor ) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index ab5946ad3ba..36dc95eea49 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -27,12 +27,13 @@ from vllm.platforms import current_platform from vllm.pooling_params import LateInteractionParams, PoolingParams from vllm.usage.usage_lib import UsageContext from vllm.utils.torch_utils import set_default_torch_num_threads -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest from vllm.v1.engine.core import EngineCore from vllm.v1.engine.core_client import ( AsyncMPClient, DPLBAsyncMPClient, EngineCoreClient, + MPClient, SyncMPClient, ) from vllm.v1.engine.utils import CoreEngineProcManager @@ -236,6 +237,30 @@ def test_dplb_non_late_interaction_still_uses_lb(): assert client.lb_engines[1][0] == 1 +def test_apply_ready_response_syncs_block_size(): + import msgspec + + client = object.__new__(MPClient) + client.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), + model_config=SimpleNamespace(max_model_len=8192), + ) + client.stats_update_address = None + + payload = msgspec.msgpack.encode( + EngineCoreReadyResponse( + max_model_len=8192, + num_gpu_blocks=100, + block_size=1056, + dp_stats_address=None, + dtype="bfloat16", + vllm_version="test", + ) + ) + client._apply_ready_response(payload) + assert client.vllm_config.cache_config.block_size == 1056 + + def loop_until_done(client: EngineCoreClient, outputs: dict): while True: engine_core_outputs = client.get_output().outputs @@ -1187,7 +1212,6 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}] mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024] mock_executor.initialize_from_config.return_value = None - mock_executor.max_concurrent_batches = 1 return mock_executor 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 040632249d3..d0a56304f2a 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 @@ -21,19 +21,20 @@ dp_ep_configs=( "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1) "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) ) +# We assume HMA enabled by default. hybrid_ssm_configs=( - "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" + "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models. - "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" + "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" # GDN (Qwen3.5) - "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" - "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--no-async-scheduling" + "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_SERVE_EXTRA_ARGS=--no-async-scheduling" ) sw_attn_configs=( # NOTE: gemma3 does not work with FlashInfer "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" # SW model - "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" - "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" + "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" + "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" ) # Select config array based on DP_EP env var @@ -50,14 +51,6 @@ else configs=("${tp_configs[@]}") fi -if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then - # Append ENABLE_HMA_FLAG=1 to each config in the selected array - echo "ENABLE_HMA_FLAG is set, appending ENABLE_HMA_FLAG=1 to each config" - for i in "${!configs[@]}"; do - configs[$i]="ENABLE_HMA_FLAG=1 ${configs[$i]}" - done -fi - run_tests() { local label=$1 local extra_args=$2 diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh index 313efc3968d..f55bd308a0a 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh @@ -11,7 +11,7 @@ SCRIPT="v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh" eagle3_config="SD_METHOD=eagle3 MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct SD_MODEL=RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3 NUM_SPEC_TOKENS=3" # MTP: Qwen3.5-0.8B-Base with hybrid SSM flags. -mtp_config="SD_METHOD=mtp MODEL_NAME=Qwen/Qwen3.5-0.8B-Base SD_MODEL=Qwen/Qwen3.5-0.8B-Base NUM_SPEC_TOKENS=1 BLOCK_SIZE=32 MAX_MODEL_LEN=4096 VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 KV_BUFFER_DEVICES=cuda" +mtp_config="SD_METHOD=mtp MODEL_NAME=Qwen/Qwen3.5-0.8B-Base SD_MODEL=Qwen/Qwen3.5-0.8B-Base NUM_SPEC_TOKENS=1 BLOCK_SIZE=32 MAX_MODEL_LEN=4096 VLLM_SSM_CONV_STATE_LAYOUT=DS KV_BUFFER_DEVICES=cuda" configs=( "$eagle3_config" 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 fc446a0e765..bde246c9b66 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -5,11 +5,6 @@ set -xe KV_BUFFER_DEVICE="cuda" # Default to cuda ATTENTION_BACKEND="" # Default to empty (use vllm default) CROSS_LAYERS_BLOCKS="False" -ENABLE_HMA_VAR="" # Default to empty (HMA disabled by default for kv connector) -# Check for ENABLE_HMA_FLAG environment variable -if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then - ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager" -fi while [[ $# -gt 0 ]]; do case $1 in @@ -37,9 +32,6 @@ echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE" if [[ -n "$ATTENTION_BACKEND" ]]; then echo "Using attention backend: $ATTENTION_BACKEND" fi -if [[ -n "$ENABLE_HMA_VAR" ]]; then - echo "HMA (Hybrid KV Cache Manager) enabled" -fi if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS" fi @@ -180,10 +172,6 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" fi - # Add HMA flag if specified - if [[ -n "$ENABLE_HMA_VAR" ]]; then - BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR" - fi FULL_CMD="$BASE_CMD" eval "$FULL_CMD &" @@ -232,10 +220,6 @@ run_tests_for_model() { BASE_CMD="${BASE_CMD} --attention-backend=$ATTENTION_BACKEND" fi - # Add HMA flag if specified - if [[ -n "$ENABLE_HMA_VAR" ]]; then - BASE_CMD="${BASE_CMD} $ENABLE_HMA_VAR" - fi # DP-EP attention mode if [[ -z "$DP_EP" ]]; then diff --git a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh index 2c5622a2f0e..bc90680a533 100755 --- a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +++ b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh @@ -27,7 +27,6 @@ # ROCM_AITER_UNIFIED_ATTN # NVIDIA options: FLASH_ATTN, FLASHINFER # VLLM_SSM_CONV_STATE_LAYOUT - SSM conv state layout (e.g. "DS" required for Mamba models) -# ENABLE_HMA_FLAG - set to 1 to enable hybrid KV cache manager # VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve set -ex @@ -85,13 +84,7 @@ if [[ -z "${ATTENTION_BACKEND:-}" ]]; then fi echo "Using attention backend: ${ATTENTION_BACKEND}" -# ── HMA & extra serve args ──────────────────────────────────────────── - -ENABLE_HMA_VAR="" -if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then - ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager" - echo "HMA (Hybrid KV Cache Manager) enabled" -fi +# ── Extra serve args ───────────────────────────────────────────────── EXTRA_SERVE_ARGS=() if [[ -n "${VLLM_SERVE_EXTRA_ARGS:-}" ]]; then @@ -258,7 +251,6 @@ run_test_for_device() { --kv-transfer-config "$kv_config" \ --speculative-config "$PREFILL_SPEC_CONFIG" \ --attention-backend $ATTENTION_BACKEND \ - ${ENABLE_HMA_VAR} \ ${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} & local SERVER_PID=$! @@ -298,7 +290,6 @@ run_test_for_device() { --kv-transfer-config "$kv_config" \ --speculative-config "$DECODE_SPEC_CONFIG" \ --attention-backend $ATTENTION_BACKEND \ - ${ENABLE_HMA_VAR} \ ${EXTRA_SERVE_ARGS[@]+"${EXTRA_SERVE_ARGS[@]}"} & local SERVER_PID=$! diff --git a/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py b/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py index c86a407ff8e..15f386f5f5a 100644 --- a/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py +++ b/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py @@ -158,6 +158,10 @@ def test_spec_decode_acceptance_length(): max_tokens=DEFAULT_OUTPUT_LEN, temperature=0.0, top_p=1.0, + # Prompts are already chat-templated (contain BOS); avoid the + # completions API prepending a second BOS, which would lower + # acceptance ~5% vs the add_special_tokens=False standalone baselines. + extra_body={"add_special_tokens": False}, ) if i < 3: text = resp.choices[0].text.strip()[:100] 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 5cc7065c1a4..20c230a4c2a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1305,3 +1305,79 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): (1, 7), ), ) + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_stale_sliding_window_block_after_prepare_store_failure( + request_runner, async_scheduling: bool +): + """Regression test: when prepare_store fails (returns None), offloading is + delayed. Meanwhile, sliding window blocks get freed and reallocated to the + same request. On retry, the stale block_id must be detected and skipped. + + Without the fix, the stale block_id would either: + - Cause a KeyError in _remove_pending_job (duplicate in + _block_id_to_pending_jobs) + - Silently offload wrong data under a wrong key + """ + block_size = 4 + # sliding_window = 8 -> window of 2 blocks + sliding_window = 8 + # Use a tight GPU block budget so freed sliding window blocks are + # immediately reused by the same request's new allocations. + num_gpu_blocks = 4 + + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window, + ), + ), + ] + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + ) + + # Request with 3 blocks of prompt. Window = 2 blocks, so block 0 is + # outside the window but won't be freed until the next allocate_slots. + runner.new_request(token_ids=[0] * block_size * 3) + + # First step: prepare_store FAILS -> offloading delayed. + # next_stored_block_idx stays at 0, block_ids[0] still holds the + # original block_id for position 0. + runner.manager.prepare_store.side_effect = lambda keys, req_context: None + runner.run(decoded_tokens=[0]) + runner.manager.prepare_store.assert_called() + + # Second step: decode more tokens -> block 3 allocated. + # allocate_slots calls remove_skipped_blocks which frees block 0 + # (it's now outside the sliding window). With num_gpu_blocks=4, + # the freed block is immediately reused for the new allocation. + # prepare_store still fails so offloading is still delayed. + runner.manager.prepare_store.side_effect = lambda keys, req_context: None + runner.run(decoded_tokens=[0] * block_size) + + # Now prepare_store succeeds. + # Without the fix, the request would try to offload the stale block_id + # at position 0 (now reused at position 3), causing a duplicate in + # sliding_window_block_ids and eventually a KeyError. + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + # block_ids=[0, ?, 3, 1]: positions 0 and 1 are zeroed (stale blocks that + # were freed by the sliding window and reallocated). Only blocks at + # positions 2 and 3 (request offsets 2, 3) are stored. + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(2, 3), + expected_flushed=(2, 3) if not async_scheduling else (), + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 2dc1071a711..22d00b0c834 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -28,6 +28,7 @@ from vllm.utils.hashing import sha256 from vllm.v1.core.kv_cache_utils import ( get_request_block_hasher, init_none_hash, + resolve_kv_cache_block_sizes, ) from vllm.v1.core.sched.async_scheduler import AsyncScheduler from vllm.v1.core.sched.scheduler import Scheduler @@ -230,13 +231,18 @@ class RequestRunner: vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks self.num_kv_groups = len(kv_cache_config.kv_cache_groups) + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) + scheduler_cls = AsyncScheduler if async_scheduling else Scheduler self.scheduler = scheduler_cls( vllm_config=vllm_config, kv_cache_config=kv_cache_config, log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), - block_size=block_size, + block_size=scheduler_block_size, + hash_block_size=hash_block_size, ) self.worker_connector = OffloadingConnector( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 6adb045277f..375aad4eeb8 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -461,6 +461,35 @@ def test_store_sending_thread_only_skips_on_no_available_handle(): assert store.batch_put_from_multi_buffers.call_count == 2 +def test_store_sending_thread_releases_pin_on_batch_is_exist_failure(): + # `batch_is_exist` raising must still decrement `stored_requests` so the + # scheduler can drop `delay_free_blocks` and release the pinned GPU blocks. + store = MagicMock() + store.batch_is_exist.side_effect = RuntimeError("mooncake down") + thread = _make_store_sending_thread(store) + + thread.add_stored_request("req-a") + with pytest.raises(RuntimeError): + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + assert thread.stored_requests["req-a"] == 0 + store.batch_put_from_multi_buffers.assert_not_called() + + +def test_store_sending_thread_releases_pin_on_batch_put_failure(): + # `batch_put_from_multi_buffers` raising is logged (not re-raised), and the + # pin must still be released through the finally block. + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.side_effect = RuntimeError("rdma error") + thread = _make_store_sending_thread(store) + + thread.add_stored_request("req-a") + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + + assert thread.stored_requests["req-a"] == 0 + + def test_store_recving_thread_reports_failed_block_ids(): store = MagicMock() store.batch_get_into_multi_buffers.return_value = [256, -5, -7] diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 6d4e6565e37..8d54353f82a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -386,8 +386,6 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): "kv_transfer_config": kv_transfer_config, "max_model_len": 2048, "max_num_seqs": 1, - # NOTE: Make sure HMA is enabled - "disable_hybrid_kv_cache_manager": False, "max_num_batched_tokens": 2048, "enable_prefix_caching": False, "block_size": block_size, diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index e4ed635b9b7..d192b04a07b 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -8,6 +8,7 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import set_random_seed from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, @@ -90,13 +91,15 @@ def test_transfer( mmap_region: SharedOffloadRegion | None = None if use_shared_memory: - cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor + cpu_page_size = round_up( + gpu_page_size_bytes * num_tensors * block_size_factor, + SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, + ) mmap_region = SharedOffloadRegion( instance_id=str(uuid.uuid4()), - total_size_bytes=num_cpu_blocks * cpu_page_size, num_blocks=num_cpu_blocks, rank=0, - num_workers=1, + kv_bytes_per_block=cpu_page_size, cpu_page_size=cpu_page_size, ) diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index b33a27ca645..f69fcf9a705 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -40,14 +40,12 @@ def _make_region( num_workers: int = 1, rank: int = 0, ) -> SharedOffloadRegion: - total_size_bytes = num_blocks * num_workers * cpu_page_size - assert total_size_bytes % PAGE_SIZE == 0 + assert cpu_page_size % PAGE_SIZE == 0 return SharedOffloadRegion( instance_id=instance_id, - total_size_bytes=total_size_bytes, num_blocks=num_blocks, rank=rank, - num_workers=num_workers, + kv_bytes_per_block=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) @@ -77,14 +75,12 @@ def _multi_region( cpu_page_size: int = PAGE_SIZE, ): """Context manager: create one SharedOffloadRegion per rank, clean up on exit.""" - total = num_blocks * num_workers * cpu_page_size regions = [ SharedOffloadRegion( instance_id=instance_id, - total_size_bytes=total, num_blocks=num_blocks, rank=rank, - num_workers=num_workers, + kv_bytes_per_block=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) for rank in range(num_workers) @@ -104,7 +100,6 @@ def _race_construct( cpu_page_size: int = PAGE_SIZE, ) -> tuple[list[SharedOffloadRegion], list[Exception]]: """Spawn num_workers threads that all race to construct SharedOffloadRegion.""" - total = num_blocks * num_workers * cpu_page_size regions: list[SharedOffloadRegion | None] = [None] * num_workers errors: list[Exception] = [] barrier = threading.Barrier(num_workers) @@ -114,10 +109,9 @@ def _race_construct( try: regions[rank] = SharedOffloadRegion( instance_id=instance_id, - total_size_bytes=total, num_blocks=num_blocks, rank=rank, - num_workers=num_workers, + kv_bytes_per_block=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) except Exception as e: @@ -134,7 +128,6 @@ def _race_construct( def _mp_race_construct_and_write( instance_id: str, - total_bytes: int, num_blocks: int, rank: int, num_workers: int, @@ -149,10 +142,9 @@ def _mp_race_construct_and_write( try: region = SharedOffloadRegion( instance_id=instance_id, - total_size_bytes=total_bytes, num_blocks=num_blocks, rank=rank, - num_workers=num_workers, + kv_bytes_per_block=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) t = region.create_next_view(cpu_page_size) @@ -309,7 +301,6 @@ def test_create_next_view_multiprocess_slots(iid): the parent verifies each slot lands at the correct interleaved offset.""" num_workers = 2 num_blocks = 4 - total_bytes = num_blocks * num_workers * PAGE_SIZE ctx = get_mp_context() done_queue = ctx.Queue() @@ -318,10 +309,9 @@ def test_create_next_view_multiprocess_slots(iid): # Parent is rank 0 (creator); child is rank 1 (joiner). region = SharedOffloadRegion( instance_id=iid, - total_size_bytes=total_bytes, num_blocks=num_blocks, rank=0, - num_workers=num_workers, + kv_bytes_per_block=num_workers * PAGE_SIZE, cpu_page_size=PAGE_SIZE, ) try: @@ -329,7 +319,6 @@ def test_create_next_view_multiprocess_slots(iid): target=_mp_race_construct_and_write, args=( iid, - total_bytes, num_blocks, 1, num_workers, @@ -464,7 +453,6 @@ def test_multiprocess_race_construct_and_write(iid): fill_value = rank+1 into their slot; parent verifies interleaved layout.""" num_workers = 4 num_blocks = 3 - total_bytes = num_blocks * num_workers * PAGE_SIZE ctx = get_mp_context() done_queue = ctx.Queue() @@ -475,7 +463,6 @@ def test_multiprocess_race_construct_and_write(iid): target=_mp_race_construct_and_write, args=( iid, - total_bytes, num_blocks, rank, num_workers, diff --git a/tests/v1/kv_offload/cpu/test_swap_blocks_triton.py b/tests/v1/kv_offload/cpu/test_swap_blocks_triton.py new file mode 100644 index 00000000000..ec14a378434 --- /dev/null +++ b/tests/v1/kv_offload/cpu/test_swap_blocks_triton.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit test for the Triton ``swap_blocks_batch`` fast-path kernel.""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.kv_offload.cpu.swap_blocks_triton import swap_blocks_batch + + +def _addrs(buffers: list[torch.Tensor]) -> torch.Tensor: + return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="Triton swap fast path requires CUDA" +) +def test_triton_swap_copies_source_bytes(): + # 8-byte-aligned, sub-threshold sizes covering 8 KiB chunk boundaries and + # odd tail-mask lengths, with enough descriptors to take the Triton path. + sizes = [8, 4096, 8192, 8200, 16384, 4088] * 8 + src = [torch.randint(256, (s,), dtype=torch.uint8, device="cuda") for s in sizes] + dst = [torch.zeros_like(s) for s in src] + sizes_t = torch.tensor(sizes, dtype=torch.int64) + + swap_blocks_batch(_addrs(src), _addrs(dst), sizes_t.clone(), bytes_per_chunk=8192) + torch.accelerator.synchronize() + + for s, t in zip(src, dst): + assert torch.equal(t, s) # kernel copied the source bytes verbatim diff --git a/tests/v1/kv_offload/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py similarity index 86% rename from tests/v1/kv_offload/test_fs_tier.py rename to tests/v1/kv_offload/tiering/test_fs_tier.py index 70ecc7943bd..ab5ed23c2dd 100644 --- a/tests/v1/kv_offload/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -8,6 +8,7 @@ The tier manager writes KV cache blocks to disk and reads them back, verifying data integrity throughout the process. """ +import mmap import os import time from unittest.mock import MagicMock @@ -26,8 +27,8 @@ from vllm.v1.kv_offload.tiering.fs.manager import ( # Helpers # --------------------------------------------------------------------------- -_BLOCK_ELEMENTS = 512 * 1024 # 2 MB per block (float32 × 512K = 2MB) -_DTYPE = torch.float32 +_BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096. +_DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") _MOCK_VLLM_CONFIG = MagicMock() @@ -90,6 +91,32 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list: return results +def _page_aligned_zero_tensor( + num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE +) -> torch.Tensor: + page_size = mmap.PAGESIZE + dtype_num_bytes = torch.tensor([], dtype=dtype).element_size() + + num_bytes = num_blocks * block_elements * dtype_num_bytes + num_bytes_aligned = num_bytes + page_size + t = torch.zeros(num_bytes_aligned, dtype=torch.uint8) + + ptr = t.data_ptr() + alignment_offset = ptr % page_size + # Move tensor to next page regardless. + shift = page_size - alignment_offset + t = t[shift : shift + num_bytes] + return t.view(dtype).view(num_blocks, block_elements) + + +def _page_aligned_rand_tensor( + num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE +) -> torch.Tensor: + rand_tensor = _page_aligned_zero_tensor(num_blocks, block_elements) + rand_tensor[:] = torch.rand(num_blocks, block_elements, dtype=dtype) + return rand_tensor + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -97,7 +124,7 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list: @pytest.fixture def fs_tier(tmp_path): - tensor = torch.zeros((4, _BLOCK_ELEMENTS), dtype=_DTYPE) + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) tier = FileSystemTierManager( offloading_spec=_MOCK_OFFLOADING_SPEC, @@ -155,7 +182,7 @@ def test_store_then_load_roundtrip(fs_tier): def test_invalid_path_raises_at_construction(): """Construction must fail immediately when the config file cannot be written.""" - tensor = torch.zeros((32, _BLOCK_ELEMENTS), dtype=_DTYPE) + tensor = _page_aligned_zero_tensor(32, _BLOCK_ELEMENTS) mock_view = memoryview(tensor.numpy()) with pytest.raises(OSError): @@ -228,7 +255,7 @@ def test_store_load_data_integrity(fs_tier): """Data written by store must be exactly recovered by load.""" tier, tensor = fs_tier # Populate tensor with random data - tensor[:] = torch.rand((4, _BLOCK_ELEMENTS), dtype=_DTYPE) + tensor[:] = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) # Store first 2 blocks num_store = 2 diff --git a/tests/v1/kv_offload/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py similarity index 96% rename from tests/v1/kv_offload/test_tiering_offloading.py rename to tests/v1/kv_offload/tiering/test_tiering_offloading.py index a64caddb2b1..5a7c11787d9 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -123,6 +123,11 @@ class TestTieringOffloadingManager: secondary_tiers=[self.secondary_tier1, self.secondary_tier2], ) + def _simulate_on_schedule_end(self): + """Simulate end of scheduler step: lifecycle flush + drain events.""" + self.manager.on_schedule_end() + list(self.manager.take_events()) + def test_basic_store_to_primary(self, manager_setup): """Test basic store operation to primary tier.""" blocks = to_keys(range(3)) @@ -185,10 +190,10 @@ class TestTieringOffloadingManager: assert block.ref_cnt == 2 # End of step 1: _maybe_process_finished_jobs() was already called by - # prepare_store() above (setting the per-step flag), so take_events() + # prepare_store() above (setting the per-step flag), so on_schedule_end() # does NOT poll get_finished_jobs() again — cascade completions remain # unprocessed until the next step. - list(self.manager.take_events()) + self._simulate_on_schedule_end() # ref_cnt still held: cascade jobs finished (sync tier) but haven't # been polled yet because the per-step guard skipped the second call. @@ -202,7 +207,7 @@ class TestTieringOffloadingManager: # End of step 2: flag was reset, so _maybe_process_finished_jobs() # runs and processes the cascade completions (complete_read → ref_cnt--) - list(self.manager.take_events()) + self._simulate_on_schedule_end() # After cascade completes, ref_cnt should be 0 for block_hash in blocks: @@ -238,10 +243,10 @@ class TestTieringOffloadingManager: assert result is None # Retry later (promotion initiated) # End of step 1: flushes deferred submit_load() calls - list(self.manager.take_events()) + self._simulate_on_schedule_end() # End of step 2: processes the completed promotion jobs - list(self.manager.take_events()) + self._simulate_on_schedule_end() # Now blocks should be in primary tier assert count_hits(self.primary_tier, blocks) == 3 @@ -271,7 +276,7 @@ class TestTieringOffloadingManager: self.manager.complete_store(blocks, _CTX, success=True) # End of step: release ref_cnt from cascade - list(self.manager.take_events()) + self._simulate_on_schedule_end() # Now try to store 2 more blocks (should trigger eviction) more_blocks = to_keys(range(5, 7)) @@ -289,7 +294,7 @@ class TestTieringOffloadingManager: # Store blocks self.manager.prepare_store(blocks, _CTX) self.manager.complete_store(blocks, _CTX, success=True) - list(self.manager.take_events()) + self._simulate_on_schedule_end() self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch) self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch) @@ -328,7 +333,7 @@ class TestTieringOffloadingManager: self.secondary_tier2.submit_store.assert_not_called() def test_lookup_batches_submit_load_per_request(self, manager_setup): - """lookup() defers submit_load until take_events(), one call per request. + """lookup() defers submit_load until on_schedule_end(), one per request. Blocks from different requests each get their own submit_load call, each carrying the correct req_context. @@ -354,7 +359,7 @@ class TestTieringOffloadingManager: self.secondary_tier1.submit_load.assert_not_called() # simulate end of step - list(self.manager.take_events()) + self._simulate_on_schedule_end() assert self.secondary_tier1.submit_load.call_count == 2 calls = self.secondary_tier1.submit_load.call_args_list @@ -389,7 +394,7 @@ class TestTieringOffloadingManager: assert result_a is None assert result_b is None - list(self.manager.take_events()) + self._simulate_on_schedule_end() # Only one submit_load call despite two lookups self.secondary_tier1.submit_load.assert_called_once() @@ -449,8 +454,7 @@ class TestTieringOffloadingManager: assert result is not None self.manager.complete_store(existing_blocks, _CTX, success=True) # Drain cascade completions - list(self.manager.take_events()) - list(self.manager.take_events()) + self._simulate_on_schedule_end() # Make tier1 request-level, tier2 stays block-level self.secondary_tier1.on_new_request = ( diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 460e0d68564..963e7423f79 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -4,6 +4,7 @@ import itertools import math from collections.abc import Generator +from types import SimpleNamespace from typing import get_args import pytest @@ -20,6 +21,7 @@ from tests.v1.sample.utils import ( from vllm import SamplingParams from vllm.config.model import LogprobsMode from vllm.distributed import cleanup_dist_env_and_memory +from vllm.exceptions import VLLMValidationError from vllm.platforms import current_platform from ...conftest import HfRunner, VllmRunner @@ -78,6 +80,14 @@ def hf_model(hf_runner) -> Generator[HfRunner, None, None]: yield hf_model +def _model_config(vocab_size: int = 10): + return SimpleNamespace( + max_logprobs=20, + logits_processors=None, + get_vocab_size=lambda: vocab_size, + ) + + def _repeat_logprob_config( test_prompts, logprob_prompt_logprob_list: BatchLogprobsSpecType, @@ -397,6 +407,27 @@ def test_max_logprobs(): runner.generate(["Hello world"], sampling_params=bad_sampling_params) +@pytest.mark.parametrize("token_ids", [[0], [0, 9]]) +def test_logprob_token_ids_validate_vocab_bounds_valid(token_ids: list[int]): + SamplingParams(logprob_token_ids=token_ids).verify( + _model_config(), + speculative_config=None, + structured_outputs_config=None, + tokenizer=None, + ) + + +@pytest.mark.parametrize("token_ids", [[-1], [10], [-35, 1873042417]]) +def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]): + with pytest.raises(VLLMValidationError, match="logprob_token_ids"): + SamplingParams(logprob_token_ids=token_ids).verify( + _model_config(), + speculative_config=None, + structured_outputs_config=None, + tokenizer=None, + ) + + def test_none_logprobs(vllm_model, example_prompts): """Engine should return `logprobs` and `prompt_logprobs` as `None` diff --git a/tests/v1/shutdown/test_forward_error.py b/tests/v1/shutdown/test_forward_error.py index eadb1abb6d5..8bc09a64cac 100644 --- a/tests/v1/shutdown/test_forward_error.py +++ b/tests/v1/shutdown/test_forward_error.py @@ -36,6 +36,7 @@ def evil_forward(self, *args, **kwargs): raise Exception("Simulated illegal memory access on Rank 0!") self.num_calls += 1 + kwargs.setdefault("intermediate_tensors", None) # required for MRV2 return self.model(*args, **kwargs) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 970e16e5279..e59905f504a 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -30,6 +30,9 @@ from vllm.v1.core.sched.output import ( NewRequestData, SchedulerOutput, ) +from vllm.v1.core.single_type_kv_cache_manager import ( + register_all_kvcache_specs, +) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -68,6 +71,9 @@ def _make_kv_cache_config( """Build a KVCacheConfig with non-empty kv_cache_tensors.""" groups = [] tensors = [] + register_all_kvcache_specs( + vllm_config=None + ) # Ensure specs are registered for tests for g in range(num_groups): layer_names = [f"layer_{g}"] groups.append( diff --git a/tests/v1/spec_decode/test_dflash_lookahead.py b/tests/v1/spec_decode/test_dflash_lookahead.py new file mode 100644 index 00000000000..d2980fb6102 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_lookahead.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from tests.v1.core.utils import create_requests +from vllm.config import ( + CacheConfig, + ModelConfig, + ParallelConfig, + SchedulerConfig, + SpeculativeConfig, + VllmConfig, +) +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.structured_output import StructuredOutputManager +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +# Matches defaults from tests/v1/spec_decode/test_eagle.py +DFLASH_TARGET_DIR = "Qwen/Qwen3-8B" +DFLASH_DRAFT_DIR = "z-lab/Qwen3-8B-DFlash-b16" + +BLOCK_SIZE = 16 +NUM_BLOCKS = 8 +NUM_SPECULATIVE_TOKENS = 3 + + +def _dflash_speculative_config(num_speculative_tokens: int) -> SpeculativeConfig: + model_config = ModelConfig( + model=DFLASH_TARGET_DIR, + runner="generate", + max_model_len=100, + trust_remote_code=True, + ) + return SpeculativeConfig( + target_model_config=model_config, + target_parallel_config=ParallelConfig(), + model=DFLASH_DRAFT_DIR, + method="dflash", + num_speculative_tokens=num_speculative_tokens, + ) + + +def _create_dflash_scheduler(num_speculative_tokens: int) -> Scheduler: + speculative_config = _dflash_speculative_config(num_speculative_tokens) + model_config = speculative_config.target_model_config + scheduler_config = SchedulerConfig( + max_num_seqs=16, + max_num_batched_tokens=8192, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + cache_config = CacheConfig( + block_size=BLOCK_SIZE, + gpu_memory_utilization=0.9, + cache_dtype="auto", + enable_prefix_caching=False, + ) + vllm_config = VllmConfig( + scheduler_config=scheduler_config, + model_config=model_config, + cache_config=cache_config, + parallel_config=ParallelConfig(), + speculative_config=speculative_config, + ) + kv_cache_config = KVCacheConfig( + num_blocks=NUM_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + cache_config.num_gpu_blocks = NUM_BLOCKS + return Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + block_size=BLOCK_SIZE, + log_stats=True, + structured_output_manager=StructuredOutputManager(vllm_config), + ) + + +def test_dflash_prefill_reserves_lookahead_blocks(): + scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS) + + assert scheduler.num_lookahead_tokens == NUM_SPECULATIVE_TOKENS + 1 + + (request,) = create_requests( + num_requests=1, + num_tokens=BLOCK_SIZE, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(request) + + output = scheduler.schedule() + + assert output.num_scheduled_tokens[request.request_id] == BLOCK_SIZE + # prefill block + one lookahead block + assert len(output.scheduled_new_reqs[0].block_ids[0]) == 2 + + +def test_dflash_first_prefill_query_window_fits_allocated_blocks(): + scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS) + + (request,) = create_requests( + num_requests=1, + num_tokens=BLOCK_SIZE, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(request) + + output = scheduler.schedule() + block_ids = output.scheduled_new_reqs[0].block_ids[0] + query_positions = range(BLOCK_SIZE, BLOCK_SIZE + scheduler.num_lookahead_tokens) + + assert all(pos // BLOCK_SIZE < len(block_ids) for pos in query_positions) + + +def test_dflash_drafter_window_reserves_bonus_token(): + # DFlash's drafter window is num_spec + 1 (the extra slot is the bonus token), + # so max_seq_len + num_spec + 1 must stay within the draft model's max len. + input_fits_in_drafter = GPUModelRunner._input_fits_in_drafter + dflash_runner = SimpleNamespace( + num_spec_tokens=NUM_SPECULATIVE_TOKENS, + effective_drafter_max_model_len=100, + speculative_config=_dflash_speculative_config(NUM_SPECULATIVE_TOKENS), + ) + # window = 4, so 96 fits (96 + 4 == 100) but 97 does not (97 + 4 == 101) + assert input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=96)) + assert not input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=97)) + assert not input_fits_in_drafter(dflash_runner, None) # no metadata + + # Other drafters don't reserve the bonus token, so 97 fits (97 + 3 == 100). + plain_runner = SimpleNamespace( + num_spec_tokens=NUM_SPECULATIVE_TOKENS, + effective_drafter_max_model_len=100, + speculative_config=SimpleNamespace(use_dflash=lambda: False), + ) + assert input_fits_in_drafter(plain_runner, SimpleNamespace(max_seq_len=97)) diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py new file mode 100644 index 00000000000..c84c3ef1c8a --- /dev/null +++ b/tests/v1/test_kv_cache_spec_registry.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, replace +from typing import Any + +import pytest +import torch + +from vllm.config import ( + CacheConfig, + DeviceConfig, + VllmConfig, +) +from vllm.v1.core.single_type_kv_cache_manager import ( + ChunkedLocalAttentionManager, + CrossAttentionManager, + FullAttentionManager, + MambaManager, + SingleTypeKVCacheManager, + SinkFullAttentionManager, + SlidingWindowManager, + register_all_kvcache_specs, +) +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + CrossAttentionSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + SinkFullAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + TQFullAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.kv_cache_spec_registry import ( + _REGISTRY_KVCACHESPEC_LIST, + KVCacheSpecRegistry, + register_kv_cache_spec, +) + + +def make_vllm_config() -> VllmConfig: + return VllmConfig( + cache_config=CacheConfig( + block_size=64, + cache_dtype="bfloat16", + ), + device_config=DeviceConfig(device="cpu"), + ) + + +vllm_config = make_vllm_config() +register_all_kvcache_specs(vllm_config) + + +@pytest.fixture(autouse=True) +def restore_kv_cache_spec_registry(): + registry = _REGISTRY_KVCACHESPEC_LIST.copy() + yield + _REGISTRY_KVCACHESPEC_LIST.clear() + _REGISTRY_KVCACHESPEC_LIST.update(registry) + + +@dataclass(frozen=True) +class _TrulyUnregisteredSpec(KVCacheSpec): + """ + A spec that inherits directly from KVCacheSpec with no registered + ancestor in the MRO. Used to test that the registry correctly raises + when no entry can be found. + """ + + @property + def page_size_bytes(self) -> int: + return self.block_size * 128 + + def max_memory_usage_bytes(self, _) -> int: + return 0 + + +spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = { + FullAttentionSpec: FullAttentionManager, + TQFullAttentionSpec: FullAttentionManager, + MLAAttentionSpec: FullAttentionManager, + HiddenStateCacheSpec: FullAttentionManager, + SlidingWindowSpec: SlidingWindowManager, + SlidingWindowMLASpec: SlidingWindowManager, + ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, + MambaSpec: MambaManager, + CrossAttentionSpec: CrossAttentionManager, + SinkFullAttentionSpec: SinkFullAttentionManager, +} + +spec_uniform_base_map: dict[type[KVCacheSpec], type[KVCacheSpec]] = { + FullAttentionSpec: FullAttentionSpec, + TQFullAttentionSpec: FullAttentionSpec, + MLAAttentionSpec: FullAttentionSpec, + HiddenStateCacheSpec: FullAttentionSpec, + SlidingWindowSpec: SlidingWindowSpec, + SlidingWindowMLASpec: SlidingWindowMLASpec, + ChunkedLocalAttentionSpec: ChunkedLocalAttentionSpec, + MambaSpec: MambaSpec, + CrossAttentionSpec: CrossAttentionSpec, + SinkFullAttentionSpec: FullAttentionSpec, +} + +spec_args_map: dict[type[KVCacheSpec], dict[str, Any]] = { + FullAttentionSpec: dict( + block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ), + TQFullAttentionSpec: dict( + block_size=64, + num_kv_heads=8, + head_size=128, + dtype=torch.bfloat16, + tq_slot_size=256, + ), + MLAAttentionSpec: dict( + block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ), + HiddenStateCacheSpec: dict( + block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 + ), + SlidingWindowSpec: dict( + block_size=64, + num_kv_heads=8, + head_size=128, + dtype=torch.bfloat16, + sliding_window=128, + ), + SlidingWindowMLASpec: dict( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=128, + ), + ChunkedLocalAttentionSpec: dict( + block_size=64, + num_kv_heads=8, + head_size=128, + dtype=torch.bfloat16, + attention_chunk_size=4, + ), + MambaSpec: dict( + block_size=64, + shapes=((2, 512), (3, 32, 32)), + dtypes=(torch.float32, torch.float32), + mamba_cache_mode="align", + num_speculative_blocks=2, + ), + CrossAttentionSpec: dict( + block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ), + SinkFullAttentionSpec: dict( + block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16, sink_len=16 + ), +} + + +def make_spec(spec_cls: type[KVCacheSpec]) -> KVCacheSpec: + return spec_cls(**spec_args_map[spec_cls]) + + +def are_uniform_specs(*specs: KVCacheSpec) -> bool: + return UniformTypeKVCacheSpecs.is_uniform_type( + {f"layer_{i}": spec for i, spec in enumerate(specs)} + ) + + +class TestKVCacheSpecRegistry: + """Test the core registry functionality.""" + + def test_builtin_kvcache_specs_registered(self): + assert set(spec_manager_map) <= set(_REGISTRY_KVCACHESPEC_LIST) + for spec_cls, manager in spec_manager_map.items(): + spec = make_spec(spec_cls) + assert KVCacheSpecRegistry.get_manager_class(spec) is manager + assert ( + KVCacheSpecRegistry.get_uniform_type_base_spec(spec) + is spec_uniform_base_map[spec_cls] + ) + + @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) + def test_custom_spec_register(self, spec_cls): + """A decorated custom spec resolves to the declared manager.""" + manager = spec_manager_map[spec_cls] + uniform_base_spec = spec_uniform_base_map[spec_cls] + + @register_kv_cache_spec( + manager_class=manager, + uniform_type_base_spec=uniform_base_spec, + ) + @dataclass(frozen=True, kw_only=True) + class _CustomSpec(spec_cls): # type: ignore[valid-type,misc] + custom_param: int = 16 + + spec = _CustomSpec(**spec_args_map[spec_cls], custom_param=100) + + assert KVCacheSpecRegistry.get_manager_class(spec) is manager + assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is uniform_base_spec + + def test_custom_spec_register_requires_manager(self): + """Invalid register decorator arguments fail early.""" + + with pytest.raises(AssertionError, match="manager_class is required"): + + @register_kv_cache_spec( + uniform_type_base_spec=FullAttentionSpec, + ) + @dataclass(frozen=True, kw_only=True) + class _CustomFullSpecWithoutManager(FullAttentionSpec): + custom_param: int = 16 + + def test_unregistered_spec_no_registered_parent_raises(self): + """ + A spec whose entire MRO contains no registered class resolves to None. + Runtime callers should use check_kv_cache_spec_registry to fail early. + Subclasses of registered specs intentionally do not fail — they inherit + their parent's manager via MRO walking. + """ + spec = _TrulyUnregisteredSpec(block_size=16) + + assert KVCacheSpecRegistry.get_manager_class(spec) is None + assert KVCacheSpecRegistry.get_uniform_type_base_spec(spec) is None + + with pytest.raises( + ValueError, match="Unsupported KV cache spec type for layer layer_0" + ): + KVCacheSpecRegistry.check_kv_cache_spec_registry({"layer_0": spec}) + + with pytest.raises(AssertionError, match="Unsupported KV cache spec type"): + UniformTypeKVCacheSpecs.is_uniform_type({"layer_0": spec}) + + def test_unregistered_subclass_inherits_parent_manager(self): + """ + An unregistered subclass of a registered spec resolves via MRO + to its parent's manager — this is intentional registry behaviour. + """ + + @dataclass(frozen=True, kw_only=True) + class _ImplicitlyInheritedSpec(FullAttentionSpec): + pass + + spec = _ImplicitlyInheritedSpec( + block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + + # MRO walk finds FullAttentionSpec → FullAttentionManager + assert KVCacheSpecRegistry.get_manager_class(spec) is FullAttentionManager + + @pytest.mark.parametrize("spec_cls", list(spec_manager_map)) + def test_builtin_specs_are_uniform_with_same_spec_type(self, spec_cls): + spec = make_spec(spec_cls) + assert are_uniform_specs(spec, replace(spec)) + + def test_full_attention_family_specs_are_uniform(self): + specs = [ + make_spec(FullAttentionSpec), + make_spec(TQFullAttentionSpec), + make_spec(MLAAttentionSpec), + make_spec(HiddenStateCacheSpec), + make_spec(SinkFullAttentionSpec), + ] + + assert are_uniform_specs(*specs) + + @pytest.mark.parametrize( + ("spec_cls", "field", "value"), + [ + (SlidingWindowSpec, "sliding_window", 256), + (SlidingWindowMLASpec, "sliding_window", 256), + (ChunkedLocalAttentionSpec, "attention_chunk_size", 8), + (MambaSpec, "num_speculative_blocks", 4), + ], + ) + def test_specs_with_type_specific_uniform_fields(self, spec_cls, field, value): + spec = make_spec(spec_cls) + changed_spec = replace(spec, **{field: value}) + + assert not are_uniform_specs(spec, changed_spec) + + @pytest.mark.parametrize( + ("left_cls", "right_cls"), + [ + (FullAttentionSpec, CrossAttentionSpec), + (FullAttentionSpec, SlidingWindowSpec), + (FullAttentionSpec, ChunkedLocalAttentionSpec), + (FullAttentionSpec, MambaSpec), + (SlidingWindowMLASpec, SlidingWindowSpec), + (ChunkedLocalAttentionSpec, SlidingWindowSpec), + (MambaSpec, CrossAttentionSpec), + ], + ) + def test_different_uniform_groups_are_not_uniform(self, left_cls, right_cls): + assert not are_uniform_specs(make_spec(left_cls), make_spec(right_cls)) + + def test_different_block_sizes_are_not_uniform(self): + spec = make_spec(FullAttentionSpec) + + assert not are_uniform_specs(spec, replace(spec, block_size=32)) + + def test_registered_custom_spec_uses_base_uniform_rule(self): + @register_kv_cache_spec( + manager_class=FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + @dataclass(frozen=True, kw_only=True) + class _CustomFullSpec(FullAttentionSpec): + custom_param: int = 16 + + custom_spec = _CustomFullSpec( + block_size=64, + num_kv_heads=8, + head_size=128, + dtype=torch.bfloat16, + ) + + assert are_uniform_specs(custom_spec, make_spec(FullAttentionSpec)) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index c2a800bd998..1db07baf93d 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -167,7 +167,9 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): events = [] runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) runner.execute_model_state = SimpleNamespace( - input_batch=SimpleNamespace(num_reqs=2), + input_batch=SimpleNamespace( + num_reqs=2, idx_mapping=torch.zeros(2, dtype=torch.int32) + ), attn_metadata=None, slot_mappings_by_layer=None, hidden_states=None, @@ -175,18 +177,19 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): finished_req_ids=set(), num_tokens_across_dp=None, ) - runner.postprocess = lambda *args, **kwargs: events.append("postprocess") - runner.eplb.step = lambda *args, **kwargs: events.append("eplb") - monkeypatch.setattr( - mrv2, - "pp_receive", - lambda *args, **kwargs: ( - torch.zeros((2, 1), dtype=torch.long), - torch.ones(2, dtype=torch.int32), - torch.zeros(2, dtype=torch.int32), - ), + runner.req_states = SimpleNamespace() + + def fake_receive(*args, **kwargs): + events.append("receive") + # all_decode_next=True, so model_state.postprocess_state is skipped. + return True + + runner.pp_handler = SimpleNamespace(receive=fake_receive) + runner.postprocess_num_computed_tokens = lambda *args, **kwargs: events.append( + "postprocess_num_computed_tokens" ) + runner.eplb.step = lambda *args, **kwargs: events.append("eplb") output = mrv2.GPUModelRunner.sample_tokens(runner, None) assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) - assert events == ["postprocess", "eplb"] + assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] diff --git a/tools/install_torchcodec_rocm.sh b/tools/install_torchcodec_rocm.sh index 6cb3b39fd66..210d7b24145 100755 --- a/tools/install_torchcodec_rocm.sh +++ b/tools/install_torchcodec_rocm.sh @@ -3,12 +3,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Script to install TorchCodec from source (required for ROCm compatibility) +# The PyPI wheel is built against upstream PyTorch and has ABI mismatches with +# ROCm's custom torch build, so we must compile from source. set -e TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}" # Pin to a specific release for reproducibility; update as needed. TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}" +# Cache directory for pre-built wheels to avoid redundant recompilation. +TORCHCODEC_WHEEL_CACHE="${TORCHCODEC_WHEEL_CACHE:-/root/.cache/torchcodec-wheels}" echo "=== TorchCodec Installation Script ===" @@ -18,9 +22,26 @@ if python3 -c "from torchcodec.decoders import VideoDecoder" 2>/dev/null; then exit 0 fi +# Try to install from cached wheel first +ARCH_TAG="${PYTORCH_ROCM_ARCH:-all}" +# Normalize arch tag (replace ; with _) for use in filename +ARCH_TAG="${ARCH_TAG//;/_}" +CACHED_WHEEL="${TORCHCODEC_WHEEL_CACHE}/torchcodec-${TORCHCODEC_BRANCH}-${ARCH_TAG}.whl" + +if [ -f "$CACHED_WHEEL" ]; then + echo "Found cached wheel: $CACHED_WHEEL" + pip install "$CACHED_WHEEL" && { + echo "Installed from cached wheel." + echo "=== TorchCodec installation complete ===" + exit 0 + } + echo "Cached wheel installation failed, rebuilding from source..." +fi + echo "TorchCodec not found. Installing from source..." -# Install system dependencies (FFmpeg + pkg-config) +# Install system dependencies (FFmpeg + pkg-config) if not already present. +# The Docker test image pre-installs these, so this is a fallback for other envs. install_system_deps() { if command -v apt-get &> /dev/null; then echo "Installing system dependencies..." @@ -56,6 +77,12 @@ export pybind11_DIR=$(python3 -c "import pybind11; print(pybind11.get_cmake_dir( export CMAKE_PREFIX_PATH="${pybind11_DIR}:${CMAKE_PREFIX_PATH}" echo "pybind11_DIR set to: $pybind11_DIR" +# Limit GPU architectures to only what this image targets. +# The default builds for all supported archs which is very slow. +if [ -n "$PYTORCH_ROCM_ARCH" ]; then + echo "Building for PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH" +fi + # Create temp directory for build BUILD_DIR=$(mktemp -d -t torchcodec-XXXXXX) echo "Building in temporary directory: $BUILD_DIR" @@ -77,9 +104,31 @@ cd torchcodec export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build" export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1 export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1 +# Use ninja for faster builds and parallelize compilation +export CMAKE_GENERATOR=Ninja +export MAX_JOBS="${MAX_JOBS:-$(nproc)}" +# Use ccache if available to speed up recompilation +if command -v ccache &> /dev/null; then + export CMAKE_C_COMPILER_LAUNCHER=ccache + export CMAKE_CXX_COMPILER_LAUNCHER=ccache +fi -echo "Building TorchCodec..." -pip install . --no-build-isolation +echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..." +pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist" + +# Install the built wheel +BUILT_WHEEL=$(ls "$BUILD_DIR/dist"/torchcodec-*.whl 2>/dev/null | head -1) +if [ -z "$BUILT_WHEEL" ]; then + echo "Error: No wheel produced" + exit 1 +fi + +pip install "$BUILT_WHEEL" + +# Cache the wheel for future runs +mkdir -p "$TORCHCODEC_WHEEL_CACHE" +cp "$BUILT_WHEEL" "$CACHED_WHEEL" +echo "Cached wheel to: $CACHED_WHEEL" # Verify installation echo "Verifying installation..." @@ -88,4 +137,4 @@ if python3 -c "from torchcodec.decoders import VideoDecoder; print('TorchCodec i else echo "Error: TorchCodec installation failed verification" exit 1 -fi \ No newline at end of file +fi diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 0df0338a648..1adad42f104 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -8,6 +8,7 @@ from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -338,7 +339,16 @@ def _xpu_mxfp8_quantize_impl( shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,) x_s = torch.empty(shape, device=x.device, dtype=torch.float32) torch.ops._C.per_token_group_fp8_quant( - x, x_q, x_s, MXFP8_BLOCK_SIZE, eps, fp8_min, fp8_max, True + x, + x_q, + x_s, + MXFP8_BLOCK_SIZE, + eps, + fp8_min, + fp8_max, + True, + False, + False, # dummy_is_scale_transposed, dummy_is_tma_aligned ) x_s = x_s.to(torch.float8_e8m0fnu) return x_q, x_s @@ -398,6 +408,312 @@ def _xpu_mxfp4_quantize_fake( return x_q, x_s +@triton.jit +def _softplus(x): + return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x) + + +@triton.jit +def _selective_scan_fwd_kernel( + # Pointers to input tensors + u_ptr, + delta_ptr, + A_ptr, + B_ptr, + C_ptr, + D_ptr, + z_ptr, + delta_bias_ptr, + # Pointers to output tensors (out aliases delta, out_z aliases z) + out_ptr, + out_z_ptr, + # SSM states + ssm_states_ptr, + # Optional pointers + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + # APC pointers + block_idx_first_ptr, + block_idx_last_ptr, + initial_state_idx_ptr, + cu_chunk_seqlen_ptr, + last_chunk_indices_ptr, + # Dimensions + batch: tl.int32, + dim: tl.int32, + seqlen: tl.int32, + dstate: tl.int32, + n_groups: tl.int32, + dim_ngroups_ratio: tl.int32, + # Strides for u (and out, since out = delta which has same layout) + u_batch_stride: tl.int64, + u_d_stride: tl.int64, + # Strides for delta + delta_batch_stride: tl.int64, + delta_d_stride: tl.int64, + # Strides for A + A_d_stride: tl.int64, + A_dstate_stride: tl.int64, + # Strides for B + B_batch_stride: tl.int64, + B_group_stride: tl.int64, + B_dstate_stride: tl.int64, + # Strides for C + C_batch_stride: tl.int64, + C_group_stride: tl.int64, + C_dstate_stride: tl.int64, + # Strides for z + z_batch_stride: tl.int64, + z_d_stride: tl.int64, + # Strides for out + out_batch_stride: tl.int64, + out_d_stride: tl.int64, + # Strides for out_z + out_z_batch_stride: tl.int64, + out_z_d_stride: tl.int64, + # Strides for ssm_states + ssm_batch_stride: tl.int64, + ssm_dim_stride: tl.int64, + ssm_dstate_stride: tl.int64, + # Cache strides + cache_indices_stride: tl.int64, + # Scalar params + null_block_id: tl.int64, + block_size: tl.int32, + # Compile-time constants + delta_softplus: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_DELTA_BIAS: tl.constexpr, + IS_VARLEN: tl.constexpr, + HAS_CACHE_INDICES: tl.constexpr, + CACHE_ENABLED: tl.constexpr, + BLOCK_DSTATE: tl.constexpr, +): + batch_idx = tl.program_id(0) + dim_idx = tl.program_id(1) + group_idx = dim_idx // dim_ngroups_ratio + + # Determine sequence boundaries + if IS_VARLEN: + seq_start = tl.load(query_start_loc_ptr + batch_idx).to(tl.int32) + seq_end = tl.load(query_start_loc_ptr + batch_idx + 1).to(tl.int32) + actual_seqlen = seq_end - seq_start + else: + seq_start = 0 + actual_seqlen = seqlen + + # Determine cache index for ssm_states + if CACHE_ENABLED: + init_state_idx = tl.load(initial_state_idx_ptr + batch_idx).to(tl.int32) + load_cache_slot = tl.load( + cache_indices_ptr + batch_idx * cache_indices_stride + init_state_idx + ).to(tl.int64) + if load_cache_slot == null_block_id: + return + elif HAS_CACHE_INDICES: + cache_index = tl.load(cache_indices_ptr + batch_idx).to(tl.int64) + if cache_index == null_block_id: + return + load_cache_slot = cache_index + else: + load_cache_slot = batch_idx.to(tl.int64) + + # Load D value + D_val = 0.0 + if HAS_D: + D_val = tl.load(D_ptr + dim_idx).to(tl.float32) + + # Load delta_bias value + delta_bias_val = 0.0 + if HAS_DELTA_BIAS: + delta_bias_val = tl.load(delta_bias_ptr + dim_idx).to(tl.float32) + + # Load A values for this dim - shape (dstate,) + dstate_offs = tl.arange(0, BLOCK_DSTATE) + dstate_mask = dstate_offs < dstate + A_vals = tl.load( + A_ptr + dim_idx * A_d_stride + dstate_offs * A_dstate_stride, + mask=dstate_mask, + other=0.0, + ).to(tl.float32) + + # Initialize state vector + state = tl.zeros((BLOCK_DSTATE,), dtype=tl.float32) + + # Load initial state if available + has_init = False + if has_initial_state_ptr is not None: + has_init = tl.load(has_initial_state_ptr + batch_idx) + if has_init: + state = tl.load( + ssm_states_ptr + + load_cache_slot * ssm_batch_stride + + dim_idx * ssm_dim_stride + + dstate_offs * ssm_dstate_stride, + mask=dstate_mask, + other=0.0, + ).to(tl.float32) + + # Compute base addresses for u and delta + if IS_VARLEN: + u_base = u_ptr + dim_idx * u_d_stride + seq_start * u_batch_stride + delta_base = ( + delta_ptr + dim_idx * delta_d_stride + seq_start * delta_batch_stride + ) + out_base = out_ptr + dim_idx * out_d_stride + seq_start * out_batch_stride + B_base = B_ptr + group_idx * B_group_stride + seq_start * B_batch_stride + C_base = C_ptr + group_idx * C_group_stride + seq_start * C_batch_stride + else: + u_base = u_ptr + batch_idx * u_batch_stride + dim_idx * u_d_stride + delta_base = ( + delta_ptr + batch_idx * delta_batch_stride + dim_idx * delta_d_stride + ) + out_base = out_ptr + batch_idx * out_batch_stride + dim_idx * out_d_stride + B_base = B_ptr + batch_idx * B_batch_stride + group_idx * B_group_stride + C_base = C_ptr + batch_idx * C_batch_stride + group_idx * C_group_stride + + if HAS_Z: + if IS_VARLEN: + z_base = z_ptr + dim_idx * z_d_stride + seq_start * z_batch_stride + out_z_base = ( + out_z_ptr + dim_idx * out_z_d_stride + seq_start * out_z_batch_stride + ) + else: + z_base = z_ptr + batch_idx * z_batch_stride + dim_idx * z_d_stride + out_z_base = ( + out_z_ptr + batch_idx * out_z_batch_stride + dim_idx * out_z_d_stride + ) + + # Determine chunk boundaries for APC mode + if CACHE_ENABLED: + last_chunk_idx = tl.load(last_chunk_indices_ptr + batch_idx).to(tl.int32) + if batch_idx == 0: + first_chunk_idx = 0 + else: + first_chunk_idx = ( + tl.load(last_chunk_indices_ptr + batch_idx - 1).to(tl.int32) + 1 + ) + n_chunks = last_chunk_idx - first_chunk_idx + 1 + first_chunk_tokens = tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + 1).to( + tl.int32 + ) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx).to(tl.int32) + block_idx_first = tl.load(block_idx_first_ptr + batch_idx).to(tl.int32) + chunk_start_offset = 0 + if n_chunks > 1 and first_chunk_tokens < block_size: + chunk_start_offset = block_size - first_chunk_tokens + current_position = block_idx_first * block_size + chunk_start_offset + else: + n_chunks = 1 + first_chunk_idx = 0 + + # Sequential scan over the sequence + tokens_processed = 0 + for chunk in range(0, n_chunks if CACHE_ENABLED else 1): + if CACHE_ENABLED: + chunk_tokens = tl.load( + cu_chunk_seqlen_ptr + first_chunk_idx + chunk + 1 + ).to(tl.int32) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + chunk).to( + tl.int32 + ) + else: + chunk_tokens = actual_seqlen + + for local_pos in range(chunk_tokens): + pos = tokens_processed + local_pos + # Load u value + u_val = tl.load(u_base + pos).to(tl.float32) + + # Load delta value + delta_val = tl.load(delta_base + pos).to(tl.float32) + + # Apply delta bias + if HAS_DELTA_BIAS: + delta_val = delta_val + delta_bias_val + + # Apply softplus + if delta_softplus: + delta_val = _softplus(delta_val) + + delta_u = delta_val * u_val + + # Compute dA = exp(delta * A) for all dstate elements + dA = tl.exp(delta_val * A_vals) + + # Load B values for this position + B_vals = tl.load( + B_base + dstate_offs * B_dstate_stride + pos, + mask=dstate_mask, + other=0.0, + ).to(tl.float32) + + # Load C values for this position + C_vals = tl.load( + C_base + dstate_offs * C_dstate_stride + pos, + mask=dstate_mask, + other=0.0, + ).to(tl.float32) + + # Update state: state = dA * state + delta * u * B + state = dA * state + delta_u * B_vals + + # Compute output: out = sum(state * C) + D * u + out_val = tl.sum(state * C_vals, axis=0) + if HAS_D: + out_val = out_val + D_val * u_val + + # Store output + tl.store(out_base + pos, out_val.to(out_ptr.dtype.element_ty)) + + if HAS_Z: + z_val = tl.load(z_base + pos).to(tl.float32) + out_z_val = out_val * z_val / (1.0 + tl.exp(-z_val)) + tl.store( + out_z_base + pos, + out_z_val.to(out_z_ptr.dtype.element_ty), + ) + + tokens_processed += chunk_tokens + + # Store intermediate state for APC mode + if CACHE_ENABLED: + if chunk == n_chunks - 1: + store_slot = tl.load( + cache_indices_ptr + + batch_idx * cache_indices_stride + + tl.load(block_idx_last_ptr + batch_idx).to(tl.int32) + ).to(tl.int64) + else: + block_idx_done = (current_position + chunk_tokens - 1) // block_size + store_slot = tl.load( + cache_indices_ptr + + batch_idx * cache_indices_stride + + block_idx_done + ).to(tl.int64) + + tl.store( + ssm_states_ptr + + store_slot * ssm_batch_stride + + dim_idx * ssm_dim_stride + + dstate_offs * ssm_dstate_stride, + state.to(ssm_states_ptr.dtype.element_ty), + mask=dstate_mask, + ) + current_position += chunk_tokens + + # Store final state for non-APC mode + if not CACHE_ENABLED: + tl.store( + ssm_states_ptr + + load_cache_slot * ssm_batch_stride + + dim_idx * ssm_dim_stride + + dstate_offs * ssm_dstate_stride, + state.to(ssm_states_ptr.dtype.element_ty), + mask=dstate_mask, + ) + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -540,6 +856,173 @@ class xpu_ops: ) return None + @staticmethod + def selective_scan_fwd( + u: torch.Tensor, + delta: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D_: torch.Tensor | None, + z_: torch.Tensor | None, + delta_bias_: torch.Tensor | None, + delta_softplus: bool, + query_start_loc: torch.Tensor | None, + cache_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + ssm_states: torch.Tensor, + null_block_id: int, + block_size: int = 1024, + block_idx_first_scheduled_token: torch.Tensor | None = None, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + cu_chunk_seqlen: torch.Tensor | None = None, + last_chunk_indices: torch.Tensor | None = None, + ) -> None: + varlen = query_start_loc is not None + batch_size = ( + (query_start_loc.shape[0] - 1) + if query_start_loc is not None + else u.shape[0] + ) + dim = u.shape[0] if varlen else u.shape[1] + total_seqlen = u.shape[1] if varlen else u.shape[2] + dstate = A.size(1) + n_groups = B.size(0) if varlen else B.size(1) + dim_ngroups_ratio = dim // n_groups + + has_z = z_ is not None + has_D = D_ is not None + has_delta_bias = delta_bias_ is not None + has_cache_indices = cache_indices is not None + cache_enabled = block_idx_first_scheduled_token is not None + + # out and out_z alias delta and z respectively + out = delta + out_z = z_ if z_ is not None else delta # won't be used if not has_z + + BLOCK_DSTATE = triton.next_power_of_2(dstate) + + # Compute strides + if varlen: + u_batch_stride = u.stride(1) + u_d_stride = u.stride(0) + delta_batch_stride = delta.stride(1) + delta_d_stride = delta.stride(0) + B_batch_stride = B.stride(2) + B_group_stride = B.stride(0) + B_dstate_stride = B.stride(1) + C_batch_stride = C.stride(2) + C_group_stride = C.stride(0) + C_dstate_stride = C.stride(1) + out_batch_stride = out.stride(1) + out_d_stride = out.stride(0) + if z_ is not None: + z_batch_stride = z_.stride(1) + z_d_stride = z_.stride(0) + out_z_batch_stride = out_z.stride(1) + out_z_d_stride = out_z.stride(0) + else: + z_batch_stride = 0 + z_d_stride = 0 + out_z_batch_stride = 0 + out_z_d_stride = 0 + else: + u_batch_stride = u.stride(0) + u_d_stride = u.stride(1) + delta_batch_stride = delta.stride(0) + delta_d_stride = delta.stride(1) + B_batch_stride = B.stride(0) + B_group_stride = B.stride(1) + B_dstate_stride = B.stride(2) + C_batch_stride = C.stride(0) + C_group_stride = C.stride(1) + C_dstate_stride = C.stride(2) + out_batch_stride = out.stride(0) + out_d_stride = out.stride(1) + if z_ is not None: + z_batch_stride = z_.stride(0) + z_d_stride = z_.stride(1) + out_z_batch_stride = out_z.stride(0) + out_z_d_stride = out_z.stride(1) + else: + z_batch_stride = 0 + z_d_stride = 0 + out_z_batch_stride = 0 + out_z_d_stride = 0 + + ssm_batch_stride = ssm_states.stride(0) + ssm_dim_stride = ssm_states.stride(1) + ssm_dstate_stride = ssm_states.stride(2) + cache_indices_stride = ( + cache_indices.stride(0) if cache_indices is not None else 0 + ) + + grid = (batch_size, dim) + _selective_scan_fwd_kernel[grid]( + u, + delta, + A, + B, + C, + D_ if has_D else u, # dummy, won't be dereferenced + z_ if has_z else u, # dummy + delta_bias_ if has_delta_bias else u, # dummy + out, + out_z, + ssm_states, + query_start_loc if varlen else u, # dummy + cache_indices if has_cache_indices else u, # dummy + has_initial_state, + # APC pointers + block_idx_first_scheduled_token if cache_enabled else u, + block_idx_last_scheduled_token if cache_enabled else u, + initial_state_idx if cache_enabled else u, + cu_chunk_seqlen if cache_enabled else u, + last_chunk_indices if cache_enabled else u, + # Dimensions + batch_size, + dim, + total_seqlen, + dstate, + n_groups, + dim_ngroups_ratio, + # Strides + u_batch_stride, + u_d_stride, + delta_batch_stride, + delta_d_stride, + A.stride(0), + A.stride(1), + B_batch_stride, + B_group_stride, + B_dstate_stride, + C_batch_stride, + C_group_stride, + C_dstate_stride, + z_batch_stride, + z_d_stride, + out_batch_stride, + out_d_stride, + out_z_batch_stride, + out_z_d_stride, + ssm_batch_stride, + ssm_dim_stride, + ssm_dstate_stride, + cache_indices_stride, + null_block_id, + block_size, + # Compile-time constants + delta_softplus=delta_softplus, + HAS_D=has_D, + HAS_Z=has_z, + HAS_DELTA_BIAS=has_delta_bias, + IS_VARLEN=varlen, + HAS_CACHE_INDICES=has_cache_indices, + CACHE_ENABLED=cache_enabled, + BLOCK_DSTATE=BLOCK_DSTATE, + ) + @staticmethod def register_ops_once() -> None: global _OPS_REGISTERED diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index af62f049ff3..59e2aa578c3 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -44,6 +44,7 @@ from vllm.lora.request import LoRARequest from vllm.lora.utils import get_adapter_absolute_path from vllm.multimodal.audio import get_audio_duration from vllm.multimodal.image import convert_image_mode +from vllm.multimodal.utils import encode_image_url, fetch_image from vllm.tokenizers import TokenizerLike from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -363,7 +364,11 @@ def lora_path_on_disk(lora_path: str) -> str: lora_tokenizer_cache: dict[int, TokenizerLike] = {} -def process_image(image: Any) -> Mapping[str, Any]: +def process_image( + image: Any, + *, + ensure_client_side_data: bool = False, +) -> Mapping[str, Any]: """ Process a single image input and return a multimedia content dictionary. @@ -380,6 +385,9 @@ def process_image(image: Any) -> Mapping[str, Any]: encoded data. - If string starts with "data:image/", treats as base64. - If string starts with "http://", "https://", or "file://", treats as URL. - Otherwise treats as local file path and prepends "file://". + - If ensure_client_side_data is True, local and HTTP(S) image references + are loaded and encoded as base64 image data URLs. Existing data:image + URLs are kept unchanged. - Returns a dictionary with the image URL or base64 data. Raises: @@ -403,6 +411,13 @@ def process_image(image: Any) -> Mapping[str, Any]: if image.startswith(("http://", "https://", "file://", "data:image/")) else f"file://{image}" ) + + if ensure_client_side_data and not image_url.startswith("data:image/"): + try: + fetched_image = fetch_image(image_url) + image_url = encode_image_url(fetched_image) + except Exception as e: + raise ValueError(f"Invalid image URL: {image_url}") from e return {"type": "image_url", "image_url": {"url": image_url}} raise ValueError( @@ -1645,6 +1660,16 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "value overrides potential output length loaded from the dataset. It is " "used only for custom dataset.", ) + custom_group.add_argument( + "--custom-ensure-client-side-data", + action="store_true", + help=( + "Ensure custom dataset media is sent as client-side data instead " + "of references. For custom_image datasets, this loads local and " + "HTTP(S) images on the benchmark client and encodes them as " + "base64 data URLs. Existing data:image URLs are kept unchanged." + ), + ) spec_bench_group = parser.add_argument_group("spec bench dataset options") spec_bench_group.add_argument( @@ -2055,6 +2080,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: tokenizer=tokenizer, output_len=args.custom_output_len, skip_chat_template=args.skip_chat_template, + chat_template_kwargs=getattr(args, "chat_template_kwargs", None), request_id_prefix=args.request_id_prefix, no_oversample=args.no_oversample, ) @@ -2075,6 +2101,9 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: tokenizer=tokenizer, output_len=args.custom_output_len, enable_multimodal_chat=args.enable_multimodal_chat, + ensure_client_side_data=getattr( + args, "custom_ensure_client_side_data", False + ), request_id_prefix=args.request_id_prefix, no_oversample=args.no_oversample, ) @@ -2381,6 +2410,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: num_requests=args.num_prompts, tokenizer=tokenizer, output_len=args.speed_bench_output_len, + chat_template_kwargs=getattr(args, "chat_template_kwargs", None), enable_multimodal_chat=args.enable_multimodal_chat, request_id_prefix=args.request_id_prefix, no_oversample=args.no_oversample, @@ -2468,6 +2498,7 @@ class CustomDataset(BenchmarkDataset): output_len: int | None = None, enable_multimodal_chat: bool = False, skip_chat_template: bool = False, + chat_template_kwargs: dict | None = None, **kwargs, ) -> list[SampleRequest]: # load all data if needed @@ -2515,6 +2546,7 @@ class CustomDataset(BenchmarkDataset): [{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False, + **(chat_template_kwargs or {}), ) prompt_len = len(tokenizer(prompt).input_ids) @@ -2627,7 +2659,12 @@ class CustomImageDataset(CustomDataset): return parts @classmethod - def _process_content_part(cls, part: dict[str, Any]) -> dict[str, Any]: + def _process_content_part( + cls, + part: dict[str, Any], + *, + ensure_client_side_data: bool = False, + ) -> dict[str, Any]: content_type = part.get("type") if content_type == "text": text = part.get("text") @@ -2638,12 +2675,22 @@ class CustomImageDataset(CustomDataset): if content_type == "image": if "image" not in part: raise ValueError("Image content parts must contain an 'image' field.") - return dict(process_image(part["image"])) + return dict( + process_image( + part["image"], + ensure_client_side_data=ensure_client_side_data, + ) + ) if content_type == "image_url": image_url = part.get("image_url") if isinstance(image_url, str): - return dict(process_image(image_url)) + return dict( + process_image( + image_url, + ensure_client_side_data=ensure_client_side_data, + ) + ) if isinstance(image_url, dict): url = image_url.get("url") @@ -2652,7 +2699,12 @@ class CustomImageDataset(CustomDataset): "Image URL content parts must contain a string 'image_url.url'." ) - processed_part = dict(process_image(url)) + processed_part = dict( + process_image( + url, + ensure_client_side_data=ensure_client_side_data, + ) + ) processed_image_url = dict(processed_part["image_url"]) processed_image_url.update( {key: value for key, value in image_url.items() if key != "url"} @@ -2671,9 +2723,17 @@ class CustomImageDataset(CustomDataset): ) @classmethod - def _process_interleaved_content(cls, content: Any) -> list[dict[str, Any]]: + def _process_interleaved_content( + cls, + content: Any, + *, + ensure_client_side_data: bool = False, + ) -> list[dict[str, Any]]: return [ - cls._process_content_part(part) + cls._process_content_part( + part, + ensure_client_side_data=ensure_client_side_data, + ) for part in cls._validate_content_parts(content) ] @@ -2682,11 +2742,23 @@ class CustomImageDataset(CustomDataset): return "".join(part["text"] for part in content if part.get("type") == "text") @staticmethod - def _process_image_files(images: Any) -> dict[str, Any] | list[dict[str, Any]]: + def _process_image_files( + images: Any, + *, + ensure_client_side_data: bool = False, + ) -> dict[str, Any] | list[dict[str, Any]]: if not isinstance(images, list) or not images: raise ValueError("'image_files' must be a non-empty list.") - mm_content = [dict(process_image(image)) for image in images] + mm_content = [ + dict( + process_image( + image, + ensure_client_side_data=ensure_client_side_data, + ) + ) + for image in images + ] if len(mm_content) == 1: return mm_content[0] @@ -2698,6 +2770,7 @@ class CustomImageDataset(CustomDataset): num_requests: int, output_len: int | None = None, enable_multimodal_chat: bool = False, + ensure_client_side_data: bool = False, request_id_prefix: str = "", no_oversample: bool = False, **kwargs, @@ -2718,9 +2791,14 @@ class CustomImageDataset(CustomDataset): break if "content" in item: - content = self._process_interleaved_content(item["content"]) + content = self._process_interleaved_content( + item["content"], + ensure_client_side_data=ensure_client_side_data, + ) text_prompt = self._get_text_from_content(content) - prompt_len = len(tokenizer(text_prompt).input_ids) + prompt_len = ( + 1 if tokenizer is None else len(tokenizer(text_prompt).input_ids) + ) prompt = ( [{"role": "user", "content": content}] if enable_multimodal_chat @@ -2741,8 +2819,11 @@ class CustomImageDataset(CustomDataset): if not isinstance(prompt, str): raise ValueError("'prompt' must be a string.") - prompt_len = len(tokenizer(prompt).input_ids) - mm_content = self._process_image_files(item["image_files"]) + prompt_len = 1 if tokenizer is None else len(tokenizer(prompt).input_ids) + mm_content = self._process_image_files( + item["image_files"], + ensure_client_side_data=ensure_client_side_data, + ) if enable_multimodal_chat: # Note: when chat is enabled the request prompt_len is no longer # accurate and we will be using request output to count the diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 2bef4b14d88..5ebc297d503 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -1609,6 +1609,15 @@ def add_cli_args(parser: argparse.ArgumentParser): "in seconds. Ready check will be skipped by default.", ) + parser.add_argument( + "--chat-template-kwargs", + type=json.loads, + default=None, + help="A JSON string of kwargs forwarded to the tokenizer's " + "apply_chat_template when a dataset renders prompts client-side " + "(e.g. custom / speed_bench). " + "Example: '{\"thinking\": true}' to enable reasoning models.", + ) parser.add_argument( "--extra-body", help="A JSON string representing extra body parameters to include " diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index e35fc5cd408..c58ce31bd29 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -70,7 +70,11 @@ class ActivationQuantPattern(VllmPatternReplacement): self.silu_and_mul_matcher = MatcherSiluAndMul() def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor: - kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs} + kwargs = { + "dtype": self.quant_dtype, + "device": current_platform.device_type, + **kwargs, + } return torch.empty(*args, **kwargs) diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 9f25a6805e9..94ae2bfcb14 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -36,14 +36,13 @@ QUANT_OPS: dict[QuantKey, OpOverload] = { kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501 kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501 + kFp8Dynamic128Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501 + kFp8Dynamic64Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501 } if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out # noqa: E501 -if current_platform.is_cuda(): - QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 - QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 SILU_MUL_OP = torch.ops._C.silu_and_mul.default diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index cc986595d43..c6a10078069 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, QuantKey, ScaleDesc, + get_fp8_min_max, kFp8Dynamic64Sym, kFp8Dynamic128Sym, kFp8DynamicTensorSym, @@ -54,19 +55,27 @@ def _rms_input_weight_dtype_match(match: pm.Match) -> bool: def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.bfloat16, device=current_platform.device_type + ) def empty_fp32(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.float32, device=current_platform.device_type + ) def empty_i32(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.int32, device=current_platform.device_type + ) def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor: - return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda") + return torch.empty( + *args, **kwargs, dtype=torch.int64, device=current_platform.device_type + ) RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default @@ -75,12 +84,11 @@ QUANT_OPS: dict[QuantKey, OpOverload] = { kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501 kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501 + kFp8Dynamic128Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501 + kFp8Dynamic64Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501 } if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out -if current_platform.is_cuda(): - QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 - QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 class FusedRMSQuantKey(NamedTuple): @@ -319,9 +327,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern): dtype=self.quant_matcher.quant_key.dtype, ) assert scale is not None - finfo = torch.finfo(self.quant_matcher.quant_key.dtype) - fp8_min = finfo.min - fp8_max = finfo.max + fp8_min, fp8_max = get_fp8_min_max() _, result, scale = auto_functionalized( self.quant_matcher.QUANT_OP, @@ -422,9 +428,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern): dtype=self.quant_matcher.quant_key.dtype, ) assert scale is not None - finfo = torch.finfo(self.quant_matcher.quant_key.dtype) - fp8_min = finfo.min - fp8_max = finfo.max + fp8_min, fp8_max = get_fp8_min_max() _, result, scale = auto_functionalized( self.quant_matcher.QUANT_OP, @@ -637,31 +641,30 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass): # Fuse rms_norm + dynamic per-token fp8 quant RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns) - # Only register group quant patterns on CUDA where the C++ op exists - if current_platform.is_cuda(): - for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]: - for has_col_major_scales in [True, False]: - for is_e8m0 in [True, False]: - for is_tma_aligned in [False, True]: - # Fuse fused_add_rms_norm + fp8 group quant - FusedAddRMSNormGroupQuantPattern( - epsilon, - FP8_DTYPE, - group_shape=group_shape, - is_e8m0=is_e8m0, - has_col_major_scales=has_col_major_scales, - is_tma_aligned=is_tma_aligned, - ).register(self.patterns) + # Only register group quant patterns on CUDA/ROCm where the C++ op exists + for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]: + for has_col_major_scales in [True, False]: + for is_e8m0 in [True, False]: + for is_tma_aligned in [False, True]: + # Fuse fused_add_rms_norm + fp8 group quant + FusedAddRMSNormGroupQuantPattern( + epsilon, + FP8_DTYPE, + group_shape=group_shape, + is_e8m0=is_e8m0, + has_col_major_scales=has_col_major_scales, + is_tma_aligned=is_tma_aligned, + ).register(self.patterns) - # Fuse rms_norm + fp8 group quant - RMSNormGroupQuantPattern( - epsilon, - FP8_DTYPE, - group_shape=group_shape, - is_e8m0=is_e8m0, - has_col_major_scales=has_col_major_scales, - is_tma_aligned=is_tma_aligned, - ).register(self.patterns) + # Fuse rms_norm + fp8 group quant + RMSNormGroupQuantPattern( + epsilon, + FP8_DTYPE, + group_shape=group_shape, + is_e8m0=is_e8m0, + has_col_major_scales=has_col_major_scales, + is_tma_aligned=is_tma_aligned, + ).register(self.patterns) self.dump_patterns(config, self.patterns) diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index e7ba3385725..03d291d4d94 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -570,9 +570,16 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): ) gated_norm_shapes: set[tuple[int, int]] = set() for layer in gdn_layers.values(): - gated_norm_shapes.add( - (layer.num_v_heads // layer.tp_size, layer.head_v_dim) + num_v_heads = getattr(layer, "num_v_heads", None) or getattr( + layer, "num_heads", None ) + head_v_dim = getattr(layer, "head_v_dim", None) or getattr( + layer, "head_dim", None + ) + + assert num_v_heads is not None and head_v_dim is not None + + gated_norm_shapes.add((num_v_heads // layer.tp_size, head_v_dim)) # Make sure fused add patterns are before simple rms norm, # as the latter is a subset of the former in torch ops. diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index fbf05f7753c..fef494ca54d 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -45,6 +45,10 @@ if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass from .fusion.collective_fusion import AsyncTPPass +if current_platform.is_xpu(): + from .fusion.act_quant_fusion import ActivationQuantFusionPass + from .fusion.rms_quant_fusion import RMSNormQuantFusionPass + from .inductor_pass import ( CustomGraphPass, InductorPass, diff --git a/vllm/config/model.py b/vllm/config/model.py index b41ab189c10..67040a423b7 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -80,16 +80,6 @@ else: logger = init_logger(__name__) - -def is_cumem_allocator_available() -> bool: - try: - from vllm.device_allocator.cumem import cumem_available - except ImportError: - return False - - return cumem_available - - RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] @@ -542,7 +532,10 @@ class ModelConfig: "Enabling cumem allocator because sleep mode requires it." ) self.enable_cumem_allocator = True - if self.enable_cumem_allocator and not is_cumem_allocator_available(): + if ( + self.enable_cumem_allocator + and not current_platform.is_cumem_allocator_available() + ): raise ValueError("cumem allocator is not supported on current platform.") hf_config = get_config( diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index fb6951ea7dd..7900c948480 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -53,7 +53,7 @@ class SchedulerConfig: In real usage, this should be set in `EngineArgs.create_engine_config`. """ - max_num_scheduled_tokens: int | None = None + max_num_scheduled_tokens: int | None = Field(default=None, ge=0) """Maximum number of tokens that the scheduler may issue in a single iteration. This is usually equal to max_num_batched_tokens, but can be smaller in cases diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f753647081c..4d80078a01f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,7 +66,13 @@ else: logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "LlamaForCausalLM", + "MistralForCausalLM", + "Qwen3ForCausalLM", + } +) class OptimizationLevel(IntEnum): @@ -487,6 +493,19 @@ class VllmConfig: ] return hash_str + @property + def max_concurrent_batches(self) -> int: + # PP requires PP-size concurrent batches to fill the pipeline. + # Async scheduling requires 2 concurrent batches to overlap. + pp_size = self.parallel_config.pipeline_parallel_size + if self.scheduler_config.async_scheduling: + if self.use_v2_model_runner: + return pp_size + 1 + # V1 Model Runner does not fully support async scheduling with PP. + if pp_size <= 1: + return 2 + return pp_size + @property def num_speculative_tokens(self) -> int: if ( diff --git a/vllm/distributed/device_communicators/cpu_communicator.py b/vllm/distributed/device_communicators/cpu_communicator.py index 067cdad7348..b8d9d6c53d5 100644 --- a/vllm/distributed/device_communicators/cpu_communicator.py +++ b/vllm/distributed/device_communicators/cpu_communicator.py @@ -32,6 +32,7 @@ class CpuCommunicator(DeviceCommunicatorBase): ( current_platform.get_cpu_architecture() == CpuArchEnum.X86 or current_platform.get_cpu_architecture() == CpuArchEnum.ARM + or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC ) and hasattr(torch.ops._C, "init_shm_manager") and (unique_name.startswith("tp") or unique_name.startswith("pp")) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index dc7e6d151a4..9482568461c 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -38,9 +38,19 @@ from vllm.utils.network_utils import ( is_valid_ipv6_address, ) -if envs.VLLM_USE_SPINLOOP_EXT: - from vllm.spinloop import spinloop +logger = init_logger(__name__) + +SPINLOOP_EXT_ENABLED = False +if envs.VLLM_USE_SPINLOOP_EXT: + try: + from vllm.spinloop import spinloop + + SPINLOOP_EXT_ENABLED = True + except ImportError: + logger.warning( + "spinloop extension could not be loaded, disabling VLLM_USE_SPINLOOP_EXT!" + ) SPINLOOP_TIMEOUT_SECONDS = 0.1 if TYPE_CHECKING: @@ -82,9 +92,6 @@ def to_bytes_big(value: int, size: int) -> bytes: return value.to_bytes(size, byteorder="big") -logger = init_logger(__name__) - - LONG_WAIT_TIME_LOG_MSG = ( "No available shared memory broadcast block found " "in %d seconds. This typically happens " @@ -552,7 +559,7 @@ class MessageQueue: written_flag = metadata_buffer[0] return not (written_flag and read_count != self.buffer.n_reader) - if envs.VLLM_USE_SPINLOOP_EXT and not check(): + if SPINLOOP_EXT_ENABLED and not check(): spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS) if not check(): @@ -673,7 +680,7 @@ class MessageQueue: written_flag = metadata_buffer[0] return not (not written_flag or read_flag) - if envs.VLLM_USE_SPINLOOP_EXT and not check(): + if SPINLOOP_EXT_ENABLED and not check(): spinloop( metadata_buffer[0 : self.local_reader_rank + 1], check, diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 28370c8e1fa..1d5f467027e 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -211,6 +211,23 @@ class ECConnectorBase(ABC): """ pass + def ensure_cache_available( + self, request: "Request", num_computed_tokens: int + ) -> bool: + """ + Ensure encoder cache items are available for the given request. + May initiate asynchronous transfers for items not yet local. + + Args: + request: the request whose multimodal features to check. + num_computed_tokens: tokens already covered by cached KV blocks. + + Returns: + True if all items are ready or no transfer is needed. + False if any items are still in transit (request should be deferred). + """ + return True + @abstractmethod def update_state_after_alloc(self, request: "Request", index: int): """ diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 92fffd22977..f10891d6cdf 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -61,25 +61,31 @@ class CpuGpuEvent: self._recorded.set() -def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: +def override_envs_for_eplb( + parallel_config: ParallelConfig, + moe_backend: str | None = None, +) -> None: """ Override environment variables for EPLB when specific conditions are met. Args: parallel_config: The parallel configuration object. + moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``). """ is_data_parallel = parallel_config.data_parallel_size > 1 is_eplb_enabled = parallel_config.enable_eplb async_eplb = parallel_config.eplb_config.use_async is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency" + is_mega_moe = moe_backend == "deep_gemm_mega_moe" is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in ( "torch_nccl", "pynccl", ) - # Override NCCL_MAX_CTAS to avoid hangs when using async EPLB with the - # DeepEP low-latency backend. + # Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange + # contends with MoE backend's cooperative-launch on GPU SMs. # + # DeepEP low-latency: # The hang happens when two ranks interleave kernel launches differently # between NCCL collectives (used by async EPLB weight exchange) and DeepEP # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries @@ -94,12 +100,14 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP # cooperative kernel to launch and complete, breaking the deadlock. # See: https://github.com/deepseek-ai/DeepEP/issues/496 + # + # DeepGEMM Mega MoE also uses cooperative launch and will cause hang even + # with sync EPLB. if ( is_data_parallel and is_eplb_enabled - and is_deepep_ll - and async_eplb and is_nccl_based_eplb_communicator + and ((is_deepep_ll and async_eplb) or is_mega_moe) ): current_value_str = os.getenv("NCCL_MAX_CTAS") @@ -108,9 +116,10 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: override_value = 8 os.environ["NCCL_MAX_CTAS"] = str(override_value) + backend = "deepep_low_latency" if is_deepep_ll else "deep_gemm_mega_moe" logger.info_once( f"EPLB: Setting NCCL_MAX_CTAS={override_value} " - "for expert parallel with NCCL-based EPLB communicator and " - "deepep_low_latency backend", + f"for expert parallel with NCCL-based EPLB communicator and " + f"cooperative MoE backend ({backend})", scope="global", ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index b16fdb7c16c..ad528140966 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -13,7 +13,6 @@ from vllm.v1.core.kv_cache_utils import ( ) from vllm.v1.core.single_type_kv_cache_manager import ( SingleTypeKVCacheManager, - spec_manager_map, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -21,6 +20,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry # Dummy placeholder hash for store_mask's template computation. _DUMMY_BLOCK_HASH = BlockHash(b"\x00" * 32) @@ -89,7 +89,10 @@ class MooncakeStoreCoordinator: ] = [] for i, g in enumerate(self.kv_cache_groups): spec = _unwrap_spec(g.kv_cache_spec) - manager_cls = spec_manager_map[type(spec)] + manager_cls = KVCacheSpecRegistry.get_manager_class(spec) + assert manager_cls is not None, ( + f"No manager registered for KVCacheSpec {spec}" + ) for existing_spec, group_ids, existing_cls in attention_groups: if existing_spec == spec: assert manager_cls is existing_cls @@ -233,7 +236,7 @@ class MooncakeStoreCoordinator: kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, - use_eagle=(0 in eagle_indices), + drop_eagle_block=(0 in eagle_indices), alignment_tokens=spec.block_size, ) num_groups = len(self.kv_cache_groups) @@ -262,9 +265,9 @@ class MooncakeStoreCoordinator: ) continue - use_eagle = idx in eagle_indices and idx not in eagle_verified + drop_eagle_block = idx in eagle_indices and idx not in eagle_verified _max_length = curr_hit_length - if use_eagle: + if drop_eagle_block: _max_length = min(curr_hit_length + spec.block_size, max_length) hashes = self.block_hashes_for_spec(block_hashes, spec) hit_blocks = manager_cls.find_longest_cache_hit( @@ -273,11 +276,11 @@ class MooncakeStoreCoordinator: kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, - use_eagle=use_eagle, + drop_eagle_block=drop_eagle_block, alignment_tokens=self.lcm_block_size, ) _new_hit_length = len(hit_blocks[0]) * spec.block_size - if use_eagle: + if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: eagle_verified.clear() 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 486c2553b6d..cd4eb5c3713 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 @@ -517,187 +517,190 @@ class KVCacheStoreSendingThread(KVTransferThread): if req_id not in self.stored_requests: self.request_queue.task_done() return - if token_len == 0: - self.dec_stored_request(req_id) - self.request_queue.task_done() - return - if self._should_skip_request(req_id): - logger.debug( - "Skipping Mooncake store for request %s while CPU/disk offloading " - "is under pressure", - req_id, - ) - self.dec_stored_request(req_id) - self.request_queue.task_done() - return - # Within each lcm region only per-spec relevant chunks are loaded - # (e.g., SWA or linear attn), so mask out irrelevant chunks - store_masks = self.coord.store_mask(token_len) - starts: list[int] = [] - ends: list[int] = [] - keys: list[str] = [] - block_hashes: list[BlockHash] = [] - group_indices: list[int] = [] - for g_idx, db in enumerate(self.token_databases): - mask = store_masks[g_idx] - for chunk_idx, (start, end, key) in enumerate( - db.process_tokens(token_len, req_meta.block_hashes) - ): - if chunk_idx >= len(mask) or not mask[chunk_idx]: - continue - starts.append(start) - ends.append(end) - keys.append(key.to_string()) - block_hashes.append(req_meta.block_hashes[chunk_idx]) - group_indices.append(g_idx) - - # Apply put_step striding for TP - sl = slice(self.tp_rank % self.put_step, None, self.put_step) - starts = starts[sl] - ends = ends[sl] - keys = keys[sl] - block_hashes = block_hashes[sl] - group_indices = group_indices[sl] - - if not keys: - self.dec_stored_request(req_id) - return - - # Check which blocks already exist (dedup) - save_exists_start = time.perf_counter() + # Decrement the in-flight counter and signal task_done() in `finally` + # so the scheduler can release the GPU blocks it pinned for this + # request (via `delay_free_blocks`) even when the store path raises. try: - exists_states = self.store.batch_is_exist(keys) - except Exception: + if token_len == 0: + return + if self._should_skip_request(req_id): + logger.debug( + "Skipping Mooncake store for request %s while CPU/disk " + "offloading is under pressure", + req_id, + ) + return + + # Within each lcm region only per-spec relevant chunks are loaded + # (e.g., SWA or linear attn), so mask out irrelevant chunks + store_masks = self.coord.store_mask(token_len) + starts: list[int] = [] + ends: list[int] = [] + keys: list[str] = [] + block_hashes: list[BlockHash] = [] + group_indices: list[int] = [] + for g_idx, db in enumerate(self.token_databases): + mask = store_masks[g_idx] + for chunk_idx, (start, end, key) in enumerate( + db.process_tokens(token_len, req_meta.block_hashes) + ): + if chunk_idx >= len(mask) or not mask[chunk_idx]: + continue + starts.append(start) + ends.append(end) + keys.append(key.to_string()) + block_hashes.append(req_meta.block_hashes[chunk_idx]) + group_indices.append(g_idx) + + # Apply put_step striding for TP + sl = slice(self.tp_rank % self.put_step, None, self.put_step) + starts = starts[sl] + ends = ends[sl] + keys = keys[sl] + block_hashes = block_hashes[sl] + group_indices = group_indices[sl] + + if not keys: + return + + # Check which blocks already exist (dedup) + save_exists_start = time.perf_counter() + try: + exists_states = self.store.batch_is_exist(keys) + except Exception: + self._record_operation( + "save_exists", + save_exists_start, + len(keys), + status="error", + num_failed_keys=len(keys), + ) + raise self._record_operation( "save_exists", save_exists_start, len(keys), - status="error", - num_failed_keys=len(keys), ) - raise - self._record_operation( - "save_exists", - save_exists_start, - len(keys), - ) - missing_indices = [i for i, exists in enumerate(exists_states) if exists != 1] + missing_indices = [ + i for i, exists in enumerate(exists_states) if exists != 1 + ] - if not missing_indices: - self.dec_stored_request(req_id) - return + if not missing_indices: + return - starts = [starts[i] for i in missing_indices] - ends = [ends[i] for i in missing_indices] - keys = [keys[i] for i in missing_indices] - block_hashes = [block_hashes[i] for i in missing_indices] - group_indices = [group_indices[i] for i in missing_indices] + starts = [starts[i] for i in missing_indices] + ends = [ends[i] for i in missing_indices] + keys = [keys[i] for i in missing_indices] + block_hashes = [block_hashes[i] for i in missing_indices] + group_indices = [group_indices[i] for i in missing_indices] - logger.debug( - "Storing KV cache for %d blocks (groups=%s) for request %s", - len(keys), - set(group_indices), - req_id, - ) - - addrs: list[list[int]] = [] - sizes: list[list[int]] = [] - stored_events: list[BlockStored] = [] - # parent_block_hash chains live within a group, not across. - prev_key_per_group: dict[int, Any] = {} - new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes] - - for idx, (s, e, g_idx) in enumerate( - zip(starts, ends, group_indices, strict=True) - ): - db = self.token_databases[g_idx] - addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx]) - addrs.append(addr) - sizes.append(size) - - if self.enable_kv_event: - token_ids = ( - req_meta.token_ids[s:e] if req_meta.token_ids is not None else None - ) - stored_event = BlockStored( - block_hashes=[new_block_hashes[idx]], - parent_block_hash=prev_key_per_group.get(g_idx), - token_ids=token_ids, - block_size=req_meta.original_block_size, - lora_id=None, - medium="cpu", - lora_name=None, - ) - stored_events.append(stored_event) - prev_key_per_group[g_idx] = new_block_hashes[idx] - - if current_event is not None: - current_event.synchronize() - - batch_bytes = _sum_batch_bytes(sizes) - put_start = time.perf_counter() - try: - res = self.store.batch_put_from_multi_buffers( - keys, - addrs, - sizes, - self.replicate_config, - ) - failed = [i for i, v in enumerate(res) if v < 0] - self._record_operation( - "save_put", - put_start, + logger.debug( + "Storing KV cache for %d blocks (groups=%s) for request %s", len(keys), - num_bytes=batch_bytes, - status="partial_failure" if failed else "ok", - num_failed_keys=len(failed), + set(group_indices), + req_id, ) - if failed: - failed_codes = set(res[i] for i in failed) - logger.warning( - "batch_put failed: %d/%d keys failed " - "(codes=%s, batch_bytes=%d, num_keys=%d), " - "first_key=%s", - len(failed), - len(keys), - failed_codes, - batch_bytes, - len(keys), - keys[0] if keys else "N/A", - ) - if ( - MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes - and not self._mark_request_skipped_for_pressure(req_id) - ): - logger.warning( - "Detected Mooncake CPU/disk offloading pressure " - "(NO_AVAILABLE_HANDLE); skipping future store " - "batches for request %s until a later store " - "batch succeeds", - req_id, + + addrs: list[list[int]] = [] + sizes: list[list[int]] = [] + stored_events: list[BlockStored] = [] + # parent_block_hash chains live within a group, not across. + prev_key_per_group: dict[int, Any] = {} + new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes] + + for idx, (s, e, g_idx) in enumerate( + zip(starts, ends, group_indices, strict=True) + ): + db = self.token_databases[g_idx] + addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx]) + addrs.append(addr) + sizes.append(size) + + if self.enable_kv_event: + token_ids = ( + req_meta.token_ids[s:e] + if req_meta.token_ids is not None + else None ) - elif self._clear_store_pressure(): - logger.info( - "Mooncake CPU/disk offloading pressure cleared after a " - "successful store batch" + stored_event = BlockStored( + block_hashes=[new_block_hashes[idx]], + parent_block_hash=prev_key_per_group.get(g_idx), + token_ids=token_ids, + block_size=req_meta.original_block_size, + lora_id=None, + medium="cpu", + lora_name=None, + ) + stored_events.append(stored_event) + prev_key_per_group[g_idx] = new_block_hashes[idx] + + if current_event is not None: + current_event.synchronize() + + batch_bytes = _sum_batch_bytes(sizes) + put_start = time.perf_counter() + try: + res = self.store.batch_put_from_multi_buffers( + keys, + addrs, + sizes, + self.replicate_config, ) - except Exception as e: - self._record_operation( - "save_put", - put_start, - len(keys), - num_bytes=batch_bytes, - status="error", - num_failed_keys=len(keys), - ) - logger.error("Failed to put key %s, error: %s", keys, e) + failed = [i for i, v in enumerate(res) if v < 0] + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="partial_failure" if failed else "ok", + num_failed_keys=len(failed), + ) + if failed: + failed_codes = set(res[i] for i in failed) + logger.warning( + "batch_put failed: %d/%d keys failed " + "(codes=%s, batch_bytes=%d, num_keys=%d), " + "first_key=%s", + len(failed), + len(keys), + failed_codes, + batch_bytes, + len(keys), + keys[0] if keys else "N/A", + ) + if ( + MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes + and not self._mark_request_skipped_for_pressure(req_id) + ): + logger.warning( + "Detected Mooncake CPU/disk offloading pressure " + "(NO_AVAILABLE_HANDLE); skipping future store " + "batches for request %s until a later store " + "batch succeeds", + req_id, + ) + elif self._clear_store_pressure(): + logger.info( + "Mooncake CPU/disk offloading pressure cleared after a " + "successful store batch" + ) + except Exception as e: + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="error", + num_failed_keys=len(keys), + ) + logger.error("Failed to put key %s, error: %s", keys, e) - if self.enable_kv_event and stored_events: - self.update_kv_event(stored_events) - - self.dec_stored_request(req_id) - self.request_queue.task_done() + if self.enable_kv_event and stored_events: + self.update_kv_event(stored_events) + finally: + self.dec_stored_request(req_id) + self.request_queue.task_done() class KVCacheStoreRecvingThread(KVTransferThread): 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 6560bcf9def..6ee827fa17e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -286,6 +286,8 @@ class OffloadingConnectorScheduler: self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} self._current_batch_jobs_to_flush: set[int] = set() + # GPU block IDs allocated in the current engine step + self._current_batch_allocated_block_ids: set[int] = set() # if GPU prefix caching is enabled, # track loaded blocks to avoid redundant loads self._blocks_being_loaded: set[OffloadKey] | None = ( @@ -589,6 +591,10 @@ class OffloadingConnectorScheduler: req_status.group_states, blocks.blocks, ): + self._current_batch_allocated_block_ids.update( + block.block_id for block in group_blocks if block.block_id != 0 + ) + gpu_block_size = group_config.gpu_block_size offloaded_block_size = group_config.offloaded_block_size offload_keys = group_state.offload_keys @@ -638,17 +644,6 @@ class OffloadingConnectorScheduler: if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL: group_state.next_stored_block_idx = num_blocks - # Fence dst blocks against finished-request pending stores. - if ( - self._block_id_to_pending_jobs - and not self._block_id_to_pending_jobs.keys().isdisjoint(dst_block_ids) - ): - self._current_batch_jobs_to_flush.update( - jid - for bid in dst_block_ids - for jid in self._block_id_to_pending_jobs.get(bid, ()) - ) - src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) dst_spec = GPULoadStoreSpec( dst_block_ids, group_sizes=group_sizes, block_indices=block_indices @@ -672,37 +667,68 @@ class OffloadingConnectorScheduler: if self._blocks_being_loaded is not None: self._blocks_being_loaded.update(keys_to_load) - def _build_store_jobs( - self, - scheduler_output: SchedulerOutput, - ) -> dict[int, TransferJob]: - block_size_factor = self.config.block_size_factor - store_jobs: dict[int, TransferJob] = {} - # iterate over both new and cached requests + def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: + """ + Update request states from the Scheduler's output. + """ + + # new_block_ids_end[req_id][i] = end of pre-existing block_ids for + # the i-th sliding window group (before this step's extend). + # Used to detect sliding window blocks that got re-allocated. + new_block_ids_end: dict[str, tuple[int, ...]] = {} + for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output): req_status = self._req_status[req_id] req_status.update_offload_keys() - req = req_status.req if preempted: for group_state in req_status.group_states: group_state.block_ids.clear() if new_block_id_groups: + if self._sliding_window_groups: + new_block_ids_end[req_id] = tuple( + len(req_status.group_states[grp_idx].block_ids) + for grp_idx in self._sliding_window_groups + ) req_status.update_block_id_groups(new_block_id_groups) - # Fence new blocks against in-flight stores. - if self._block_id_to_pending_jobs: - new_blocks_flat = [ - bid for new_blocks in new_block_id_groups for bid in new_blocks - ] - if not self._block_id_to_pending_jobs.keys().isdisjoint( - new_blocks_flat - ): - self._current_batch_jobs_to_flush.update( - jid - for bid in new_blocks_flat - for jid in self._block_id_to_pending_jobs.get(bid, ()) - ) + for new_blocks in new_block_id_groups: + for bid in new_blocks: + if bid != 0: + self._current_batch_allocated_block_ids.add(bid) + + # Zero out stale block_ids in sliding window groups' pending-store + # positions. Only sliding window groups can have stale entries (blocks + # freed by remove_skipped_blocks then reallocated). Only positions in + # [next_stored_block_idx * bsf, end) need checking where end is the + # pre-extend length: earlier positions were already offloaded, later + # ones are fresh allocations from this step. + if self._sliding_window_groups and self._current_batch_allocated_block_ids: + block_size_factor = self.config.block_size_factor + for req_id, req_status in self._req_status.items(): + ends = new_block_ids_end.get(req_id) + for i, grp_idx in enumerate(self._sliding_window_groups): + group_state = req_status.group_states[grp_idx] + start = group_state.next_stored_block_idx * block_size_factor + end = ends[i] if ends is not None else len(group_state.block_ids) + for j in range(start, end): + if ( + group_state.block_ids[j] + in self._current_batch_allocated_block_ids + ): + group_state.block_ids[j] = 0 + + def _build_store_jobs( + self, + scheduler_output: SchedulerOutput, + ) -> dict[int, TransferJob]: + block_size_factor = self.config.block_size_factor + store_jobs: dict[int, TransferJob] = {} + for req_id in scheduler_output.num_scheduled_tokens: + req_status = self._req_status.get(req_id) + if req_status is None: + continue + req = req_status.req num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens @@ -735,11 +761,8 @@ class OffloadingConnectorScheduler: # For each block to offload, take the last corresponding GPU block. # e.g. if block size factor is 3 and GPU block IDs are # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8. - # We will use these GPU blocks to determine if the block needs - # offloading, or (if the GPU block ID is 0) this block should - # be skipped due to sliding window attention / SSM. - # We know that if a block is skipped, then all the previous blocks - # are skipped as well. This is why we take the last of each block. + # A block_id of 0 means either a sliding window / SSM skip + # or a stale entry that was zeroed out — skip it either way. offload_block_ids = group_state.block_ids[ start_block_idx * block_size_factor + block_size_factor @@ -814,10 +837,8 @@ class OffloadingConnectorScheduler: for i in range(block_size_factor): block_id = block_ids[gpu_block_idx + i] if block_id == 0: - # skipped blocks cannot appear after non-skipped blocks - assert start_gpu_block_idx is None continue - elif start_gpu_block_idx is None: + if start_gpu_block_idx is None: start_gpu_block_idx = gpu_block_idx + i src_block_ids.append(block_id) num_group_blocks += 1 @@ -875,6 +896,10 @@ class OffloadingConnectorScheduler: def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> KVConnectorMetadata: + self._update_req_states(scheduler_output) + self.manager.on_schedule_end() + + # Flush jobs for preempted requests. for req_id in scheduler_output.preempted_req_ids or (): req_status = self._req_status.get(req_id) if req_status is None or not req_status.transfer_jobs: @@ -883,6 +908,20 @@ class OffloadingConnectorScheduler: assert self._jobs[any_jid].is_store self._current_batch_jobs_to_flush.update(req_status.transfer_jobs) + # Flush jobs that contain re-allocated blocks. + if ( + self._block_id_to_pending_jobs + and not self._block_id_to_pending_jobs.keys().isdisjoint( + self._current_batch_allocated_block_ids + ) + ): + self._current_batch_jobs_to_flush.update( + jid + for bid in self._current_batch_allocated_block_ids + if bid in self._block_id_to_pending_jobs + for jid in self._block_id_to_pending_jobs[bid] + ) + # If all tracked requests are finished, flush all pending jobs # (both store and load) - there might not be a future scheduler # step to trigger their completion. @@ -898,6 +937,7 @@ class OffloadingConnectorScheduler: ) self._current_batch_load_jobs = {} self._current_batch_jobs_to_flush = set() + self._current_batch_allocated_block_ids = set() return meta def update_connector_output(self, connector_output: KVConnectorOutput): @@ -1013,6 +1053,7 @@ class OffloadingConnectorScheduler: # reset_cache cannot be called in the middle of a schedule step assert not self._current_batch_load_jobs assert not self._current_batch_jobs_to_flush + assert not self._current_batch_allocated_block_ids # Flush all in-flight jobs self._current_batch_jobs_to_flush.update(self._jobs.keys()) diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 712167c601c..331e0684e32 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -359,6 +359,9 @@ class GroupCoordinator: assert self_cpu_group is not None assert self_device_group is not None + self.group_ranks = group_ranks + self.torch_distributed_backend = torch_distributed_backend + self.cpu_group = self_cpu_group self.device_group = self_device_group @@ -406,6 +409,22 @@ class GroupCoordinator: and getattr(self.device_communicator, "supports_tensor_dict", False) ) + def make_sibling_device_group(self, group_desc: str | None = None) -> ProcessGroup: + """Create a new device-side ProcessGroup with the same per-rank membership + as this coordinator's `device_group`, but backed by a distinct communicator. + 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`. + """ + 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 + ) + if self.rank in ranks: + sibling = pg + assert sibling is not None + return sibling + def create_mq_broadcaster( self, writer_rank=0, external_writer_handle=None, blocking=True ): diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 3ebc171173e..279f3625345 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -65,7 +65,7 @@ class AnthropicContentBlock(BaseModel): class AnthropicMessage(BaseModel): """Message structure""" - role: Literal["user", "assistant"] + role: Literal["user", "assistant", "system"] content: str | list[AnthropicContentBlock] diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 915cee59f98..2bdec6f4ec3 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -143,23 +143,36 @@ class AnthropicServingMessages(OpenAIServingChat): openai_messages: list[dict[str, Any]], ) -> None: """Convert Anthropic system message to OpenAI format""" - if not anthropic_request.system: - return + system_parts: list[str] = [] - if isinstance(anthropic_request.system, str): - openai_messages.append( - {"role": "system", "content": anthropic_request.system} - ) - else: - system_prompt = "" - for block in anthropic_request.system: - if block.type == "text" and block.text: - # Strip Claude Code's attribution header which contains - # a per-request hash that defeats prefix caching. - if block.text.startswith("x-anthropic-billing-header"): - continue - system_prompt += block.text - openai_messages.append({"role": "system", "content": system_prompt}) + # Top-level system field + if anthropic_request.system: + if isinstance(anthropic_request.system, str): + system_parts.append(anthropic_request.system) + else: + for block in anthropic_request.system: + if block.type == "text" and block.text: + # Strip Claude Code's attribution header which contains + # a per-request hash that defeats prefix caching. + if block.text.startswith("x-anthropic-billing-header"): + continue + system_parts.append(block.text) + + # System messages embedded inside the messages array + for msg in anthropic_request.messages: + if msg.role != "system": + continue + if isinstance(msg.content, str): + system_parts.append(msg.content) + else: + for block in msg.content: + if block.type == "text" and block.text: + if block.text.startswith("x-anthropic-billing-header"): + continue + system_parts.append(block.text) + + if system_parts: + openai_messages.append({"role": "system", "content": "".join(system_parts)}) @classmethod def _convert_messages( @@ -167,6 +180,9 @@ class AnthropicServingMessages(OpenAIServingChat): ) -> None: """Convert Anthropic messages to OpenAI format""" for msg in messages: + if msg.role == "system": + continue + openai_msg: dict[str, Any] = {"role": msg.role} # type: ignore if isinstance(msg.content, str): diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index 35256bc647d..52fc881aff8 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -58,6 +58,7 @@ from vllm.renderers.embed_utils import ( safe_load_prompt_embeds, safe_load_prompt_embeds_async, ) +from vllm.transformers_utils.processor import get_video_processor_cls_name from vllm.utils import random_uuid from vllm.utils.collection_utils import is_list_of from vllm.utils.import_utils import LazyLoader @@ -577,6 +578,10 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]): def mm_processor(self): return self.mm_registry.create_processor(self.model_config) + @property + def video_processor_name(self) -> str | None: + return get_video_processor_cls_name(self.model_config) + def add(self, modality: ModalityStr, item: _T) -> str | None: """ Add a multi-modal item to the current prompt and returns the @@ -1025,7 +1030,14 @@ class MultiModalContentParser(BaseMultiModalContentParser): return self.parse_audio(audio_url, uuid) def parse_video(self, video_url: str | None, uuid: str | None = None) -> None: - video = self._connector.fetch_video(video_url=video_url) if video_url else None + video = ( + self._connector.fetch_video( + video_url=video_url, + video_processor=self._tracker.video_processor_name, + ) + if video_url + else None + ) placeholder = self._tracker.add("video", (video, uuid)) self._add_placeholder("video", placeholder) @@ -1205,7 +1217,12 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser): async def _video_with_uuid_async(self, video_url: str | None, uuid: str | None): video = ( - await self._connector.fetch_video_async(video_url) if video_url else None + await self._connector.fetch_video_async( + video_url, + video_processor=self._tracker.video_processor_name, + ) + if video_url + else None ) return video, uuid @@ -1837,7 +1854,8 @@ def _postprocess_messages(messages: list[ConversationMessage]) -> None: # if arguments is None or empty string, set to {} if content := function.get("arguments"): if not isinstance(content, (dict, list)): - function["arguments"] = json.loads(content) + parsed = json.loads(content) + function["arguments"] = parsed if parsed is not None else {} else: function["arguments"] = {} diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 5455f1ca427..892f9d82d70 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -43,9 +43,7 @@ from vllm.entrypoints.openai.server_utils import ( validation_exception_handler, ) from vllm.entrypoints.sagemaker.api_router import sagemaker_standards_bootstrap -from vllm.entrypoints.serve.elastic_ep.middleware import ( - ScalingMiddleware, -) +from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization from vllm.entrypoints.utils import ( @@ -195,6 +193,11 @@ def build_app( register_sagemaker_api_router(app, supported_tasks, model_config) + if envs.VLLM_SERVER_DEV_MODE: + from vllm.entrypoints.serve import register_vllm_dev_api_routers + + register_vllm_dev_api_routers(app) + if "generate" in supported_tasks: from vllm.entrypoints.generate.api_router import ( register_generate_api_routers, @@ -208,12 +211,6 @@ def build_app( attach_disagg_router(app) - from vllm.entrypoints.serve.rlhf.api_router import ( - attach_router as attach_rlhf_router, - ) - - attach_rlhf_router(app) - from vllm.entrypoints.serve.elastic_ep.api_router import ( attach_router as elastic_ep_attach_router, ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 0be220fff77..184ace56805 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -30,6 +30,8 @@ from vllm.entrypoints.openai.engine.protocol import ( StructuralTagResponseFormat, ToolCall, UsageInfo, + validate_structural_tag_response_format, + validate_structured_outputs_structural_tag, ) from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger @@ -671,6 +673,9 @@ class ChatCompletionRequest(OpenAIBaseModel): parameter="response_format", ) + if rf_type == "structural_tag": + validate_structural_tag_response_format(response_format) + return data @model_validator(mode="before") @@ -754,6 +759,7 @@ class ChatCompletionRequest(OpenAIBaseModel): "You can only either use constraints for structured outputs " "or tools, not both.", ) + validate_structured_outputs_structural_tag(structured_outputs_kwargs) return data @model_validator(mode="before") @@ -979,6 +985,16 @@ class BatchChatCompletionRequest(OpenAIBaseModel): "Batch chat completions do not support beam search. " "Please set `use_beam_search` to False." ) + response_format = data.get("response_format") + rf_type = ( + response_format.get("type") + if isinstance(response_format, dict) + else getattr(response_format, "type", None) + ) + if rf_type == "structural_tag": + validate_structural_tag_response_format(response_format) + if (structured_outputs := data.get("structured_outputs")) is not None: + validate_structured_outputs_structural_tag(structured_outputs) n = data.get("n", 1) if n is not None and n != 1: raise ValueError( diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 92ffc141548..a378fb79d3b 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -3,7 +3,6 @@ import asyncio import io -import json import time from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence @@ -40,9 +39,7 @@ from vllm.entrypoints.openai.chat_completion.stream_harmony import ( extract_harmony_streaming_delta, ) from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, DeltaMessage, - DeltaToolCall, ErrorResponse, FunctionCall, PromptTokenUsageInfo, @@ -65,7 +62,7 @@ from vllm.entrypoints.utils import get_max_tokens, should_include_usage from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.logprobs import Logprob -from vllm.outputs import CompletionOutput, RequestOutput +from vllm.outputs import RequestOutput from vllm.parser import ParserManager from vllm.parser.abstract_parser import Parser from vllm.reasoning import ReasoningParser @@ -360,6 +357,14 @@ class OpenAIServingChat(OpenAIServing): assert len(generators) == 1 (result_generator,) = generators + parser: Parser | None = None + if self.parser_cls is not None: + parser = self.parser_cls( + tokenizer, + request.tools, + chat_template_kwargs=chat_template_kwargs, + ) + if request.stream: return self.chat_completion_stream_generator( request, @@ -381,7 +386,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, - reasoning_parser, + parser, ) def get_chat_request_role(self, request: ChatCompletionRequest) -> str: @@ -715,6 +720,7 @@ class OpenAIServingChat(OpenAIServing): delta_token_ids=as_list(output.token_ids), request=request, prompt_token_ids=res.prompt_token_ids, + finished=output.finish_reason is not None, ) if delta_message and delta_message.tool_calls: tools_streamed[i] = True @@ -805,81 +811,13 @@ class OpenAIServingChat(OpenAIServing): # finish_reason='error' indicates a retryable error self._raise_if_error(output.finish_reason, request_id) - # check to make sure we haven't "forgotten" to stream - # any tokens that were generated but previously - # matched by partial json parsing - # only happens if we are NOT using structured outputs - index = 0 - auto_tools_called = False - if tool_parser: - auto_tools_called = len(tool_parser.prev_tool_call_arr) > 0 - index = ( - len(tool_parser.prev_tool_call_arr) - 1 - if auto_tools_called - else 0 - ) - should_check = ( - self._should_check_for_unstreamed_tool_arg_tokens( - delta_message, output - ) - ) - # only check if there are any tool calls - # detected by partial parsing - if should_check and tool_parser and auto_tools_called: - latest_delta_len = 0 - if ( - isinstance( - delta_message.tool_calls[0].function, - DeltaFunctionCall, - ) - ) and isinstance( - delta_message.tool_calls[0].function.arguments, str - ): - latest_delta_len = len( - delta_message.tool_calls[0].function.arguments - ) - - # get the expected call based on partial JSON - # parsing which "autocompletes" the JSON. - # Tool parsers (e.g. Qwen3Coder) store - # arguments as a JSON string in - # prev_tool_call_arr. Calling json.dumps() - # on an already-serialized string would - # double-serialize it (e.g. '{"k":1}' becomes - # '"{\\"k\\":1}"'), which then causes the - # replace() below to fail and append the - # entire double-serialized string as a - # spurious final delta. - args = tool_parser.prev_tool_call_arr[index].get( - "arguments", {} - ) - if isinstance(args, str): - expected_call = args - else: - expected_call = json.dumps(args, ensure_ascii=False) - - # get what we've streamed so far for arguments - # for the current tool - actual_call = tool_parser.streamed_args_for_tool[index] - if latest_delta_len > 0: - actual_call = actual_call[:-latest_delta_len] - - # check to see if there's anything left to stream - remaining_call = expected_call.replace(actual_call, "", 1) - # set that as a delta message - delta_message = self._create_remaining_args_delta( - delta_message, remaining_call, index - ) - # Send the finish response for each request.n only once # In OpenAI's API, when a tool is called, the # finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. - if ( - auto_tools_called - or (tools_streamed[i] and not tool_choice_function_name) - or (self.use_harmony and harmony_tools_streamed[i]) + if (tools_streamed[i] and not tool_choice_function_name) or ( + self.use_harmony and harmony_tools_streamed[i] ): finish_reason_ = "tool_calls" else: @@ -1004,7 +942,7 @@ class OpenAIServingChat(OpenAIServing): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, - reasoning_parser: ReasoningParser | None = None, + parser: Parser | None = None, ) -> ErrorResponse | ChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -1113,28 +1051,20 @@ class OpenAIServingChat(OpenAIServing): choices.append(choice_data) continue - if reasoning_parser: - # If the reasoning parser is enabled, - # tool calls are extracted exclusively from the content. - reasoning, content = reasoning_parser.extract_reasoning( - output.text, request=request + if parser is not None: + reasoning, content, tool_calls = parser.parse( + output.text, + request, + enable_auto_tools=self.enable_auto_tools, ) if not request.include_reasoning: reasoning = None else: reasoning = None content = output.text + tool_calls = [] auto_tools_called = False - # if auto tools are not enabled, and a named tool choice using - # outlines is not being used - tool_calls, content = self._parse_tool_calls_from_content( - request=request, - tokenizer=tokenizer, - content=content, - enable_auto_tools=self.enable_auto_tools, - tool_parser_cls=self.tool_parser, - ) if is_mistral_tokenizer(tokenizer): from vllm.tool_parsers.mistral_tool_parser import MistralToolCall @@ -1535,56 +1465,3 @@ class OpenAIServingChat(OpenAIServing): and self.enable_auto_tools and request.tool_choice in ["auto", None] ) - - def _should_check_for_unstreamed_tool_arg_tokens( - self, - delta_message: DeltaMessage | None, - output: CompletionOutput, - ) -> bool: - """ - Check to see if we should check for unstreamed tool arguments tokens. - This is only applicable when auto tool parsing is enabled, the delta - is a tool call with arguments. - """ - - return bool( - # if there is a delta message that includes tool calls which - # include a function that has arguments - output.finish_reason is not None - and self.enable_auto_tools - and self.tool_parser - and delta_message - and delta_message.tool_calls - and delta_message.tool_calls[0] - and delta_message.tool_calls[0].function - and delta_message.tool_calls[0].function.arguments is not None - ) - - @staticmethod - def _create_remaining_args_delta( - delta_message: DeltaMessage, - remaining_call: str, - index: int, - ) -> DeltaMessage: - """ - Create a delta message for remaining tool arguments, preserving - id/type/name from the original delta. - """ - original_tc = next( - (tc for tc in delta_message.tool_calls if tc.index == index), - None, - ) - original_fn = original_tc.function if original_tc else None - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=index, - id=original_tc.id if original_tc else None, - type=original_tc.type if original_tc else None, - function=DeltaFunctionCall( - name=original_fn.name if original_fn else None, - arguments=remaining_call, - ), - ) - ] - ) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index a6c3f9c93dc..30a4f20084e 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -18,6 +18,8 @@ from vllm.entrypoints.openai.engine.protocol import ( StreamOptions, StructuralTagResponseFormat, UsageInfo, + validate_structural_tag_response_format, + validate_structured_outputs_structural_tag, ) from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger @@ -370,6 +372,9 @@ class CompletionRequest(OpenAIBaseModel): parameter="response_format", ) + if rf_type == "structural_tag": + validate_structural_tag_response_format(response_format) + return data @model_validator(mode="before") @@ -397,6 +402,7 @@ class CompletionRequest(OpenAIBaseModel): "outputs ('json', 'regex' or 'choice').", parameter="structured_outputs", ) + validate_structured_outputs_structural_tag(structured_outputs_kwargs) return data @model_validator(mode="before") diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 890af0300ef..434888df9ef 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -17,6 +17,7 @@ from pydantic import ( ) from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.utils import random_uuid from vllm.utils.import_utils import resolve_obj_by_qualname @@ -158,6 +159,80 @@ AnyResponseFormat: TypeAlias = ( ) +def validate_structural_tag_response_format( + response_format: AnyStructuralTagResponseFormat | dict[str, Any], +) -> None: + """Validate structural tags before they are sent to the engine. + + Engine-side validation reports malformed structural tags as generation + failures. OpenAI request parsing should classify them as bad requests. + """ + import json + + from pydantic import TypeAdapter, ValidationError + + if isinstance(response_format, dict): + try: + response_format = TypeAdapter( + AnyStructuralTagResponseFormat + ).validate_python(response_format) + except ValidationError as exc: + raise VLLMValidationError( + "Invalid response_format structural_tag specification.", + parameter="response_format", + ) from exc + + try: + payload = json.dumps(response_format.model_dump(by_alias=True)) + validate_structural_tag_payload(payload, parameter="response_format") + except (TypeError, ValueError) as exc: + raise VLLMValidationError( + "Invalid response_format structural_tag specification.", + parameter="response_format", + ) from exc + + +def validate_structural_tag_payload(payload: Any, *, parameter: str) -> None: + from vllm.sampling_params import SamplingParams, StructuredOutputsParams + from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar + + if isinstance(payload, str) and not payload: + raise VLLMValidationError( + f"Invalid {parameter} structural_tag specification.", + parameter=parameter, + ) + + try: + validate_xgrammar_grammar( + SamplingParams( + structured_outputs=StructuredOutputsParams(structural_tag=payload) + ) + ) + except (TypeError, ValueError) as exc: + raise VLLMValidationError( + f"Invalid {parameter} structural_tag specification.", + parameter=parameter, + ) from exc + + +def validate_structured_outputs_structural_tag( + structured_outputs: Any, +) -> None: + from vllm.sampling_params import StructuredOutputsParams + + if isinstance(structured_outputs, StructuredOutputsParams): + structural_tag = structured_outputs.structural_tag + elif isinstance(structured_outputs, dict): + structural_tag = structured_outputs.get("structural_tag") + else: + return + if structural_tag is not None: + validate_structural_tag_payload( + structural_tag, + parameter="structured_outputs", + ) + + class StreamOptions(OpenAIBaseModel): include_usage: bool | None = False continuous_usage_stats: bool | None = False diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index ff67575fcc6..61b2656bac0 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import contextlib import json import time from collections.abc import Awaitable, Mapping @@ -9,8 +8,7 @@ from http import HTTPStatus from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeVar from fastapi import Request -from openai.types.responses import ToolChoiceFunction -from pydantic import ConfigDict, TypeAdapter, ValidationError +from pydantic import ConfigDict from starlette.datastructures import Headers import vllm.envs as envs @@ -21,7 +19,6 @@ from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.chat_completion.protocol import ( BatchChatCompletionRequest, - ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponse, ) @@ -31,8 +28,6 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, - FunctionCall, - FunctionDefinition, GenerationError, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels @@ -61,14 +56,12 @@ from vllm.renderers.inputs.preprocess import ( ) from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers import ToolParser from vllm.tracing import ( contains_trace_headers, extract_trace_headers, log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.mistral import is_mistral_tool_parser logger = init_logger(__name__) @@ -451,124 +444,6 @@ class OpenAIServing(BeamSearchOnlineMixin): exc_info=True, ) - @staticmethod - def _parse_tool_calls_from_content( - request: ResponsesRequest | ChatCompletionRequest, - tokenizer: TokenizerLike | None, - enable_auto_tools: bool, - tool_parser_cls: type[ToolParser] | None, - content: str | None = None, - ) -> tuple[list[FunctionCall] | None, str | None]: - # When the Mistral grammar factory injected structured outputs, - # let the parser handle the output. - use_mistral_tool_parser = ( - isinstance(request, ChatCompletionRequest) - and is_mistral_tool_parser(tool_parser_cls) - and request._grammar_from_tool_parser - ) - - function_calls = list[FunctionCall]() - if ( - not use_mistral_tool_parser - and request.tool_choice - and isinstance(request.tool_choice, ToolChoiceFunction) - ): - # Forced Function Call (Responses API) - if content is None: - return [], None - function_calls.append( - FunctionCall(name=request.tool_choice.name, arguments=content) - ) - content = None # Clear content since tool is called. - elif ( - not use_mistral_tool_parser - and request.tool_choice - and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) - and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named) - ): - # Named function with standard JSON-based parsing - if content is None: - return [], None - function_calls.append( - FunctionCall(name=request.tool_choice.function.name, arguments=content) - ) - content = None # Clear content since tool is called. - elif ( - not use_mistral_tool_parser - and request.tool_choice == "required" - and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named) - ): - # "required" with standard JSON-based parsing - tool_calls = [] - with contextlib.suppress(ValidationError): - content = content or "" - tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json( - content - ) - for tool_call in tool_calls: - function_calls.append( - FunctionCall( - name=tool_call.name, - arguments=json.dumps(tool_call.parameters, ensure_ascii=False), - ) - ) - content = None # Clear content since tool is called. - elif tool_parser_cls and ( - use_mistral_tool_parser - or ( - enable_auto_tools - and ( - request.tool_choice == "auto" - or request.tool_choice is None - or ( - not tool_parser_cls.supports_required_and_named - and request.tools - and ( - request.tool_choice == "required" - or isinstance( - request.tool_choice, - ChatCompletionNamedToolChoiceParam, - ) - ) - ) - ) - ) - ): - # Automatic Tool Call Parsing (also used as fallback for - # required/named when supports_required_and_named=False) - if tokenizer is None: - raise ValueError( - "Tokenizer not available when `skip_tokenizer_init=True`" - ) - - try: - tool_parser = tool_parser_cls(tokenizer, request.tools) - except RuntimeError as e: - logger.exception("Error in tool parser creation.") - raise e - tool_call_info = tool_parser.extract_tool_calls( - content if content is not None else "", - request=request, # type: ignore - ) - if tool_call_info is not None and tool_call_info.tools_called: - # extract_tool_calls() returns a list of tool calls. - function_calls.extend( - FunctionCall( - id=tool_call.id, - name=tool_call.function.name, - arguments=tool_call.function.arguments, - ) - for tool_call in tool_call_info.tool_calls - ) - content = tool_call_info.content - if content and content.strip() == "": - content = None - else: - # No tool calls. - return None, content - - return function_calls, content - @staticmethod def _get_decoded_token( logprob: Logprob, diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index 347752c912c..504d30f69d2 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -194,10 +194,16 @@ class OpenAIServingModels: lora_request.lora_name, lora_request.lora_path ) ) in str(e): - raise LoRAAdapterNotFoundError( - lora_request.lora_name, lora_request.lora_path - ) from e - raise + return create_error_response( + LoRAAdapterNotFoundError( + lora_request.lora_name, lora_request.lora_path + ) + ) + return create_error_response( + message=str(e), + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) self.lora_requests[lora_name] = lora_request logger.info( diff --git a/vllm/entrypoints/openai/parser/responses_parser.py b/vllm/entrypoints/openai/parser/responses_parser.py index 1868a31ca28..809b601fd21 100644 --- a/vllm/entrypoints/openai/parser/responses_parser.py +++ b/vllm/entrypoints/openai/parser/responses_parser.py @@ -10,10 +10,6 @@ from openai.types.responses.response_function_tool_call_output_item import ( 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.response_reasoning_item import ( - Content, - ResponseReasoningItem, -) from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption from vllm.entrypoints.constants import MCP_PREFIX @@ -22,9 +18,8 @@ from vllm.entrypoints.openai.responses.protocol import ( ResponsesRequest, ) from vllm.outputs import CompletionOutput -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.parser.abstract_parser import Parser from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.utils import random_uuid logger = logging.getLogger(__name__) @@ -37,12 +32,13 @@ class ResponsesParser: self, *, tokenizer: TokenizerLike, - reasoning_parser_cls: type[ReasoningParser], + parser_cls: type[Parser] | None, response_messages: list[ResponseInputOutputItem], request: ResponsesRequest, - tool_parser_cls: type[ToolParser] | None, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, + enable_auto_tools: bool = False, + tool_call_id_type: str = "random", ): self.response_messages: list[ResponseInputOutputItem] = ( # TODO: initial messages may not be properly typed @@ -52,17 +48,22 @@ class ResponsesParser: self.tokenizer = tokenizer self.request = request - self.reasoning_parser_instance = reasoning_parser_cls( - tokenizer, - chat_template_kwargs=_effective_chat_template_kwargs( + self.parser_instance: Parser | None = None + if parser_cls is not None: + chat_template_kwargs = _effective_chat_template_kwargs( request, chat_template=chat_template, chat_template_content_format=chat_template_content_format, - ), - ) - self.tool_parser_instance = None - if tool_parser_cls is not None: - self.tool_parser_instance = tool_parser_cls(tokenizer, request.tools) + ) + + self.parser_instance = parser_cls( + tokenizer, + tools=request.tools, + chat_template_kwargs=chat_template_kwargs, + ) + + self.enable_auto_tools = enable_auto_tools + self.tool_call_id_type = tool_call_id_type # Store the last finish_reason to determine response status self.finish_reason: str | None = None @@ -71,66 +72,34 @@ class ResponsesParser: # Store the finish_reason from the output self.finish_reason = output.finish_reason - reasoning, content = self.reasoning_parser_instance.extract_reasoning( - output.text, request=self.request - ) - if reasoning: - self.response_messages.append( - ResponseReasoningItem( - type="reasoning", - id=f"rs_{random_uuid()}", - summary=[], - content=[ - Content( - type="reasoning_text", - text=reasoning, - ) - ], - ) + if self.parser_instance is not None: + output_items = self.parser_instance.extract_response_outputs( + model_output=output.text, + model_output_token_ids=output.token_ids, + request=self.request, + enable_auto_tools=self.enable_auto_tools, + tool_call_id_type=self.tool_call_id_type, ) - - function_calls: list[ResponseFunctionToolCall] = [] - if self.tool_parser_instance is not None: - tool_call_info = self.tool_parser_instance.extract_tool_calls( - content if content is not None else "", - request=self.request, # type: ignore - ) - if tool_call_info is not None and tool_call_info.tools_called: - # extract_tool_calls() returns a list of tool calls. - function_calls.extend( - ResponseFunctionToolCall( - id=f"fc_{random_uuid()}", - call_id=f"call_{random_uuid()}", - type="function_call", + self.response_messages.extend(output_items) + else: + # No parser configured, treat entire output as text content + if output.text: + self.response_messages.append( + ResponseOutputMessage( + type="message", + id=f"msg_{random_uuid()}", status="completed", - name=tool_call.function.name, - arguments=tool_call.function.arguments, + role="assistant", + content=[ + ResponseOutputText( + annotations=[], # TODO + type="output_text", + text=output.text, + logprobs=None, # TODO + ) + ], ) - for tool_call in tool_call_info.tool_calls ) - content = tool_call_info.content - if content and content.strip() == "": - content = None - - if content: - self.response_messages.append( - ResponseOutputMessage( - type="message", - id=f"msg_{random_uuid()}", - status="completed", - role="assistant", - content=[ - ResponseOutputText( - annotations=[], # TODO - type="output_text", - text=content, - logprobs=None, # TODO - ) - ], - ) - ) - if len(function_calls) > 0: - self.response_messages.extend(function_calls) return self @@ -169,27 +138,29 @@ class ResponsesParser: def get_responses_parser_for_simple_context( *, tokenizer: TokenizerLike, - reasoning_parser_cls: type[ReasoningParser], + parser_cls: type[Parser] | None, response_messages: list[ResponseInputOutputItem], request: ResponsesRequest, - tool_parser_cls, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, + enable_auto_tools: bool = False, + tool_call_id_type: str = "random", ) -> ResponsesParser: """Factory function to create a ResponsesParser with - optional reasoning parser. + optional unified parser. Returns: ResponsesParser instance configured with the provided parser """ return ResponsesParser( tokenizer=tokenizer, - reasoning_parser_cls=reasoning_parser_cls, + parser_cls=parser_cls, response_messages=response_messages, request=request, - tool_parser_cls=tool_parser_cls, chat_template=chat_template, chat_template_content_format=chat_template_content_format, + enable_auto_tools=enable_auto_tools, + tool_call_id_type=tool_call_id_type, ) diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 644dc8cfaaa..62de02ef826 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -41,9 +41,8 @@ from vllm.entrypoints.openai.responses.protocol import ( ) from vllm.entrypoints.openai.responses.utils import construct_tool_dicts from vllm.outputs import RequestOutput -from vllm.reasoning.abs_reasoning_parsers import ReasoningParser +from vllm.parser.abstract_parser import Parser from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ToolParser from vllm.utils import random_uuid if TYPE_CHECKING: @@ -272,12 +271,13 @@ class ParsableContext(ConversationContext): *, response_messages: list[ResponseInputOutputItem], tokenizer: TokenizerLike, - reasoning_parser_cls: type[ReasoningParser] | None, + parser_cls: type[Parser] | None, request: ResponsesRequest, available_tools: list[str] | None, - tool_parser_cls: type[ToolParser] | None, chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, + enable_auto_tools: bool = False, + tool_call_id_type: str = "random", ): self.num_prompt_tokens = 0 self.num_output_tokens = 0 @@ -286,19 +286,17 @@ class ParsableContext(ConversationContext): # not implemented yet for ParsableContext self.all_turn_metrics: list[TurnMetrics] = [] - if reasoning_parser_cls is None: - raise ValueError("reasoning_parser_cls must be provided.") - self.parser = get_responses_parser_for_simple_context( tokenizer=tokenizer, - reasoning_parser_cls=reasoning_parser_cls, + parser_cls=parser_cls, response_messages=response_messages, request=request, - tool_parser_cls=tool_parser_cls, chat_template=chat_template, chat_template_content_format=chat_template_content_format, + enable_auto_tools=enable_auto_tools, + tool_call_id_type=tool_call_id_type, ) - self.tool_parser_cls = tool_parser_cls + self.parser_cls = parser_cls self.request = request self.available_tools = available_tools or [] diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 7da04b3994d..eee02707a97 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -460,16 +460,13 @@ class OpenAIServingResponses(OpenAIServing): context = ParsableContext( response_messages=messages, tokenizer=tokenizer, - reasoning_parser_cls=self.parser.reasoning_parser_cls - if self.parser - else None, + parser_cls=self.parser, request=request, - tool_parser_cls=self.parser.tool_parser_cls - if self.parser - else None, available_tools=available_tools, chat_template=self.chat_template, chat_template_content_format=self.chat_template_content_format, + enable_auto_tools=self.enable_auto_tools, + tool_call_id_type=self.tool_call_id_type, ) else: context = SimpleContext() @@ -708,7 +705,7 @@ class OpenAIServingResponses(OpenAIServing): context.request, context.parser.response_messages, context.tool_dicts, - context.tool_parser_cls, + context.parser_cls.tool_parser_cls if context.parser_cls else None, context.chat_template, context.chat_template_content_format, ) @@ -1411,6 +1408,7 @@ class OpenAIServingResponses(OpenAIServing): delta_token_ids=delta_token_ids, request=request, prompt_token_ids=ctx.last_output.prompt_token_ids, + finished=output.finish_reason is not None, ) else: delta_message = DeltaMessage(content=output.text) diff --git a/vllm/entrypoints/serve/__init__.py b/vllm/entrypoints/serve/__init__.py index 8233d3324d6..57491d45f63 100644 --- a/vllm/entrypoints/serve/__init__.py +++ b/vllm/entrypoints/serve/__init__.py @@ -3,18 +3,15 @@ from fastapi import FastAPI -import vllm.envs as envs from vllm.logger import init_logger logger = init_logger(__name__) def register_vllm_serve_api_routers(app: FastAPI): - if envs.VLLM_SERVER_DEV_MODE: - logger.warning( - "SECURITY WARNING: Development endpoints are enabled! " - "This should NOT be used in production!" - ) + from .instrumentator import register_instrumentator_api_routers + + register_instrumentator_api_routers(app) from vllm.entrypoints.serve.lora.api_router import ( attach_router as attach_lora_router, @@ -28,30 +25,37 @@ def register_vllm_serve_api_routers(app: FastAPI): attach_profile_router(app) - from vllm.entrypoints.serve.sleep.api_router import ( - attach_router as attach_sleep_router, - ) - - attach_sleep_router(app) - - from vllm.entrypoints.serve.rpc.api_router import ( - attach_router as attach_rpc_router, - ) - - attach_rpc_router(app) - - from vllm.entrypoints.serve.cache.api_router import ( - attach_router as attach_cache_router, - ) - - attach_cache_router(app) - from vllm.entrypoints.serve.tokenize.api_router import ( attach_router as attach_tokenize_router, ) attach_tokenize_router(app) - from .instrumentator import register_instrumentator_api_routers - register_instrumentator_api_routers(app) +def register_vllm_dev_api_routers(app: FastAPI): + logger.warning( + "SECURITY WARNING: Development endpoints are enabled! " + "This should NOT be used in production!" + ) + + from .dev.cache.api_router import attach_router as attach_cache_router + + attach_cache_router(app) + + from .dev.rlhf.api_router import attach_router as attach_rlhf_router + + attach_rlhf_router(app) + + from .dev.rpc.api_router import attach_router as attach_rpc_router + + attach_rpc_router(app) + + from .dev.server_info.api_router import ( + attach_router as attach_server_info_router, + ) + + attach_server_info_router(app) + + from .dev.sleep.api_router import attach_router as attach_sleep_router + + attach_sleep_router(app) diff --git a/vllm/entrypoints/serve/rlhf/__init__.py b/vllm/entrypoints/serve/dev/__init__.py similarity index 100% rename from vllm/entrypoints/serve/rlhf/__init__.py rename to vllm/entrypoints/serve/dev/__init__.py diff --git a/vllm/entrypoints/serve/rpc/__init__.py b/vllm/entrypoints/serve/dev/cache/__init__.py similarity index 100% rename from vllm/entrypoints/serve/rpc/__init__.py rename to vllm/entrypoints/serve/dev/cache/__init__.py diff --git a/vllm/entrypoints/serve/cache/api_router.py b/vllm/entrypoints/serve/dev/cache/api_router.py similarity index 96% rename from vllm/entrypoints/serve/cache/api_router.py rename to vllm/entrypoints/serve/dev/cache/api_router.py index 10015f02caa..c274717c0a8 100644 --- a/vllm/entrypoints/serve/cache/api_router.py +++ b/vllm/entrypoints/serve/dev/cache/api_router.py @@ -5,7 +5,6 @@ from fastapi import APIRouter, FastAPI, Query, Request from fastapi.responses import Response -import vllm.envs as envs from vllm.engine.protocol import EngineClient from vllm.logger import init_logger @@ -67,6 +66,4 @@ async def reset_encoder_cache(raw_request: Request): def attach_router(app: FastAPI): - if not envs.VLLM_SERVER_DEV_MODE: - return app.include_router(router) diff --git a/vllm/entrypoints/serve/sleep/__init__.py b/vllm/entrypoints/serve/dev/rlhf/__init__.py similarity index 100% rename from vllm/entrypoints/serve/sleep/__init__.py rename to vllm/entrypoints/serve/dev/rlhf/__init__.py diff --git a/vllm/entrypoints/serve/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py similarity index 98% rename from vllm/entrypoints/serve/rlhf/api_router.py rename to vllm/entrypoints/serve/dev/rlhf/api_router.py index dcae3889dc7..6237de87769 100644 --- a/vllm/entrypoints/serve/rlhf/api_router.py +++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py @@ -8,7 +8,6 @@ from typing import Annotated from fastapi import APIRouter, FastAPI, HTTPException, Query, Request from fastapi.responses import JSONResponse -import vllm.envs as envs from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, WeightTransferUpdateRequest, @@ -186,6 +185,4 @@ async def get_world_size( def attach_router(app: FastAPI): - if not envs.VLLM_SERVER_DEV_MODE: - return app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/rpc/__init__.py b/vllm/entrypoints/serve/dev/rpc/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/serve/rpc/api_router.py b/vllm/entrypoints/serve/dev/rpc/api_router.py similarity index 95% rename from vllm/entrypoints/serve/rpc/api_router.py rename to vllm/entrypoints/serve/dev/rpc/api_router.py index 54f582c408d..99b904c2f63 100644 --- a/vllm/entrypoints/serve/rpc/api_router.py +++ b/vllm/entrypoints/serve/dev/rpc/api_router.py @@ -8,7 +8,6 @@ from typing import Any from fastapi import APIRouter, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response -import vllm.envs as envs from vllm.engine.protocol import EngineClient from vllm.logger import init_logger @@ -56,6 +55,4 @@ async def collective_rpc(raw_request: Request): def attach_router(app: FastAPI): - if not envs.VLLM_SERVER_DEV_MODE: - return app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/server_info/__init__.py b/vllm/entrypoints/serve/dev/server_info/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/serve/instrumentator/server_info.py b/vllm/entrypoints/serve/dev/server_info/api_router.py similarity index 93% rename from vllm/entrypoints/serve/instrumentator/server_info.py rename to vllm/entrypoints/serve/dev/server_info/api_router.py index 60967c5a66a..64b7cdeb2fb 100644 --- a/vllm/entrypoints/serve/instrumentator/server_info.py +++ b/vllm/entrypoints/serve/dev/server_info/api_router.py @@ -7,7 +7,7 @@ import functools from typing import Annotated, Literal import pydantic -from fastapi import APIRouter, Query, Request +from fastapi import APIRouter, FastAPI, Query, Request from fastapi.responses import JSONResponse import vllm.envs as envs @@ -57,3 +57,7 @@ async def show_server_info( "system_env": await asyncio.to_thread(_get_system_env_info_cached), } return JSONResponse(content=server_info) + + +def attach_router(app: FastAPI): + app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/sleep/__init__.py b/vllm/entrypoints/serve/dev/sleep/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/entrypoints/serve/sleep/api_router.py b/vllm/entrypoints/serve/dev/sleep/api_router.py similarity index 95% rename from vllm/entrypoints/serve/sleep/api_router.py rename to vllm/entrypoints/serve/dev/sleep/api_router.py index 46fa1c3f43f..0861c867732 100644 --- a/vllm/entrypoints/serve/sleep/api_router.py +++ b/vllm/entrypoints/serve/dev/sleep/api_router.py @@ -5,7 +5,6 @@ from fastapi import APIRouter, FastAPI, Request from fastapi.responses import JSONResponse, Response -import vllm.envs as envs from vllm.engine.protocol import EngineClient from vllm.logger import init_logger @@ -50,7 +49,4 @@ async def is_sleeping(raw_request: Request): def attach_router(app: FastAPI): - if not envs.VLLM_SERVER_DEV_MODE: - return - app.include_router(router) diff --git a/vllm/entrypoints/serve/instrumentator/__init__.py b/vllm/entrypoints/serve/instrumentator/__init__.py index 8abce02325a..c987394ad03 100644 --- a/vllm/entrypoints/serve/instrumentator/__init__.py +++ b/vllm/entrypoints/serve/instrumentator/__init__.py @@ -3,8 +3,6 @@ from fastapi import FastAPI -from vllm import envs - def register_instrumentator_api_routers(app: FastAPI): from .basic import router as basic_router @@ -22,8 +20,3 @@ def register_instrumentator_api_routers(app: FastAPI): from .offline_docs import attach_router as offline_docs_attach_router offline_docs_attach_router(app) - - if envs.VLLM_SERVER_DEV_MODE: - from .server_info import router as server_info_router - - app.include_router(server_info_router) diff --git a/vllm/envs.py b/vllm/envs.py index c12e3cae247..dc11fbd224d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -95,7 +95,6 @@ if TYPE_CHECKING: CMAKE_BUILD_TYPE: Literal["Debug", "Release", "RelWithDebInfo"] | None = None VERBOSE: bool = False VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False - VLLM_RPC_TIMEOUT: int = 10000 # ms VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds VLLM_MAX_N_SEQUENCES: int = 16384 VLLM_PLUGINS: list[str] | None = None @@ -1015,9 +1014,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv( "VLLM_TEST_FORCE_LOAD_FORMAT", "dummy" ), - # Time in ms for the zmq client to wait for a response from the backend - # server for simple data operations - "VLLM_RPC_TIMEOUT": lambda: int(os.getenv("VLLM_RPC_TIMEOUT", "10000")), # Timeout in seconds for keeping HTTP connections alive in API server "VLLM_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int( os.environ.get("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5") diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 37e5b8e1d54..39d2e86d3c3 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -161,6 +161,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import ( TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( + XPUFp8BlockScaledMMKernel, XPUFP8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.zentorch import ( @@ -317,6 +318,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ CPUFp8BlockScaledMMKernel, ], PlatformEnum.XPU: [ + XPUFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], } diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py index 68528bbd488..17900c75058 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py @@ -51,13 +51,6 @@ class XPUwNa16LinearKernel(MPLinearKernel): "XPUwNa16, supported sizes are multiples of 32", ) - if c.partition_weight_shape[1] % 32 != 0: - return ( - False, - f"Output size ({c.partition_weight_shape[1]}) not supported by " - "XPUWNA16, supported sizes are multiples of 32", - ) - return True, None def process_weights_after_loading(self, layer: torch.nn.Module): diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 9bd644b6299..39f9abd460e 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( from vllm.model_executor.kernels.linear.scaled_mm.triton import ( TritonInt8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( + XPUFp8BlockScaledMMKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.zentorch import ( ZentorchInt8ScaledMMLinearKernel, ) @@ -64,4 +67,5 @@ __all__ = [ "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", + "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index 7003e727bfa..78dad872958 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -160,7 +160,7 @@ class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): @classmethod def is_supported(cls, compute_capability=None): if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): - return False, "only cuda-like and xpu devices are supported." + return False, "only CUDA-alike and XPU devices are supported." return True, None def apply_block_scaled_mm( diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index 0e4ead39219..670a021ef0c 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -5,10 +5,6 @@ from collections.abc import Sequence import torch -from vllm.model_executor.kernels.linear import ( # noqa: E501 - FP8ScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticChannelSym, kFp8StaticTensorSym, @@ -16,6 +12,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform +from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel +from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig + class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): @classmethod @@ -84,3 +83,38 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): output_shape: list, ) -> torch.Tensor: pass + + +class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + 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) + replace_parameter(layer, scale_attr, scale.data.t().contiguous()) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + # Weight is [N, K]. Use .t() to create a [K, N] view without copying. + return torch.ops._xpu_C.fp8_gemm( + A, + B.t(), + self.config.out_dtype, + As, + Bs, + torch.Tensor(), + ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b87b87f136d..430947235e9 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -102,23 +102,26 @@ def _quant_flags_to_group_shape( class RoutingMethodType(IntEnum): # Default: Softmax -> TopK Default = (0,) - # Renormalize: TopK -> Softmax/Sigmoid + # Renormalize: TopK -> Softmax Renormalize = (1,) # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups # -> Top8 experts from the Top4 groups DeepSeekV3 = (2,) # Llama4: Top1 -> Sigmoid Llama4 = (3,) - # RenormalizeNaive: Softmax/Sigmoid -> TopK -> Renormalize + # RenormalizeNaive: Softmax -> TopK -> Renormalize RenormalizeNaive = (4,) # TopK: TopK (no softmax) TopK = (5,) # SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K) SigmoidRenorm = (6,) # MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize + # (routeScale=1.0, epsilon=1e-20) MiniMax2 = (7,) + # Sigmoid: Sigmoid -> TopK (no renormalization) + Sigmoid = (8,) # Unspecified - Unspecified = (8,) + Unspecified = (9,) # other routing types (not passed to FlashInfer kernels) # Deepseek V4 -> sqrtsoftplus + Bias + Normalize DeepseekV4 = (100,) @@ -132,6 +135,7 @@ def get_routing_method_type( renormalize: bool, num_expert_group: int | None, has_e_score_bias: bool, + routed_scaling_factor: float | None = 1.0, ) -> RoutingMethodType: if scoring_func == "sqrtsoftplus": # DeepSeek V4 uses sqrtsoftplus routing with optional routing bias @@ -142,20 +146,21 @@ def get_routing_method_type( return RoutingMethodType.Unspecified if has_e_score_bias: - if (num_expert_group or 0) > 0 and scoring_func == "sigmoid": - return RoutingMethodType.DeepSeekV3 - elif scoring_func == "sigmoid": - return RoutingMethodType.MiniMax2 + if scoring_func == "sigmoid": + if not renormalize: + return RoutingMethodType.Unspecified + if (num_expert_group or 0) > 0: + return RoutingMethodType.DeepSeekV3 + if routed_scaling_factor in (None, 1.0): + return RoutingMethodType.MiniMax2 + return RoutingMethodType.Unspecified else: return RoutingMethodType.Unspecified if scoring_func == "sigmoid": - if top_k == 1: - return RoutingMethodType.Llama4 - elif renormalize: + if renormalize: return RoutingMethodType.SigmoidRenorm - else: - return RoutingMethodType.Unspecified + return RoutingMethodType.Sigmoid if scoring_func == "softmax": if renormalize: diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 9192b6a9b7e..d49270122a7 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -53,6 +53,10 @@ _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = { MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x), MoEActivation.SWIGLUOAI: _swigluoai_forward_native, MoEActivation.GELU: _gelu_and_mul, + MoEActivation.GELU_TANH: ( + lambda x: F.gelu(x[..., : x.shape[-1] // 2], approximate="tanh") + * x[..., x.shape[-1] // 2 :] + ), } diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index feb49d260e1..d8570049af2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py @@ -322,6 +322,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular): return activation in [ MoEActivation.SILU, MoEActivation.GELU, + MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, ] @@ -719,10 +720,12 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular): return activation in [ MoEActivation.SILU, MoEActivation.GELU, + MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUSTEP, MoEActivation.SILU_NO_MUL, MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, MoEActivation.RELU2_NO_MUL, ] 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 6481434f2e7..38200d9d090 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 @@ -54,6 +54,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): self.out_dtype = moe_config.in_dtype self.num_local_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank + # FC2 input scale tensor bound in process_weights_after_loading: the + # calibrated (now-zeroed) a2_gscale for static-quant checkpoints, or + # a synthesized uniform-1.0 tensor for W4A16 checkpoints that lack + # one. Holding it on the instance keeps apply() alloc-free. + self._fc2_input_scale: 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 @@ -86,6 +91,18 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): # its own per-block dynamic scale. if self.a2_gscale is not None: self.a2_gscale.fill_(1.0) + self._fc2_input_scale = self.a2_gscale + else: + # W4A16 NVFP4 checkpoints have no calibrated a2_gscale; b12x + # performs dynamic per-block FC2-input quantization, so a uniform + # 1.0 scale per expert is equivalent to the bake-in above for + # static-quant checkpoints. Allocate once here so apply() stays + # alloc-free. + self._fc2_input_scale = torch.ones( + self.num_local_experts, + device=layer.w13_weight.device, + dtype=torch.float32, + ) # Precompute MMA-layout views of the weight scale factors once here # rather than recomputing on every forward pass. @@ -131,7 +148,13 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) + # b12x performs in-kernel BF16->FP4 activation quant, so W4A16 + # NVFP4 checkpoints (activation_key=None, e.g. mixed-precision + # compressed-tensors layouts) are runtime-compatible. + return (weight_key, activation_key) in ( + (kNvfp4Static, kNvfp4Dynamic), + (kNvfp4Static, None), + ) @staticmethod def _supports_activation(activation: MoEActivation) -> bool: @@ -198,8 +221,8 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): assert self.g1_alphas is not None and self.g2_alphas is not None, ( "g1_alphas and g2_alphas must not be None for FlashInferB12xExperts" ) - assert self.a2_gscale is not None, ( - "a2_gscale must not be None for FlashInferB12xExperts" + assert self._fc2_input_scale is not None, ( + "_fc2_input_scale must be set by process_weights_after_loading" ) top_k = topk_ids.shape[1] @@ -211,7 +234,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): w1_weight=w1, w1_weight_sf=self.w1_sf_mma, w1_alpha=self.g1_alphas, - fc2_input_scale=self.a2_gscale, + fc2_input_scale=self._fc2_input_scale, w2_weight=w2, w2_weight_sf=self.w2_sf_mma, w2_alpha=self.g2_alphas, diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 1d0cf91d427..64c68018f36 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -45,6 +45,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt8Static, kMxfp4Static, kMxfp8Static, @@ -611,6 +613,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): kInt4Static, kInt8Static, kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, ] return weight_key in SUPPORTED_W 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 43126195205..9230fea6e5c 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 @@ -257,7 +257,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - return router_logits_dtype != torch.float32 + return router_logits_dtype in [torch.bfloat16, torch.float32] @staticmethod def _supports_routing_method( @@ -279,6 +279,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.SigmoidRenorm, + RoutingMethodType.Sigmoid, RoutingMethodType.MiniMax2, RoutingMethodType.Simulated, ] @@ -290,6 +291,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.SigmoidRenorm, + RoutingMethodType.Sigmoid, RoutingMethodType.MiniMax2, RoutingMethodType.Simulated, ] @@ -317,9 +319,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout assert not apply_router_weight_on_input - assert activation == MoEActivation.SILU + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + activation_type = activation_to_flashinfer_int(activation) assert self.topk <= global_num_experts - assert self.topk <= 10 assert global_num_experts % 4 == 0 assert self.quant_config.block_shape in [[128, 128], [1, 32]] # Kernel expects #experts <= #threads 512 @@ -333,13 +335,19 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit use_shuffled_weight = True weight_layout = WeightLayout.MajorK hidden_states_scale = a1q_scale + # FlashInfer expects None for non-grouped MXFP8 routing configs. + n_group = num_expert_group or None + selected_topk_group = topk_group or None else: + assert self.topk <= 10 fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 use_shuffled_weight = True weight_layout = WeightLayout.BlockMajorK hidden_states_scale = a1q_scale.t().contiguous() + n_group = num_expert_group or 0 + selected_topk_group = topk_group or 0 - return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( + kwargs = dict( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, @@ -350,8 +358,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit gemm2_weights_scale=self.quant_config.w2_scale, num_experts=global_num_experts, top_k=self.topk, - n_group=(num_expert_group or 0), - topk_group=(topk_group or 0), + n_group=n_group, + topk_group=selected_topk_group, intermediate_size=self.intermediate_size_per_partition, local_expert_offset=self.ep_rank * self.local_num_experts, local_num_experts=self.local_num_experts, @@ -361,6 +369,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit weight_layout=weight_layout, fp8_quantization_type=fp8_quant_type, ) + if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL: + kwargs["activation_type"] = activation_type + return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs) def _apply_per_tensor( self, @@ -393,11 +404,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit else: assert not apply_router_weight_on_input - # Currently FI requires bfloat16 routing bias. - # https://github.com/flashinfer-ai/flashinfer/issues/2909 - if e_score_correction_bias is not None: - e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16) - out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, 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 5ee023aa27c..cbfabce502e 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 @@ -360,7 +360,7 @@ class TrtLlmNvFp4ExpertsMonolithic( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - return router_logits_dtype != torch.float32 + return router_logits_dtype in [torch.bfloat16, torch.float32] def apply( self, @@ -393,11 +393,6 @@ class TrtLlmNvFp4ExpertsMonolithic( and self.routing_method_type != RoutingMethodType.Llama4 ) - # Currently FI requires bfloat16 routing bias. - # https://github.com/flashinfer-ai/flashinfer/issues/2909 - if e_score_correction_bias is not None: - e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16) - output1_scale_gate_scalar = self.quant_config.g1_alphas # Invoke kernel. 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 9de7a6ba119..6ad60d62e97 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_moe_permute_scales, marlin_permute_bias, moe_awq_to_marlin_zero_points, + moe_packed_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -496,6 +497,23 @@ def _process_weights_marlin( marlin_w2_scales ) + # --- Permute zero points --- + if w13_qzeros is not None and w2_qzeros is not None: + w13_qzeros = moe_packed_to_marlin_zero_points( + w13_qzeros, + size_k=w13_qzeros.shape[1], + size_n=w13_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + w2_qzeros = moe_packed_to_marlin_zero_points( + w2_qzeros, + size_k=w2_qzeros.shape[1], + size_n=w2_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + # --- Permute bias --- if w13_bias is not None: w13_bias_out = marlin_permute_bias(w13_bias) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 4a9b190335d..8e4012d3ec8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -68,6 +68,12 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac UnquantizedMoeBackend.BATCHED_TRITON, ] + # On Hopper (SM90), the FlashInfer unquantized MoE kernels are slower + # than Triton, so prefer Triton by default. + if current_platform.is_device_capability_family(90): + _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM) + _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) + # HACK: Qwen3.5 has crash with FLASHINFER_CUTLASS BF16 if DEP. # Updating the oracle querying logic is out of the scope of this # PR. Need to fix the kernel or update structure in follow up. 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 31a75c860d3..cd9aff83536 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 @@ -326,6 +326,7 @@ class FusedTopKBiasRouter(BaseRouter): renormalize=self.renormalize, num_expert_group=None, has_e_score_bias=True, + routed_scaling_factor=self.routed_scaling_factor, ) def _compute_routing( diff --git a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py index 6f792b46a0a..ac95de346e5 100644 --- a/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py @@ -283,6 +283,7 @@ class GroupedTopKRouter(BaseRouter): renormalize=self.renormalize, num_expert_group=self.num_expert_group, has_e_score_bias=self.e_score_correction_bias is not None, + routed_scaling_factor=self.routed_scaling_factor, ) def _compute_routing( diff --git a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py index 54f0fa4fb0a..0c477322e99 100644 --- a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py +++ b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py @@ -63,6 +63,7 @@ class ZeroExpertRouter(BaseRouter): renormalize=self.renormalize, num_expert_group=None, has_e_score_bias=True, + routed_scaling_factor=self.routed_scaling_factor, ) def _compute_routing( diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 92fc6442ced..7a0d50c74e3 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -177,7 +177,7 @@ def _resolve_gdn_prefill_backend( return backend, "triton" head_k_dim = getattr( - vllm_config.model_config.hf_config, "linear_key_head_dim", None + vllm_config.model_config.hf_text_config, "linear_key_head_dim", None ) supports_flashinfer = False @@ -218,7 +218,7 @@ def _log_gdn_backend_decision( ) -> None: """Log the GDN prefill backend choice in the attention-selector style.""" head_k_dim = getattr( - vllm_config.model_config.hf_config, "linear_key_head_dim", None + vllm_config.model_config.hf_text_config, "linear_key_head_dim", None ) chosen = { "flashinfer": "FlashInfer", diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float16.json new file mode 100644 index 00000000000..fdf38cdf042 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float16.json @@ -0,0 +1,87 @@ +{ + "triton_version": "3.6.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "128": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "512": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "4096": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 32, + "num_warps": 2 + }, + "12288": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "16384": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "24576": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "32768": { + "BLOCK_SIZE_M": 32, + "num_warps": 2 + }, + "49152": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "65536": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "98304": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "131072": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "196608": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + }, + "262144": { + "BLOCK_SIZE_M": 16, + "num_warps": 2 + } +} \ 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=NVIDIA_H200,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float32.json new file mode 100644 index 00000000000..82bdff70134 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float32.json @@ -0,0 +1,87 @@ +{ + "triton_version": "3.6.0", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 4 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "12288": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "24576": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "49152": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "65536": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "98304": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "196608": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "262144": { + "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=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float16.json new file mode 100644 index 00000000000..6be92a4bc28 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float16.json @@ -0,0 +1,87 @@ +{ + "triton_version": "3.6.0", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 4 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "128": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "2048": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "4096": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "8192": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "12288": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "24576": { + "BLOCK_SIZE_M": 32, + "num_warps": 4 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "49152": { + "BLOCK_SIZE_M": 32, + "num_warps": 2 + }, + "65536": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "98304": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "196608": { + "BLOCK_SIZE_M": 32, + "num_warps": 1 + }, + "262144": { + "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=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float32.json new file mode 100644 index 00000000000..7b55fab1add --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float32.json @@ -0,0 +1,87 @@ +{ + "triton_version": "3.6.0", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "64": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "128": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 8 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "2048": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "4096": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 8 + }, + "12288": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "24576": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 4, + "num_warps": 4 + }, + "49152": { + "BLOCK_SIZE_M": 16, + "num_warps": 4 + }, + "65536": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "98304": { + "BLOCK_SIZE_M": 16, + "num_warps": 1 + }, + "131072": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "196608": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "262144": { + "BLOCK_SIZE_M": 64, + "num_warps": 4 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 166bd43bbdd..e5ef487ee9b 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -12,11 +12,6 @@ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_torch, causal_conv1d_update_torch, ) -from vllm.model_executor.layers.mamba.ops.cpu.recurrent_gated_delta_rule import ( - chunk_gated_delta_rule, - gdn_gating, - recurrent_gated_delta_rule, -) from vllm.utils.torch_utils import ( LayerNameType, _resolve_layer_name, @@ -55,88 +50,91 @@ def cpu_gdn_attention_core( attn_metadata_i.spec_sequence_masks is None and attn_metadata_i.num_accepted_tokens is None ), "speculative decode not supported in CPU GDN attention." - - if torch.cpu._is_amx_tile_supported(): - return cpu_gdn_attention_core_amx( - mixed_qkv, - b, - a, - core_attn_out, - attn_metadata_i, - layer, - ) + assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16." state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor query_start_loc = attn_metadata_i.non_spec_query_start_loc assert state_indices_tensor is not None assert query_start_loc is not None - # [num_allocated_slots, conv_dim, kernel - 1] + is_amx = torch.cpu._is_amx_tile_supported() + conv_state = layer.kv_cache[0] - if not is_conv_state_dim_first(): - conv_state = conv_state.transpose(-1, -2) + if is_amx: + # AMX causal conv requires [num_allocated_slots, kernel - 1, conv_dim]. + if is_conv_state_dim_first(): + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") + conv_state = conv_state.transpose(1, 2) + else: + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + conv_weights = layer.conv1d.weight.view( + layer.conv1d.weight.size(0), layer.conv1d.weight.size(2) + ) # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim] ssm_state = layer.kv_cache[1] + mixed_qkv = mixed_qkv.contiguous() + a = a.contiguous() + b = b.contiguous() + + num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size() + ssm_state = ssm_state.view( + num_allocated_slots, + head_num, + k_dim, + v_dim, + ) num_decodes = attn_metadata_i.num_decodes num_decode_tokens = attn_metadata_i.num_decode_tokens num_prefills = attn_metadata_i.num_prefills num_prefill_tokens = attn_metadata_i.num_prefill_tokens - conv_weights = layer.conv1d.weight.view( - layer.conv1d.weight.size(0), layer.conv1d.weight.size(2) - ) - # all decode requests (batched) if num_decodes > 0: decode_mixed_qkv = mixed_qkv[:num_decode_tokens] decode_b = b[:num_decode_tokens] decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] - decode_conv_state = conv_state[decode_state_indices].contiguous() + if is_amx: + decode_mixed_qkv = ops.causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_states=conv_state, + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + silu_activation=layer.activation == "silu", + conv_state_indices=decode_state_indices, + is_vnni=True, + ) + else: + decode_conv_state = conv_state[decode_state_indices].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - # [B, dim] -> [B, dim, 1] - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, - weight=conv_weights, - bias=layer.conv1d.bias, - activation=layer.activation, - ).squeeze(-1) - conv_state[decode_state_indices] = decode_conv_state + decode_mixed_qkv = causal_conv1d_update_torch( + # [B, dim] -> [B, dim, 1] + x=decode_mixed_qkv.unsqueeze(-1), + conv_state=decode_conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + ).squeeze(-1) + conv_state[decode_state_indices] = decode_conv_state query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) - # [1, L, H, D] -> [B, 1, H, D] for batched decode - query = query.transpose(0, 1).contiguous() - key = key.transpose(0, 1).contiguous() - value = value.transpose(0, 1).contiguous() - - g, beta_output = gdn_gating( + attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( A_log=layer.A_log, + dt_bias=layer.dt_bias, + q=query, + k=key, + v=value, a=decode_a, b=decode_b, - dt_bias=layer.dt_bias, - ) - if g.ndim == 2: - g = g.unsqueeze(1) - beta_output = beta_output.unsqueeze(1) - - initial_state = ssm_state[decode_state_indices].contiguous() - attn_out, last_recurrent_state = recurrent_gated_delta_rule( - query=query, - key=key, - value=value, - g=g, - beta=beta_output, - initial_state=initial_state, - scale=None, + initial_state_source=ssm_state, + initial_state_indices=decode_state_indices, + cu_seqlens=query_start_loc[: num_decodes + 1], use_qk_l2norm_in_kernel=True, ) - ssm_state[decode_state_indices] = last_recurrent_state.to( - ssm_state.dtype - ).contiguous() core_attn_out[:num_decode_tokens] = attn_out.squeeze(1) # all prefill requests: (varlen) currently naively loops over sequences @@ -160,154 +158,29 @@ def cpu_gdn_attention_core( num_decodes : num_decodes + num_prefills ] - prefill_mixed_qkv = causal_conv1d_torch( - x=prefill_mixed_qkv.transpose(0, 1), - weight=conv_weights, - bias=layer.conv1d.bias, - conv_states=conv_state, - query_start_loc=prefill_query_start_loc, - cache_indices=prefill_state_indices, - has_initial_state=prefill_has_initial_state, - activation=layer.activation, - ).transpose(0, 1) - - query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) - g, beta = gdn_gating(layer.A_log, prefill_a, prefill_b, layer.dt_bias) - if g.ndim == 2: - g = g.unsqueeze(0) - beta = beta.unsqueeze(0) - - initial_state = ssm_state[prefill_state_indices].contiguous() - initial_state[~prefill_has_initial_state, ...] = 0 - attn_out, last_recurrent_state = chunk_gated_delta_rule( - q=query, - k=key, - v=value, - g=g, - beta=beta, - scale=None, - initial_state=initial_state, - cu_seqlens=prefill_query_start_loc, - use_qk_l2norm_in_kernel=True, - ) - ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) - core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) - - -def cpu_gdn_attention_core_fake( - mixed_qkv: torch.Tensor, - b: torch.Tensor, - a: torch.Tensor, - core_attn_out: torch.Tensor, - layer_name: LayerNameType, -) -> None: - """Fake implementation for torch.compile.""" - return - - -def cpu_gdn_attention_core_amx( - mixed_qkv: torch.Tensor, - b: torch.Tensor, - a: torch.Tensor, - core_attn_out: torch.Tensor, - attn_metadata_i: GDNAttentionMetadata, - layer: torch.nn.Module, -): - state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor - query_start_loc = attn_metadata_i.non_spec_query_start_loc - assert state_indices_tensor is not None - assert query_start_loc is not None - - # [num_allocated_slots, kernel - 1, conv_dim] - conv_state = layer.kv_cache[0] - if is_conv_state_dim_first(): - raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.") - # reshape to [num_allocated_slots, conv_dim, kernel - 1] - conv_state_t = conv_state.transpose(1, 2) - - # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim] - ssm_state = layer.kv_cache[1] - # rehape to [num_allocated_slots, num_v_heads / tp_size, k_dim, v_dim] - num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size() - ssm_state = ssm_state.view( - num_allocated_slots, - head_num, - k_dim, - v_dim, - ) - - mixed_qkv = mixed_qkv.contiguous() - a = a.contiguous() - b = b.contiguous() - - num_decodes = attn_metadata_i.num_decodes - num_decode_tokens = attn_metadata_i.num_decode_tokens - num_prefills = attn_metadata_i.num_prefills - num_prefill_tokens = attn_metadata_i.num_prefill_tokens - - if num_decodes > 0: - decode_mixed_qkv = mixed_qkv[:num_decode_tokens] - decode_b = b[:num_decode_tokens] - decode_a = a[:num_decode_tokens] - decode_state_indices = state_indices_tensor[:num_decodes] - - decode_mixed_qkv = ops.causal_conv1d_update_cpu( - x=decode_mixed_qkv, - conv_states=conv_state_t, - weight=layer.conv1d.weight, - bias=layer.conv1d.bias, - silu_activation=layer.activation == "silu", - conv_state_indices=decode_state_indices, - is_vnni=True, - ) - - query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) - attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( - A_log=layer.A_log, - dt_bias=layer.dt_bias, - q=query, - k=key, - v=value, - a=decode_a, - b=decode_b, - initial_state_source=ssm_state, - initial_state_indices=decode_state_indices, - cu_seqlens=query_start_loc[: num_decodes + 1], - use_qk_l2norm_in_kernel=True, - ) - core_attn_out[:num_decode_tokens] = attn_out.squeeze(1) - - if num_prefills > 0: - has_initial_state = attn_metadata_i.has_initial_state - assert has_initial_state is not None - - prefill_token_start = num_decode_tokens - prefill_token_end = prefill_token_start + num_prefill_tokens - prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end] - prefill_b = b[prefill_token_start:prefill_token_end] - prefill_a = a[prefill_token_start:prefill_token_end] - prefill_state_indices = state_indices_tensor[ - num_decodes : num_decodes + num_prefills - ] - prefill_query_start_loc = ( - query_start_loc[num_decodes : num_decodes + num_prefills + 1] - - num_decode_tokens - ) - prefill_has_initial_state = has_initial_state[ - num_decodes : num_decodes + num_prefills - ] - - prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( - x=prefill_mixed_qkv.transpose(0, 1), - weight=layer.conv1d.weight, - bias=layer.conv1d.bias, - conv_states=conv_state_t, - query_start_loc=prefill_query_start_loc, - cache_indices=prefill_state_indices, - has_initial_state=prefill_has_initial_state, - silu_activation=layer.activation == "silu", - is_vnni=True, - ).transpose(0, 1) + if is_amx: + prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu( + x=prefill_mixed_qkv.transpose(0, 1), + weight=layer.conv1d.weight, + bias=layer.conv1d.bias, + conv_states=conv_state, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + silu_activation=layer.activation == "silu", + is_vnni=True, + ).transpose(0, 1) + else: + prefill_mixed_qkv = causal_conv1d_torch( + x=prefill_mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_state, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + activation=layer.activation, + ).transpose(0, 1) query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) g, beta = ops.fused_gdn_gating_cpu( @@ -334,6 +207,17 @@ def cpu_gdn_attention_core_amx( core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) +def cpu_gdn_attention_core_fake( + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + layer_name: LayerNameType, +) -> None: + """Fake implementation for torch.compile.""" + return + + def register_cpu_gdn_attention_ops() -> None: global _CPU_GDN_ATTENTION_OPS_REGISTERED if _CPU_GDN_ATTENTION_OPS_REGISTERED: diff --git a/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py b/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py deleted file mode 100644 index 30fca3423a3..00000000000 --- a/vllm/model_executor/layers/mamba/ops/cpu/recurrent_gated_delta_rule.py +++ /dev/null @@ -1,223 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -import torch -import torch.nn.functional as F - - -def l2norm( - x: torch.Tensor, - dim: int = -1, - eps: float = 1e-6, -) -> torch.Tensor: - inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) - return x * inv_norm - - -def recurrent_gated_delta_rule( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - initial_state: torch.Tensor, - scale: float | None = None, - use_qk_l2norm_in_kernel: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - initial_dtype = query.dtype - if use_qk_l2norm_in_kernel: - query = l2norm(query, dim=-1, eps=1e-6) - key = l2norm(key, dim=-1, eps=1e-6) - - if query.shape[2] != value.shape[2]: - repeat_factor = value.shape[2] // query.shape[2] - query = query.repeat_interleave(repeat_factor, dim=2) - key = key.repeat_interleave(repeat_factor, dim=2) - - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - - batch_size, num_heads, sequence_length, _ = key.shape - v_head_dim = value.shape[-1] - if scale is None: - scale = 1 / (query.shape[-1] ** 0.5) - query = query * scale - - core_attn_out = torch.empty( - batch_size, - num_heads, - sequence_length, - v_head_dim, - dtype=value.dtype, - ) - last_recurrent_state = initial_state.to(value) - - for token_idx in range(sequence_length): - q_t = query[:, :, token_idx] - k_t = key[:, :, token_idx] - v_t = value[:, :, token_idx] - g_t = g[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1) - beta_t = beta[:, :, token_idx].unsqueeze(-1) - - last_recurrent_state = last_recurrent_state * g_t - kv_mem = (last_recurrent_state * k_t.unsqueeze(-2)).sum(dim=-1) - delta = (v_t - kv_mem) * beta_t - last_recurrent_state = last_recurrent_state + delta.unsqueeze( - -1 - ) * k_t.unsqueeze(-2) - core_attn_out[:, :, token_idx] = (last_recurrent_state * q_t.unsqueeze(-2)).sum( - dim=-1 - ) - - core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) - return core_attn_out, last_recurrent_state - - -def gdn_gating( - A_log: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - dt_bias: torch.Tensor, - beta: float = 1.0, - threshold: float = 20.0, -) -> tuple[torch.Tensor, torch.Tensor]: - softplus_x = F.softplus(a.float() + dt_bias.float(), beta=beta, threshold=threshold) - g = -torch.exp(A_log.float()) * softplus_x - beta_output = torch.sigmoid(b.float()).to(dtype=b.dtype) - return g, beta_output - - -def chunk_gated_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - *, - initial_state: torch.Tensor, - scale: float | None = None, - cu_seqlens: torch.Tensor, - use_qk_l2norm_in_kernel: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - output = torch.empty_like(v) - state_dtype = initial_state.dtype - chunk_size = 128 - sequence_bounds = [ - ( - seq_idx, - int(cu_seqlens[seq_idx].item()), - int(cu_seqlens[seq_idx + 1].item()), - ) - for seq_idx in range(len(cu_seqlens) - 1) - ] - chunk_eye = torch.eye(chunk_size, dtype=torch.float32) - num_sequences = len(sequence_bounds) - num_value_heads = v.shape[2] - value_head_dim = v.shape[3] - key_head_dim = k.shape[3] - final_state = torch.empty( - (num_sequences, num_value_heads, value_head_dim, key_head_dim), - dtype=state_dtype, - ) - - for seq_idx, begin, end in sequence_bounds: - q_seq = q[:, begin:end] - k_seq = k[:, begin:end] - v_seq = v[:, begin:end] - g_seq = g[:, begin:end] - beta_seq = beta[:, begin:end] - - initial_dtype = q_seq.dtype - if use_qk_l2norm_in_kernel: - q_seq = l2norm(q_seq, dim=-1, eps=1e-6) - k_seq = l2norm(k_seq, dim=-1, eps=1e-6) - - num_qk_heads = q_seq.shape[2] - num_value_heads = v_seq.shape[2] - if num_qk_heads != num_value_heads: - repeat_factor = num_value_heads // num_qk_heads - q_seq = q_seq.repeat_interleave(repeat_factor, dim=2) - k_seq = k_seq.repeat_interleave(repeat_factor, dim=2) - - q_seq, k_seq, v_seq, beta_seq, g_seq = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (q_seq, k_seq, v_seq, beta_seq, g_seq) - ] - seq_batch_size, num_heads, seq_len, qk_head_dim = q_seq.shape - value_head_dim = v_seq.shape[-1] - - if scale is None: - scale = 1 / (qk_head_dim**0.5) - - q_seq = q_seq * scale - - seq_state = initial_state[seq_idx : seq_idx + 1].to(v_seq) - seq_output = torch.empty( - seq_batch_size, - num_heads, - seq_len, - value_head_dim, - dtype=v_seq.dtype, - ) - - for chunk_start in range(0, seq_len, chunk_size): - chunk_end = min(chunk_start + chunk_size, seq_len) - q_chunk = q_seq[:, :, chunk_start:chunk_end] - k_chunk = k_seq[:, :, chunk_start:chunk_end] - v_chunk = v_seq[:, :, chunk_start:chunk_end] - beta_chunk = beta_seq[:, :, chunk_start:chunk_end] - g_chunk = g_seq[:, :, chunk_start:chunk_end] - chunk_len = chunk_end - chunk_start - - cum_g = g_chunk.cumsum(dim=-1) - exp_cum_g = cum_g.exp() - decay = (cum_g.unsqueeze(-1) - cum_g.unsqueeze(-2)).exp() - - interaction = (k_chunk * beta_chunk.unsqueeze(-1)) @ k_chunk.transpose( - -1, -2 - ) - interaction = torch.tril(interaction * decay, diagonal=-1) - system = interaction + chunk_eye[:chunk_len, :chunk_len] - - solved_values = torch.linalg.solve_triangular( - system, - v_chunk * beta_chunk.unsqueeze(-1), - upper=False, - ) - solved_keys = torch.linalg.solve_triangular( - system, - (k_chunk * beta_chunk.unsqueeze(-1)) * exp_cum_g.unsqueeze(-1), - upper=False, - ) - - incoming_memory = torch.einsum("bhvk,bhck->bhcv", seq_state, solved_keys) - transformed_values = solved_values - incoming_memory - - # Each chunk contributes both from the incoming recurrent state and - # from its own in-chunk interactions. - inter_chunk = torch.einsum( - "bhvk,bhck->bhcv", - seq_state, - q_chunk * exp_cum_g.unsqueeze(-1), - ) - intra_chunk = torch.tril((q_chunk @ k_chunk.transpose(-1, -2)) * decay) - seq_output[:, :, chunk_start:chunk_end] = ( - inter_chunk + intra_chunk @ transformed_values - ) - - # Carry the recurrent state forward to the next chunk boundary. - end_decay = (cum_g[:, :, -1:] - cum_g).exp().unsqueeze(-1) - decayed_keys = k_chunk * end_decay - seq_state = seq_state * exp_cum_g[:, :, -1, None, None] + torch.einsum( - "bhcv,bhck->bhvk", transformed_values, decayed_keys - ) - - output[0, begin:end].copy_( - seq_output.transpose(1, 2).contiguous().to(initial_dtype).squeeze(0) - ) - final_state[seq_idx].copy_(seq_state.squeeze(0).to(state_dtype).contiguous()) - - return output, final_state diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 2aef3337577..8c5a6355803 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -21,6 +21,9 @@ from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +if current_platform.is_xpu(): + from vllm._xpu_ops import xpu_ops + logger = init_logger(__name__) TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0")) @@ -790,28 +793,52 @@ def selective_scan_fn( if C.dim() == 2 and query_start_loc is not None: C = C.unsqueeze(0) - ops.selective_scan_fwd( - u, - delta, - A, - B, - C, - D, - z, - delta_bias, - delta_softplus, - query_start_loc, - cache_indices, - has_initial_state, - ssm_states, - null_block_id, - block_size, - block_idx_first_scheduled_token, - block_idx_last_scheduled_token, - initial_state_idx, - cu_chunk_seqlen, - last_chunk_indices, - ) + if current_platform.is_xpu(): + xpu_ops.selective_scan_fwd( + u, + delta, + A, + B, + C, + D, + z, + delta_bias, + delta_softplus, + query_start_loc, + cache_indices, + has_initial_state, + ssm_states, + null_block_id, + block_size, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + cu_chunk_seqlen, + last_chunk_indices, + ) + else: + ops.selective_scan_fwd( + u, + delta, + A, + B, + C, + D, + z, + delta_bias, + delta_softplus, + query_start_loc, + cache_indices, + has_initial_state, + ssm_states, + null_block_id, + block_size, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + cu_chunk_seqlen, + last_chunk_indices, + ) if z is None: return delta # output written inplace to delta diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index 02267b8f682..23aa3210179 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import cached_property from typing import Any, Union import torch @@ -168,6 +169,12 @@ class BitsAndBytesConfig(QuantizationConfig): return None +class BitsAndBytesWeightParameter(torch.nn.Parameter): + @cached_property + def dtype(self) -> torch.dtype: + return torch.get_default_dtype() + + def is_layer_skipped_bnb(prefix: str, llm_int8_skip_modules: list[str]): # Split the prefix into its dot-separated components components = prefix.split(".") @@ -246,7 +253,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase): "The input size is not aligned with the quantized weight shape." ) - qweight = torch.nn.Parameter( + qweight = BitsAndBytesWeightParameter( torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8), requires_grad=False, ) 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 2d629d73edd..2a98d444afd 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 @@ -29,6 +29,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso ) from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa WNA16_SUPPORTED_TYPES_MAP, + WNA16_ZP_SUPPORTED_TYPES_MAP, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, @@ -56,9 +57,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.weight_quant = weight_quant self.input_quant = input_quant - assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" - ) + self.symmetric = weight_quant.symmetric # Extract properties from weight_quant self.num_bits = weight_quant.num_bits self.packed_factor = 32 // weight_quant.num_bits @@ -66,7 +65,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): self.group_size = weight_quant.group_size self.actorder = weight_quant.actorder - self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + self.quant_type = ( + WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + if self.symmetric + else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits] + ) + self.marlin_input_dtype = get_marlin_input_dtype(layer_name) if self.num_bits == 4: @@ -82,7 +86,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now." ) - weight_key = QuantKey(self.quant_type, scale) + weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) # Select WNA16 MoE backend via oracle. self.wna16_backend, self.experts_cls = select_wna16_moe_backend( @@ -103,15 +107,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): Get the shape of the weight based on the weight name, number of experts hidden size, intermediate size per partition, number of groups for w2, and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 - for weight scales. + for weight scales/zero_points. """ - if weight_name == "w13_scale": + if weight_name in ("w13_scale", "w13_zp"): assert num_groups_w13 is not None, ( - "num_groups_w13 must be provided for weight scales" + "num_groups_w13 must be provided for weight scales/zero_points" ) - if weight_name == "w2_scale": + if weight_name in ("w2_scale", "w2_zp"): assert num_groups_w2 is not None, ( - "num_groups_w2 must be provided for weight scales" + "num_groups_w2 must be provided for weight scales/zero_points" ) w13_num_shards = 2 if self.moe.is_act_and_mul else 1 is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM @@ -140,6 +144,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w13_num_shards * intermediate_size_per_partition, ), }, + "w13_zp": { + "Marlin": ( + num_experts, + num_groups_w13, + w13_num_shards + * intermediate_size_per_partition + // self.packed_factor, + ), + }, "w2_weight": { "Flashinfer": ( num_experts, @@ -156,6 +169,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): "Flashinfer": (num_experts, hidden_size, num_groups_w2), "Marlin": (num_experts, num_groups_w2, hidden_size), }, + "w2_zp": { + "Marlin": ( + num_experts, + num_groups_w2, + hidden_size // self.packed_factor, + ), + }, } backend_key = "Flashinfer" if is_flashinfer else "Marlin" return shape_map[weight_name][backend_key] @@ -263,6 +283,39 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w2_scale, extra_weight_attrs) set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + if not self.symmetric: + w13_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w13_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w13=num_groups_w13, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_zero_point", w13_zp) + set_weight_attrs(w13_zp, extra_weight_attrs) + + w2_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w2_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w2=num_groups_w2, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_zero_point", w2_zp) + set_weight_attrs(w2_zp, extra_weight_attrs) + w2_weight_shape = torch.nn.Parameter( torch.empty(num_experts, 2), requires_grad=False ) @@ -334,8 +387,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_g_idx_processed, w13_g_idx_sort_indices, w2_g_idx_sort_indices, - _, # w13_qzeros - _, # w2_qzeros + w13_qzeros, + w2_qzeros, w13_input_global_scale, w2_input_global_scale, _, # w13_bias @@ -351,6 +404,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): 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), ) # Replace common parameters @@ -359,6 +414,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) + if not self.symmetric: + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) + # Marlin-specific parameters (not needed for Flashinfer) if not is_flashinfer: replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) @@ -417,6 +476,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_scale=layer.w2_weight_scale, group_size=self.group_size, num_bits=self.num_bits, + w1_zp=getattr(layer, "w13_weight_zero_point", None), + w2_zp=getattr(layer, "w2_weight_zero_point", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 12bb07a4022..e4d27efe370 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -810,8 +810,8 @@ class HummingMoEMethod(FusedMoEMethodBase): param = torch.nn.Parameter(tensor, requires_grad=False) setattr(layer, name, param) - layer.weight_schemas[sublayer_name] = weight_schema - layer.input_schemas[sublayer_name] = input_schema + layer.weight_schemas[sublayer_name] = weight_schema + layer.input_schemas[sublayer_name] = input_schema # force requant (origin quant setting -> fp16/bf16 -> new_quant setting) assert isinstance(weight_schema, HummingWeightSchema) @@ -865,6 +865,7 @@ class HummingMoEMethod(FusedMoEMethodBase): # use moe modular experts: HummingIndexedExperts | HummingGroupedExperts + layer.ensure_moe_quant_config_init() assert self.moe_quant_config is not None if get_humming_moe_gemm_type() == "indexed": experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config) diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index 35e0b4533f4..d7fa6cf2633 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -158,11 +158,6 @@ class QuantFP8(CustomOp): if use_aiter_per_token_quant: return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale) - # Fallback to native implementation for group quantization. - if self.is_group_quant: - assert scale is None, "Dynamic group quantization does not use scale" - return self._quantize_group_native(x) - # Fallback to CUDA implementation return self.forward_cuda(x, scale, scale_ub) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 471febab044..ee4b455ddc4 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -13,7 +13,6 @@ from vllm.model_executor.layers.fused_moe import ( RoutedExperts, SharedExperts, ) -from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, int4_w4a16_moe_quant_config, @@ -367,16 +366,13 @@ class MoeWNA16Method(FusedMoEMethodBase): ) -> torch.Tensor: from vllm.model_executor.layers.fused_moe import fused_experts - assert layer.activation == MoEActivation.SILU, ( - f"Only SiLU activation is supported, not {layer.activation}." - ) - return fused_experts( x, layer.w13_qweight, layer.w2_qweight, topk_weights=topk_weights, topk_ids=topk_ids, + activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 8b20c13a97f..71442fb1add 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -575,7 +575,9 @@ def per_token_group_quant_fp8( # prefer CUDA/XPU kernel if available # TODO(bnell): this causes some fp8 moe test to fail. - if current_platform.is_cuda() and x.is_contiguous(): + if ( + current_platform.is_cuda_alike() or current_platform.is_xpu() + ) and x.is_contiguous(): torch.ops._C.per_token_group_fp8_quant( x, x_q, @@ -590,12 +592,6 @@ def per_token_group_quant_fp8( ) return x_q, x_s - if current_platform.is_xpu() and x.is_contiguous(): - torch.ops._C.per_token_group_fp8_quant( - x, x_q, x_s, group_size, eps, fp8_min, fp8_max, use_ue8m0 - ) - return x_q, x_s - # TRITON FALLBACK M = x.numel() // group_size N = group_size @@ -670,8 +666,7 @@ def per_token_group_quant_fp8_packed_for_deepgemm( ) assert x.stride(-1) == 1, "`x` groups must be contiguous" - finfo = torch.finfo(dtype) - fp8_min, fp8_max = finfo.min, finfo.max + fp8_min, fp8_max = get_fp8_min_max() # compute DeepGEMM-style packed scale tensor shape. hidden_dim = x.shape[-1] @@ -687,10 +682,10 @@ def per_token_group_quant_fp8_packed_for_deepgemm( dtype=torch.int32, ) - # CUDA kernel path only (DeepGEMM + E8M0 is CUDA-specific). - assert current_platform.is_cuda(), ( - "per_token_group_quant_fp8_packed_for_deepgemm is only valid on CUDA " - "platforms using DeepGEMM." + # Native kernel (libtorch stable); used with DeepGEMM on CUDA and + # available on ROCm for the same packed UE8M0 scale layout. + assert current_platform.is_cuda_alike(), ( + "per_token_group_quant_fp8_packed_for_deepgemm requires a CUDA or ROCm GPU." ) x_contiguous = x.contiguous() diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index a98e29ffd57..eac6b11b219 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -235,8 +235,8 @@ def per_token_group_quant_int8( device=x.device, dtype=torch.float32, ) - # prefer CUDA kernel if available - if current_platform.is_cuda(): + # Prefer native stable kernel on CUDA/ROCm when available. + if current_platform.is_cuda_alike(): torch.ops._C.per_token_group_quant_int8( x, x_q, x_s, group_size, eps, float(int8_min), float(int8_max) ) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index eca04eed74b..19f2605dc48 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -429,6 +429,30 @@ def maybe_warn_marlin_atomic_add(device, dtype): ) +def moe_packed_to_marlin_zero_points( + q_zp_packed: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +): + """Convert compressed-tensors packed zero points to Marlin format. + + Unlike AWQ, compressed-tensors uses standard bit packing without + interleaving, so we just unpack and apply Marlin permutation directly. + """ + num_experts = q_zp_packed.shape[0] + output = torch.empty( + (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]), + device=q_zp_packed.device, + dtype=q_zp_packed.dtype, + ) + for e in range(num_experts): + q_zp = unpack_cols(q_zp_packed[e], num_bits, size_k, size_n) + output[e] = marlin_zero_points(q_zp, size_k, size_n, num_bits, is_a_8bit) + return output + + def maybe_warn_marlin_atomic_add_env(): if torch.compiler.is_dynamo_compiling(): return diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 947e428ca82..ba1016a4fb9 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -181,6 +181,13 @@ kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True) kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32)) kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True) +kInt4StaticAsym = QuantKey( + scalar_types.uint4, scale=kInt4StaticGroupScale, symmetric=False +) +kInt4Static32Asym = QuantKey( + scalar_types.uint4, scale=kInt4Static32GroupScale, symmetric=False +) + kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True) kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True) diff --git a/vllm/model_executor/model_loader/bitsandbytes_loader.py b/vllm/model_executor/model_loader/bitsandbytes_loader.py index 4d9ae267e08..bc2504b09c5 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -744,6 +744,29 @@ class BitsAndBytesModelLoader(BaseModelLoader): stacked_quant_state_dict[quant_param_name][shard_index] = quant_state_dict[ non_stacked_param_name ] + + # repeat k_proj for v_proj for k_eq_v models (e.g. Gemma4) + config = getattr(model, "config", None) + if config is not None: + text_config = config.get_text_config() + if getattr(text_config, "attention_k_eq_v", False): + shard_packed = { + name + for name, subs in self.modules_mapping.packed_mapping.items() + if len(subs) == 3 + } + for param_name, shards in stacked_quant_state_dict.items(): + is_target = ( + isinstance(shards, dict) + and len(shards) == 2 + and any( + param_name.endswith(f"{p}.weight") for p in shard_packed + ) + ) + if is_target: + assert 1 in shards and 2 not in shards + shards[2] = shards[1] + return stacked_quant_state_dict def _bind_quant_states_to_params( diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index b546040b741..8f593ab640c 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -16,7 +16,7 @@ reason about temporal order. import math from collections.abc import Iterable, Mapping, Sequence -from typing import Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, Literal import numpy as np import torch @@ -41,6 +41,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.transformers.utils import recursive_replace_linear from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -71,6 +72,7 @@ from .interfaces import ( SupportsLoRA, SupportsMultiModal, SupportsPP, + SupportsQuant, ) from .utils import ( AutoWeightsLoader, @@ -79,6 +81,9 @@ from .utils import ( maybe_prefix, ) +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization import QuantizationConfig + logger = init_logger(__name__) # Video constants — match transformers Gemma4VideoProcessor defaults. @@ -514,6 +519,25 @@ class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): + def _apply_hf_processor_text_only( + self, + prompt_text: str, + tokenization_kwargs: Mapping[str, object], + ) -> list[int]: + # Bypass the HF processor and tokenize directly. The HF + # processor expands multimodal placeholders (<|video|>, etc.) + # via get_text_with_replacements, which raises StopIteration + # when the prompt contains placeholders without matching data. + # The text-only path only needs token IDs, so the tokenizer + # alone is sufficient. + processor = self.info.get_hf_processor() + text_inputs = processor.tokenizer([prompt_text], **tokenization_kwargs) + input_ids = text_inputs["input_ids"] + if not isinstance(input_ids, list): + input_ids = input_ids.tolist() + (prompt_ids,) = input_ids + return prompt_ids + def _call_hf_processor( self, prompt: str, @@ -872,6 +896,9 @@ class Gemma4MultimodalEmbedder(nn.Module): self, multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig, text_config: Gemma4TextConfig, + *, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() @@ -895,6 +922,8 @@ class Gemma4MultimodalEmbedder(nn.Module): embedding_dim, self.text_hidden_size, bias=False, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "embedding_projection"), ) def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: @@ -917,6 +946,7 @@ class Gemma4MultimodalEmbedder(nn.Module): class Gemma4ForConditionalGeneration( nn.Module, SupportsMultiModal, + SupportsQuant, SupportsPP, SupportsLoRA, SupportsEagle3, @@ -936,11 +966,14 @@ class Gemma4ForConditionalGeneration( # Maps checkpoint prefixes to vLLM module paths. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ - "model.embed_audio.": "embed_audio.", - "model.embed_vision.": "embed_vision.", - "model.language_model.": "language_model.model.", - "model.vision_tower.": "vision_tower.", + # vision tower + "model.vision_tower": "vision_tower", + "model.embed_vision": "embed_vision", + # audio tower "model.audio_tower.": "audio_tower.", + "model.embed_audio.": "embed_audio.", + # backbone + "model.language_model.": "language_model.model.", "lm_head.": "language_model.lm_head.", "model": "language_model.model", } @@ -959,7 +992,15 @@ class Gemma4ForConditionalGeneration( with self._mark_tower_model(vllm_config, {"image", "video"}): self.vision_tower = AutoModel.from_config(config=config.vision_config) self.embed_vision = Gemma4MultimodalEmbedder( - config.vision_config, config.text_config + config.vision_config, + config.text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "embed_vision"), + ) + recursive_replace_linear( + self.vision_tower, + quant_config, + prefix=maybe_prefix(prefix, "vision_tower"), ) # ---- Audio tower (variants with audio_config) ---- @@ -972,7 +1013,15 @@ class Gemma4ForConditionalGeneration( # position embeddings, softcap, gradient_clipping). self.audio_tower.post_init() self.embed_audio = Gemma4MultimodalEmbedder( - config.audio_config, config.text_config + config.audio_config, + config.text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "embed_audio"), + ) + recursive_replace_linear( + self.audio_tower, + quant_config, + prefix=maybe_prefix(prefix, "audio_tower"), ) else: self.audio_tower = None @@ -1153,6 +1202,7 @@ class Gemma4ForConditionalGeneration( vt = self.vision_tower vision_cfg = self.config.vision_config pooling_k2 = vision_cfg.pooling_kernel_size**2 + target_dtype = self.language_model.model.embed_tokens.weight.dtype # Concurrent requests with different image resolutions may # arrive as a list of per-image tensors, while same-resolution @@ -1193,7 +1243,11 @@ class Gemma4ForConditionalGeneration( ) pad_tensor = (pp_tensor == -1).all(dim=-1) - inputs_embeds = vt.patch_embedder(pv_tensor, pp_tensor, pad_tensor) + inputs_embeds = vt.patch_embedder( + pv_tensor, + pp_tensor, + pad_tensor, + ).to(target_dtype) encoder_outputs = vt.encoder( inputs_embeds=inputs_embeds, attention_mask=~pad_tensor, @@ -1230,7 +1284,9 @@ class Gemma4ForConditionalGeneration( all_valid_states[orig_idx] = valid_states valid_lens[orig_idx] = valid_states.shape[0] - target_dtype = self.embed_vision.embedding_projection.weight.dtype + # Use embed_tokens dtype as compute dtype; embedding_projection.weight + # may be uint8 under BnB 4-bit, which would corrupt the cast. + target_dtype = self.language_model.model.embed_tokens.weight.dtype # Project all images in a single batched call. flat_valid_states = torch.cat(all_valid_states, dim=0).to(target_dtype) @@ -1273,7 +1329,7 @@ class Gemma4ForConditionalGeneration( vt = self.vision_tower vision_cfg = self.config.vision_config pooling_k2 = vision_cfg.pooling_kernel_size**2 - target_dtype = self.embed_vision.embedding_projection.weight.dtype + target_dtype = self.language_model.model.embed_tokens.weight.dtype if isinstance(frame_counts, torch.Tensor): fc_list = frame_counts.tolist() @@ -1301,7 +1357,11 @@ class Gemma4ForConditionalGeneration( pp_chunk = pixel_position_ids[i : i + max_batch_size] pad_chunk = padding_positions[i : i + max_batch_size] - inputs_embeds = vt.patch_embedder(pv_chunk, pp_chunk, pad_chunk) + inputs_embeds = vt.patch_embedder( + pv_chunk, + pp_chunk, + pad_chunk, + ).to(target_dtype) encoder_outputs = vt.encoder( inputs_embeds=inputs_embeds, attention_mask=~pad_chunk, diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index 9251b14728e..a8786f677ba 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -26,7 +26,7 @@ import os from collections.abc import Callable, Iterable, Mapping, Sequence -from typing import Annotated, Any, Literal, TypeAlias +from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias import torch from torch import nn @@ -75,6 +75,9 @@ from .utils import AutoWeightsLoader, cast_overflow_tensors, maybe_prefix CPU_DEVICE = torch.device("cpu") +if TYPE_CHECKING: + from vllm.transformers_utils.processors.minicpmo import MiniCPMOProcessor + if os.getenv("USE_FLAGOS") == "1": import flag_gems @@ -173,9 +176,64 @@ MiniCPMOAudioInputs: TypeAlias = ( def _minicpmo_field_config(hf_inputs: Mapping[str, torch.Tensor]): + audio_features = hf_inputs.get("audio_features") + audio_feature_lens = hf_inputs.get("audio_feature_lens") + + # For multi-chunk audio (>30s), audio_features has one item per chunk + # (total_chunks) while audio_feature_lens has one item per audio (N). + # Use flat to group audio_features by audio so both fields + # share the same batch size (N). + audio_features_cfg = MultiModalFieldConfig.batched("audio") + + if audio_features is not None and audio_feature_lens is not None: + num_features = ( + len(audio_features) + if isinstance(audio_features, (list, tuple)) + else audio_features.shape[0] + ) + num_audios = ( + len(audio_feature_lens) + if isinstance(audio_feature_lens, (list, tuple)) + else audio_feature_lens.shape[0] + ) + + if num_features > num_audios: + # Compute the number of chunks belonging to each audio + chunks_per_audio: list[int] = [] + for lens in audio_feature_lens: + if isinstance(lens, torch.Tensor): + chunks_per_audio.append(lens.numel()) + else: + chunks_per_audio.append(1) + + # When audio_feature_lens is padded (e.g. from batched HF + # processor output), numel() over-counts. Fall back to + # counting non-zero entries so the sizes sum to num_features. + if sum(chunks_per_audio) != num_features: + chunks_per_audio = [] + for lens in audio_feature_lens: + if isinstance(lens, torch.Tensor): + n = int((lens != 0).sum()) + chunks_per_audio.append(max(n, 1)) + else: + chunks_per_audio.append(1) + + # Use flat (not flat_from_sizes) because audio_features + # is list[Tensor] with variable-length chunks (post-unpad). + slice_idxs = [0] + for n in chunks_per_audio: + slice_idxs.append(slice_idxs[-1] + n) + audio_features_cfg = MultiModalFieldConfig.flat( + "audio", + [ + slice(slice_idxs[i], slice_idxs[i + 1]) + for i in range(len(chunks_per_audio)) + ], + ) + return dict( **_minicpmv_field_config(hf_inputs), - audio_features=MultiModalFieldConfig.batched("audio"), + audio_features=audio_features_cfg, audio_feature_lens=MultiModalFieldConfig.batched("audio"), audio_embeds=MultiModalFieldConfig.batched("audio"), ) @@ -215,6 +273,38 @@ class MiniCPMOMultiModalDataParser(MiniCPMVMultiModalDataParser): class MiniCPMOProcessingInfo(MiniCPMVProcessingInfo): audio_pattern = "()" + def get_hf_processor(self, **kwargs: object) -> "MiniCPMOProcessor": + """Get vendored MiniCPMOProcessor for multimodal (image+audio) inputs. + + Creates a vendored processor that reuses the HF image processor, + feature extractor, and tokenizer; applies the correct audio pooling + configuration; and converts numpy arrays in the image processor to + lists for serialization compatibility. The returned processor is + compatible with Transformers v5. + """ + import numpy as np + + hf_processor = self.ctx.get_hf_processor(**kwargs) + + from vllm.transformers_utils.processors.minicpmo import MiniCPMOProcessor + + # Create vendored processor with correct configuration + vendored_processor = MiniCPMOProcessor( + image_processor=hf_processor.image_processor, + feature_extractor=hf_processor.feature_extractor, + tokenizer=hf_processor.tokenizer, + pool_step=self.get_default_audio_pool_step(), + ) + + # Convert numpy arrays in image processor to lists for serialization + image_processor = vendored_processor.image_processor + for attr in ("mean", "std"): + val = getattr(image_processor, attr, None) + if val is not None and isinstance(val, np.ndarray): + setattr(image_processor, attr, val.tolist()) + + return vendored_processor + def get_data_parser(self): return MiniCPMOMultiModalDataParser( target_sr=self.get_default_audio_sampling_rate(), @@ -364,11 +454,23 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing # Avoid padding since we need the output for each audio to be # independent of other audios for the cache to work correctly + # Flatten audio_feature_lens (list of tensors of any + # dimensionality, one per audio, each containing per-chunk + # lengths) into a flat list of integer lengths so there is + # one length per chunk, matching the first dimension of + # audio_features. Using flatten() handles 0-D, 1-D, and + # higher-dimensional tensors uniformly. + flat_feature_lens: list[int] = [] + for lens in audio_inputs["audio_feature_lens"]: + if isinstance(lens, torch.Tensor): + flat_feature_lens.extend(lens.flatten().tolist()) + else: + flat_feature_lens.append(int(lens)) unpadded_audio_features = [ - feat[:, :feature_len] - for feat, feature_len in zip( + feat[:, :length] + for feat, length in zip( audio_inputs["audio_features"], - audio_inputs["audio_feature_lens"], + flat_feature_lens, ) ] audio_inputs["audio_features"] = unpadded_audio_features diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index af5f5651bbf..001329b1762 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -545,6 +545,14 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo): def get_hf_processor(self, **kwargs: object): hf_processor = self.ctx.get_hf_processor(**kwargs) + from vllm.transformers_utils.processors.minicpmv import MiniCPMVProcessor + + vendored_processor = MiniCPMVProcessor( + image_processor=hf_processor.image_processor, + tokenizer=hf_processor.tokenizer, + ) + hf_processor = vendored_processor + # NumPy arrays are considered as Iterable but not Sequence in # https://github.com/huggingface/transformers/blob/main/src/transformers/image_transforms.py#L428 image_processor = hf_processor.image_processor # type: ignore diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index 04d6de28efd..dbf0a084f78 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.linear import ( ReplicatedLinear, RowParallelLinear, ) +from vllm.model_executor.models.utils import maybe_prefix from vllm.transformers_utils.config import is_rope_parameters_nested if TYPE_CHECKING: @@ -227,6 +228,34 @@ def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm: return RMSNorm(**kwargs) +def recursive_replace_linear( + model: nn.Module, + quant_config: "QuantizationConfig | None", + prefix: str = "", +): + """Recursively replace linear modules in the model as needed.""" + + def _recursive_replace(module: nn.Module, prefix: str): + for child_name, child_module in module.named_children(): + new_module = child_module + qual_name = maybe_prefix(prefix, child_name) + # Replace modules as needed + if isinstance(child_module, nn.Linear): + style = "replicate" + new_module = replace_linear_class( + child_module, + style, + quant_config, + prefix=qual_name, + ) + else: + _recursive_replace(child_module, prefix=qual_name) + if new_module is not child_module: + setattr(module, child_name, new_module) + + _recursive_replace(model, prefix=prefix) + + def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module): logger.debug("%s: %s -> %s", name, old_module, new_module) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 28836a2b143..fb724fbe2f1 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -30,7 +30,6 @@ from vllm.model_executor.layers.mhc import ( MHCPreOp, ) from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -48,9 +47,9 @@ from vllm.model_executor.models.utils import ( ) from vllm.models.deepseek_v4.attention import ( DeepseekV4Indexer, - DeepseekV4MLAModules, - DeepseekV4MultiHeadLatentAttentionWrapper, + DeepseekV4MLA, ) +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.import_utils import has_tilelang @@ -314,26 +313,13 @@ class DeepseekV4Attention(nn.Module): self.rope_parameters = config.rope_scaling - # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it) - rope_parameters = config.rope_parameters - rope_parameters["rope_theta"] = ( - config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta - ) - if config.rope_parameters["rope_type"] != "default": - config.rope_parameters["rope_type"] = ( - "deepseek_yarn" - if config.rope_parameters.get("apply_yarn_scaling", True) - else "deepseek_llama_scaling" - ) - rope_parameters["mscale"] = 0 # Disable mscale - rope_parameters["mscale_all_dim"] = 0 # Disable mscale - rope_parameters["is_deepseek_v4"] = True - rope_parameters["rope_dim"] = self.rope_head_dim - self.rotary_emb = get_rope( - self.head_dim, - max_position=self.max_position_embeddings, - rope_parameters=rope_parameters, - is_neox_style=False, + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + compress_ratio=self.compress_ratio, ) self.indexer = None @@ -351,7 +337,17 @@ class DeepseekV4Attention(nn.Module): prefix=f"{prefix}.indexer", ) - mla_modules = DeepseekV4MLAModules( + self.mla_attn = DeepseekV4MLA( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, vllm_config=vllm_config, fused_wqa_wkv=self.fused_wqa_wkv, q_norm=self.q_norm, @@ -365,19 +361,6 @@ class DeepseekV4Attention(nn.Module): indexer_rotary_emb=self.rotary_emb, topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, - ) - self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper( - hidden_size=self.hidden_size, - num_heads=self.n_local_heads, - head_dim=self.head_dim, - scale=self.softmax_scale, - qk_nope_head_dim=self.nope_head_dim, - qk_rope_head_dim=self.rope_head_dim, - v_head_dim=self.head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.head_dim, - o_lora_rank=self.o_lora_rank, - mla_modules=mla_modules, window_size=self.window_size, compress_ratio=self.compress_ratio, cache_config=vllm_config.cache_config, @@ -618,7 +601,7 @@ class DeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps # Three aux streams: one per non-default input GEMM in - # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute + # DeepseekV4MLA.attn_gemm_parallel_execute # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. # Disable them on ROCm because of hang issues. diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 4fae5dc0529..55cb3d94ba6 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -5,7 +5,6 @@ DeepseekV4 MLA Attention Layer """ from collections.abc import Callable -from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast import torch @@ -38,9 +37,8 @@ from vllm.config import ( get_current_vllm_config, ) from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.forward_context import ForwardContext, get_forward_context +from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization import QuantizationConfig @@ -90,46 +88,7 @@ def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]": return DeepseekV4FlashMLASparseImpl -@dataclass -class DeepseekV4MLAModules: - """Modules used in DeepseekV4 MLA.""" - - vllm_config: VllmConfig - fused_wqa_wkv: torch.nn.Module - q_norm: torch.nn.Module - wq_b: torch.nn.Module - kv_norm: torch.nn.Module - wo_a: torch.nn.Module - wo_b: torch.nn.Module - attn_sink: torch.nn.Module - rotary_emb: torch.nn.Module - indexer: torch.nn.Module | None - indexer_rotary_emb: torch.nn.Module - topk_indices_buffer: torch.Tensor | None - aux_stream_list: list[torch.cuda.Stream] | None = None - - -# --8<-- [start:multi_head_latent_attention] -@PluggableLayer.register("deepseek_v4_multi_head_latent_attention") -class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): - """Pluggable MLA layer which allows OOT backends to add - custom implementations of the outer MLA layer (including rope & o_proj). - Note that currently oot platforms can still use CustomOp.register_oot to - replace MLA layer entirely, although we use PluggableLayer to register - this layer now. - - This class takes positions and hidden_states as input. - The input tensors can either contain prefill tokens or decode tokens. - The class does the following: - - 1. MLA Preprocess. - 2. Perform multi-head attention to prefill tokens and - multi-query attention to decode tokens separately. - 3. Return the output tensor. - """ - - # --8<-- [end:multi_head_latent_attention] - +class DeepseekV4MLA(nn.Module): def __init__( self, hidden_size: int, @@ -142,7 +101,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): q_lora_rank: int | None, kv_lora_rank: int, o_lora_rank: int | None, - mla_modules: DeepseekV4MLAModules, + vllm_config: VllmConfig, + fused_wqa_wkv: torch.nn.Module, + q_norm: torch.nn.Module, + wq_b: torch.nn.Module, + kv_norm: torch.nn.Module, + wo_a: torch.nn.Module, + wo_b: torch.nn.Module, + attn_sink: torch.nn.Module, + rotary_emb: torch.nn.Module, + indexer: torch.nn.Module | None, + indexer_rotary_emb: torch.nn.Module, + topk_indices_buffer: torch.Tensor | None, + aux_stream_list: list[torch.cuda.Stream] | None, window_size: int, compress_ratio: int | None, cache_config: CacheConfig | None = None, @@ -162,7 +133,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): self.prefix = prefix # Extract config from vllm_config - config = mla_modules.vllm_config.model_config.hf_config + config = vllm_config.model_config.hf_config tp_size = get_tensor_model_parallel_world_size() # DeepseekV4-specific attributes (num_heads is already TP-adjusted) @@ -173,12 +144,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): self.o_lora_rank = config.o_lora_rank # Store projection modules - self.fused_wqa_wkv = mla_modules.fused_wqa_wkv - self.q_norm = mla_modules.q_norm - self.wq_b = mla_modules.wq_b + self.fused_wqa_wkv = fused_wqa_wkv + self.q_norm = q_norm + self.wq_b = wq_b - self.kv_norm = mla_modules.kv_norm - self.wo_a = mla_modules.wo_a + self.kv_norm = kv_norm + self.wo_a = wo_a self._wo_a_act_quant = QuantFP8( static=False, @@ -188,7 +159,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): # Bypass packed-for-deepgemm path — we need FP32 scales (not packed # INT32) so fp8_einsum can handle layout transform internally. self._wo_a_act_quant.use_deep_gemm_supported = False - self.wo_b = mla_modules.wo_b + self.wo_b = wo_b # Pick fp8_einsum recipe based on GPU arch: # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128 @@ -198,11 +169,11 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) self._tma_aligned_scales = cap.major >= 10 - self.rotary_emb = mla_modules.rotary_emb - self.indexer_rotary_emb = mla_modules.indexer_rotary_emb - self.topk_indices_buffer = mla_modules.topk_indices_buffer + self.rotary_emb = rotary_emb + self.indexer_rotary_emb = indexer_rotary_emb + self.topk_indices_buffer = topk_indices_buffer - self.indexer = mla_modules.indexer + self.indexer = indexer # Per-head RMS normalization for Q (no learnable weights) self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False) @@ -216,7 +187,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): ) # Will be None on ROCm for now. - self.aux_stream_list = mla_modules.aux_stream_list + self.aux_stream_list = aux_stream_list # [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. @@ -243,7 +214,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): window_size=self.window_size, head_bytes=head_bytes, swa_cache_layer=self.swa_cache_layer, - attn_sink=mla_modules.attn_sink, # already padded with -inf + attn_sink=attn_sink, # already padded with -inf cache_config=cache_config, quant_config=quant_config, prefix=prefix, @@ -253,21 +224,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): # Mirror the inner layer's padded head count (single source of truth). self.padded_heads = self.mla_attn.padded_heads - # Register this layer in the compilation config's static forward context - # This allows the custom op to retrieve the layer during execution - compilation_config = mla_modules.vllm_config.compilation_config - # HACK - self.layer_name = prefix + ".deepseek_v4_multi_head_latent_attention" - if self.layer_name in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {self.layer_name}") - compilation_config.static_forward_context[self.layer_name] = self - # Create the compressor for layers with compress_ratio > 1; after # creating the DeepseekV4MLAAttention layer to get its cache. self.compressor = None if self.compress_ratio > 1: self.compressor = DeepseekCompressor( - vllm_config=mla_modules.vllm_config, + vllm_config=vllm_config, compress_ratio=self.compress_ratio, hidden_size=self.hidden_size, head_dim=self.head_dim, @@ -291,15 +253,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): device=hidden_states.device, ) - # @eager_break_during_capture: this is where the breakable - # cudagraph capture breaks (the attention op runs eagerly between - # captured graph segments). - deepseek_v4_attention( - hidden_states, - positions, - o_padded, - self.layer_name, - ) + # attention_impl is wrapped with @eager_break_during_capture: this is + # where the breakable cudagraph capture breaks (the attention op runs + # eagerly between captured graph segments). + self.attention_impl(hidden_states, positions, o_padded) o = o_padded[:, : self.n_local_heads, :] # Keep ROCm on the BF16 reference wo_a path util kernel ready. @@ -405,6 +362,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): return qr_kv, kv_score, indexer_kv_score, indexer_weights + @eager_break_during_capture def attention_impl( self, hidden_states: torch.Tensor, @@ -541,18 +499,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): ) -@eager_break_during_capture -def deepseek_v4_attention( - hidden_states: torch.Tensor, - positions: torch.Tensor, - out: torch.Tensor, - layer_name: str, -) -> None: - forward_context: ForwardContext = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - self.attention_impl(hidden_states, positions, out) - - class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): def __init__( self, diff --git a/vllm/models/deepseek_v4/common/rope.py b/vllm/models/deepseek_v4/common/rope.py new file mode 100644 index 00000000000..44ae3286eb2 --- /dev/null +++ b/vllm/models/deepseek_v4/common/rope.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepseekV4 rotary embedding initialization.""" + +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding + + +def build_deepseek_v4_rope( + config, + *, + head_dim: int, + rope_head_dim: int, + max_position_embeddings: int, + compress_ratio: int, +) -> RotaryEmbedding: + rope_parameters = config.rope_parameters + rope_parameters["rope_theta"] = ( + config.compress_rope_theta if compress_ratio > 1 else config.rope_theta + ) + if rope_parameters["rope_type"] != "default": + rope_parameters["rope_type"] = ( + "deepseek_yarn" + if rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + rope_parameters["mscale"] = 0 # Disable mscale + rope_parameters["mscale_all_dim"] = 0 # Disable mscale + rope_parameters["is_deepseek_v4"] = True + rope_parameters["rope_dim"] = rope_head_dim + return get_rope( + head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 30a7e6e747f..547048ab58f 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, MutableSequence, Sequence from itertools import islice import regex as re @@ -15,6 +15,7 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, @@ -23,6 +24,9 @@ from vllm.model_executor.kernels.mhc.tilelang import ( ) from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( fused_topk_bias, ) @@ -35,13 +39,12 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, @@ -54,9 +57,9 @@ from vllm.model_executor.models.utils import ( from vllm.model_executor.utils import set_weight_attrs from vllm.models.deepseek_v4.attention import ( DeepseekV4Indexer, - DeepseekV4MLAModules, - DeepseekV4MultiHeadLatentAttentionWrapper, + DeepseekV4MLA, ) +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.sequence import IntermediateTensors @@ -145,6 +148,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): hidden_size: int, intermediate_size: int, prefix: str = "", + num_logical_experts: int | None = None, ): super().__init__() self.prefix = prefix @@ -157,6 +161,12 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.intermediate_size = intermediate_size self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.num_logical_experts = ( + num_logical_experts if num_logical_experts is not None else num_experts + ) + + self.eplb_state = EplbLayerState() + weight_attrs = {"weight_loader": self.weight_loader} self.w13_weight = nn.Parameter( torch.zeros( @@ -207,10 +217,22 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None - def _map_global_expert_id(self, expert_id: int) -> int: - if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: - return -1 - return expert_id - self.experts_start_idx + # Register in the static forward context so the custom-op wrapper + # can look up this module by name from within a torch.compile graph. + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _map_global_expert_id(self, expert_id: int) -> list[int]: + """Return local (per-rank) slot offsets where logical expert + `expert_id` should land on this rank. + """ + physical_ids: list[int] = [] + for p in range(self.experts_start_idx, self.experts_end_idx): + if p % self.num_logical_experts == expert_id: + physical_ids.append(p - self.experts_start_idx) + return physical_ids def weight_loader( self, @@ -221,30 +243,38 @@ class DeepseekV4MegaMoEExperts(nn.Module): expert_id: int, return_success: bool = False, ) -> bool | None: - local_expert_id = self._map_global_expert_id(expert_id) - if local_expert_id == -1: + local_expert_ids = self._map_global_expert_id(expert_id) + if not local_expert_ids: return False if return_success else None - expert_data = param.data[local_expert_id] - if shard_id in ("w1", "w3"): - if "w13_" not in weight_name: - return False if return_success else None - shard_offset = 0 if shard_id == "w1" else self.intermediate_size - expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) - elif shard_id == "w2": - if "w2_" not in weight_name: - return False if return_success else None - else: - raise ValueError(f"Unsupported expert shard id: {shard_id}") + loaded_any = False + for local_expert_id in local_expert_ids: + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if "w13_" not in weight_name: + continue + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow( + 0, shard_offset, self.intermediate_size + ) + elif shard_id == "w2": + if "w2_" not in weight_name: + continue + else: + raise ValueError(f"Unsupported expert shard id: {shard_id}") - if expert_data.shape != loaded_weight.shape: - raise ValueError( - f"DeepSeek V4 MegaMoE expert weight shape mismatch for " - f"{weight_name}: parameter shard {tuple(expert_data.shape)} " - f"vs checkpoint {tuple(loaded_weight.shape)}" - ) - expert_data.copy_(loaded_weight) - return True if return_success else None + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{weight_name}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + loaded_any = True + + if return_success: + return loaded_any + return None @staticmethod def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: @@ -265,7 +295,9 @@ class DeepseekV4MegaMoEExperts(nn.Module): return self._check_runtime_supported() - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() w13_scale = deep_gemm.transform_sf_into_required_layout( self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), @@ -299,7 +331,9 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.w2_weight_scale = None def get_symm_buffer(self): - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() group = get_ep_group().device_group device = torch.accelerator.current_device_index() @@ -325,6 +359,52 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._symm_buffer_cache[key] = symm_buffer return symm_buffer + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, + ) + + def get_expert_weights(self) -> list[torch.Tensor]: + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + + def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: + """Return a (num_local_experts, -1) view with contiguous memory layout.""" + assert t.shape[0] == self.num_local_experts + if t.is_contiguous(): + return t.view(self.num_local_experts, -1) + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: + # scales have shape (E, M, N) with memory layout (E, N, M) + back = torch.transpose(t, 1, 2) + assert back.is_contiguous() + return back.view(self.num_local_experts, -1) + + raise AssertionError( + f"DSv4 EPLB {name}: non-contiguous expert tensor with " + f"unexpected layout shape={tuple(t.shape)} " + f"stride={tuple(t.stride())} dtype={t.dtype}" + ) + + return [ + _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), + _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), + _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), + _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), + ] + + def update_expert_map(self) -> None: + pass + def forward( self, hidden_states: torch.Tensor, @@ -340,29 +420,28 @@ class DeepseekV4MegaMoEExperts(nn.Module): f"but the symmetric buffer was sized for {self.max_num_tokens}." ) y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - self._run_mega_moe( - hidden_states, - topk_weights, - topk_ids, - y, - activation_clamp, - fast_math, - ) - return y - def _run_mega_moe( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - y: torch.Tensor, - activation_clamp: float | None, - fast_math: bool, - ) -> None: - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] + + # EPLB: map logical expert IDs to physical replicas and record load. + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + ) + prepare_megamoe_inputs( hidden_states, topk_weights, @@ -387,6 +466,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): activation_clamp=activation_clamp, fast_math=fast_math, ) + return y DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] @@ -494,17 +574,33 @@ class DeepseekV4MoE(nn.Module): self.ep_group = get_ep_group() self.ep_size = self.ep_group.world_size self.ep_rank = self.ep_group.rank_in_group - assert config.n_routed_experts % self.ep_size == 0 - self.n_local_experts = config.n_routed_experts // self.ep_size - self.experts_start_idx = self.ep_rank * self.n_local_experts - self.experts_end_idx = self.experts_start_idx + self.n_local_experts + eplb_config = vllm_config.parallel_config.eplb_config + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_routed_experts = config.n_routed_experts + self.n_shared_experts = config.n_shared_experts or 0 + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + assert self.n_physical_experts % self.ep_size == 0, ( + f"n_physical_experts={self.n_physical_experts} must be divisible by " + f"ep_size={self.ep_size}. Adjust num_redundant_experts." + ) + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + self.n_local_experts = self.n_local_physical_experts + self.experts_start_idx = self.physical_expert_start + self.experts_end_idx = self.physical_expert_end self.experts = DeepseekV4MegaMoEExperts( vllm_config, - num_experts=config.n_routed_experts, - num_local_experts=self.n_local_experts, - experts_start_idx=self.experts_start_idx, + num_experts=self.n_physical_experts, + num_local_experts=self.n_local_physical_experts, + experts_start_idx=self.physical_expert_start, + num_logical_experts=self.n_logical_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, @@ -524,6 +620,14 @@ class DeepseekV4MoE(nn.Module): self.experts_start_idx = self.tp_rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.n_redundant_experts = 0 + self.n_shared_experts = config.n_shared_experts or 0 + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_local_physical_experts = self.n_local_experts + self.physical_expert_start = self.experts_start_idx + self.physical_expert_end = self.experts_end_idx + self.experts = FusedMoE( shared_experts=self.shared_experts, gate=self.gate, @@ -697,26 +801,13 @@ class DeepseekV4Attention(nn.Module): self.rope_parameters = config.rope_scaling - # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it) - rope_parameters = config.rope_parameters - rope_parameters["rope_theta"] = ( - config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta - ) - if config.rope_parameters["rope_type"] != "default": - config.rope_parameters["rope_type"] = ( - "deepseek_yarn" - if config.rope_parameters.get("apply_yarn_scaling", True) - else "deepseek_llama_scaling" - ) - rope_parameters["mscale"] = 0 # Disable mscale - rope_parameters["mscale_all_dim"] = 0 # Disable mscale - rope_parameters["is_deepseek_v4"] = True - rope_parameters["rope_dim"] = self.rope_head_dim - self.rotary_emb = get_rope( - self.head_dim, - max_position=self.max_position_embeddings, - rope_parameters=rope_parameters, - is_neox_style=False, + # Initialize rotary embedding BEFORE DeepseekV4MLA (which needs it) + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + compress_ratio=self.compress_ratio, ) self.indexer = None @@ -741,7 +832,17 @@ class DeepseekV4Attention(nn.Module): aux_stream=indexer_aux_stream, ) - mla_modules = DeepseekV4MLAModules( + self.mla_attn = DeepseekV4MLA( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, vllm_config=vllm_config, fused_wqa_wkv=self.fused_wqa_wkv, q_norm=self.q_norm, @@ -755,19 +856,6 @@ class DeepseekV4Attention(nn.Module): indexer_rotary_emb=self.rotary_emb, topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, - ) - self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper( - hidden_size=self.hidden_size, - num_heads=self.n_local_heads, - head_dim=self.head_dim, - scale=self.softmax_scale, - qk_nope_head_dim=self.nope_head_dim, - qk_rope_head_dim=self.rope_head_dim, - v_head_dim=self.head_dim, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.head_dim, - o_lora_rank=self.o_lora_rank, - mla_modules=mla_modules, window_size=self.window_size, compress_ratio=self.compress_ratio, cache_config=vllm_config.cache_config, @@ -955,7 +1043,7 @@ class DeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps # Three aux streams: one per non-default input GEMM in - # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute + # DeepseekV4MLA.attn_gemm_parallel_execute # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. aux_stream_list = [torch.cuda.Stream() for _ in range(3)] @@ -1259,7 +1347,44 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: ) -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): +class DeepseekV4MixtureOfExperts(MixtureOfExperts): + moe_mlp_layers: list["DeepseekV4MoE"] + + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None: + if example_moe is None: + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_shared_experts = 0 + self.num_redundant_experts = 0 + return + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_shared_experts = example_moe.n_shared_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + 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 moe in self.moe_mlp_layers: + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts): model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. @@ -1291,6 +1416,28 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP): self.model.make_empty_intermediate_tensors ) + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = [] + self.num_expert_groups = getattr(self.config, "n_group", 1) + self.num_moe_layers = self.config.num_hidden_layers + self.moe_layers: list[nn.Module] = [] + self.moe_mlp_layers: list[DeepseekV4MoE] = [] + example_moe: DeepseekV4MoE | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if not isinstance(layer, DeepseekV4DecoderLayer): + continue + if isinstance(layer.ffn, DeepseekV4MoE): + example_moe = layer.ffn + self.moe_mlp_layers.append(layer.ffn) + self.moe_layers.append(layer.ffn.experts) + + self.num_moe_layers = len(self.moe_layers) + self.extract_moe_parameters(example_moe) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index 1e6c2ed829d..ed16ca6d3b5 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -370,11 +370,11 @@ class SparseAttnCompressNormRopeStoreC4Kernel: inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) for pair in cutlass.range_constexpr(self.elems_per_lane // 2): elem = const_expr(pair * 2) - y0 = cute.arch.fmin( + y0 = cutlass.min( cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), Float32(self.fp8_max), ) - y1 = cute.arch.fmin( + y1 = cutlass.min( cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), Float32(self.fp8_max), ) @@ -1026,11 +1026,11 @@ class SparseAttnNormRopeStoreKernel: bits = _recast_val(scale_raw, Uint32) ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - y0 = cute.arch.fmin( + y0 = cutlass.min( cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), Float32(self.fp8_max), ) - y1 = cute.arch.fmin( + y1 = cutlass.min( cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), Float32(self.fp8_max), ) diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index babc4c742a3..312239ad3fd 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -22,6 +22,7 @@ from urllib3.util import Url, parse_url import vllm.envs as envs from vllm.connections import HTTPConnection, global_http_connection from vllm.logger import init_logger +from vllm.multimodal.video import get_video_loader_backend_for_processor from vllm.utils.registry import ExtensionManager from .audio import AudioEmbeddingMediaIO, AudioMediaIO @@ -452,6 +453,7 @@ class MediaConnector: video_url: str, *, image_mode: str = "RGB", + video_processor: str | None = None, ) -> tuple[npt.NDArray, dict[str, Any]]: """ Load video from an HTTP or base64 data URL. @@ -459,7 +461,12 @@ class MediaConnector: image_io = ImageMediaIO( image_mode=image_mode, **self.media_io_kwargs.get("image", {}) ) - video_io = VideoMediaIO(image_io, **self.media_io_kwargs.get("video", {})) + video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) + if "video_backend" not in video_io_kwargs and ( + video_backend := get_video_loader_backend_for_processor(video_processor) + ): + video_io_kwargs["video_backend"] = video_backend + video_io = VideoMediaIO(image_io, **video_io_kwargs) return self.load_from_url( video_url, @@ -472,6 +479,7 @@ class MediaConnector: video_url: str, *, image_mode: str = "RGB", + video_processor: str | None = None, ) -> tuple[npt.NDArray, dict[str, Any]]: """ Asynchronously load video from an HTTP or base64 data URL. @@ -481,7 +489,12 @@ class MediaConnector: image_io = ImageMediaIO( image_mode=image_mode, **self.media_io_kwargs.get("image", {}) ) - video_io = VideoMediaIO(image_io, **self.media_io_kwargs.get("video", {})) + video_io_kwargs = dict(self.media_io_kwargs.get("video", {})) + if "video_backend" not in video_io_kwargs and ( + video_backend := get_video_loader_backend_for_processor(video_processor) + ): + video_io_kwargs["video_backend"] = video_backend + video_io = VideoMediaIO(image_io, **video_io_kwargs) return await self.load_from_url_async( video_url, diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 2d321cb67b4..56c65dc8ea7 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import bisect import mimetypes from collections import defaultdict from collections.abc import Generator, Sequence @@ -18,6 +19,7 @@ from vllm.utils.import_utils import LazyLoader from .hasher import MultiModalHasher from .inputs import ( BatchedTensorInputs, + MultiModalFeatureSpec, MultiModalFieldElem, MultiModalKwargsItem, MultiModalSharedField, @@ -109,6 +111,29 @@ def encode_video_url( return f"data:{mimetype};base64,{video_b64}" +def get_mm_features_in_window( + mm_features: list[MultiModalFeatureSpec], + start: int, + end: int, +) -> tuple[int, int]: + """Return (lo, hi) indices for features overlapping [start, end). + + Assumes mm_features are sorted by offset and non-overlapping, so + offset + length is also sorted. + """ + lo = bisect.bisect_left( + mm_features, + start + 1, + key=lambda f: f.mm_position.offset + f.mm_position.length, + ) + hi = bisect.bisect_left( + mm_features, + end, + key=lambda f: f.mm_position.offset, + ) + return lo, hi + + def argsort_mm_positions( mm_positions: MultiModalPlaceholders, ) -> list[tuple[str, int]]: diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index df296437fe6..1324e79c5c9 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -28,6 +28,60 @@ except ImportError: logger = init_logger(__name__) +class VideoLoaderRegistry(ExtensionManager): + def __init__(self) -> None: + super().__init__() + self.processor2backend: dict[str, str] = {} + + @staticmethod + def _normalize_registered_video_processors( + video_processor: str | tuple[str, ...] | None, + ) -> tuple[str, ...]: + if video_processor is None: + return () + + if isinstance(video_processor, str): + return (video_processor,) + + if all(isinstance(processor, str) for processor in video_processor): + return video_processor + + raise TypeError( + "video_processor must be a class name or a tuple of class names" + ) + + def register( + self, + name: str, + *, + video_processor: str | tuple[str, ...] | None = None, + ): + processors = self._normalize_registered_video_processors(video_processor) + + def wrap(cls_to_register): + self.name2class[name] = cls_to_register + for processor_name in processors: + self.processor2backend[processor_name] = name + return cls_to_register + + return wrap + + def get_backend_for_video_processor( + self, + video_processor: str | None, + ) -> str | None: + if video_processor is None: + return None + + return self.processor2backend.get(video_processor) + + +def get_video_loader_backend_for_processor( + video_processor: str | None, +) -> str | None: + return VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(video_processor) + + def resize_video(frames: npt.NDArray, size: tuple[int, int]) -> npt.NDArray: num_frames, _, _, channels = frames.shape new_height, new_width = size @@ -113,7 +167,7 @@ class VideoLoader: } -VIDEO_LOADER_REGISTRY = ExtensionManager() +VIDEO_LOADER_REGISTRY = VideoLoaderRegistry() class OpenCVVideoBackendMixin: @@ -550,7 +604,10 @@ class VideoBackend(VideoLoader, OpenCVVideoBackendMixin, PyAVVideoBackendMixin): ) -@VIDEO_LOADER_REGISTRY.register("opencv_dynamic") +@VIDEO_LOADER_REGISTRY.register( + "opencv_dynamic", + video_processor="Glm4vVideoProcessor", +) class DynamicVideoBackend(VideoBackend): """Duration-aware dynamic-sampling video backend. @@ -639,8 +696,11 @@ class DynamicVideoBackend(VideoBackend): ) -@VIDEO_LOADER_REGISTRY.register("glm4_6v") -class GLM4_6VVideoBackend(VideoBackend): +@VIDEO_LOADER_REGISTRY.register( + "glmga", + video_processor="GlmgaVideoProcessor", +) +class GLMGAVideoBackend(VideoBackend): @classmethod def _prepare_source(cls, source: VideoSourceMetadata) -> VideoSourceMetadata: # Estimate duration from frame count and fps when the container @@ -740,7 +800,10 @@ class GLM4_6VVideoBackend(VideoBackend): return frames, metadata -@VIDEO_LOADER_REGISTRY.register("molmo2") +@VIDEO_LOADER_REGISTRY.register( + "molmo2", + video_processor="Molmo2VideoProcessor", +) class Molmo2VideoBackend(VideoLoader, OpenCVVideoBackendMixin): @classmethod def get_candidate_target_fps( diff --git a/vllm/parser/__init__.py b/vllm/parser/__init__.py index dc256daaa7e..de815b2e1fd 100644 --- a/vllm/parser/__init__.py +++ b/vllm/parser/__init__.py @@ -4,7 +4,6 @@ from vllm.parser.abstract_parser import ( DelegatingParser, Parser, - _WrappedParser, ) from vllm.parser.parser_manager import ParserManager @@ -12,21 +11,4 @@ __all__ = [ "Parser", "DelegatingParser", "ParserManager", - "_WrappedParser", ] - -_PARSERS_TO_REGISTER = { - "minimax_m2": ( # name - "minimax_m2_parser", # filename - "MiniMaxM2Parser", # class_name - ), -} - - -def register_lazy_parsers(): - for name, (file_name, class_name) in _PARSERS_TO_REGISTER.items(): - module_path = f"vllm.parser.{file_name}" - ParserManager.register_lazy_module(name, module_path, class_name) - - -register_lazy_parsers() diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 2a13f138607..9e4d1830b4d 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -37,13 +37,13 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.reasoning.abs_reasoning_parsers import ReasoningParser from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser from vllm.tool_parsers.streaming import ( extract_named_tool_call_streaming, extract_required_tool_call_streaming, ) -from vllm.tool_parsers.utils import Tool from vllm.utils import random_uuid +from vllm.utils.mistral import is_mistral_tool_parser logger = init_logger(__name__) @@ -90,19 +90,25 @@ class Parser: reasoning_parser_cls: type[ReasoningParser] | None = None tool_parser_cls: type[ToolParser] | None = None - def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): - """ - Initialize the Parser. - - Args: - tokenizer: The tokenizer used by the model. This is required for - token-based parsing operations. - """ + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + *args, + **kwargs, + ): self.model_tokenizer = tokenizer self._reasoning_parser: ReasoningParser | None = None self._tool_parser: ToolParser | None = None self._stream_state = StreamState() + if self.__class__.reasoning_parser_cls is not None: + self._reasoning_parser = self.__class__.reasoning_parser_cls( + tokenizer, *args, **kwargs + ) + if self.__class__.tool_parser_cls is not None: + self._tool_parser = self.__class__.tool_parser_cls(tokenizer, tools) + @cached_property def vocab(self) -> dict[str, int]: """Get the vocabulary mapping from tokens to IDs.""" @@ -313,6 +319,24 @@ class Parser: A DeltaMessage with tool_calls field, or None. """ + @abstractmethod + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + """Parse a complete model output, extracting reasoning and tool calls. + + Args: + model_output: The complete model-generated string. + request: The request object used to generate the output. + enable_auto_tools: Whether to enable automatic tool call parsing. + + Returns: + A tuple of (reasoning, content, tool_calls). + """ + @abstractmethod def parse_delta( self, @@ -320,6 +344,8 @@ class Parser: delta_token_ids: list[int], request: ChatCompletionRequest | ResponsesRequest, prompt_token_ids: list[int] | None = None, + *, + finished: bool, ) -> DeltaMessage | None: """Parse a single streaming delta, orchestrating reasoning then tool call extraction via internal stream state. @@ -510,6 +536,99 @@ class DelegatingParser(Parser): # No tool calls return [], content + def _extract_tool_calls( + self, + content: str | None, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + ) -> tuple[list[FunctionCall] | None, str | None]: + tool_parser = self._tool_parser + if tool_parser is None: + return [], content + + # When the Mistral grammar factory injected structured outputs, + # let the parser handle the output. + use_mistral_tool_parser = ( + is_mistral_tool_parser(type(tool_parser)) + and isinstance(request, ChatCompletionRequest) + and request._grammar_from_tool_parser + ) + + supports_required_and_named = tool_parser.supports_required_and_named + is_named_tool_choice = request.tool_choice and isinstance( + request.tool_choice, + (ToolChoiceFunction, ChatCompletionNamedToolChoiceParam), + ) + is_required_tool_choice = request.tool_choice == "required" + is_auto_tool_choice = enable_auto_tools and ( + request.tool_choice == "auto" + or request.tool_choice is None + or ( + not supports_required_and_named + and (is_named_tool_choice or is_required_tool_choice) + ) + ) + + tool_calls = list[FunctionCall]() + if ( + is_named_tool_choice + and supports_required_and_named + and not use_mistral_tool_parser + ): + if content is None: + return [], None + tool_calls.append( + FunctionCall( + name=self._get_function_name(request), + arguments=content, + ) + ) + content = None + elif ( + is_required_tool_choice + and supports_required_and_named + and not use_mistral_tool_parser + ): + # "required" with standard JSON-based parsing + parsed_calls = [] + with contextlib.suppress(ValidationError): + content = content or "" + parsed_calls = TypeAdapter(list[FunctionDefinition]).validate_json( + content + ) + for tc in parsed_calls: + tool_calls.append( + FunctionCall( + name=tc.name, + arguments=json.dumps(tc.parameters, ensure_ascii=False), + ) + ) + content = None + elif is_auto_tool_choice or use_mistral_tool_parser: + # Automatic Tool Call Parsing (also used as fallback for + # required/named when supports_required_and_named=False) + tool_call_info = tool_parser.extract_tool_calls( + content if content is not None else "", + request=request, # type: ignore + ) + if tool_call_info is not None and tool_call_info.tools_called: + tool_calls.extend( + FunctionCall( + id=tc.id, + name=tc.function.name, + arguments=tc.function.arguments, + ) + for tc in tool_call_info.tool_calls + ) + content = tool_call_info.content + if content and content.strip() == "": + content = None + else: + # No tool calls. + return None, content + + return tool_calls, content + def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: @@ -656,12 +775,43 @@ class DelegatingParser(Parser): return False return state.reasoning_ended + def _append_unstreamed_tool_args( + self, + delta_message: DeltaMessage | None, + ) -> None: + """Append parsed-but-unstreamed tool-call arguments to *delta_message*.""" + if ( + self._tool_parser is not None + and delta_message + and delta_message.tool_calls + and (last_tc := delta_message.tool_calls[-1]).function + ): + last_tc.function.arguments = ( + last_tc.function.arguments or "" + ) + self._tool_parser.get_remaining_unstreamed_args() + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + reasoning, content = self.extract_reasoning(model_output, request) + tool_calls, content = self._extract_tool_calls( + content=content, + request=request, + enable_auto_tools=enable_auto_tools, + ) + return reasoning, content, tool_calls + def parse_delta( self, delta_text: str, delta_token_ids: list[int], request: ChatCompletionRequest | ResponsesRequest, prompt_token_ids: list[int] | None = None, + *, + finished: bool, ) -> DeltaMessage | None: state = self._stream_state @@ -745,35 +895,8 @@ class DelegatingParser(Parser): state.previous_text = current_text state.previous_token_ids = current_token_ids + + if finished: + self._append_unstreamed_tool_args(delta_message) + return delta_message - - -class _WrappedParser(DelegatingParser): - """ - A DelegatingParser subclass that instantiates parsers from class attributes. - - This class is used to dynamically create a parser that wraps individual - ReasoningParser and ToolParser classes. The class attributes - `reasoning_parser_cls` and `tool_parser_cls` should be set before - instantiation. - - Usage: - _WrappedParser.reasoning_parser_cls = MyReasoningParser - _WrappedParser.tool_parser_cls = MyToolParser - parser = _WrappedParser(tokenizer) - """ - - reasoning_parser_cls: type[ReasoningParser] | None = None - tool_parser_cls: type[ToolParser] | None = None - - def __init__( - self, tokenizer: TokenizerLike, tools: list[Tool] | None = None, **kwargs - ): - super().__init__(tokenizer) - # Instantiate the underlying parsers from class attributes - if self.__class__.reasoning_parser_cls is not None: - self._reasoning_parser = self.__class__.reasoning_parser_cls( - tokenizer, **kwargs - ) - if self.__class__.tool_parser_cls is not None: - self._tool_parser = self.__class__.tool_parser_cls(tokenizer, tools) diff --git a/vllm/parser/minimax_m2_parser.py b/vllm/parser/minimax_m2_parser.py deleted file mode 100644 index 34aaa726844..00000000000 --- a/vllm/parser/minimax_m2_parser.py +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -MiniMax M2 Parser - A unified parser for MiniMax M2 models. - -This parser combines the existing MiniMaxM2ReasoningParser and -MinimaxM2ToolParser into a single unified interface by delegating -to those implementations. -""" - -from vllm.logger import init_logger -from vllm.parser.abstract_parser import DelegatingParser -from vllm.reasoning.minimax_m2_reasoning_parser import MiniMaxM2ReasoningParser -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, -) -from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser - -logger = init_logger(__name__) - - -class MiniMaxM2Parser(DelegatingParser): - """ - Unified parser for MiniMax M2 models that handles both reasoning - extraction and tool call parsing. - - This parser delegates to the existing implementations: - - MiniMaxM2ReasoningParser for reasoning extraction - - MinimaxM2ToolParser for tool call parsing - - MiniMax M2 models have two special behaviors: - 1. Reasoning: They don't generate start token, only end - token. All content before is reasoning, content after is the - actual response. - 2. Tool Calls: They use ... tags - with ... and ... - syntax. - """ - - # Class-level parser classes for compatibility - reasoning_parser_cls = MiniMaxM2ReasoningParser - tool_parser_cls = MinimaxM2ToolParser - - def __init__( - self, - tokenizer: TokenizerLike, - tools: list[Tool] | None = None, - *args, - **kwargs, - ): - super().__init__(tokenizer, *args, **kwargs) - - # Initialize the underlying parsers - self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer, *args, **kwargs) - self._tool_parser = MinimaxM2ToolParser(tokenizer, tools) - - logger.debug( - "vLLM Successfully initialized parser %s!", self.__class__.__name__ - ) diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index f8bded62d59..7afd39d4fea 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -3,14 +3,9 @@ from __future__ import annotations -import importlib -import os -from collections.abc import Callable from typing import TYPE_CHECKING from vllm.logger import init_logger -from vllm.utils.collection_utils import is_list_of -from vllm.utils.import_utils import import_from_path if TYPE_CHECKING: from vllm.parser.abstract_parser import Parser @@ -22,170 +17,10 @@ logger = init_logger(__name__) class ParserManager: """ - Central registry for Parser implementations. - - Supports two registration modes: - - Eager registration via `register_module` - - Lazy registration via `register_lazy_module` + Provides a unified Parser by composing individual reasoning and tool + parsers from their respective registries. """ - parsers: dict[str, type[Parser]] = {} - lazy_parsers: dict[str, tuple[str, str]] = {} # name -> (module_path, class_name) - - @classmethod - def get_parser_internal(cls, name: str) -> type[Parser]: - """ - Retrieve a registered or lazily registered Parser class. - - Args: - name: The registered name of the parser. - - Returns: - The Parser class. - - Raises: - KeyError: If no parser is found under the given name. - """ - if name in cls.parsers: - return cls.parsers[name] - - if name in cls.lazy_parsers: - return cls._load_lazy_parser(name) - - registered = ", ".join(cls.list_registered()) - raise KeyError(f"Parser '{name}' not found. Available parsers: {registered}") - - @classmethod - def _load_lazy_parser(cls, name: str) -> type[Parser]: - """Import and register a lazily loaded parser.""" - from vllm.parser.abstract_parser import Parser - - module_path, class_name = cls.lazy_parsers[name] - try: - mod = importlib.import_module(module_path) - parser_cls = getattr(mod, class_name) - if not issubclass(parser_cls, Parser): - raise TypeError( - f"{class_name} in {module_path} is not a Parser subclass." - ) - cls.parsers[name] = parser_cls # cache - return parser_cls - except Exception as e: - logger.exception( - "Failed to import lazy parser '%s' from %s: %s", - name, - module_path, - e, - ) - raise - - @classmethod - def _register_module( - cls, - module: type[Parser], - module_name: str | list[str] | None = None, - force: bool = True, - ) -> None: - """Register a Parser class immediately.""" - from vllm.parser.abstract_parser import Parser - - if not issubclass(module, Parser): - raise TypeError( - f"module must be subclass of Parser, but got {type(module)}" - ) - - if module_name is None: - module_names = [module.__name__] - elif isinstance(module_name, str): - module_names = [module_name] - elif is_list_of(module_name, str): - module_names = module_name - else: - raise TypeError("module_name must be str, list[str], or None.") - - for name in module_names: - if not force and name in cls.parsers: - existed = cls.parsers[name] - raise KeyError(f"{name} is already registered at {existed.__module__}") - cls.parsers[name] = module - - @classmethod - def register_lazy_module(cls, name: str, module_path: str, class_name: str) -> None: - """ - Register a lazy module mapping for delayed import. - - Example: - ParserManager.register_lazy_module( - name="minimax_m2", - module_path="vllm.parser.minimax_m2_parser", - class_name="MiniMaxM2Parser", - ) - """ - cls.lazy_parsers[name] = (module_path, class_name) - - @classmethod - def register_module( - cls, - name: str | list[str] | None = None, - force: bool = True, - module: type[Parser] | None = None, - ) -> type[Parser] | Callable[[type[Parser]], type[Parser]]: - """ - Register a Parser class. - - Can be used as a decorator or called directly. - - Usage: - @ParserManager.register_module("my_parser") - class MyParser(Parser): - ... - - Or: - ParserManager.register_module(module=MyParser) - """ - if not isinstance(force, bool): - raise TypeError(f"force must be a boolean, but got {type(force)}") - - # Immediate registration - if module is not None: - cls._register_module(module=module, module_name=name, force=force) - return module - - # Decorator usage - def _decorator(obj: type[Parser]) -> type[Parser]: - module_path = obj.__module__ - class_name = obj.__name__ - - if isinstance(name, str): - names = [name] - elif name is not None and is_list_of(name, str): - names = name - else: - names = [class_name] - - for n in names: - cls.lazy_parsers[n] = (module_path, class_name) - - return obj - - return _decorator - - @classmethod - def list_registered(cls) -> list[str]: - """Return names of all registered parsers.""" - return sorted(set(cls.parsers.keys()) | set(cls.lazy_parsers.keys())) - - @classmethod - def import_parser(cls, plugin_path: str) -> None: - """Import a user-defined parser from an arbitrary path.""" - module_name = os.path.splitext(os.path.basename(plugin_path))[0] - try: - import_from_path(module_name, plugin_path) - except Exception: - logger.exception( - "Failed to load module '%s' from %s.", module_name, plugin_path - ) - @classmethod def get_tool_parser( cls, @@ -246,12 +81,10 @@ class ParserManager: model_name: str | None = None, ) -> type[Parser] | None: """ - Get a unified Parser that handles both reasoning and tool parsing. + Get a Parser that handles both reasoning and tool parsing. - This method checks if a unified Parser exists that can handle both - reasoning extraction and tool call parsing. If no unified parser - exists, it creates a DelegatingParser that wraps the individual - reasoning and tool parsers. + Composes individual reasoning and tool parsers into a single + DelegatingParser subclass. Args: tool_parser_name: The name of the tool parser. @@ -262,37 +95,9 @@ class ParserManager: Returns: A Parser class, or None if neither parser is specified. """ - from vllm.parser.abstract_parser import _WrappedParser - if not tool_parser_name and not reasoning_parser_name: return None - # Strategy 1: If both names match, check for a unified parser with that name - if tool_parser_name and tool_parser_name == reasoning_parser_name: - try: - parser = cls.get_parser_internal(tool_parser_name) - logger.info( - "Using unified parser '%s' for both reasoning and tool parsing.", - tool_parser_name, - ) - return parser - except KeyError: - pass # No unified parser with this name - - # Strategy 2: Check for parser with either name - for name in [tool_parser_name, reasoning_parser_name]: - if name: - try: - parser = cls.get_parser_internal(name) - logger.info( - "Using unified parser '%s' for reasoning and tool parsing.", - name, - ) - return parser - except KeyError: - pass - - # Strategy 3: Create a DelegatingParser with the individual parser classes reasoning_parser_cls = cls.get_reasoning_parser(reasoning_parser_name) tool_parser_cls = cls.get_tool_parser( tool_parser_name, enable_auto_tools, model_name @@ -301,8 +106,13 @@ class ParserManager: if reasoning_parser_cls is None and tool_parser_cls is None: return None - # Set the class-level attributes on the imported _WrappedParser - _WrappedParser.reasoning_parser_cls = reasoning_parser_cls - _WrappedParser.tool_parser_cls = tool_parser_cls + from vllm.parser.abstract_parser import DelegatingParser - return _WrappedParser + r_cls = reasoning_parser_cls + t_cls = tool_parser_cls + + class _Parser(DelegatingParser): + reasoning_parser_cls = r_cls + tool_parser_cls = t_cls + + return _Parser diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index abec3bbe886..b357c5798bf 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -199,6 +199,14 @@ class Platform: # all ROCm platforms for now. return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM) + def is_cumem_allocator_available(self) -> bool: + try: + from vllm.device_allocator.cumem import cumem_available + except ImportError: + return False + + return cumem_available + @classmethod def get_pass_manager_cls(cls) -> str: """ @@ -690,6 +698,13 @@ class Platform: mamba_padding_pct, ) + @classmethod + def register_custom_kv_cache_specs(cls, vllm_config: "VllmConfig") -> None: + """ + Register custom KVCacheSpec class on current platform. + """ + pass + @classmethod def verify_model_arch(cls, model_arch: str) -> None: """ diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index c2be7ff03ab..5947bff9b08 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -110,6 +110,13 @@ class XPUPlatform(Platform): dtype: torch.dtype, backend: "AttentionBackendEnum | None" = None, ) -> "AttentionBackendEnum": + if dtype == torch.float32: + logger.warning_once( + "Flash Attention on XPU does not support float32 dtype. " + "Falling back to Triton Attention backend for vit attention." + ) + return AttentionBackendEnum.TRITON_ATTN + if backend is not None: assert backend in cls.get_supported_vit_attn_backends(), ( f"Backend {backend} is not supported for vit attention. " @@ -197,24 +204,25 @@ class XPUPlatform(Platform): ) # Disable fusion passes not yet supported on XPU. + from vllm.config.compilation import CompilationMode + pass_config = compilation_config.pass_config fusion_passes_to_disable = { "enable_sp": "Sequence parallelism", "fuse_gemm_comms": "Async TP", "fuse_allreduce_rms": "AllReduce + RMSNorm fusion", - "fuse_norm_quant": "RMSNorm + quant fusion", - "fuse_act_quant": "Activation + quant fusion", "fuse_attn_quant": "Attention + quant fusion", "fuse_act_padding": "Activation + padding fusion", "fuse_rope_kvcache": "RoPE + KV cache fusion", } - for flag, feature_name in fusion_passes_to_disable.items(): - if getattr(pass_config, flag): - logger.warning( - "Feature %r is not yet supported on XPU and will be disabled.", - feature_name, - ) - setattr(pass_config, flag, False) + if compilation_config.mode != CompilationMode.NONE: + for flag, feature_name in fusion_passes_to_disable.items(): + if getattr(pass_config, flag): + logger.warning( + "Feature %r is not yet supported on XPU and will be disabled.", + feature_name, + ) + setattr(pass_config, flag, False) # check and update parallel config parallel_config = vllm_config.parallel_config diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index e796607722a..e57d0586aa0 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -450,6 +450,66 @@ def _detect_content_format( return "openai" +@lru_cache(maxsize=32) +def _detect_developer_role_support(chat_template: str) -> bool: + return '"developer"' in chat_template or "'developer'" in chat_template + + +def _convert_developer_to_system( + conversation: list[ConversationMessage], +) -> list[ConversationMessage]: + converted: list[ConversationMessage] = [] + for msg in conversation: + if msg["role"] == "developer": + new_msg = dict(msg) + new_msg["role"] = "system" + new_msg.pop("tools", None) + converted.append(new_msg) # type: ignore[arg-type] + else: + converted.append(msg) + return converted + + +def _consolidate_system_messages( + conversation: list[ConversationMessage], +) -> list[ConversationMessage]: + """Merge all system messages into one at position 0. + + Some chat templates (e.g. Qwen 3.6) require the system message to be the + very first message. After developer-to-system conversion, system messages + may appear at non-first positions; this merges them into a single message. + """ + system_contents: list[str] = [] + non_system: list[ConversationMessage] = [] + needs_consolidation = False + for i, msg in enumerate(conversation): + if msg["role"] == "system": + if i > 0 or system_contents: + needs_consolidation = True + content = msg.get("content", "") + if isinstance(content, list): + parts = [] + for part in content: + if isinstance(part, dict) and "text" in part: + parts.append(part["text"]) + elif isinstance(part, str): + parts.append(part) + content = "\n".join(parts) + if content: + system_contents.append(content) + else: + non_system.append(msg) + + if not needs_consolidation: + return conversation + + merged: ConversationMessage = { + "role": "system", + "content": "\n\n".join(system_contents), + } + return [merged, *non_system] + + def _resolve_chat_template_content_format( chat_template: str | None, tools: list[dict[str, Any]] | None, @@ -653,7 +713,15 @@ def safe_apply_chat_template( "allowed, so you must provide a chat template if the tokenizer " "does not define one." ) - + if any( + msg["role"] == "developer" for msg in conversation + ) and not _detect_developer_role_support(chat_template): + conversation = _convert_developer_to_system(conversation) + conversation = _consolidate_system_messages(conversation) + logger.info_once( + "Chat template does not support the 'developer' message role. " + "Converting developer messages to 'system' role.", + ) resolved_kwargs = resolve_chat_template_kwargs( tokenizer=tokenizer, chat_template=chat_template, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 6e0be9dbff5..6beb1423ce2 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -737,6 +737,20 @@ class SamplingParams( parameter="logprob_token_ids", value=n, ) + vocab_size = model_config.get_vocab_size() + invalid_token_ids = [ + token_id + for token_id in self.logprob_token_ids + if token_id < 0 or token_id >= vocab_size + ] + if invalid_token_ids: + raise VLLMValidationError( + f"token_id(s) {invalid_token_ids} in logprob_token_ids " + f"contain out-of-vocab token ids. Vocabulary size: " + f"{vocab_size}", + parameter="logprob_token_ids", + value=invalid_token_ids, + ) if self.logprobs is not None and self.logprobs != n: raise VLLMValidationError( f"When both logprobs and logprob_token_ids are set, " diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index c3438082a72..94543b82350 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -79,6 +79,25 @@ class ToolParser: else: self.tools = [] + def get_remaining_unstreamed_args(self) -> str: + """Return tool call arguments parsed but not yet streamed.""" + if not self.prev_tool_call_arr: + return "" + index = len(self.prev_tool_call_arr) - 1 + args = self.prev_tool_call_arr[index].get("arguments", {}) + if isinstance(args, str): + expected = args + else: + expected = json.dumps(args, ensure_ascii=False) + actual = ( + self.streamed_args_for_tool[index] + if index < len(self.streamed_args_for_tool) + else "" + ) + if expected.startswith(actual): + return expected[len(actual) :] + return "" + @cached_property def vocab(self) -> dict[str, int]: # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab diff --git a/vllm/tool_parsers/glm4_moe_tool_parser.py b/vllm/tool_parsers/glm4_moe_tool_parser.py index 1779896e5b6..213a774535b 100644 --- a/vllm/tool_parsers/glm4_moe_tool_parser.py +++ b/vllm/tool_parsers/glm4_moe_tool_parser.py @@ -11,7 +11,6 @@ The fix streams string values incrementally as they arrive, providing a true streaming experience for long content. """ -import ast import json from collections.abc import Sequence from typing import Any @@ -42,6 +41,7 @@ from vllm.tool_parsers.utils import ( extract_types_from_schema, find_tool_properties, partial_tag_overlap, + safe_literal_eval, ) logger = init_logger(__name__) @@ -110,7 +110,7 @@ class Glm4MoeModelToolParser(ToolParser): pass try: - return ast.literal_eval(value) + return safe_literal_eval(value) except (ValueError, SyntaxError): pass diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py index 496deb4f2d5..619be5e9cc2 100644 --- a/vllm/tool_parsers/hy_v3_tool_parser.py +++ b/vllm/tool_parsers/hy_v3_tool_parser.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from collections.abc import Sequence from typing import Any @@ -27,6 +26,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) +from vllm.tool_parsers.utils import safe_literal_eval logger = init_logger(__name__) @@ -183,13 +183,13 @@ class HYV3ToolParser(ToolParser): @staticmethod def _deserialize(value: str) -> Any: - """Deserialize a string value using json.loads then ast.literal_eval.""" + """Deserialize a string value using json.loads then safe_literal_eval.""" try: return json.loads(value) except Exception: pass try: - return ast.literal_eval(value) + return safe_literal_eval(value) except Exception: pass return value diff --git a/vllm/tool_parsers/minicpm5xml_tool_parser.py b/vllm/tool_parsers/minicpm5xml_tool_parser.py index fed9677411b..a5b5252415c 100644 --- a/vllm/tool_parsers/minicpm5xml_tool_parser.py +++ b/vllm/tool_parsers/minicpm5xml_tool_parser.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from collections.abc import Sequence from typing import Any @@ -28,7 +27,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) -from vllm.tool_parsers.utils import partial_tag_overlap +from vllm.tool_parsers.utils import partial_tag_overlap, safe_literal_eval from vllm.utils import random_uuid logger = init_logger(__name__) @@ -116,7 +115,7 @@ def _parse_arguments(json_value: str) -> tuple[Any, bool]: try: parsed_value = json.loads(json_value) except json.JSONDecodeError: - parsed_value = ast.literal_eval(json_value) + parsed_value = safe_literal_eval(json_value) return parsed_value, True except Exception: return json_value, False diff --git a/vllm/tool_parsers/poolside_v1_tool_parser.py b/vllm/tool_parsers/poolside_v1_tool_parser.py index f14b4736291..e515e1ce637 100644 --- a/vllm/tool_parsers/poolside_v1_tool_parser.py +++ b/vllm/tool_parsers/poolside_v1_tool_parser.py @@ -11,7 +11,6 @@ The fix streams string values incrementally as they arrive, providing a true streaming experience for long content. """ -import ast import json from collections.abc import Sequence from typing import Any @@ -41,6 +40,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) +from vllm.tool_parsers.utils import safe_literal_eval logger = init_logger(__name__) @@ -106,7 +106,7 @@ class PoolsideV1ToolParser(ToolParser): pass try: - return ast.literal_eval(value) + return safe_literal_eval(value) except (ValueError, SyntaxError): pass diff --git a/vllm/tool_parsers/qwen3xml_tool_parser.py b/vllm/tool_parsers/qwen3xml_tool_parser.py index d5b87ea074e..e5d2b896e00 100644 --- a/vllm/tool_parsers/qwen3xml_tool_parser.py +++ b/vllm/tool_parsers/qwen3xml_tool_parser.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from collections.abc import Sequence from typing import Any @@ -26,7 +25,7 @@ from vllm.tool_parsers.abstract_tool_parser import ( Tool, ToolParser, ) -from vllm.tool_parsers.utils import find_tool_properties +from vllm.tool_parsers.utils import find_tool_properties, safe_literal_eval logger = init_logger(__name__) @@ -824,7 +823,7 @@ class StreamingXMLToolCallParser: try: parsed_value = json.loads(raw_for_parse) except json.JSONDecodeError: - parsed_value = ast.literal_eval(raw_for_parse) + parsed_value = safe_literal_eval(raw_for_parse) output_arguments = json.dumps(parsed_value, ensure_ascii=False) except Exception: # Fallback: output as string as-is diff --git a/vllm/tool_parsers/step3p5_tool_parser.py b/vllm/tool_parsers/step3p5_tool_parser.py index b46f899ce2c..8a48e2686e5 100644 --- a/vllm/tool_parsers/step3p5_tool_parser.py +++ b/vllm/tool_parsers/step3p5_tool_parser.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import ast import json from collections.abc import Sequence from typing import Any @@ -23,6 +22,7 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser +from vllm.tool_parsers.utils import safe_literal_eval logger = init_logger(__name__) @@ -1016,7 +1016,7 @@ class StreamingXMLToolCallParser: raw_for_parse = raw_text + "\n" else: raw_for_parse = raw_text - parsed_value = ast.literal_eval(raw_for_parse) + parsed_value = safe_literal_eval(raw_for_parse) output_arguments = json.dumps(parsed_value, ensure_ascii=False) except Exception: # Fallback: output as string as-is diff --git a/vllm/tool_parsers/utils.py b/vllm/tool_parsers/utils.py index 1c7830b320f..6ee107433c5 100644 --- a/vllm/tool_parsers/utils.py +++ b/vllm/tool_parsers/utils.py @@ -3,6 +3,7 @@ import ast import json +import warnings from json import JSONDecodeError, JSONDecoder from typing import Any, TypeAlias @@ -31,6 +32,12 @@ Tool: TypeAlias = ChatCompletionToolsParam | ResponsesTool logger = init_logger(__name__) +def safe_literal_eval(text: str): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SyntaxWarning) + return ast.literal_eval(text) + + def partial_tag_overlap(text: str, tag: str) -> int: """Length of the longest prefix of *tag* that matches a suffix of *text*. diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index 0e241f6abfd..ec01f65d774 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -161,6 +161,44 @@ def get_processor_cls_name_from_config( return None +def get_video_processor_cls_name_from_config( + processor_name: str, + revision: str | None = "main", +) -> str | None: + processor_name = convert_model_repo_to_path(processor_name) + config_file = [ + "video_preprocessor_config.json", + "preprocessor_config.json", + ] + for file in config_file: + config = get_hf_file_to_dict(file, processor_name, revision=revision) + if config and "video_processor_type" in config: + return config["video_processor_type"] + return None + + +_cached_get_video_processor_cls_name = lru_cache( + get_video_processor_cls_name_from_config +) + + +def get_video_processor_cls_name( + model_config: "ModelConfig", +) -> str | None: + if is_gguf(model_config.model): + assert not is_gguf(model_config.tokenizer), ( + "For multimodal GGUF models, the original tokenizer " + "should be used to correctly load video processor metadata." + ) + model = model_config.tokenizer + revision = model_config.tokenizer_revision + else: + model = model_config.model + revision = model_config.revision + + return _cached_get_video_processor_cls_name(model, revision=revision) + + def get_processor( processor_name: str, *args: Any, diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index ba2872f8927..b53dd87d608 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -29,6 +29,8 @@ __all__ = [ "KimiAudioProcessor", "KimiK25Processor", "MiMoOmniProcessor", + "MiniCPMOProcessor", + "MiniCPMVProcessor", "MistralCommonPixtralProcessor", "MistralCommonVoxtralProcessor", "NanoNemotronVLProcessor", @@ -61,6 +63,8 @@ _CLASS_TO_MODULE: dict[str, str] = { "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio", "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25", "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni", + "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo", + "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv", "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral", "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral", "Moondream3Processor": "vllm.transformers_utils.processors.moondream3", diff --git a/vllm/transformers_utils/processors/minicpmo.py b/vllm/transformers_utils/processors/minicpmo.py new file mode 100644 index 00000000000..3059b8bac99 --- /dev/null +++ b/vllm/transformers_utils/processors/minicpmo.py @@ -0,0 +1,603 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# coding=utf-8 +# Copyright 2025 The OpenBMB Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for MiniCPMO. +""" + +import math +from typing import Literal, TypeAlias + +import numpy as np +import regex +import torch +import torchaudio +from transformers.image_processing_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.utils import TensorType + +MiniCPMOBatchFeature: TypeAlias = BatchFeature + + +class MiniCPMOProcessor(ProcessorMixin): + r""" + Constructs a MiniCPMV processor which wraps a MiniCPMV image + processor and a MiniCPMV tokenizer into a single processor. + + [`MiniCPMVProcessor`] offers all the functionalities of + [`MiniCPMVImageProcessor`] and [`LlamaTokenizerWrapper`]. See the + [`~MiniCPMVProcessor.__call__`] and [`~MiniCPMVProcessor.decode`] + for more information. + + Args: + image_processor ([`MiniCPMVImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`LlamaTokenizerWrapper`], *optional*): + The tokenizer is a required input. + """ + + attributes = ["image_processor", "feature_extractor", "tokenizer"] + feature_extractor_class = "WhisperFeatureExtractor" + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__( + self, + image_processor=None, + feature_extractor=None, + tokenizer=None, + pool_step=2, + ): + super().__init__(image_processor, feature_extractor, tokenizer) + self.version = image_processor.version + self.pool_step = pool_step + + def _safe_get_token_id(self, attr_name, default_token_str): + """Get token ID safely, with fallback to default.""" + val = getattr(self.tokenizer, attr_name, None) + if val is None: + val = self.tokenizer.convert_tokens_to_ids(default_token_str) + if val is None: + return -1 + return val + + def _safe_get_token_str(self, attr_name, default_token_str): + """Get token string safely, with fallback to default.""" + return getattr(self.tokenizer, attr_name, default_token_str) + + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + images: ImageInput = None, + audios: np.ndarray | list[np.ndarray] | list[list[np.ndarray]] = None, + audio_parts: list | None = None, + max_length: int | None = None, + do_pad: bool | None = True, + max_slice_nums: int | None = None, + use_image_id: bool = True, + chunk_input: bool = False, + return_tensors: str | TensorType | None = TensorType.PYTORCH, + sampling_rate: int | None = 16000, + **kwargs, + ) -> MiniCPMOBatchFeature: + if images is not None: + image_inputs = self.image_processor( + images, + do_pad=do_pad, + max_slice_nums=max_slice_nums, + return_tensors=return_tensors, + ) + else: + image_inputs = None + + if audios is not None: + audio_features, audio_feature_lens, audio_phs = self.audio_feature_extract( + audios, audio_parts, chunk_input, sampling_rate + ) + else: + audio_features, audio_feature_lens, audio_phs = [], [], [] + + model_inputs = self._convert_omni_to_inputs( + image_inputs, + audio_phs, + text, + max_slice_nums=max_slice_nums, + use_image_id=use_image_id, + max_length=max_length, + **kwargs, + ) + + model_inputs["audio_features"] = audio_features + model_inputs["audio_feature_lens"] = audio_feature_lens + + return MiniCPMOBatchFeature(data={**model_inputs}) + + def get_audio_placeholder(self, audio_lens, chunk_input, chunk_length): + pool_step = self.pool_step + feature_lens = math.ceil(audio_lens / self.feature_extractor.hop_length) + + feature_lens = (feature_lens - 1) // 2 + 1 + output_lens = (feature_lens - pool_step) // pool_step + 1 + + audio_start = getattr(self.tokenizer, "audio_start", "") + + if chunk_input: + fbank_feat_in_chunk = int(chunk_length * 100) + cnn_feat_in_chunk = (fbank_feat_in_chunk - 1) // 2 + 1 + audio_embeds_in_chunk = (cnn_feat_in_chunk - pool_step) // pool_step + 1 + num_audio_chunks = ( + output_lens + audio_embeds_in_chunk - 1 + ) // audio_embeds_in_chunk + + place_holders = "" + total_unk_len = 0 + for _ in range(num_audio_chunks): + unk_len = min(audio_embeds_in_chunk, output_lens - total_unk_len) + place_holders += audio_start + "" * unk_len + audio_end + total_unk_len += unk_len + audio_placeholder = place_holders + else: + audio_placeholder = audio_start + "" * output_lens + audio_end + + return audio_placeholder + + def audio_feature_extract( + self, + audios: np.ndarray | list[np.ndarray] | list[list[np.ndarray]], + audio_parts: list | None = None, + chunk_input: bool | None = False, + sampling_rate: int | None = None, + chunk_length: int | None = 1, + **kwargs, + ): + if isinstance(audios, np.ndarray): + audios_list = [[audios]] + elif isinstance(audios[0], np.ndarray): + audios_list = [audios] + else: + audios_list = audios + + if audio_parts is not None: + assert len(audio_parts) == len(audios_list) + for parts, audios in zip(audio_parts, audios_list): + assert len(parts) == len(audios) + + audio_feature_lens_list = [] + audio_ph_list = [] + + audio_features_all = [] + + # audio placeholder not dependent on audio_parts + for audios in audios_list: + if audios: + audio_ph_list.append( + [ + self.get_audio_placeholder(len(a), chunk_input, chunk_length) + for a in audios + ] + ) + else: + audio_ph_list.append([]) + + for idx, audios in enumerate(audios_list): + if audio_parts is not None: + # same audio part merge + audio_part = audio_parts[idx] + merge_audio = [] + cur_audio = [] + for aid, (part, audio) in enumerate(zip(audio_part, audios)): + if aid == 0 or audio_part[aid] == audio_part[aid - 1]: + cur_audio.append(audio) + else: + merge_audio.append(np.hstack(cur_audio)) + cur_audio = [audio] + if cur_audio: + merge_audio.append(np.hstack(cur_audio)) + + else: + merge_audio = audios + + audio_feature_lens = [] + + # If the audio exceeds 30 seconds, split it into chunks every 30 seconds. + final_merge_audio = [] + max_audio_inp_len = 30 * (sampling_rate or 16000) + for audio in merge_audio: + if len(audio) <= max_audio_inp_len: + final_merge_audio.append(audio) + else: + for i in range(math.ceil(len(audio) / max_audio_inp_len)): + final_merge_audio.append( + audio[i * max_audio_inp_len : (i + 1) * max_audio_inp_len] + ) + + if audios: + audio_inputs = self.feature_extractor( + final_merge_audio, + sampling_rate=sampling_rate, + return_attention_mask=True, + padding="max_length", + return_tensors="pt", + **kwargs, + ) + audio_feature = audio_inputs["input_features"] + actual_lens = audio_inputs["attention_mask"].sum(dim=1) + + for feat, lens in zip(audio_feature, actual_lens): + audio_features_all.append(feat[:, :lens]) + audio_feature_lens.append(lens) + + audio_feature_lens = torch.hstack(audio_feature_lens) + audio_feature_lens_list.append(audio_feature_lens) + else: + audio_feature_lens_list.append([]) + + if audio_features_all: + audio_features = [i.permute(1, 0) for i in audio_features_all] + audio_features = torch.nn.utils.rnn.pad_sequence( + audio_features, batch_first=True, padding_value=0.0 + ).permute(0, 2, 1) + else: + audio_features = [] + + return audio_features, audio_feature_lens_list, audio_ph_list + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode + # with CLIP->Llama + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's + [`~PreTrainedTokenizer.batch_decode`]. Please refer to the + docstring of this method for more information. + """ + output_ids = args[0] + result_text = [] + for result in output_ids: + result = result[result != 0] + if len(result) > 0 and result[0] == self.tokenizer.bos_id: + result = result[1:] + if len(result) > 0 and result[-1] == self.tokenizer.eos_id: + result = result[:-1] + result_text.append( + self.tokenizer.decode(result, *args[1:], **kwargs).strip() + ) + return result_text + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode + # with CLIP->Llama + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's + [`~PreTrainedTokenizer.decode`]. Please refer to the docstring + of this method for more information. + """ + result = args[0] + result = result[result != 0] + if len(result) > 0 and result[0] == self.tokenizer.bos_id: + result = result[1:] + if len(result) > 0 and ( + result[-1] == self.tokenizer.eos_id + or ( + hasattr(self.tokenizer, "eot_id") + and result[-1] == self.tokenizer.eot_id + ) + ): + result = result[:-1] + return self.tokenizer.decode(result, *args[1:], **kwargs).strip() + + def _convert(self, input_str, max_inp_length: int | None = None, **kwargs): + input_ids = self.tokenizer.encode(input_str, **kwargs) + if max_inp_length is not None: + input_ids = input_ids[:max_inp_length] + input_ids = torch.tensor(input_ids, dtype=torch.int32) + + ## image bound + start_cond = (input_ids == self.tokenizer.im_start_id) | ( + input_ids == self.tokenizer.slice_start_id + ) + end_cond = (input_ids == self.tokenizer.im_end_id) | ( + input_ids == self.tokenizer.slice_end_id + ) + + image_start_idx = torch.where(start_cond)[0] + image_start_idx += 1 + image_end_idx = torch.where(end_cond)[0] + + assert len(image_start_idx) == len(image_end_idx), ( + f"The number of image start tokens ({len(image_start_idx)}) " + f"and end tokens ({len(image_end_idx)}) must match." + ) + + image_bounds = torch.hstack( + [ + image_start_idx.unsqueeze(-1), + image_end_idx.unsqueeze(-1), + ] + ) + + ## audio bound + audio_start_idx = torch.where(input_ids == self.tokenizer.audio_start_id)[0] + audio_end_idx = torch.where(input_ids == self.tokenizer.audio_end_id)[0] + assert len(audio_start_idx) == len(audio_end_idx) + audio_bounds = torch.hstack( + [(audio_start_idx + 1).unsqueeze(-1), audio_end_idx.unsqueeze(-1)] + ) + + spk_start_idx = torch.where(input_ids == self.tokenizer.spk_start_id)[0] + spk_end_idx = torch.where(input_ids == self.tokenizer.spk_end_id)[0] + assert len(spk_start_idx) == len(spk_end_idx) + spk_bounds = torch.hstack( + [(spk_start_idx + 1).unsqueeze(-1), spk_end_idx.unsqueeze(-1)] + ) + + return input_ids, image_bounds, audio_bounds, spk_bounds + + def _convert_omni_to_inputs( + self, + images, + audio_phs, + texts: str | list[str], + truncation=None, + max_length=None, + max_slice_nums=None, + use_image_id=None, + return_tensors=None, + **kwargs, + ): + if images is None and audio_phs is None: + model_inputs = self.tokenizer( + texts, + return_tensors=return_tensors, + truncation=truncation, + max_length=max_length, + **kwargs, + ) + return MiniCPMOBatchFeature(data={**model_inputs}) + + image_tag = "(./)" + image_pattern = r"\(./\)" + audio_tag = "()" + audio_pattern = r"\(\)" + split_pattern = rf"({image_pattern}|{audio_pattern})" + + if isinstance(texts, str): + texts = [texts] + + bs = len(texts) + if images is not None: + images, image_sizes, tgt_sizes = ( + images["pixel_values"], + images["image_sizes"], + images["tgt_sizes"], + ) + else: + images, image_sizes, tgt_sizes = [[]] * bs, [[]] * bs, [[]] * bs + + input_ids_list = [] + image_bounds_list = [] + audio_bounds_list = [] + spk_bounds_list = [] + + for index, text in enumerate(texts): + text_chunks = regex.split(split_pattern, text) + + image_tags = regex.findall(image_pattern, text) + audio_tags = regex.findall(audio_pattern, text) + + if image_tags: + assert images is not None + assert len(image_tags) == len(image_sizes[index]) + if audio_tags: + assert audio_phs is not None + assert len(audio_tags) == len(audio_phs[index]) + + image_id = 0 + audio_id = 0 + for i, chunk in enumerate(text_chunks): + if chunk == image_tag: + image_placeholder = ( + self.image_processor.get_slice_image_placeholder( + image_sizes[index][image_id], + image_id, + max_slice_nums, + use_image_id, + ) + ) + image_id += 1 + text_chunks[i] = image_placeholder + elif chunk == audio_tag: + audio_placeholder = audio_phs[index][audio_id] + audio_id += 1 + text_chunks[i] = audio_placeholder + + final_text = "".join(text_chunks) + input_ids, image_bounds, audio_bounds, spk_bounds = self._convert( + final_text, max_length, **kwargs + ) + + input_ids_list.append(input_ids) + image_bounds_list.append(image_bounds) + audio_bounds_list.append(audio_bounds) + spk_bounds_list.append(spk_bounds) + + padded_input_ids, padding_lengths = self.pad( + input_ids_list, padding_side="left" + ) + attention_mask = torch.ones_like(padded_input_ids, dtype=torch.bool) + for i, length in enumerate(padding_lengths): + image_bounds_list[i] = image_bounds_list[i] + length + audio_bounds_list[i] = audio_bounds_list[i] + length + spk_bounds_list[i] = spk_bounds_list[i] + length + attention_mask[i, :length] = False + + data = { + "input_ids": padded_input_ids, + "attention_mask": attention_mask, + "pixel_values": images, + "image_sizes": image_sizes, + "image_bound": image_bounds_list, + "tgt_sizes": tgt_sizes, + "audio_bounds": audio_bounds_list, + "spk_bounds": spk_bounds_list, + } + + return data + + @property + # Copied from + # transformers.models.clip.processing_clip.CLIPProcessor.model_input_names + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + feature_extractor_input_names = self.feature_extractor.model_input_names + return list( + dict.fromkeys( + tokenizer_input_names + + image_processor_input_names + + feature_extractor_input_names + ) + ) + + def pad( + self, + inputs, + max_length=None, + padding_value=0, + padding_side="left", + ): + if not inputs: + return torch.empty(0), [] + + items = [] + if isinstance(inputs[0], list): + assert isinstance(inputs[0][0], torch.Tensor) + for it in inputs: + for tr in it: + items.append(tr) + else: + assert isinstance(inputs[0], torch.Tensor) + items = inputs + + batch_size = len(items) + shape = items[0].shape + dim = len(shape) + assert dim <= 2 + if max_length is None: + max_length = 0 + max_length = max(max_length, max(item.shape[-1] for item in items)) + min_length = min(item.shape[-1] for item in items) + dtype = items[0].dtype + + if dim == 0: + return torch.stack([item for item in items], dim=0), [0] + elif dim == 1: + if max_length == min_length: + return ( + torch.stack([item for item in items], dim=0), + [0] * batch_size, + ) + tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value + else: + tensor = ( + torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype) + + padding_value + ) + + padding_length = [] + for i, item in enumerate(items): + if dim == 1: + if padding_side == "left": + tensor[i, -len(item) :] = item.clone() + else: + tensor[i, : len(item)] = item.clone() + elif dim == 2: + if padding_side == "left": + tensor[i, -len(item) :, :] = item.clone() + else: + tensor[i, : len(item), :] = item.clone() + padding_length.append(tensor.shape[-1] - len(item)) + + return tensor, padding_length + + +class MelSpectrogramFeatures(torch.nn.Module): + def __init__( + self, + sample_rate=24000, + n_fft=1024, + hop_length=256, + n_mels=100, + padding: Literal["center", "same"] = "center", + ): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.mel_spec = torchaudio.transforms.MelSpectrogram( + sample_rate=sample_rate, + n_fft=n_fft, + hop_length=hop_length, + n_mels=n_mels, + center=padding == "center", + power=1, + ) + + def __call__(self, audio: torch.Tensor) -> torch.Tensor: + """ + audio: Tensor([num_channels, num_samples]) + """ + return super().__call__(audio) + + def forward(self, audio: torch.Tensor) -> torch.Tensor: + """ + audio: Tensor([num_channels, num_samples]) + """ + mel: torch.Tensor = self.mel_spec(audio) + features = torch.log(torch.clip(mel, min=1e-5)) + return features + + +class ChatTTSProcessor: + def __init__(self, text_tokenizer): + self.audio_processor = MelSpectrogramFeatures() + self.text_tokenizer = text_tokenizer + + def __call__(self, text_list, audio_list): + assert len(text_list) == len(audio_list) + input_ids_varlen = [] + for text in text_list: + input_ids_ = self.text_tokenizer.encode( + text, return_tensors="pt", add_special_tokens=False + ) # [1, seq_len] + input_ids_ = input_ids_.squeeze(0) # [seq_len] + input_ids_varlen.append(input_ids_) + + audio_features_varlen = [] + for audio in audio_list: + assert audio.shape.__len__() == 1 # [seq_len] + try: + mel = self.audio_processor(audio) # [100(num_mel_bins), seq_len_mel] + except Exception as e: + raise e + audio_features_varlen.append(mel) + + return { + "tts_input_ids_varlen": input_ids_varlen, # return List[Tensor] + "tts_input_features_varlen": audio_features_varlen, # return List[Tensor] + } diff --git a/vllm/transformers_utils/processors/minicpmv.py b/vllm/transformers_utils/processors/minicpmv.py new file mode 100644 index 00000000000..cc0dee8dacd --- /dev/null +++ b/vllm/transformers_utils/processors/minicpmv.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2024 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Processor class for MiniCPMV. +""" + +from typing import TypeAlias + +import regex +import torch +from transformers.image_processing_utils import BatchFeature +from transformers.image_utils import ImageInput +from transformers.processing_utils import ProcessorMixin +from transformers.tokenization_utils_base import ( + PaddingStrategy, + PreTokenizedInput, + TextInput, + TruncationStrategy, +) +from transformers.utils import TensorType + +MiniCPMVBatchFeature: TypeAlias = BatchFeature + + +class MiniCPMVProcessor(ProcessorMixin): + r""" + Constructs a MiniCPMV processor which wraps a MiniCPMV image + processor and a MiniCPMV tokenizer into a single processor. + + [`MiniCPMVProcessor`] offers all the functionalities of + [`MiniCPMVImageProcessor`] and [`LlamaTokenizerWrapper`]. See the + [`~MiniCPMVProcessor.__call__`] and [`~MiniCPMVProcessor.decode`] + for more information. + + Args: + image_processor ([`MiniCPMVImageProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`LlamaTokenizerWrapper`], *optional*): + The tokenizer is a required input. + """ + + attributes = ["image_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__(self, image_processor=None, tokenizer=None): + super().__init__(image_processor, tokenizer) + self.version = image_processor.version + + def __call__( + self, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput], + images: ImageInput = None, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = None, + max_length: int | None = None, + do_pad: bool | None = True, + return_tensors: str | TensorType | None = TensorType.PYTORCH, + ) -> MiniCPMVBatchFeature: + """Run the vendored MiniCPMV processor on a (text, images) pair. + + Only single-sample input is currently supported; batched input is + coming soon. ``images`` is forwarded to the underlying image + processor and ``text`` is tokenized with image placeholders + replaced by the appropriate slice tokens. Returns a + ``MiniCPMVBatchFeature`` with at minimum ``input_ids`` and (when + images are provided) ``pixel_values``, ``image_sizes``, + ``image_bound`` and ``tgt_sizes``. + """ + if images is not None: + image_inputs = self.image_processor( + images, do_pad=do_pad, return_tensors=return_tensors + ) + else: + image_inputs = {} + return self._convert_images_texts_to_inputs( + image_inputs, text, max_length=max_length + ) + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor + # .batch_decode with CLIP->Llama + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's + [`~PreTrainedTokenizer.batch_decode`]. Please refer to the + docstring of this method for more information. + """ + output_ids = args[0] + result_text = [] + + bos_id = getattr( + self.tokenizer, + "bos_token_id", + getattr(self.tokenizer, "bos_id", 1), + ) + eos_id = getattr( + self.tokenizer, + "eos_token_id", + getattr(self.tokenizer, "eos_id", 2), + ) + + for result in output_ids: + result = result[result != 0] + if len(result) > 0 and result[0] == bos_id: + result = result[1:] + if len(result) > 0 and result[-1] == eos_id: + result = result[:-1] + result_text.append( + self.tokenizer.decode(result, *args[1:], **kwargs).strip() + ) + return result_text + + # Copied from transformers.models.clip.processing_clip.CLIPProcessor + # .decode with CLIP->Llama + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LlamaTokenizerFast's + [`~PreTrainedTokenizer.decode`]. Please refer to the docstring + of this method for more information. + """ + result = args[0] + result = result[result != 0] + + bos_id = getattr( + self.tokenizer, + "bos_token_id", + getattr(self.tokenizer, "bos_id", 1), + ) + eos_id = getattr( + self.tokenizer, + "eos_token_id", + getattr(self.tokenizer, "eos_id", 2), + ) + eot_id = getattr(self.tokenizer, "eot_id", None) + + if len(result) > 0 and result[0] == bos_id: + result = result[1:] + if len(result) > 0 and ( + result[-1] == eos_id or (eot_id is not None and result[-1] == eot_id) + ): + result = result[:-1] + return self.tokenizer.decode(result, *args[1:], **kwargs).strip() + + 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: + input_ids = self.tokenizer.encode(input_str) + else: + bos_id = getattr( + self.tokenizer, + "bos_token_id", + getattr(self.tokenizer, "bos_id", 1), + ) + input_ids = [bos_id] + self.tokenizer.encode(input_str) + + if max_inp_length is not None: + input_ids = input_ids[:max_inp_length] + input_ids = torch.tensor(input_ids, dtype=torch.int32) + + im_start_id = getattr( + self.tokenizer, + "im_start_id", + self.tokenizer.convert_tokens_to_ids(""), + ) + im_end_id = getattr( + self.tokenizer, + "im_end_id", + self.tokenizer.convert_tokens_to_ids(""), + ) + + image_start_tokens = torch.where(input_ids == im_start_id)[0] + image_start_tokens += 1 + image_end_tokens = torch.where(input_ids == im_end_id)[0] + assert len(image_start_tokens) == len(image_end_tokens), ( + f"The number of image start tokens ({len(image_start_tokens)}) " + f"and end tokens ({len(image_end_tokens)}) must match." + ) + image_bounds = torch.hstack( + [ + image_start_tokens.unsqueeze(-1), + image_end_tokens.unsqueeze(-1), + ] + ) + return input_ids.unsqueeze(0), image_bounds + + def _convert_images_texts_to_inputs( + self, + images, + texts, + do_pad=False, + truncation=None, + max_length=None, + return_tensors=None, + ): + if not len(images): + model_inputs = self.tokenizer( + texts, + return_tensors=return_tensors, + padding=do_pad, + truncation=truncation, + max_length=max_length, + ) + return MiniCPMVBatchFeature(data={**model_inputs}) + + pattern = "(./)" + images_val = images["pixel_values"] + image_sizes = images["image_sizes"] + tgt_sizes = images["tgt_sizes"] + + image_tags = regex.findall(pattern, texts) + assert len(image_tags) == len(image_sizes[0]) + text_chunks = texts.split(pattern) + final_texts = "" + for i in range(len(image_tags)): + placeholder = self.image_processor.get_slice_image_placeholder( + image_sizes[0][i] + ) + final_texts = final_texts + text_chunks[i] + placeholder + final_texts += text_chunks[-1] + input_ids, image_bounds = self._convert(final_texts, max_length) + return MiniCPMVBatchFeature( + data={ + "input_ids": input_ids, + "pixel_values": images_val, + "image_sizes": image_sizes, + "image_bound": [image_bounds], + "tgt_sizes": tgt_sizes, + } + ) + + @property + # Copied from + # transformers.models.clip.processing_clip.CLIPProcessor.model_input_names + 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 pad( + self, + orig_items, + key, + max_length=None, + padding_value=0, + padding_side="left", + ): + if not orig_items: + return torch.empty(0) + + items = [] + if isinstance(orig_items[0][key], list): + assert isinstance(orig_items[0][key][0], torch.Tensor) + for it in orig_items: + for tr in it[key]: + items.append({key: tr}) + else: + assert isinstance(orig_items[0][key], torch.Tensor) + items = orig_items + + batch_size = len(items) + shape = items[0][key].shape + dim = len(shape) + assert dim <= 3 + if max_length is None: + max_length = 0 + max_length = max(max_length, max(item[key].shape[-1] for item in items)) + min_length = min(item[key].shape[-1] for item in items) + dtype = items[0][key].dtype + + if dim == 1: + return torch.cat([item[key] for item in items], dim=0) + elif dim == 2: + if max_length == min_length: + return torch.cat([item[key] for item in items], dim=0) + tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value + else: + tensor = ( + torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype) + + padding_value + ) + + for i, item in enumerate(items): + tensor_to_pad = item[key] + if tensor_to_pad.shape[0] != 1: + raise ValueError( + f"Expected leading batch size of 1 for padding, " + f"but got shape {tensor_to_pad.shape}" + ) + squeezed = tensor_to_pad.squeeze(0) + if dim == 2: + if padding_side == "left": + tensor[i, -squeezed.shape[0] :] = squeezed.clone() + else: + tensor[i, : squeezed.shape[0]] = squeezed.clone() + elif dim == 3: + if padding_side == "left": + tensor[i, -squeezed.shape[0] :, :] = squeezed.clone() + else: + tensor[i, : squeezed.shape[0], :] = squeezed.clone() + + return tensor diff --git a/vllm/transformers_utils/processors/voxtral.py b/vllm/transformers_utils/processors/voxtral.py index 93f97729134..829bab2d415 100644 --- a/vllm/transformers_utils/processors/voxtral.py +++ b/vllm/transformers_utils/processors/voxtral.py @@ -44,7 +44,7 @@ class MistralCommonFeatureExtractor: if not self.audio_encoder.audio_config.is_streaming: audio = self.audio_encoder.pad(audio, self.sampling_rate) - audios_processed.append(torch.tensor(audio)) + audios_processed.append(torch.from_numpy(audio)) return BatchFeature( {"audio_arrays": audios_processed}, tensor_type=return_tensors diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 6b89f5c3320..4252ce87754 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -140,6 +140,7 @@ _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None +@functools.cache def _import_deep_gemm(): """Import the deep_gemm module. diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index e008e17d806..c37b3b6c70c 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -404,7 +404,7 @@ def _has_module(module_name: str) -> bool: if importlib.util.find_spec(module_name) is None: return False importlib.import_module(module_name) - except ImportError: + except Exception: logger.warning( "Module %s was found but failed to import", module_name, exc_info=True ) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 005975c4775..3519691a3c5 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -93,7 +93,7 @@ class CPUAttentionBackend(AttentionBackend): head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: - return 2, num_blocks, num_kv_heads, block_size, head_size + return num_blocks, num_kv_heads, block_size, 2 * head_size @classmethod def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": @@ -308,7 +308,7 @@ class CPUAttentionBackendImpl(AttentionImpl): key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = - [2, num_blocks, num_kv_heads, block_size, head_size] + [num_blocks, num_kv_heads, block_size, 2 * head_size] attn_metadata: Metadata for attention. Returns: shape = [num_tokens, num_heads * head_size] @@ -338,8 +338,12 @@ class CPUAttentionBackendImpl(AttentionImpl): ) # For decoder and cross-attention, use KV cache, size are - # [num_blocks, num_kv_heads, block_size, head_size] - key_cache, value_cache = kv_cache.unbind(0) + # [num_blocks, num_kv_heads, block_size, 2 * head_size] + # Make a view [num_blocks, num_kv_heads, block_size * 2, head_size] + # Then slice KV at dim 2 + num_blocks, num_kv_heads, block_size, _ = kv_cache.size() + kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1)) + key_cache, value_cache = kv_cache.chunk(2, dim=2) # key and value may be None in the case of cross attention. They are # calculated once based on the output from the encoder and then cached diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index a81c5742c1b..83e3072546f 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -623,6 +623,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4" if self.is_kvcache_nvfp4: + # trtllm-gen FP4 FMHA kernels only exist for sm100f (sm_100/sm_103). + # Fail fast at init rather than crashing on the first request. + if not current_platform.is_device_capability_family(100): + raise ValueError( + "--kv-cache-dtype nvfp4 requires sm100f, " + "please try a different dtype or remove" + ) # For NVFP4, kv_cache_dtype stays as the string "nvfp4" # which is passed to FlashInferImpl self.kv_cache_dtype = self.cache_dtype diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index 6f1a2560d2c..d479239e3b1 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -71,6 +71,8 @@ class EncoderCacheManager: # mm_hash of mm_data => ids of requests that reference the mm_data self.cached: dict[str, set[str]] = {} + # request_id => set of input_ids cached for that request + self.request_cached_ids: dict[str, set[int]] = {} # mm_hash of mm_data => num_encoder_embeds of the mm_data self.freeable: OrderedDict[str, int] = OrderedDict() @@ -83,6 +85,7 @@ class EncoderCacheManager: Called when model weights are updated to invalidate stale embeddings. """ self.cached.clear() + self.request_cached_ids.clear() self.freeable.clear() self.freed.clear() self.num_free_slots = self.cache_size @@ -114,6 +117,7 @@ class EncoderCacheManager: self.num_freeable_slots -= num_encoder_embeds self.cached[mm_hash].add(request.request_id) + self.request_cached_ids.setdefault(request.request_id, set()).add(input_id) return True def can_allocate( @@ -201,22 +205,13 @@ class EncoderCacheManager: assert self.num_freeable_slots >= num_encoder_embeds self.cached[mm_hash].add(request_id) + self.request_cached_ids.setdefault(request_id, set()).add(input_id) self.num_free_slots -= num_encoder_embeds self.num_freeable_slots -= num_encoder_embeds def get_cached_input_ids(self, request: Request) -> set[int]: - """Get all cached multimodal input IDs for a request. - - Returns the set of input IDs whose `mm_hash` exists in the cache map. - This includes entries that are currently unreferenced (and thus present - in `freeable`); for such entries, freeing for this request will be a - no-op. - """ - return { - input_id - for input_id in range(len(request.mm_features)) - if request.mm_features[input_id].identifier in self.cached - } + """Get all cached multimodal input IDs for a request.""" + return self.request_cached_ids.get(request.request_id, set()) def free_encoder_input(self, request: Request, input_id: int) -> None: """Free the request's reference to the encoder input (`mm_data`) @@ -230,6 +225,12 @@ class EncoderCacheManager: """ req_id = request.request_id mm_hash = request.mm_features[input_id].identifier + # Always clean up request_cached_ids, even if the mm_hash was + # already evicted from cache (e.g. by can_allocate). + if req_id in self.request_cached_ids: + self.request_cached_ids[req_id].discard(input_id) + if not self.request_cached_ids[req_id]: + del self.request_cached_ids[req_id] # The mm_hash not in cache or the req_id set is empty if not self.cached.get(mm_hash, None): return @@ -248,8 +249,7 @@ class EncoderCacheManager: Typically called when a request is finished, cancelled, or aborted. """ - input_ids = self.get_cached_input_ids(request) - for input_id in input_ids: + for input_id in list(self.get_cached_input_ids(request)): self.free_encoder_input(request, input_id) def get_freed_mm_hashes(self) -> list[str]: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c5e8953745a..387f1a1e335 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod from collections.abc import Sequence -from math import lcm +from typing import NamedTuple from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector @@ -40,12 +40,20 @@ class KVCacheCoordinator(ABC): enable_kv_cache_events: bool, dcp_world_size: int, pcp_world_size: int, + scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, ): self.kv_cache_config = kv_cache_config self.max_model_len = max_model_len self.enable_caching = enable_caching + # The scheduling granularity (LCM of all group block sizes), must be a multiple + # of the hash_block_size and the block size of each group. + assert scheduler_block_size % hash_block_size == 0 and all( + scheduler_block_size % g.kv_cache_spec.block_size == 0 + for g in kv_cache_config.kv_cache_groups + ) + self.scheduler_block_size = scheduler_block_size self.block_pool = BlockPool( num_gpu_blocks=kv_cache_config.num_blocks, @@ -73,6 +81,7 @@ class KVCacheCoordinator(ABC): kv_cache_group_id=i, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=self.scheduler_block_size, ) for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups) ) @@ -290,6 +299,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): enable_kv_cache_events: bool, dcp_world_size: int, pcp_world_size: int, + scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, ): @@ -302,6 +312,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator): enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) @@ -338,6 +349,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): enable_kv_cache_events: bool, dcp_world_size: int, pcp_world_size: int, + scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, ): @@ -350,6 +362,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) @@ -369,6 +382,8 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): assert len(self.kv_cache_config.kv_cache_groups) == 1, ( "UnitaryKVCacheCoordinator assumes only one kv cache group" ) + # Single group; useless but just set ``use_eagle`` for consistency regardless. + self.single_type_managers[0].use_eagle = 0 in self.eagle_group_ids def find_longest_cache_hit( self, @@ -381,7 +396,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): kv_cache_group_ids=[0], block_pool=self.block_pool, kv_cache_spec=self.kv_cache_spec, - use_eagle=0 in self.eagle_group_ids, + drop_eagle_block=0 in self.eagle_group_ids, alignment_tokens=self.block_size, dcp_world_size=self.dcp_world_size, pcp_world_size=self.pcp_world_size, @@ -389,6 +404,21 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): return hit_blocks, len(hit_blocks[0]) * self.block_size +class SpecGroup(NamedTuple): + """KV cache groups that share one spec, batched together for a single + cache-hit lookup. + + ``use_eagle`` is True iff any member group is an EAGLE/MTP group. Members + sharing a spec are cached and looked up jointly, so the EAGLE last-block drop + is necessarily decided for the whole spec group. + """ + + spec: KVCacheSpec + group_ids: list[int] + manager_cls: type[SingleTypeKVCacheManager] + use_eagle: bool + + class HybridKVCacheCoordinator(KVCacheCoordinator): """ KV cache coordinator for hybrid models with multiple KV cache types, and @@ -405,6 +435,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): enable_kv_cache_events: bool, dcp_world_size: int, pcp_world_size: int, + scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, ): @@ -417,6 +448,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) @@ -438,66 +470,63 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): Groups KV cache groups by their spec type for efficient batch processing during cache hit lookup. """ - attention_groups: list[ - tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]] - ] = [] - + self.attention_groups: list[SpecGroup] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): manager_cls = self.single_type_managers[i].__class__ spec = g.kv_cache_spec + use_eagle = i in self.eagle_group_ids # Try to find an existing group with the same spec - for existing_spec, group_ids, existing_cls in attention_groups: - if existing_spec == spec: - assert manager_cls is existing_cls, ( + for idx, group in enumerate(self.attention_groups): + if group.spec == spec: + assert manager_cls is group.manager_cls, ( "Expected same manager class for identical KV cache specs." ) - group_ids.append(i) + group.group_ids.append(i) + if use_eagle and not group.use_eagle: + self.attention_groups[idx] = group._replace(use_eagle=True) break else: - attention_groups.append((spec, [i], manager_cls)) + self.attention_groups.append( + SpecGroup(spec, [i], manager_cls, use_eagle) + ) - assert len(attention_groups) > 1, ( + assert len(self.attention_groups) > 1, ( "HybridKVCacheCoordinator requires at least two attention groups." ) # Put full attention first: its efficient left-to-right scan provides # a tighter initial bound, reducing work for subsequent groups. - self.attention_groups = sorted( - attention_groups, - key=lambda x: not isinstance(x[0], FullAttentionSpec), + self.attention_groups.sort( + key=lambda g: not isinstance(g.spec, FullAttentionSpec) ) - # The LCM of the block sizes of all attention types. - # The cache hit length must be a multiple of the LCM of the block sizes - # to make sure the cache hit length is a multiple of the block size of - # each attention type. Requiring this because we don't support partial - # block cache hit yet. - block_sizes = [spec.block_size for spec, _, _ in attention_groups] - self.lcm_block_size = lcm(*block_sizes) - - # Attention-group indices (into ``self.attention_groups``) that - # contain at least one EAGLE/MTP KV cache group. - self.eagle_attn_group_indices: set[int] = { - i - for i, (_, group_ids, _) in enumerate(self.attention_groups) - if any(gid in self.eagle_group_ids for gid in group_ids) - } + # Propagate the eagle bit to each manager (default to ``use_eagle=False``). + for group in self.attention_groups: + if group.use_eagle: + for gid in group.group_ids: + self.single_type_managers[gid].use_eagle = True def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: # Cache hits in this coordinator are always a multiple of - # ``lcm_block_size`` tokens (see ``find_longest_cache_hit``). Within an - # aligned region, SWA groups only consult a subset of blocks per - # ``lcm_block_size``-segment so the unused blocks also stay out of the - # prefix-cache hash map. - num_computed_tokens = ( - num_computed_tokens // self.lcm_block_size * self.lcm_block_size + # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). + # Within an aligned region, SWA groups may only consult a subset of blocks + # per ``scheduler_block_size``-segment so the unused blocks also stay + # out of the prefix-cache hash map. + aligned_num_computed_tokens = ( + num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size ) for manager in self.single_type_managers: + num_tokens_to_cache = aligned_num_computed_tokens + # EAGLE groups match one block past each aligned boundary and drop + # it, so make that lookahead block eligible to be cached. + if manager.use_eagle and aligned_num_computed_tokens > 0: + num_tokens_to_cache = min( + num_computed_tokens, + aligned_num_computed_tokens + manager.block_size, + ) manager.cache_blocks( - request, - num_computed_tokens, - alignment_tokens=self.lcm_block_size, + request, num_tokens_to_cache, alignment_tokens=self.scheduler_block_size ) def find_longest_cache_hit( @@ -537,7 +566,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): # Simple hybrid (1 full attn + 1 other): one iteration suffices. # Full attn is always first if it exists. is_simple_hybrid = len(self.attention_groups) == 2 and isinstance( - self.attention_groups[0][0], FullAttentionSpec + self.attention_groups[0].spec, FullAttentionSpec ) # Attention-group indices whose EAGLE drop is verified at the current @@ -548,7 +577,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): while True: curr_hit_length = hit_length - for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): + for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate( + self.attention_groups + ): cached_blocks = hit_blocks_by_group[group_ids[0]] if isinstance(spec, FullAttentionSpec) and cached_blocks is not None: # Full attention is downward-closed: we only need to look @@ -559,12 +590,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): ) continue - use_eagle = ( - idx in self.eagle_attn_group_indices and idx not in eagle_verified - ) + drop_eagle_block = use_eagle and idx not in eagle_verified _max_length = curr_hit_length - if use_eagle: + if drop_eagle_block: # Eagle needs to match one more block and then pop the last. _max_length = min( curr_hit_length + spec.block_size, max_cache_hit_length @@ -575,11 +604,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, - use_eagle=use_eagle, - alignment_tokens=self.lcm_block_size, + drop_eagle_block=drop_eagle_block, + alignment_tokens=self.scheduler_block_size, ) _new_hit_length = len(hit_blocks[0]) * spec.block_size - if use_eagle: + if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: # length shrunk; invalidate previous eagle verifications @@ -595,10 +624,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): break # Truncate full attention blocks to final hit_length (if present) - spec, group_ids, _ = self.attention_groups[0] - if isinstance(spec, FullAttentionSpec): - num_blocks = hit_length // spec.block_size - for group_id in group_ids: + first_group = self.attention_groups[0] + if isinstance(first_group.spec, FullAttentionSpec): + num_blocks = hit_length // first_group.spec.block_size + for group_id in first_group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: del blks[num_blocks:] @@ -616,6 +645,7 @@ def get_kv_cache_coordinator( enable_kv_cache_events: bool, dcp_world_size: int, pcp_world_size: int, + scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, ) -> KVCacheCoordinator: @@ -628,6 +658,7 @@ def get_kv_cache_coordinator( enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) @@ -641,6 +672,7 @@ def get_kv_cache_coordinator( enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) @@ -653,6 +685,7 @@ def get_kv_cache_coordinator( enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9359d8843a9..d98520da95f 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -112,6 +112,7 @@ class KVCacheManager: self, kv_cache_config: KVCacheConfig, max_model_len: int, + scheduler_block_size: int, hash_block_size: int, max_num_batched_tokens: int | None = None, enable_caching: bool = True, @@ -147,6 +148,7 @@ class KVCacheManager: enable_kv_cache_events=enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=self.metrics_collector, ) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 7f3a5e4fdf3..cfa79f077a1 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request from vllm.v1.utils import tensor_data @@ -1991,6 +1992,9 @@ def get_kv_cache_configs( "across workers. This is not supported yet." ) + # Check if the KV cache specs are registered correctly. + # This is to prevent that some layers are initialized with unregistered specs. + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) # Get global KV cache groups. This also handles spec unification for # hybrid models when disable_hybrid_kv_cache_manager is enabled. # After this call, merged_kv_cache_specs may be modified in-place. diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index cb61bcabd3e..2fd22f4c0cb 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -14,6 +14,7 @@ class AsyncScheduler(Scheduler): super().__init__(*args, **kwargs) # reusable read-only placeholder list for speculative decoding. self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens + self.pp_size = self.parallel_config.pipeline_parallel_size def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: super()._update_after_schedule(scheduler_output) @@ -34,6 +35,11 @@ class AsyncScheduler(Scheduler): # We will update the actual spec token ids in the worker process. request.spec_token_ids = self._spec_token_placeholders + if self.use_v2_model_runner: + # Set the next step index in which this request is eligible to be + # scheduled for decode (for PP microbatching). + request.next_decode_eligible_step = self.current_step + self.pp_size + def _update_request_with_output( self, request: Request, new_token_ids: list[int] ) -> tuple[list[int], bool]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 73d3dcb4b65..c39e80c24eb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget +from vllm.multimodal.utils import get_mm_features_in_window from vllm.v1.core.encoder_cache_manager import ( EncoderCacheManager, EncoderDecoderCacheManager, @@ -103,7 +104,7 @@ class Scheduler(SchedulerInterface): self.max_num_running_reqs = self.scheduler_config.max_num_seqs self.max_num_scheduled_tokens = ( self.scheduler_config.max_num_scheduled_tokens - if self.scheduler_config.max_num_scheduled_tokens + if self.scheduler_config.max_num_scheduled_tokens is not None else self.scheduler_config.max_num_batched_tokens ) self.max_model_len = vllm_config.model_config.max_model_len @@ -237,6 +238,7 @@ class Scheduler(SchedulerInterface): enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, pcp_world_size=self.pcp_world_size, + scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, ) @@ -247,6 +249,9 @@ class Scheduler(SchedulerInterface): self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = vllm_config.use_v2_model_runner + # Scheduler iteration counter. Drives the V2+PP+async decode-throttle + # cadence (`next_decode_eligible_step`). + self.current_step = 0 self.scheduler_reserve_full_isl = ( self.scheduler_config.scheduler_reserve_full_isl ) @@ -332,6 +337,7 @@ class Scheduler(SchedulerInterface): return num_new_tokens def schedule(self) -> SchedulerOutput: + self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: # There's no "decoding phase" nor "prefill phase" in the scheduler. # Each request just has the num_computed_tokens and @@ -387,6 +393,12 @@ class Scheduler(SchedulerInterface): req_index += 1 continue + if self.current_step < request.next_decode_eligible_step: + # V2+PP+async: enforce `pp_size` steps between same-req decodes + # to match worker-side sampled-tokens broadcast slot ring cadence. + req_index += 1 + continue + num_new_tokens = ( request.num_tokens_with_spec + request.num_output_placeholders @@ -628,6 +640,18 @@ class Scheduler(SchedulerInterface): ) assert num_computed_tokens <= request.num_tokens + # Skip request with pending mm encoding prefetches + if ( + self.ec_connector is not None + and request.mm_features + and not self.ec_connector.ensure_cache_available( + request, num_computed_tokens + ) + ): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + # Track first scheduled prefill, not post-preemption repeat prefills if request.prefill_stats is not None: assert num_computed_tokens <= request.num_prompt_tokens @@ -1140,22 +1164,23 @@ class Scheduler(SchedulerInterface): # trackers for accounting at the encoder input level. mm_hashes_to_schedule = set() num_embeds_to_schedule = 0 - for i, mm_feature in enumerate(mm_features): + + lo, hi = get_mm_features_in_window( + mm_features, + start=num_computed_tokens, + end=num_computed_tokens + num_new_tokens + shift_computed_tokens, + ) + # For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always. + if self.is_encoder_decoder: + lo = 0 + + for i in range(lo, hi): + mm_feature = mm_features[i] start_pos = mm_feature.mm_position.offset num_encoder_tokens = mm_feature.mm_position.length num_encoder_embeds = mm_feature.mm_position.get_num_embeds() item_identifier = mm_feature.identifier - # The encoder output is needed if the two ranges overlap: - # [num_computed_tokens, num_computed_tokens + num_new_tokens) and - # [start_pos, start_pos + num_encoder_tokens) - if ( - start_pos - >= num_computed_tokens + num_new_tokens + shift_computed_tokens - ): - # The encoder input is not needed in this step. - break - if self.is_encoder_decoder and num_computed_tokens > 0: assert start_pos == 0, ( "Encoder input should be processed at the beginning of " @@ -1171,10 +1196,6 @@ class Scheduler(SchedulerInterface): # decoder tokens (num_computed_tokens > 0), then we know we # already calculated encoder inputs and can skip here. continue - elif start_pos + num_encoder_tokens <= num_computed_tokens: - # The encoder input is already computed and stored - # in the decoder's KV cache. - continue if not self.is_encoder_decoder: # We are not using the encoder cache for encoder-decoder models, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index cd000dc849e..281b79639db 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -25,6 +25,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, TQFullAttentionSpec, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -40,6 +41,7 @@ class SingleTypeKVCacheManager(ABC): block_pool: BlockPool, enable_caching: bool, kv_cache_group_id: int, + scheduler_block_size: int, dcp_world_size: int = 1, pcp_world_size: int = 1, max_admission_blocks_per_request: int | None = None, @@ -50,6 +52,8 @@ class SingleTypeKVCacheManager(ABC): kv_cache_spec: The kv_cache_spec for this manager. block_pool: The block pool. 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``. 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, @@ -57,6 +61,7 @@ class SingleTypeKVCacheManager(ABC): correct for full-attention-style specs that hold every block until the request finishes. """ + self.scheduler_block_size = scheduler_block_size self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size @@ -82,6 +87,12 @@ class SingleTypeKVCacheManager(ABC): self.kv_cache_group_id = kv_cache_group_id self._null_block = block_pool.null_block + # Whether this group's prefix-cache hits drop the EAGLE/MTP lookahead + # block. Only consulted by managers whose hit logic is sparse within an + # aligned segment (SWA). Initialized lazily by the coordinator after + # determining the attention groups. + self.use_eagle = False + @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) @@ -313,8 +324,12 @@ class SingleTypeKVCacheManager(ABC): if alignment_tokens is None or alignment_tokens <= self.block_size: block_mask = None else: - block_mask = self._cache_block_mask( - num_cached_blocks, num_full_blocks, alignment_tokens + block_mask = self.reachable_block_mask( + num_cached_blocks, + num_full_blocks, + alignment_tokens, + self.kv_cache_spec, + self.use_eagle, ) self.block_pool.cache_full_blocks( request=request, @@ -328,11 +343,14 @@ class SingleTypeKVCacheManager(ABC): self.num_cached_block[request.request_id] = num_full_blocks - def _cache_block_mask( - self, - num_cached_blocks: int, - num_full_blocks: int, + @classmethod + def reachable_block_mask( + cls, + start_block: int, + num_blocks: int, alignment_tokens: int, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, ) -> list[bool] | None: """Per-block mask for ``cache_full_blocks``. ``None`` means cache every (non-null) block — the default for full attention. @@ -385,7 +403,7 @@ class SingleTypeKVCacheManager(ABC): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -405,7 +423,10 @@ class SingleTypeKVCacheManager(ABC): kv_cache_group_ids: The ids of the kv cache groups. block_pool: The block pool. kv_cache_spec: The kv cache spec. - use_eagle: Whether to use eagle. + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP. + Always False for non-EAGLE/MTP groups, but can be False for EAGLE/MTP + groups too if the last block is already dropped (e.g., in a + convergence loop in `find_longest_cache_hit`). alignment_tokens: The returned cache hit length (in tokens) should be a multiple of this value (in tokens). By default, it should be set to the block_size. @@ -495,7 +516,7 @@ class FullAttentionManager(SingleTypeKVCacheManager): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -524,7 +545,7 @@ class FullAttentionManager(SingleTypeKVCacheManager): computed.append(cached) else: break - if use_eagle and computed_blocks[0]: + if drop_eagle_block and computed_blocks[0]: # Need to drop the last matched block if eagle is enabled. for computed in computed_blocks: computed.pop() @@ -552,6 +573,19 @@ class SlidingWindowManager(SingleTypeKVCacheManager): super().__init__(kv_cache_spec, **kwargs) self.sliding_window = kv_cache_spec.sliding_window + @classmethod + def _contiguous_blocks_for_hit( + cls, window_size: int, block_size: int, use_eagle: bool + ) -> int: + blocks = cdiv(window_size - 1, block_size) + if use_eagle: + # Need to drop the last matched block if eagle is enabled. For + # sliding window layer, we achieve this by increasing the number of + # contiguous blocks needed for prefix cache hit by one and dropping + # the last matched block. + blocks += 1 + return blocks + @classmethod def find_longest_cache_hit( cls, @@ -560,7 +594,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -571,17 +605,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager): assert dcp_world_size == 1, "DCP not support sliding window attn now." assert pcp_world_size == 1, "PCP not support sliding window attn now." - # The number of contiguous blocks needed for prefix cache hit. - # -1 since the input token itself is also included in the window - sliding_window_contiguous_blocks = cdiv( - kv_cache_spec.sliding_window - 1, kv_cache_spec.block_size + # The number of contiguous blocks needed for a prefix cache hit. + sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit( + kv_cache_spec.sliding_window, kv_cache_spec.block_size, drop_eagle_block ) - if use_eagle: - # Need to drop the last matched block if eagle is enabled. For - # sliding window layer, we achieve this by increasing the number of - # contiguous blocks needed for prefix cache hit by one and dropping - # the last matched block. - sliding_window_contiguous_blocks += 1 # TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to # optimize the time complexity from O(max_num_blocks) to @@ -604,7 +631,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager): # Skip prefix matching check if the block is not aligned with # `alignment_tokens`. if num_contiguous_blocks == 0 and block_size != alignment_tokens: - post_pop_blocks = i if use_eagle else i + 1 + post_pop_blocks = i if drop_eagle_block else i + 1 if (post_pop_blocks * block_size) % alignment_tokens != 0: continue # Add the cached block to the computed blocks. @@ -632,7 +659,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager): ): for computed in computed_blocks: computed.pop() - if use_eagle and computed_blocks[0]: + if drop_eagle_block and computed_blocks[0]: for computed in computed_blocks: computed.pop() # Re-align after eagle pop: the pop may break the alignment @@ -646,17 +673,33 @@ class SlidingWindowManager(SingleTypeKVCacheManager): computed.pop() return computed_blocks - def _cache_block_mask( - self, num_cached_blocks: int, num_full_blocks: int, alignment_tokens: int + @classmethod + def reachable_block_mask( + cls, + start_block: int, + num_blocks: int, + alignment_tokens: int, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, ) -> list[bool] | None: - assert alignment_tokens > self.block_size - per_segment = alignment_tokens // self.block_size - tail = cdiv(self.sliding_window - 1, self.block_size) - if tail >= per_segment: + assert alignment_tokens > kv_cache_spec.block_size + assert isinstance(kv_cache_spec, SlidingWindowSpec) + per_segment = alignment_tokens // kv_cache_spec.block_size + need = cls._contiguous_blocks_for_hit( + window_size=kv_cache_spec.sliding_window, + block_size=kv_cache_spec.block_size, + use_eagle=use_eagle, + ) + if need >= per_segment: return None - skip = per_segment - tail + # The matched run's right edge sits on the aligned boundary block when + # EAGLE peeks one block past it (shift=1), otherwise on the last block + # before the boundary (shift=0). A block is reachable iff it falls in + # the ``need``-wide run ending at some boundary's right edge. + shift = 1 if use_eagle else 0 return [ - i % per_segment >= skip for i in range(num_cached_blocks, num_full_blocks) + i >= shift and (i - shift) % per_segment >= per_segment - need + for i in range(start_block, num_blocks) ] def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: @@ -710,7 +753,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -741,7 +784,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): kv_cache_group_ids: The ids of the kv cache groups. block_pool: The block pool. kv_cache_spec: The kv cache spec. - use_eagle: Whether to use eagle. + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP. dcp_world_size: The world size of decode context parallelism. pcp_world_size: The world size of prefill context parallelism. alignment_tokens: The returned cache hit length (in tokens) should @@ -754,7 +797,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): "ChunkedLocalAttentionManager can only be used for " "chunked local attention groups" ) - assert use_eagle is False, ( + assert drop_eagle_block is False, ( "Hybrid KV cache is not supported for " + "eagle + chunked local attention." ) assert dcp_world_size == 1, "DCP not support chunked local attn now." @@ -870,7 +913,7 @@ class MambaManager(SingleTypeKVCacheManager): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -1164,7 +1207,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager): kv_cache_group_ids: list[int], block_pool: BlockPool, kv_cache_spec: KVCacheSpec, - use_eagle: bool, + drop_eagle_block: bool, alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, @@ -1205,27 +1248,30 @@ class SinkFullAttentionManager(FullAttentionManager): self.sink_blocks = self.block_pool.free_block_queue.popleft_n(num_sink_block) -spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = { - FullAttentionSpec: FullAttentionManager, - TQFullAttentionSpec: FullAttentionManager, - MLAAttentionSpec: FullAttentionManager, - HiddenStateCacheSpec: FullAttentionManager, - SlidingWindowSpec: SlidingWindowManager, - SlidingWindowMLASpec: SlidingWindowManager, - ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager, - MambaSpec: MambaManager, - CrossAttentionSpec: CrossAttentionManager, - SinkFullAttentionSpec: SinkFullAttentionManager, -} - - def get_manager_for_kv_cache_spec( kv_cache_spec: KVCacheSpec, max_num_batched_tokens: int, max_model_len: int, **kwargs, ) -> SingleTypeKVCacheManager: - manager_class = spec_manager_map[type(kv_cache_spec)] + """ + Get the appropriate manager for a given KVCacheSpec. + + Uses the KVCacheSpecRegistry to look up the manager class, supporting + both built-in and custom specs registered via @register_kv_cache_spec + and KVCacheSpecRegistry.register. + + Args: + kv_cache_spec: The KVCacheSpec instance + max_num_batched_tokens: The maximum number of tokens in a batch + max_model_len: The maximum context length the model could serve + Returns: + An instance of the appropriate SingleTypeKVCacheManager subclass + """ + manager_class = KVCacheSpecRegistry.get_manager_class(kv_cache_spec) + assert manager_class is not None, ( + f"No manager registered for KVCacheSpec {type(kv_cache_spec)}" + ) # SlidingWindow / ChunkedLocalAttention managers recycle blocks across # chunks; the runtime admission cap must match the recycling-aware bound # the startup pool sizer uses (single source of truth: the spec method). @@ -1238,3 +1284,64 @@ def get_manager_for_kv_cache_spec( ) manager = manager_class(kv_cache_spec, **kwargs) return manager + + +def register_all_kvcache_specs(vllm_config): + """Built-in spec registration""" + KVCacheSpecRegistry.register( + FullAttentionSpec, + FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + + KVCacheSpecRegistry.register( + SlidingWindowSpec, + SlidingWindowManager, + uniform_type_base_spec=SlidingWindowSpec, + ) + KVCacheSpecRegistry.register( + SlidingWindowMLASpec, + SlidingWindowManager, + uniform_type_base_spec=SlidingWindowMLASpec, + ) + + KVCacheSpecRegistry.register( + MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec + ) + KVCacheSpecRegistry.register( + ChunkedLocalAttentionSpec, + ChunkedLocalAttentionManager, + uniform_type_base_spec=ChunkedLocalAttentionSpec, + ) + KVCacheSpecRegistry.register( + CrossAttentionSpec, + CrossAttentionManager, + uniform_type_base_spec=CrossAttentionSpec, + ) + + # FullAttentionSpec subclasses — grouped with FullAttentionSpec + KVCacheSpecRegistry.register( + TQFullAttentionSpec, + FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + KVCacheSpecRegistry.register( + MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec + ) + # NOTE(Mengqing): HiddenStateCacheSpec won't take part in + # grouping, thus the uniform_type_base_spec is just a + # placeholder. + KVCacheSpecRegistry.register( + HiddenStateCacheSpec, + FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + KVCacheSpecRegistry.register( + SinkFullAttentionSpec, + SinkFullAttentionManager, + uniform_type_base_spec=FullAttentionSpec, + ) + + from vllm.platforms import current_platform + + current_platform.register_custom_kv_cache_specs(vllm_config) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index aa1756bf682..848f530ce33 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -74,6 +74,7 @@ class EngineCoreReadyResponse: max_model_len: int num_gpu_blocks: int + block_size: int dp_stats_address: str | None dtype: str vllm_version: str diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index c21a4de5d30..b12aa9d0505 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import ( ) from vllm.v1.core.sched.interface import PauseState, SchedulerInterface from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs from vllm.v1.engine import ( EEP_NOTIFICATION_CALL_ID, EEPNotificationType, @@ -188,7 +189,7 @@ class EngineCore: # Batch queue for scheduled batches. This enables us to asynchronously # schedule and execute batches, and is required by pipeline parallelism # to eliminate pipeline bubbles. - self.batch_queue_size = self.model_executor.max_concurrent_batches + self.batch_queue_size = vllm_config.max_concurrent_batches self.batch_queue: ( deque[tuple[Future[ModelRunnerOutput], SchedulerOutput, Future[Any]]] | None ) = None @@ -235,6 +236,9 @@ class EngineCore: def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: start = time.time() + # register all kvcache specs in enginecore process. + register_all_kvcache_specs(vllm_config) + # Get all kv cache needed by the model kv_cache_specs = self.model_executor.get_kv_cache_specs() @@ -534,14 +538,12 @@ class EngineCore: if not deferred_scheduler_output: # Add this step's future to the queue. batch_queue.appendleft((future, scheduler_output, exec_future)) - if ( - model_executed - and len(batch_queue) < self.batch_queue_size - and not batch_queue[-1][0].done() + if len(batch_queue) < self.batch_queue_size and ( + model_executed or self.scheduler.has_requests() ): # Don't block on next worker response unless the queue is full # or there are no more requests to schedule. - return None, True + return None, model_executed elif not batch_queue: # Queue is empty. We should not reach here since this method should @@ -1462,6 +1464,7 @@ class EngineCoreProc(EngineCore): ready_response = EngineCoreReadyResponse( max_model_len=self.vllm_config.model_config.max_model_len, num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, + block_size=self.vllm_config.cache_config.block_size, dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index c26380e6e15..14257b020ee 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -713,6 +713,10 @@ class MPClient(EngineCoreClient): num_gpu_blocks += response.num_gpu_blocks vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks + # Sync block_size: may be enlarged by _align_hybrid_block_size in the + # worker for hybrid Mamba models. + vllm_config.cache_config.block_size = response.block_size + # In external DP LB mode, the coordinator address that the # front-end procs connect to is obtained by each engine via it's # initial handshake with the rank 0 front-end. diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index e68c0283f57..7beef598e27 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -253,10 +253,6 @@ class Executor(ABC): output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") return output[0] - @property - def max_concurrent_batches(self) -> int: - return 1 - def profile(self, is_start: bool = True, profile_prefix: str | None = None): self.collective_rpc("profile", args=(is_start, profile_prefix)) diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index 449fbdd9736..c5766c923c8 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -15,7 +15,7 @@ from concurrent.futures import Future, InvalidStateError from contextlib import suppress from dataclasses import dataclass from enum import Enum, auto -from functools import cached_property, partial +from functools import partial from multiprocessing.connection import Connection from multiprocessing.process import BaseProcess from multiprocessing.synchronize import Lock as LockType @@ -472,12 +472,6 @@ class MultiprocExecutor(Executor): self.collective_rpc("check_health", timeout=10) return - @cached_property - def max_concurrent_batches(self) -> int: - # PP requires PP-size concurrent batches to fill the pipeline. - pp_size = self.parallel_config.pipeline_parallel_size - return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size - def _get_output_rank(self) -> int: # Only returns ModelRunnerOutput from TP rank=0 and PP rank=-1 # (the first TP worker of the last PP stage). diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index cfeebb5e09d..749e59e04c2 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -96,14 +96,6 @@ class RayDistributedExecutor(Executor): self.scheduler_output: SchedulerOutput | None = None - @property - def max_concurrent_batches(self) -> int: - """Ray distributed executor supports pipeline parallelism, - meaning that it allows PP size batches to be executed concurrently. - """ - pp_size = self.parallel_config.pipeline_parallel_size - return 2 if pp_size <= 1 and self.scheduler_config.async_scheduling else pp_size - def shutdown(self) -> None: if logger: # Somehow logger can be None here. diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py index c3be3300fd3..dd04b718d67 100644 --- a/vllm/v1/executor/uniproc_executor.py +++ b/vllm/v1/executor/uniproc_executor.py @@ -3,7 +3,6 @@ import os from collections.abc import Callable from concurrent.futures import Future -from functools import cached_property from multiprocessing import Lock from typing import Any @@ -77,10 +76,6 @@ class UniProcExecutor(Executor): local_rank = int(device_info[1]) if len(device_info) > 1 else 0 return distributed_init_method, 0, local_rank - @cached_property - def max_concurrent_batches(self) -> int: - return 2 if self.scheduler_config.async_scheduling else 1 - def collective_rpc( # type: ignore[override] self, method: str | Callable, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 31ee89bc72a..3bbfba1a0fe 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -17,6 +17,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry if TYPE_CHECKING: from vllm.config import VllmConfig @@ -139,6 +140,21 @@ class KVCacheSpec: ) return copy.deepcopy(specs[0]) + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + """ + Whether this KVCacheSpec is uniform with all specs of all layers. + """ + uniform_type_base_spec = KVCacheSpecRegistry.get_uniform_type_base_spec(self) + assert uniform_type_base_spec is not None, ( + f"Unsupported KV cache spec type: {type(self)}. " + "Please register it using @register_kv_cache_spec decorator." + ) + return all( + isinstance(spec, uniform_type_base_spec) for spec in kv_cache_specs.values() + ) + @dataclass(frozen=True, kw_only=True) class AttentionSpec(KVCacheSpec): @@ -430,6 +446,15 @@ class ChunkedLocalAttentionSpec(AttentionSpec): ) return max_blocks * self.page_size_bytes + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + return all( + isinstance(spec, ChunkedLocalAttentionSpec) + and spec.attention_chunk_size == self.attention_chunk_size + for spec in kv_cache_specs.values() + ) + @dataclass(frozen=True, kw_only=True) class SlidingWindowSpec(AttentionSpec): @@ -493,6 +518,15 @@ class SlidingWindowSpec(AttentionSpec): ) return max_blocks * self.page_size_bytes + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + return all( + isinstance(spec, SlidingWindowSpec) + and spec.sliding_window == self.sliding_window + for spec in kv_cache_specs.values() + ) + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): @@ -558,6 +592,15 @@ class SlidingWindowMLASpec(SlidingWindowSpec): model_version=model_version_set.pop(), ) + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + return all( + isinstance(spec, SlidingWindowMLASpec) + and spec.sliding_window == self.sliding_window + for spec in kv_cache_specs.values() + ) + @dataclass(frozen=True) class MambaSpec(KVCacheSpec): @@ -590,6 +633,15 @@ class MambaSpec(KVCacheSpec): else: return self.page_size_bytes * (1 + self.num_speculative_blocks) + def is_uniform_with_collection( + self, kv_cache_specs: dict[str, KVCacheSpec] + ) -> bool: + return all( + isinstance(spec, MambaSpec) + and spec.num_speculative_blocks == self.num_speculative_blocks + for spec in kv_cache_specs.values() + ) + @dataclass(frozen=True) class EncoderOnlyAttentionSpec(AttentionSpec): @@ -689,53 +741,16 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: """ Whether all layers have the same type of KV cache spec. + + Uses the registry to determine grouping base classes, so custom specs + that inherit from FullAttentionSpec are treated as full attention. """ block_sizes = set(spec.block_size for spec in kv_cache_specs.values()) if len(block_sizes) > 1: # Different block sizes, not uniform. return False - one_spec = next(iter(kv_cache_specs.values())) - # NOTE: Check subclasses before parent classes since isinstance() - # returns True for subclasses. - if isinstance(one_spec, SlidingWindowMLASpec): - # SlidingWindowMLASpec is uniform if all specs are SlidingWindowMLASpec - # with the same sliding_window size. - return all( - isinstance(spec, SlidingWindowMLASpec) - and spec.sliding_window == one_spec.sliding_window - for spec in kv_cache_specs.values() - ) - elif isinstance(one_spec, FullAttentionSpec): - return all( - isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values() - ) - elif isinstance(one_spec, CrossAttentionSpec): - return all( - isinstance(spec, CrossAttentionSpec) for spec in kv_cache_specs.values() - ) - elif isinstance(one_spec, SlidingWindowSpec): - return all( - isinstance(spec, SlidingWindowSpec) - and spec.sliding_window == one_spec.sliding_window - for spec in kv_cache_specs.values() - ) - elif isinstance(one_spec, ChunkedLocalAttentionSpec): - return all( - isinstance(spec, ChunkedLocalAttentionSpec) - and spec.attention_chunk_size == one_spec.attention_chunk_size - for spec in kv_cache_specs.values() - ) - elif isinstance(one_spec, MambaSpec): - return all( - isinstance(spec, MambaSpec) - and spec.num_speculative_blocks == one_spec.num_speculative_blocks - for spec in kv_cache_specs.values() - ) - else: - # NOTE(Chen): Please add new branches for new KV cache spec types. - raise NotImplementedError( - f"Unsupported KV cache spec type: {type(one_spec)}" - ) + first_spec = next(iter(kv_cache_specs.values())) + return first_spec.is_uniform_with_collection(kv_cache_specs) @classmethod def from_specs(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> Self | None: diff --git a/vllm/v1/kv_cache_spec_registry.py b/vllm/v1/kv_cache_spec_registry.py new file mode 100644 index 00000000000..816a6862dae --- /dev/null +++ b/vllm/v1/kv_cache_spec_registry.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Registry for KVCacheSpec types and their associated managers. + +This module provides a pluggable architecture for registering custom KVCacheSpec +subclasses without modifying vLLM core code. Out-of-tree platforms can define +custom specs and managers by using the @register_kv_cache_spec decorator. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager + from vllm.v1.kv_cache_interface import KVCacheSpec + + +@dataclass(frozen=True) +class KVCacheSpecMetadata: + """Metadata for a registered KVCacheSpec.""" + + kvcache_spec_cls: type["KVCacheSpec"] + manager_class: type["SingleTypeKVCacheManager"] + # The base spec class for grouping compatibility checks. + # KVCacheSpecs with the same uniform_type_base_spec will be + # grouped into one kvcache group + uniform_type_base_spec: type["KVCacheSpec"] + + +_REGISTRY_KVCACHESPEC_LIST: dict[type["KVCacheSpec"], KVCacheSpecMetadata] = {} + + +class KVCacheSpecRegistry: + """Global registry for KVCacheSpec types and their associated managers.""" + + @classmethod + def _ensure_registered(cls, vllm_config=None) -> None: + """ + Run full KVCacheSpec registration if the registration is not done. + """ + if _REGISTRY_KVCACHESPEC_LIST: + return + + if vllm_config is None: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + + # lazy import to avoid circular dependency + from vllm.v1.core.single_type_kv_cache_manager import ( + register_all_kvcache_specs, + ) + + register_all_kvcache_specs(vllm_config) + + @classmethod + def register( + cls, + kvcache_spec_cls: type["KVCacheSpec"], + manager_class: type["SingleTypeKVCacheManager"] | None = None, + uniform_type_base_spec: type["KVCacheSpec"] | None = None, + ) -> None: + """ + Register a KVCacheSpec class with its manager and base spec. + + Args: + kvcache_spec_cls: The KVCacheSpec subclass to register + manager_class: The SingleTypeKVCacheManager to use for this spec + uniform_type_base_spec: The base spec class for grouping compatibility. + instead of being grouped to different kvcache group, `kvcache_spec_cls` + and `uniform_type_base_spec` will be trated as uniform type. + If None, defaults to kvcache_spec_cls itself (for built-in base specs). + """ + assert manager_class is not None, "manager_class is required" + if uniform_type_base_spec is None: + uniform_type_base_spec = kvcache_spec_cls + assert issubclass(kvcache_spec_cls, uniform_type_base_spec), ( + f"{kvcache_spec_cls.__name__} must inherit from its declared " + f"uniform_type_base_spec {uniform_type_base_spec.__name__}." + ) + + if kvcache_spec_cls in _REGISTRY_KVCACHESPEC_LIST: + registered_spec = _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] + is_same_registration = ( + manager_class == registered_spec.manager_class + and uniform_type_base_spec == registered_spec.uniform_type_base_spec + ) + assert is_same_registration, ( + f"Conflicting registration for KVCacheSpec " + f": {kvcache_spec_cls.__name__}" + ) + + _REGISTRY_KVCACHESPEC_LIST[kvcache_spec_cls] = KVCacheSpecMetadata( + kvcache_spec_cls=kvcache_spec_cls, + manager_class=manager_class, + uniform_type_base_spec=uniform_type_base_spec, + ) + + @classmethod + def get_manager_class( + cls, kvcache_spec: "KVCacheSpec" + ) -> type["SingleTypeKVCacheManager"] | None: + """ + Get the single type kvcache manager class for a given kvcache spec instance. + + Args: + kvcache_spec: A KVCacheSpec instance + + Returns: + The SingleTypeKVCacheManager class to use for this kvcache_spec + """ + cls._ensure_registered() + kvcache_spec_cls = type(kvcache_spec) + + # Walk up the MRO to find a registered base class + for base in kvcache_spec_cls.__mro__: + if base in _REGISTRY_KVCACHESPEC_LIST: + return _REGISTRY_KVCACHESPEC_LIST[base].manager_class + + return None + + @classmethod + def get_uniform_type_base_spec( + cls, kvcache_spec: "KVCacheSpec" + ) -> type["KVCacheSpec"] | None: + """ + Get the base kvcache spec class for grouping compatibility checks. + KVCacheSpecs with uniform_type_base_spec will be trated as one group. + + Args: + kvcache_spec: A KVCacheSpec instance + + Returns: + The base KVCacheSpec class for checking uniform type kvcache specs + """ + cls._ensure_registered() + kvcache_spec_cls = type(kvcache_spec) + + # Walk up the MRO to find a registered base spec + for base in kvcache_spec_cls.__mro__: + if base in _REGISTRY_KVCACHESPEC_LIST: + return _REGISTRY_KVCACHESPEC_LIST[base].uniform_type_base_spec + + return None + + @classmethod + def check_kv_cache_spec_registry( + cls, kv_cache_spec: dict[str, "KVCacheSpec"] + ) -> None: + """ + Check if the KVCacheSpecs of each layer are registered as expected. + """ + cls._ensure_registered() + for layer_name, spec in kv_cache_spec.items(): + # use raise instead of assert to make it effective in production environment + if cls.get_uniform_type_base_spec(spec) is None: + raise ValueError( + f"Unsupported KV cache spec type for layer {layer_name}: " + f"{type(spec)}. Please register it using " + f"@register_kv_cache_spec decorator." + ) + if cls.get_manager_class(spec) is None: + raise ValueError( + f"No manager found for KV cache spec type for layer " + f"{layer_name}: {type(spec)}. Please register it using " + f"@register_kv_cache_spec decorator." + ) + + +def register_kv_cache_spec( + manager_class: type["SingleTypeKVCacheManager"] | None = None, + uniform_type_base_spec: type["KVCacheSpec"] | None = None, +): + """ + Decorator to register a custom KVCacheSpec class. + + Args: + manager_class: The SingleTypeKVCacheManager to use for this spec. + Required for all registered specs. + uniform_type_base_spec: The base spec class for uniform type kv cache specs + compatibility. If None, the spec is treated as a new base + type. + + Examples: + - Register a new specs: + @register_kv_cache_spec( + manager_class=FullAttentionManager, + uniform_type_base_spec=FullAttentionSpec + ) + @dataclass(frozen=True, kw_only=True) + class CustomFullAttentionSpec(FullAttentionSpec): + pass + """ + + def decorator(kvcache_spec_cls: type["KVCacheSpec"]) -> type["KVCacheSpec"]: + KVCacheSpecRegistry.register( + kvcache_spec_cls=kvcache_spec_cls, + manager_class=manager_class, + uniform_type_base_spec=uniform_type_base_spec, + ) + return kvcache_spec_cls + + return decorator diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index ec168de8917..5f798f41eac 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, 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 @@ -256,6 +257,14 @@ class OffloadingManager(ABC): """ return () + def on_schedule_end(self) -> None: + """Called once at the end of each scheduler step. + + Managers may override this to flush deferred work accumulated + during the step (e.g., batched promotions). + """ + return + def reset_cache(self) -> None: """Evict all tracked blocks and reset internal state.""" return @@ -311,6 +320,7 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec): self.block_indices: Sequence[int] = block_indices @staticmethod + @override def medium() -> str: return "GPU" diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py index cf5b2b39dd6..42f576bb705 100644 --- a/vllm/v1/kv_offload/cpu/common.py +++ b/vllm/v1/kv_offload/cpu/common.py @@ -1,5 +1,7 @@ # 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 @@ -9,5 +11,6 @@ class CPULoadStoreSpec(BlockIDsLoadStoreSpec): """ @staticmethod + @override def medium() -> str: return "CPU" diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 119778368ca..4fbda71d9ed 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -1,14 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools import time from collections import deque from dataclasses import dataclass import numpy as np import torch +from typing_extensions import override from vllm import _custom_ops as ops from vllm.logger import init_logger +from vllm.triton_utils import HAS_TRITON, triton from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.kv_offload.base import ( @@ -18,6 +21,10 @@ from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, ) from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +from vllm.v1.kv_offload.cpu.swap_blocks_triton import ( + THRESHOLD_BYTES, + swap_blocks_batch, +) from vllm.v1.kv_offload.worker.worker import ( OffloadingHandler, TransferResult, @@ -27,6 +34,30 @@ from vllm.v1.kv_offload.worker.worker import ( logger = init_logger(__name__) +def _select_swap_blocks_fn( + kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], + gpu_to_cpu: bool, +): + """Resolve the swap_blocks function for a handler at init time.""" + # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton. + if gpu_to_cpu: + return ops.swap_blocks_batch + # Fall back to the C++ DMA path on platforms where Triton isn't usable + # (e.g. ROCm builds without Triton). + if not HAS_TRITON: + return ops.swap_blocks_batch + page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g] + # Triton wins only on small, 8-byte-aligned payloads. + if ( + not page_sizes + or max(page_sizes) >= THRESHOLD_BYTES + or any(s % 8 for s in page_sizes) + ): + return ops.swap_blocks_batch + chunk = min(triton.next_power_of_2(max(page_sizes)), 8192) + return functools.partial(swap_blocks_batch, bytes_per_chunk=chunk) + + @dataclass class Transfer: job_id: int @@ -34,6 +65,9 @@ class Transfer: start_event: torch.Event end_event: torch.Event num_bytes: int + batch_src: torch.Tensor + batch_dst: torch.Tensor + batch_sizes: torch.Tensor def compute_sub_block_ptrs( @@ -108,6 +142,17 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None: region.is_pinned = True +def _new_descriptor_buffers( + num_copy_ops: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + pin = is_pin_memory_available() + return ( + torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin), + torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin), + torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin), + ) + + class SingleDirectionOffloadingHandler(OffloadingHandler): """ SingleDirectionOffloadingHandler handles transfers for a single direction, @@ -161,6 +206,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): ) self.gpu_to_cpu: bool = gpu_to_cpu self.kv_cache_groups_data_refs = kv_cache_groups_data_refs + self._swap_blocks_batch = _select_swap_blocks_fn( + kv_cache_groups_data_refs, gpu_to_cpu + ) # GPU blocks may be smaller # cpu_page_size = gpu_page_size * block_size_factor. @@ -178,7 +226,10 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): self._stream_pool: list[torch.cuda.Stream] = [] # list of CUDA events available for re-use self._event_pool: list[torch.Event] = [] + # list of pinned descriptor buffer sets available for re-use + self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + @override def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool: src_spec, dst_spec = transfer_spec assert isinstance(src_spec, BlockIDsLoadStoreSpec) @@ -227,9 +278,21 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): ): num_copy_ops += group_size * len(group_data_refs) - all_src = np.empty(num_copy_ops, dtype=np.int64) - all_dst = np.empty(num_copy_ops, dtype=np.int64) - all_sizes = np.empty(num_copy_ops, dtype=np.int64) + # reuse a pooled buffer set, growing it if this transfer needs more room + batch_src, batch_dst, batch_sizes = ( + self._buffer_pool.pop() + if self._buffer_pool + else _new_descriptor_buffers(num_copy_ops) + ) + if batch_src.numel() < num_copy_ops: + batch_src, batch_dst, batch_sizes = _new_descriptor_buffers(num_copy_ops) + + src = batch_src[:num_copy_ops] + dst = batch_dst[:num_copy_ops] + sizes = batch_sizes[:num_copy_ops] + all_src = src.numpy() + all_dst = dst.numpy() + all_sizes = sizes.numpy() src_offset = 0 dst_offset = 0 @@ -292,10 +355,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): assert dst_offset == num_dst_blocks assert op_idx == num_copy_ops - batch_src = torch.from_numpy(all_src) - batch_dst = torch.from_numpy(all_dst) - batch_sizes = torch.from_numpy(all_sizes) - stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream() start_event = ( self._event_pool.pop() @@ -326,10 +385,10 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): with torch.cuda.stream(stream): start_event.record(stream) if num_copy_ops > 0: - ops.swap_blocks_batch( - batch_src, - batch_dst, - batch_sizes, + self._swap_blocks_batch( + src, + dst, + sizes, is_src_access_order_any=is_src_access_order_any, ) end_event.record(stream) @@ -342,12 +401,16 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): start_event=start_event, end_event=end_event, num_bytes=num_transfer_bytes, + batch_src=batch_src, + batch_dst=batch_dst, + batch_sizes=batch_sizes, ) ) # success return True + @override def get_finished(self) -> list[TransferResult]: results: list[TransferResult] = [] while self._transfers and self._transfers[0].end_event.query(): @@ -367,15 +430,20 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): self._stream_pool.append(transfer.stream) self._event_pool.append(transfer.end_event) self._event_pool.append(transfer.start_event) + self._buffer_pool.append( + (transfer.batch_src, transfer.batch_dst, transfer.batch_sizes) + ) del self._transfer_events[transfer.job_id] return results + @override def wait(self, job_ids: set[int]): for job_id in job_ids: event = self._transfer_events.get(job_id) if event is not None: event.synchronize() + @override def shutdown(self) -> None: while self._transfers: transfer = self._transfers.popleft() @@ -383,6 +451,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): self._transfer_events.clear() self._stream_pool.clear() self._event_pool.clear() + self._buffer_pool.clear() self.src_tensors.clear() self.dst_tensors.clear() if self._mmap_region is not None: diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 4700e250c5d..a1d3a30ebb1 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -4,6 +4,8 @@ from collections import OrderedDict from collections.abc import Collection, Iterable from typing import Literal +from typing_extensions import override + from vllm.v1.kv_offload.base import ( LoadStoreSpec, OffloadingEvent, @@ -95,9 +97,11 @@ class CPUOffloadingManager(OffloadingManager): # --- OffloadingManager interface --- + @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return RequestOffloadingContext() + @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: if self.counts is not None: if key in self.counts: @@ -114,6 +118,7 @@ class CPUOffloadingManager(OffloadingManager): return None # write in-flight; caller should retry return True + @override def prepare_load( self, keys: Collection[OffloadKey], @@ -128,9 +133,11 @@ class CPUOffloadingManager(OffloadingManager): blocks.append(block) return self._get_load_store_spec(keys, blocks) + @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: self._policy.touch(keys) + @override def complete_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> None: @@ -140,6 +147,7 @@ class CPUOffloadingManager(OffloadingManager): assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0" block.ref_cnt -= 1 + @override def prepare_store( self, keys: Collection[OffloadKey], @@ -197,6 +205,7 @@ class CPUOffloadingManager(OffloadingManager): evicted_keys=to_evict, ) + @override def complete_store( self, keys: Collection[OffloadKey], @@ -227,6 +236,7 @@ class CPUOffloadingManager(OffloadingManager): ) ) + @override def reset_cache(self) -> None: # Clear ALL blocks unconditionally. The scheduler's _stale_job_threshold # guarantees that complete_load / complete_store are never called for @@ -238,6 +248,7 @@ class CPUOffloadingManager(OffloadingManager): self._free_list.clear() self._num_allocated_blocks = 0 + @override def take_events(self) -> Iterable[OffloadingEvent]: if self.events is not None: yield from self.events diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index 5b01815c2d7..7d22e518654 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -3,6 +3,8 @@ from collections import OrderedDict from collections.abc import Iterable +from typing_extensions import override + from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -54,18 +56,22 @@ class ARCCachePolicy(CachePolicy): self.b1: OrderedDict[OffloadKey, None] = OrderedDict() self.b2: OrderedDict[OffloadKey, None] = OrderedDict() + @override def get(self, key: OffloadKey) -> BlockStatus | None: return self.t1.get(key) or self.t2.get(key) + @override def insert(self, key: OffloadKey, block: BlockStatus) -> None: self.t1[key] = block self.b1.pop(key, None) self.b2.pop(key, None) + @override def remove(self, key: OffloadKey) -> None: if self.t1.pop(key, None) is None: self.t2.pop(key, None) + @override def touch(self, keys: Iterable[OffloadKey]) -> None: for key in reversed(list(keys)): if key in self.t1: @@ -94,6 +100,7 @@ class ARCCachePolicy(CachePolicy): # move to MRU position (end) to keep it fresh in the ghost list self.b2.move_to_end(key) + @override def clear(self) -> None: self.t1.clear() self.t2.clear() @@ -101,6 +108,7 @@ class ARCCachePolicy(CachePolicy): self.b2.clear() self.target_t1_size = 0.0 + @override def evict( self, n: int, protected: set[OffloadKey] ) -> list[tuple[OffloadKey, BlockStatus]] | None: diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index 51680d8bcc5..75fbc6015e1 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -3,6 +3,8 @@ from collections import OrderedDict from collections.abc import Iterable +from typing_extensions import override + from vllm.v1.kv_offload.base import OffloadKey from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -14,23 +16,29 @@ class LRUCachePolicy(CachePolicy): # cache_capacity unused by LRU but accepted for a uniform constructor self.blocks: OrderedDict[OffloadKey, BlockStatus] = OrderedDict() + @override def get(self, key: OffloadKey) -> BlockStatus | None: return self.blocks.get(key) + @override def insert(self, key: OffloadKey, block: BlockStatus) -> None: self.blocks[key] = block + @override def remove(self, key: OffloadKey) -> None: del self.blocks[key] + @override def touch(self, keys: Iterable[OffloadKey]) -> None: for key in reversed(list(keys)): if key in self.blocks: self.blocks.move_to_end(key) + @override def clear(self) -> None: self.blocks.clear() + @override def evict( self, n: int, protected: set[OffloadKey] ) -> list[tuple[OffloadKey, BlockStatus]] | None: diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 1166b44fc7e..b9b415f12d1 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -35,24 +35,26 @@ class SharedOffloadRegion: File path: /dev/shm/vllm_offload_{instance_id}.mmap """ + BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE + def __init__( self, instance_id: str, - total_size_bytes: int, num_blocks: int, rank: int | None, - num_workers: int, + kv_bytes_per_block: int, cpu_page_size: int, ) -> None: self.page_size = mmap.PAGESIZE + assert kv_bytes_per_block % self.page_size == 0 + + self.num_blocks = num_blocks + self._row_stride = kv_bytes_per_block + self.total_size_bytes = self.num_blocks * self._row_stride - self.total_size_bytes = total_size_bytes self.mmap_path = f"/dev/shm/vllm_offload_{instance_id}.mmap" self._creator = False # set True only if this worker creates the file - self.num_blocks = num_blocks self.rank = rank - # interleaved-layout stride: one row = all workers' data for one block - self._row_stride = cpu_page_size * num_workers if rank is not None: # byte offset to this worker's first slot within each block row self._worker_offset = rank * cpu_page_size diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 6d17d5317f1..8791ff5d391 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -2,8 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterator +from typing_extensions import override + from vllm.config import VllmConfig from vllm.platforms import current_platform +from vllm.utils.math_utils import round_up from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.kv_offload.base import ( CanonicalKVCaches, @@ -19,6 +22,8 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler class CPUOffloadingSpec(OffloadingSpec): + BLOCK_SIZE_ALIGNMENT = 1 + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): super().__init__(vllm_config, kv_cache_config) @@ -28,26 +33,34 @@ class CPUOffloadingSpec(OffloadingSpec): "cpu_bytes_to_use must be specified in kv_connector_extra_config" ) - # calculate kv_bytes_per_offloaded_block + world_size = vllm_config.parallel_config.world_size + self.num_blocks = 0 + self.kv_bytes_per_offloaded_block = 0 + self.cpu_page_size_per_worker = 0 assert kv_cache_config is not None - if kv_cache_config.num_blocks > 0: + if kv_cache_config.num_blocks > 0 and world_size > 0: total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors) kv_bytes_per_block = ( total_gpu_kv_bytes // kv_cache_config.num_blocks - ) * vllm_config.parallel_config.world_size - else: - kv_bytes_per_block = 0 + ) * world_size + kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor - kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor - self.num_blocks = ( - int(cpu_bytes_to_use) // kv_bytes_per_offloaded_block - if kv_bytes_per_offloaded_block > 0 - else 0 - ) - world_size = vllm_config.parallel_config.world_size - self.cpu_page_size_per_worker: int = ( - kv_bytes_per_offloaded_block // world_size if world_size > 0 else 0 - ) + # calculate cpu_page_size_per_worker + self.cpu_page_size_per_worker = kv_bytes_per_offloaded_block // world_size + + # calculate num_blocks + aligned_kv_bytes_per_offloaded_block = round_up( + kv_bytes_per_offloaded_block, self.BLOCK_SIZE_ALIGNMENT + ) + self.num_blocks = ( + int(cpu_bytes_to_use) // aligned_kv_bytes_per_offloaded_block + ) + + # Expose aligned_kv_bytes_per_offloaded_block as + # kv_bytes_per_offloaded_block. Note that this might contain + # some padding. i.e. each offloaded block is of the form, + # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** | + self.kv_bytes_per_offloaded_block = aligned_kv_bytes_per_offloaded_block # scheduler-side self._manager: OffloadingManager | None = None @@ -57,6 +70,7 @@ class CPUOffloadingSpec(OffloadingSpec): self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru") + @override def get_manager(self) -> OffloadingManager: if not self._manager: kv_events_config = self.vllm_config.kv_events_config @@ -88,6 +102,7 @@ class CPUOffloadingSpec(OffloadingSpec): num_cpu_blocks=self.num_blocks, ) + @override def get_handlers( self, kv_caches: CanonicalKVCaches ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]: diff --git a/vllm/v1/kv_offload/cpu/swap_blocks_triton.py b/vllm/v1/kv_offload/cpu/swap_blocks_triton.py new file mode 100644 index 00000000000..77d9028d739 --- /dev/null +++ b/vllm/v1/kv_offload/cpu/swap_blocks_triton.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernel + tuned constants for the ``swap_blocks_batch`` fast path.""" + +from __future__ import annotations + +import torch + +from vllm import _custom_ops as ops +from vllm.triton_utils import tl, triton + +# Constants tuned empirically on H100 (PCIe Gen5): +# NUM_SMS - smallest SM slice within 5% of peak bandwidth at the +# 8-32 KB block sizes that matter in practice +# THRESHOLD_BYTES - max payload per descriptor where Triton beats DMA; above +# this the C++ cuMemcpyBatchAsync path takes the lead +# MIN_N - minimum batch size where Triton's per-launch cost is +# amortized; below this DMA wins +NUM_SMS = 12 +THRESHOLD_BYTES = 28 * 1024 +MIN_N = 16 + + +@triton.jit +def _swap_blocks_kernel( + src_addrs, + dst_addrs, + sizes, + n_jobs, # type: ignore[name-defined] + BYTES_PER_CHUNK: tl.constexpr, # type: ignore[name-defined] +): + pid = tl.program_id(0) + num_progs = tl.num_programs(0) + WORDS_PER_CHUNK: tl.constexpr = BYTES_PER_CHUNK // 8 + offsets = tl.arange(0, WORDS_PER_CHUNK) + job = pid + while job < n_jobs: + src = tl.load(src_addrs + job).to(tl.pointer_type(tl.int64)) + dst = tl.load(dst_addrs + job).to(tl.pointer_type(tl.int64)) + words = tl.load(sizes + job) // 8 + for start in range(0, words, WORDS_PER_CHUNK): + idx = start + offsets + mask = idx < words + data = tl.load(src + idx, mask=mask, other=0) + tl.store(dst + idx, data, mask=mask) + job += num_progs + + +def swap_blocks_batch( + src_addrs: torch.Tensor, + dst_addrs: torch.Tensor, + sizes: torch.Tensor, + is_src_access_order_any: bool = False, + *, + bytes_per_chunk: int, +) -> None: + """Triton implementation of ``swap_blocks_batch`` for small CPU->GPU batches.""" + n = src_addrs.numel() + # Too few descriptors to amortize Triton's launch cost. + if n < MIN_N: + ops.swap_blocks_batch( + src_addrs, + dst_addrs, + sizes, + is_src_access_order_any=is_src_access_order_any, + ) + return + _swap_blocks_kernel[(min(NUM_SMS, n),)]( + src_addrs.to("cuda", non_blocking=True), + dst_addrs.to("cuda", non_blocking=True), + sizes.to("cuda", non_blocking=True), + n, + BYTES_PER_CHUNK=bytes_per_chunk, + ) diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index b824eec1c0a..d4f0cefe5eb 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -185,6 +185,14 @@ class SecondaryTierManager(ABC): """ return + def on_schedule_end(self) -> None: + """Called once at the end of each scheduler step. + + Secondary tiers may override this for per-step cleanup or + deferred work submission. + """ + return + def shutdown(self) -> None: """Release resources held by this tier (threads, connections, etc.).""" return diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index 29ec627ca62..caf1d2c71b4 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -13,6 +13,8 @@ import logging from collections.abc import Iterable from typing import TYPE_CHECKING +from typing_extensions import override + from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext from vllm.v1.kv_offload.tiering.base import ( JobMetadata, @@ -64,6 +66,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): # Completed jobs waiting to be retrieved by get_finished_jobs() self.completed_jobs: list[JobResult] = [] + @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ Check whether a block exists in this secondary tier. @@ -77,6 +80,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): """ return key in self.blocks + @override def submit_store(self, job_metadata: JobMetadata) -> None: """ Submit a job to store blocks from primary tier to this tier. @@ -96,6 +100,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.blocks[key] = True self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + @override def submit_load(self, job_metadata: JobMetadata) -> None: """ Submit a job to load blocks from this tier to primary tier. @@ -120,6 +125,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) + @override def get_finished_jobs(self) -> Iterable[JobResult]: """ Poll for finished jobs. @@ -132,6 +138,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs = [] return result + @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return RequestOffloadingContext() diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 1d37983d558..a33de02f43d 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -21,6 +21,8 @@ import os from collections.abc import Iterable from typing import TYPE_CHECKING +from typing_extensions import override + from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.file_mapper import FileMapper @@ -50,6 +52,14 @@ class FileSystemTierManager(SecondaryTierManager): submit_store / submit_load are non-blocking: they enqueue tasks and return. get_finished_jobs() polls job completion and returns completed JobResults. + Cross-process sharing: + In order to enable KV cache sharing between multiple vLLM instances + using the same ``root_dir`` (e.g., via a shared PVC) the environment + variable ``PYTHONHASHSEED`` must be set to the same fixed value + (e.g., "0") on all instances. Without this, each process initializes + ``NONE_HASH`` (the chain-hash seed for block content hashes) with + random bytes, producing different block filenames for identical token + content. """ def __init__( @@ -101,14 +111,17 @@ class FileSystemTierManager(SecondaryTierManager): thread_name_prefix="vllm_kv_py_fs", ) + @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return RequestOffloadingContext() + @override def lookup( self, key: OffloadKey, req_context: ReqContext | None = None ) -> bool | None: return os.path.exists(self.file_mapper.get_file_name(key)) + @override def submit_store(self, job_metadata: JobMetadata) -> None: tasks = ( functools.partial( @@ -122,6 +135,7 @@ class FileSystemTierManager(SecondaryTierManager): ) self._pool.enqueue_store(job_metadata.job_id, len(job_metadata.keys), tasks) + @override def submit_load(self, job_metadata: JobMetadata) -> None: tasks = ( functools.partial( @@ -135,6 +149,7 @@ class FileSystemTierManager(SecondaryTierManager): ) self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks) + @override def get_finished_jobs(self) -> Iterable[JobResult]: """ Collect completed jobs from the finished-jobs queue. @@ -144,6 +159,7 @@ class FileSystemTierManager(SecondaryTierManager): for job_id, success in self._pool.get_finished() ) + @override def shutdown(self) -> None: """ Release resources held by this tier. diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index ae8121b4aa1..cb8de749ec7 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -25,6 +25,7 @@ from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field import numpy as np +from typing_extensions import override from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( @@ -100,6 +101,7 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): """ return self._kv_memoryview + @override def shutdown(self) -> None: super().shutdown() self._kv_memoryview.release() @@ -150,7 +152,7 @@ class TieringOffloadingManager(OffloadingManager): self._transfer_jobs: dict[JobId, JobMetadata] = {} # Pending promotion requests accumulated during lookup() calls; flushed - # as one batched submit_load() per (tier, request) in take_events(). + # as one batched submit_load() per (tier, request) in on_schedule_end(). # Outer key: tier. Inner key: req_context.req_id — the same ReqContext # object is reused for all block lookups of a given request per engine step. self._pending_load_submissions: dict[ @@ -158,7 +160,7 @@ class TieringOffloadingManager(OffloadingManager): ] = {} # Gate for once-per-step execution of _maybe_process_finished_jobs(). - # Reset at the end of each step in take_events(). + # Reset at the end of each step in on_schedule_end(). self._processed_jobs_this_step: bool = False # Per-request set of secondary tiers that requested REQUEST_LEVEL @@ -180,7 +182,7 @@ class TieringOffloadingManager(OffloadingManager): Guarded by _processed_jobs_this_step: the first call in an engine step does the actual polling; subsequent calls are no-ops. The flag is reset - in take_events() at the end of each step. + in on_schedule_end() at the end of each step. """ if self._processed_jobs_this_step: return @@ -222,6 +224,7 @@ class TieringOffloadingManager(OffloadingManager): job_metadata.keys, job_metadata.req_context ) + @override def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: """ Check whether a single block is offloaded and ready. @@ -301,8 +304,8 @@ class TieringOffloadingManager(OffloadingManager): store_spec = primary_write_result.store_spec assert isinstance(store_spec, CPULoadStoreSpec) - # Defer submit_load to take_events(). Group by (tier, request) so each - # request's blocks are submitted as one batched job per tier. + # Defer submit_load to on_schedule_end(). Group by (tier, request) so + # each request's blocks are submitted as one batched job per tier. tier_pending = self._pending_load_submissions.setdefault(tier, {}) ctx_id = req_context.req_id if ctx_id not in tier_pending: @@ -317,8 +320,8 @@ class TieringOffloadingManager(OffloadingManager): def _flush_pending_promotions(self) -> None: """Submit one batched submit_load() per (tier, request). - Called from take_events() at the end of each engine step, flushing - all promotion requests deferred during lookup(). + Called from on_schedule_end() at the end of each scheduler step, + flushing all promotion requests deferred during lookup(). """ if not self._pending_load_submissions: return @@ -338,6 +341,7 @@ class TieringOffloadingManager(OffloadingManager): self._pending_load_submissions.clear() + @override def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: @@ -362,6 +366,7 @@ class TieringOffloadingManager(OffloadingManager): return self.primary_tier.prepare_load(keys, req_context) + @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as recently used in all tiers. @@ -374,6 +379,7 @@ class TieringOffloadingManager(OffloadingManager): for tier in self.secondary_tiers: tier.touch(keys, req_context) + @override def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext): """ Mark blocks as done loading from primary tier to GPU. @@ -387,6 +393,7 @@ class TieringOffloadingManager(OffloadingManager): """ self.primary_tier.complete_load(keys, req_context) + @override def prepare_store( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: @@ -470,6 +477,7 @@ class TieringOffloadingManager(OffloadingManager): self._transfer_jobs[job_id] = job_metadata tier.submit_store(job_metadata) + @override def complete_store( self, keys: Collection[OffloadKey], @@ -528,6 +536,7 @@ class TieringOffloadingManager(OffloadingManager): # Note: The async transfers are now in flight. Their completion is # tracked via get_finished_jobs() / _maybe_process_finished_jobs(). + @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: """ Query each secondary tier for its offload policy preference. @@ -547,44 +556,41 @@ class TieringOffloadingManager(OffloadingManager): ) return RequestOffloadingContext(policy=policy) + @override def on_request_finished(self, req_context: ReqContext) -> None: self.primary_tier.on_request_finished(req_context) for tier in self.secondary_tiers: tier.on_request_finished(req_context) self._request_level_tiers.pop(req_context.req_id, None) - def take_events(self) -> Iterable[OffloadingEvent]: - """ - End-of-step hook: flush deferred work, yield events, reset per-step state. + @override + def on_schedule_end(self) -> None: + """End-of-schedule hook: process finished jobs, flush deferred + promotions, and reset the per-step gate. - Called once per engine step from Scheduler.update_from_output() → - connector.take_events(). Ensures _maybe_process_finished_jobs() has run - at least once this step, flushes pending promotions, yields collected - events, and resets the per-step flag. + Called once per scheduler step from + OffloadingConnectorScheduler.build_connector_meta(). + """ + self._maybe_process_finished_jobs() + self._processed_jobs_this_step = False + self._flush_pending_promotions() + for tier in self.secondary_tiers: + tier.on_schedule_end() + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + """Yield offloading events collected since the last call. Yields: New OffloadingEvents collected since the last call. """ - # TODO: Move _flush_pending_promotions() to a dedicated end_of_batch() - # hook once one exists. For now, take_events() serves as the flush - # point under the assumption that it is called at the end of each - # engine step (Scheduler.update_from_output() → connector.take_events()). - # When the dedicated hook is added, update tests that rely on - # take_events() to signal end of step. - - self._maybe_process_finished_jobs() - - self._flush_pending_promotions() - - # Reset the per-step gate so next step's first call does real work. - self._processed_jobs_this_step = False - if self.events is not None: yield from self.events self.events.clear() yield from self.primary_tier.take_events() + @override def shutdown(self) -> None: """Shutdown all tiers and release resources.""" for tier in self.secondary_tiers: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index ced8a7fc654..a4ea46e08eb 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -63,6 +63,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec): memory and must transfer data through the primary tier. """ + BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): super().__init__(vllm_config, kv_cache_config) # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it @@ -96,21 +98,16 @@ class TieringOffloadingSpec(CPUOffloadingSpec): # Create scheduler-side SharedOffloadRegion (rank=None) so the # primary tier can eagerly create a memoryview over _base. - world_size = self.vllm_config.parallel_config.world_size scheduler_mmap = SharedOffloadRegion( instance_id=self.vllm_config.instance_id, - total_size_bytes=self.cpu_page_size_per_worker - * world_size - * self.num_blocks, num_blocks=self.num_blocks, rank=None, - num_workers=world_size, + kv_bytes_per_block=self.kv_bytes_per_offloaded_block, cpu_page_size=self.cpu_page_size_per_worker, ) self._scheduler_mmap = scheduler_mmap # Create primary tier (CPU-based) - assert len(self.gpu_block_size) == 1 primary_tier = CPUPrimaryTierOffloadingManager( num_blocks=self.num_blocks, cache_policy=self.eviction_policy, # type: ignore[arg-type] @@ -166,16 +163,12 @@ class TieringOffloadingSpec(CPUOffloadingSpec): @override def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers: - world_size = self.vllm_config.parallel_config.world_size rank = torch.accelerator.current_device_index() worker_mmap = SharedOffloadRegion( instance_id=self.vllm_config.instance_id, - total_size_bytes=self.cpu_page_size_per_worker - * world_size - * self.num_blocks, num_blocks=self.num_blocks, rank=rank, - num_workers=world_size, + kv_bytes_per_block=self.kv_bytes_per_offloaded_block, cpu_page_size=self.cpu_page_size_per_worker, ) return CpuGpuOffloadingHandlers( diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 26cc82fc4a6..44246e70a8b 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -141,6 +141,10 @@ class Request: self.num_output_placeholders = 0 self.async_tokens_to_discard = 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 + self.spec_token_ids: list[int] = [] self.num_computed_tokens = 0 self.cache_salt: str | None = cache_salt diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index f98b12a379d..66806ab8a9b 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -19,6 +19,69 @@ if HAS_TRITON: logger = init_logger(__name__) +_FLASHINFER_MIN_VERSION = "0.2.3" + + +def flashinfer_sampler_supported() -> bool: + """Decide whether FlashInfer's top-p/top-k sampler can be used. + + Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER`` + is 0, when the platform isn't CUDA, when the GPU's compute capability is + unsupported, or when the installed flashinfer is missing or too old. Raises + ``RuntimeError`` if the user explicitly opted in via the env var but + FlashInfer is unavailable. + + Note: callers must additionally ensure ``logprobs_mode`` doesn't require + post-top-k/top-p logits/logprobs for any request whose logprobs will be + returned in this step, since FlashInfer doesn't expose those. + """ + if not current_platform.is_cuda(): + return False + if not envs.VLLM_USE_FLASHINFER_SAMPLER: + logger.info_once( + "FlashInfer top-p/top-k sampling disabled via " + "VLLM_USE_FLASHINFER_SAMPLER=0." + ) + return False + from vllm.v1.attention.backends.flashinfer import FlashInferBackend + + capability = current_platform.get_device_capability() + assert capability is not None + unsupported_reason: str | None = None + if not FlashInferBackend.supports_compute_capability(capability): + unsupported_reason = ( + f"unsupported compute capability {capability.as_version_str()}" + ) + else: + try: + import flashinfer + + if version.parse(flashinfer.__version__) < version.parse( + _FLASHINFER_MIN_VERSION + ): + unsupported_reason = ( + f"flashinfer {flashinfer.__version__} is too old " + f"(>={_FLASHINFER_MIN_VERSION} required)" + ) + except ImportError: + unsupported_reason = "flashinfer is not installed" + + if unsupported_reason is None: + logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global") + return True + if envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"): + raise RuntimeError( + f"FlashInfer top-p/top-k sampling unavailable: {unsupported_reason}. " + "Unset VLLM_USE_FLASHINFER_SAMPLER=1." + ) + logger.warning_once( + "FlashInfer top-p/top-k sampling unavailable: %s; falling back. " + "Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.", + unsupported_reason, + ) + return False + + class TopKTopPSampler(nn.Module): """ Module that performs optional top-k and top-p filtering followed by @@ -30,49 +93,16 @@ class TopKTopPSampler(nn.Module): def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs") -> None: super().__init__() self.logprobs_mode = logprobs_mode - # flashinfer optimization does not apply if intermediate - # logprobs/logits after top_k/top_p need to be returned - if ( - logprobs_mode not in ("processed_logits", "processed_logprobs") - and current_platform.is_cuda() - ): - if envs.VLLM_USE_FLASHINFER_SAMPLER: - from vllm.v1.attention.backends.flashinfer import FlashInferBackend - - capability = current_platform.get_device_capability() - assert capability is not None - if FlashInferBackend.supports_compute_capability(capability): - logger.info_once( - "Using FlashInfer for top-p & top-k sampling.", - scope="global", - ) - self.forward = self.forward_cuda - elif envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"): - # User explicitly opted in but the GPU can't run FlashInfer. - capability_str = capability.as_version_str() - raise RuntimeError( - "FlashInfer does not support compute capability " - f"{capability_str}, unset VLLM_USE_FLASHINFER_SAMPLER=1." - ) - else: - # Default-on path; hardware can't run FlashInfer → - # quietly fall back to the PyTorch-native sampler - # instead of failing server startup. - logger.warning_once( - "FlashInfer top-p/top-k sampling not supported on " - "compute capability %s; falling back to PyTorch-native " - "sampler. Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.", - capability.as_version_str(), - ) - self.forward = self.forward_native - else: - # User explicitly set VLLM_USE_FLASHINFER_SAMPLER=0. - logger.info_once( - "FlashInfer top-p/top-k sampling disabled via " - "VLLM_USE_FLASHINFER_SAMPLER=0; using PyTorch-native sampler." - ) - self.forward = self.forward_native - + if current_platform.is_cuda(): + # FlashInfer doesn't expose post-top-k/top-p logits/logprobs, + # so it can't be used when the configured mode requires them. + can_use_flashinfer = ( + logprobs_mode not in ("processed_logits", "processed_logprobs") + and flashinfer_sampler_supported() + ) + self.forward = ( + self.forward_cuda if can_use_flashinfer else self.forward_native + ) elif current_platform.is_cpu(): arch = current_platform.get_cpu_architecture() # Fall back to native implementation for POWERPC and RISCV. @@ -417,7 +447,7 @@ def flashinfer_sample( logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None, - generators: dict[int, torch.Generator], + generators: dict[int, torch.Generator] = {}, # noqa ) -> torch.Tensor: """Sample from the logits using FlashInfer. @@ -431,11 +461,6 @@ def flashinfer_sample( """ import flashinfer - if version.parse(flashinfer.__version__) < version.parse("0.2.3"): - raise ImportError( - "FlashInfer version >= 0.2.3 required for top-k and top-p sampling. " - ) - assert not (k is None and p is None) if k is None: # Top-p only. diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 4b17831dfc9..bfe6fd6ae52 100644 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -186,7 +186,7 @@ def _topk_topp_kernel( # max so the search converges to -inf (no masking). min_logit = tl.minimum(min_logit, max_logit) - # Second passes: Ternary search for pivots + # Second passes: Ternary search for pivot num_iters = 0 k_pivot = float("inf") k_pivots_num = tl.zeros((), dtype=tl.uint32) @@ -279,6 +279,8 @@ def _topk_topp_kernel( num_iters += 1 if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: k_pivot = (max_range + min_range) / 2.0 + min_larger = min_larger_0 + num_min_larger = num_min_larger_0 found_pivot = 1 else: # If top-k outlier gathering failed, search whole logit space @@ -286,12 +288,12 @@ def _topk_topp_kernel( min_range = min_logit found_pivot = 0 while found_pivot == 0: - k_pivot_0 = (max_range - min_range) * 1.0 / 4.0 + min_range + k_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range k_pivots_num_0 = tl.zeros((), dtype=tl.uint32) min_larger_0 = float("inf") num_min_larger_0 = tl.zeros((), dtype=tl.uint32) - k_pivot_1 = (max_range - min_range) * 2.0 / 4.0 + min_range + k_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range k_pivots_num_1 = tl.zeros((), dtype=tl.uint32) min_larger_1 = float("inf") num_min_larger_1 = tl.zeros((), dtype=tl.uint32) @@ -359,6 +361,8 @@ def _topk_topp_kernel( num_iters += 1 if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9: k_pivot = (max_range + min_range) / 2.0 + min_larger = min_larger_0 + num_min_larger = num_min_larger_0 found_pivot = 1 duplicate_logit = min_larger @@ -520,17 +524,15 @@ def _topk_topp_kernel( # Fifth passes: Search for p_pivot found_pivot = 0 while found_pivot == 0: - p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range + p_pivot_0 = (max_range - min_range) * 0.5 + min_range p_pivots_sum_0 = 0.0 min_larger_0 = 1.0 num_min_larger_0 = tl.zeros((), dtype=tl.uint32) - p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range - p_pivots_sum_1 = 0.0 - min_larger_1 = 1.0 - num_min_larger_1 = tl.zeros((), dtype=tl.uint32) - - # First pass: Calculate p_pivots_sum and min_larger + # Single fused pass: compute p_pivots_sum, + # min_larger, and num_min_larger together. + # See _update_min_larger_stats for the + # tile-level merge logic. for i in range(0, search_iters): offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( 0, BLOCK_SIZE_TRUNC @@ -540,52 +542,20 @@ def _topk_topp_kernel( BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 ) - p_pivots_sum_0 += tl.sum( - probs_blk * (probs_blk > p_pivot_0) - ) - masked_larger_0 = tl.where( - probs_blk > p_pivot_0, probs_blk, 1.0 - ) - min_larger_0 = tl.minimum( - min_larger_0, tl.min(masked_larger_0) + above_0 = probs_blk > p_pivot_0 + p_pivots_sum_0 += tl.sum(probs_blk * above_0) + + min_larger_0, num_min_larger_0 = ( + _update_min_larger_stats( + probs_blk, + above_0, + min_larger_0, + num_min_larger_0, + 1.0, + ) ) - p_pivots_sum_1 += tl.sum( - probs_blk * (probs_blk > p_pivot_1) - ) - masked_larger_1 = tl.where( - probs_blk > p_pivot_1, probs_blk, 1.0 - ) - min_larger_1 = tl.minimum( - min_larger_1, tl.min(masked_larger_1) - ) - - # Second pass: Calculate num_min_larger - for i in range(0, search_iters): - offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( - 0, BLOCK_SIZE_TRUNC - ) - mask_n_2 = offs_n < search_range - probs_blk = tl.load( - BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 - ) - - num_min_larger_0 += tl.sum( - tl.abs(probs_blk - min_larger_0) < 1e-9 - ) - num_min_larger_1 += tl.sum( - tl.abs(probs_blk - min_larger_1) < 1e-9 - ) - - # Check if any of the pivots satisfy termination condition - if p_pivots_sum_1 >= p and ( - p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p - ): - p_pivot = p_pivot_1 - min_larger_prob = min_larger_1 - num_min_larger = num_min_larger_1 - p_pivots_sum = p_pivots_sum_1 - found_pivot = 1 + # Check if the pivot satisfies termination condition if p_pivots_sum_0 >= p and ( p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p ): @@ -596,19 +566,17 @@ def _topk_topp_kernel( found_pivot = 1 # Update range - if p_pivots_sum_1 > p: - min_range = p_pivot_1 - elif p_pivots_sum_0 > p: + if p_pivots_sum_0 > p: min_range = p_pivot_0 - - if p_pivots_sum_0 < p: + elif p_pivots_sum_0 < p: max_range = p_pivot_0 - elif p_pivots_sum_1 < p: - max_range = p_pivot_1 num_iters += 1 if (max_range - min_range) < 1e-9 or num_iters >= 18: p_pivot = (max_range + min_range) / 2.0 + min_larger_prob = min_larger_0 + num_min_larger = num_min_larger_0 + p_pivots_sum = p_pivots_sum_0 found_pivot = 1 duplicate_logit = ( @@ -725,17 +693,15 @@ def _topk_topp_kernel( found_pivot = 0 while found_pivot == 0: - p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range + p_pivot_0 = (max_range - min_range) * 0.5 + min_range p_pivots_sum_0 = 0.0 min_larger_0 = 1.0 num_min_larger_0 = tl.zeros((), dtype=tl.uint32) - p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range - p_pivots_sum_1 = 0.0 - min_larger_1 = 1.0 - num_min_larger_1 = tl.zeros((), dtype=tl.uint32) - - # First pass: Calculate p_pivots_sum and min_larger + # Single fused pass: compute p_pivots_sum, + # min_larger, and num_min_larger together. + # See _update_min_larger_stats for the + # tile-level merge logic. for i in range(0, search_iters): offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( 0, BLOCK_SIZE_TRUNC @@ -745,53 +711,18 @@ def _topk_topp_kernel( BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 ) - p_pivots_sum_0 += tl.sum( - probs_blk * (probs_blk > p_pivot_0) - ) - masked_larger_0 = tl.where( - probs_blk > p_pivot_0, probs_blk, 1.0 - ) - min_larger_0 = tl.minimum( - min_larger_0, tl.min(masked_larger_0) + above_0 = probs_blk > p_pivot_0 + p_pivots_sum_0 += tl.sum(probs_blk * above_0) + + min_larger_0, num_min_larger_0 = _update_min_larger_stats( + probs_blk, + above_0, + min_larger_0, + num_min_larger_0, + 1.0, ) - p_pivots_sum_1 += tl.sum( - probs_blk * (probs_blk > p_pivot_1) - ) - masked_larger_1 = tl.where( - probs_blk > p_pivot_1, probs_blk, 1.0 - ) - min_larger_1 = tl.minimum( - min_larger_1, tl.min(masked_larger_1) - ) - - # Second pass: Calculate num_min_larger - for i in range(0, search_iters): - offs_n = i * BLOCK_SIZE_TRUNC + tl.arange( - 0, BLOCK_SIZE_TRUNC - ) - mask_n_2 = offs_n < search_range - probs_blk = tl.load( - BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0 - ) - - num_min_larger_0 += tl.sum( - tl.abs(probs_blk - min_larger_0) < 1e-9 - ) - num_min_larger_1 += tl.sum( - tl.abs(probs_blk - min_larger_1) < 1e-9 - ) - - # Check if any of the pivots satisfy termination condition - if ( - p_pivots_sum_1 >= p - and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p - ): - p_pivot = p_pivot_1 - min_larger_prob = min_larger_1 - num_min_larger = num_min_larger_1 - p_pivots_sum = p_pivots_sum_1 - found_pivot = 1 + # Check if the pivot satisfies termination condition if ( p_pivots_sum_0 >= p and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p @@ -803,19 +734,17 @@ def _topk_topp_kernel( found_pivot = 1 # Update range - if p_pivots_sum_1 > p: - min_range = p_pivot_1 - elif p_pivots_sum_0 > p: + if p_pivots_sum_0 > p: min_range = p_pivot_0 - - if p_pivots_sum_0 < p: + elif p_pivots_sum_0 < p: max_range = p_pivot_0 - elif p_pivots_sum_1 < p: - max_range = p_pivot_1 num_iters += 1 if (max_range - min_range) < 1e-9 or num_iters >= 18: p_pivot = (max_range + min_range) / 2.0 + min_larger_prob = min_larger_0 + num_min_larger = num_min_larger_0 + p_pivots_sum = p_pivots_sum_0 found_pivot = 1 else: # Re-populate the buffer with full softmax probabilities @@ -832,17 +761,15 @@ def _topk_topp_kernel( found_pivot = 0 while found_pivot == 0: - p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range + p_pivot_0 = (max_range - min_range) * 0.5 + min_range p_pivots_sum_0 = 0.0 min_larger_0 = 1.0 num_min_larger_0 = tl.zeros((), dtype=tl.uint32) - p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range - p_pivots_sum_1 = 0.0 - min_larger_1 = 1.0 - num_min_larger_1 = tl.zeros((), dtype=tl.uint32) - - # First pass: Calculate p_pivots_sum and min_larger + # Single fused pass: compute p_pivots_sum, + # min_larger, and num_min_larger together. + # See _update_min_larger_stats for the + # tile-level merge logic. for i in range(0, NUM_TILES): offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask_n = offs_n < VOCAB_SIZE @@ -850,51 +777,18 @@ def _topk_topp_kernel( BUFFER_ROW + offs_n, mask=mask_n, other=0.0 ) - p_pivots_sum_0 += tl.sum( - probs_blk * (probs_blk > p_pivot_0) - ) - masked_larger_0 = tl.where( - probs_blk > p_pivot_0, probs_blk, 1.0 - ) - min_larger_0 = tl.minimum( - min_larger_0, tl.min(masked_larger_0) + above_0 = probs_blk > p_pivot_0 + p_pivots_sum_0 += tl.sum(probs_blk * above_0) + + min_larger_0, num_min_larger_0 = _update_min_larger_stats( + probs_blk, + above_0, + min_larger_0, + num_min_larger_0, + 1.0, ) - p_pivots_sum_1 += tl.sum( - probs_blk * (probs_blk > p_pivot_1) - ) - masked_larger_1 = tl.where( - probs_blk > p_pivot_1, probs_blk, 1.0 - ) - min_larger_1 = tl.minimum( - min_larger_1, tl.min(masked_larger_1) - ) - - # Second pass: Calculate num_min_larger - for i in range(0, NUM_TILES): - offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask_n = offs_n < VOCAB_SIZE - probs_blk = tl.load( - BUFFER_ROW + offs_n, mask=mask_n, other=0.0 - ) - - num_min_larger_0 += tl.sum( - tl.abs(probs_blk - min_larger_0) < 1e-9 - ) - num_min_larger_1 += tl.sum( - tl.abs(probs_blk - min_larger_1) < 1e-9 - ) - - # Check if any of the pivots satisfy termination condition - if ( - p_pivots_sum_1 >= p - and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p - ): - p_pivot = p_pivot_1 - min_larger_prob = min_larger_1 - num_min_larger = num_min_larger_1 - p_pivots_sum = p_pivots_sum_1 - found_pivot = 1 + # Check if the pivot satisfies termination condition if ( p_pivots_sum_0 >= p and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p @@ -906,22 +800,20 @@ def _topk_topp_kernel( found_pivot = 1 # Update range - if p_pivots_sum_1 > p: - min_range = p_pivot_1 - elif p_pivots_sum_0 > p: + if p_pivots_sum_0 > p: min_range = p_pivot_0 - - if p_pivots_sum_0 < p: + elif p_pivots_sum_0 < p: max_range = p_pivot_0 - elif p_pivots_sum_1 < p: - max_range = p_pivot_1 num_iters += 1 if (max_range - min_range) < 1e-9 or num_iters >= 18: p_pivot = (max_range + min_range) / 2.0 + min_larger_prob = min_larger_0 + num_min_larger = num_min_larger_0 + p_pivots_sum = p_pivots_sum_0 found_pivot = 1 - duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_logit + duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_sample num_duplicate_logit = num_min_larger num_keep = num_duplicate_logit - tl.cast( (p_pivots_sum - p) / min_larger_prob, tl.uint32 @@ -952,9 +844,7 @@ def _topk_topp_kernel( tl.abs(logits_blk - duplicate_logit) < 1e-9 ) & mask_n duplicate_count = tl.cumsum(duplicate_mask) + num_kept - duplicate_keep_mask = ( - duplicate_count <= num_duplicate_logit - ) & duplicate_mask + duplicate_keep_mask = (duplicate_count <= num_keep) & duplicate_mask duplicate_remove_mask = duplicate_mask & ~duplicate_keep_mask num_kept += tl.sum(duplicate_keep_mask) keep_mask = keep_mask & (~duplicate_remove_mask) diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index ca5e2b66e03..8789e6afdc4 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any import torch +from vllm.platforms import current_platform from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.sample.logits_processor.interface import ( BatchUpdate, @@ -511,14 +512,31 @@ class ThinkingBudgetStateHolder: if active_indices_cpu: device = logits.device - active_indices = async_tensor_h2d( - active_indices_cpu, dtype=torch.long, device=device - ) - force_tokens = async_tensor_h2d( - force_tokens_cpu, dtype=torch.long, device=device - ) - # Avoid CPU->GPU sync. - fill = logits.new_full((len(active_indices_cpu),), 1e9) - logits.index_put_((active_indices, force_tokens), fill) + if current_platform.is_rocm() and logits.is_contiguous(): + # Flattened index_fill avoids ROCm faults seen with 2-D + # advanced-indexing writes on the thinking-budget path. + vocab_size = logits.shape[1] + flat_indices_cpu = [ + row * vocab_size + token + for row, token in zip(active_indices_cpu, force_tokens_cpu) + ] + flat_indices = async_tensor_h2d( + flat_indices_cpu, dtype=torch.long, device=device + ) + logits.view(-1).index_fill_(0, flat_indices, 1e9) + elif current_platform.is_rocm(): + fill = logits.new_tensor(1e9) + for row, token in zip(active_indices_cpu, force_tokens_cpu): + logits[row, token] = fill + else: + active_indices = async_tensor_h2d( + active_indices_cpu, dtype=torch.long, device=device + ) + force_tokens = async_tensor_h2d( + force_tokens_cpu, dtype=torch.long, device=device + ) + # Avoid CPU->GPU sync. + fill = logits.new_full((len(active_indices_cpu),), 1e9) + logits.index_put_((active_indices, force_tokens), fill) return logits diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 24b6a178ce9..f61c4320dff 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -127,6 +127,7 @@ class SimpleCPUOffloadScheduler: enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size, + scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, ) self.cpu_block_pool: BlockPool = self.cpu_coordinator.block_pool diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index ae9e8d58b8b..583fd78ced0 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -153,6 +153,7 @@ class EncoderCudaGraphManager: ) self.budget_graphs: dict[int, BudgetGraphMetadata] = {} + self.graph_pool: Any | None = None self.graph_hits = 0 self.graph_misses = 0 self.log_stats_interval = 100 @@ -183,9 +184,16 @@ class EncoderCudaGraphManager: """Check if a modality is supported by this manager.""" return modality in self.config.modalities - def capture(self): + def clear(self) -> None: + """Release captured encoder CUDA graphs and the manager-local pool.""" + self.budget_graphs.clear() + self.graph_pool = None + + def capture(self, graph_pool: Any): """Capture CUDA graphs for all token budgets.""" - for token_budget in self.token_budgets: + self.graph_pool = graph_pool + + for token_budget in sorted(self.token_budgets, reverse=True): self._capture_budget_graph(token_budget) logger.info( @@ -193,6 +201,9 @@ class EncoderCudaGraphManager: len(self.budget_graphs), ) + def get_num_graphs_to_capture(self) -> int: + return len(self.token_budgets) + def _capture_budget_graph(self, token_budget: int): """Capture CUDA graph for a single token budget.""" logger.debug( @@ -218,7 +229,7 @@ class EncoderCudaGraphManager: output_buffer = torch.empty_like(output) graph = torch.cuda.CUDAGraph() - with torch.inference_mode(), torch.cuda.graph(graph): + with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool): output = self.model.encoder_cudagraph_forward({**values}) output_buffer.copy_(output) diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 5963790a779..e4497de43a7 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -13,6 +13,15 @@ from vllm.utils.torch_utils import ( get_accelerator_view_from_cpu_tensor, ) +# Default round-robin depth for the UVA buffer pools. Must be >= the number of +# concurrent in-flight steps (engine batch_queue_size). +_DEFAULT_MAX_CONCURRENCY = 2 + + +def set_default_max_concurrency(n: int) -> None: + global _DEFAULT_MAX_CONCURRENCY + _DEFAULT_MAX_CONCURRENCY = max(2, n) + def async_copy_to_gpu( x: torch.Tensor | np.ndarray, @@ -47,8 +56,10 @@ class UvaBufferPool: self, size: int | Sequence[int], dtype: torch.dtype, - max_concurrency: int = 2, + max_concurrency: int | None = None, ): + if max_concurrency is None: + max_concurrency = _DEFAULT_MAX_CONCURRENCY self.size = size self.dtype = dtype self.max_concurrency = max_concurrency @@ -80,7 +91,10 @@ class UvaBufferPool: class UvaBackedTensor: def __init__( - self, size: int | Sequence[int], dtype: torch.dtype, max_concurrency: int = 2 + self, + size: int | Sequence[int], + dtype: torch.dtype, + max_concurrency: int | None = None, ): self.dtype = dtype @@ -104,9 +118,11 @@ class StagedWriteTensor: size: int | Sequence[int], dtype: torch.dtype, device: torch.device, - max_concurrency: int = 2, + max_concurrency: int | None = None, uva_instead_of_gpu: bool = False, ): + if max_concurrency is None: + max_concurrency = _DEFAULT_MAX_CONCURRENCY supported_dtypes = [torch.int32, torch.int64, torch.float32] if dtype not in supported_dtypes: raise ValueError( diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 384f1192435..0648de29859 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -299,7 +299,6 @@ class CudaGraphManager: self.breakable_cg_runner = BreakableCUDAGraphWrapper( model, self.vllm_config ) - self.breakable_cg_runner.graph_pool = self.pool def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any: if not self.use_breakable_cg: diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index b253d7d8c06..f905d09e45f 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -67,9 +67,18 @@ class InputBatch: seq_lens_cpu_upper_bound: torch.Tensor # [num_reqs] dcp_local_seq_lens: torch.Tensor | None - # [num_reqs] CPU bool array. + # [num_reqs] + num_computed_tokens_np: np.ndarray + # [num_reqs] + prefill_len_np: np.ndarray + # [num_reqs] + num_computed_prefill_tokens_np: np.ndarray + # [num_reqs] CPU bool array == (num_computed_prefill_tokens_np < prefill_len_np). is_prefilling_np: np.ndarray + # [num_reqs] only populated when pipeline parallelism is enabled. + max_seq_len_np: np.ndarray | None + # [num_tokens_after_padding] input_ids: torch.Tensor # [num_tokens_after_padding] @@ -148,7 +157,11 @@ class InputBatch: seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=None, + num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32), + prefill_len_np=np.zeros(num_reqs, dtype=np.int32), + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), + max_seq_len_np=None, input_ids=input_ids, positions=positions, logits_indices=logits_indices, @@ -438,6 +451,9 @@ def _post_update_kernel( ): req_id = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_id) + if req_state_idx < 0: + # Filter rows with negative index entries. + return total_len = tl.load(total_len_ptr + req_state_idx) num_sampled = tl.load(num_sampled_ptr + req_id) @@ -464,18 +480,22 @@ def _post_update_kernel( count = tl.load(token_ptr) tl.store(token_ptr, count + 1) - query_start = tl.load(query_start_loc_ptr + req_id) - query_end = tl.load(query_start_loc_ptr + req_id + 1) - query_len = query_end - query_start + if query_start_loc_ptr is None: + query_len = 0 + else: + query_start = tl.load(query_start_loc_ptr + req_id) + query_end = tl.load(query_start_loc_ptr + req_id + 1) + query_len = query_end - query_start num_rejected = tl.load(num_rejected_ptr + req_id) - num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) - num_computed += query_len - num_rejected - tl.store(num_computed_tokens_ptr + req_state_idx, num_computed) + computed_delta = query_len - num_rejected + if computed_delta != 0: + num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + tl.store(num_computed_tokens_ptr + req_state_idx, num_computed + computed_delta) def post_update( - # [num_reqs] + # [num_reqs] batch_idx -> req_state_idx; negative index means skip. idx_mapping: torch.Tensor, # [max_num_reqs] num_computed_tokens: torch.Tensor, @@ -490,7 +510,7 @@ def post_update( # [num_reqs] num_rejected: torch.Tensor, # [num_reqs + 1] - query_start_loc: torch.Tensor, + query_start_loc: torch.Tensor | None, # [max_num_reqs, max_model_len] all_token_ids: torch.Tensor, # [max_num_reqs] @@ -516,7 +536,7 @@ def post_update( @triton.jit -def _post_update_pool_kernel( +def _post_update_num_computed_tokens_kernel( idx_mapping_ptr, num_computed_tokens_ptr, query_start_loc_ptr, @@ -531,7 +551,7 @@ def _post_update_pool_kernel( tl.store(num_computed_tokens_ptr + req_state_idx, num_computed + query_len) -def post_update_pool( +def post_update_num_computed_tokens( # [num_reqs] idx_mapping: torch.Tensor, # [max_num_reqs] @@ -540,7 +560,7 @@ def post_update_pool( query_start_loc: torch.Tensor, ) -> None: num_reqs = idx_mapping.shape[0] - _post_update_pool_kernel[(num_reqs,)]( + _post_update_num_computed_tokens_kernel[(num_reqs,)]( idx_mapping, num_computed_tokens, query_start_loc, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 519395e90b2..2e3133822fd 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -47,6 +47,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib +from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec @@ -60,7 +61,10 @@ from vllm.v1.worker.gpu.attn_utils import ( init_kv_cache, ) from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.v1.worker.gpu.buffer_utils import ( + async_copy_to_gpu, + set_default_max_concurrency, +) from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -76,7 +80,7 @@ from vllm.v1.worker.gpu.input_batch import ( expand_idx_mapping, get_num_sampled_and_rejected, post_update, - post_update_pool, + post_update_num_computed_tokens, prepare_pos_seq_lens, prepare_prefill_inputs, ) @@ -89,7 +93,7 @@ from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states import init_model_state from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner -from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive +from vllm.v1.worker.gpu.pp_utils import PPHandler from vllm.v1.worker.gpu.sample.output import SamplerOutput from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker from vllm.v1.worker.gpu.sample.sampler import Sampler @@ -103,6 +107,7 @@ from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin +from vllm.v1.worker.utils import KVBlockZeroer logger = init_logger(__name__) @@ -129,6 +134,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.cache_config.cache_dtype ] + # Lazily initialized in _init_kv_zero_meta() when the KV cache needs + # zeroing (e.g. hybrid models with fp8 KV cache). + self.kv_block_zeroer: KVBlockZeroer | None = None + self.vocab_size = self.model_config.get_vocab_size() self.max_model_len = self.model_config.max_model_len self.max_num_tokens = self.scheduler_config.max_num_batched_tokens @@ -143,6 +152,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.is_first_pp_rank = get_pp_group().is_first_rank self.is_last_pp_rank = get_pp_group().is_last_rank + # Size the UVA buffer pools to the max number of concurrent in-flight + # steps. Must run before any pooled buffer is constructed + set_default_max_concurrency(vllm_config.max_concurrent_batches) + + # PP broadcast/recv helper. Runs the collective on a side stream. + self.pp_handler: PPHandler | None = None + # Persistent buffer for intermediate tensors (non-first PP ranks). self.intermediate_tensors: IntermediateTensors | None = None @@ -204,6 +220,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): device=self.device, ) + if self.use_pp: + self.pp_handler = PPHandler( + max_num_reqs=self.max_num_reqs, + num_speculative_steps=self.num_speculative_steps, + device=self.device, + ) + self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None self.prompt_logprobs_worker: PromptLogprobsWorker | None = None @@ -338,6 +361,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.reset_encoder_cache() self.reset_mm_cache() + def apply_sparse_weight_patches(self, *args, **kwargs) -> None: + # TODO: Use full version instead of import when fully migrated to v2 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 + + GPUModelRunnerV1.apply_sparse_weight_patches(self, *args, **kwargs) # type: ignore[arg-type] + def update_config(self, *args, **kwargs) -> None: # TODO(Wentao): Use full version instead of import when fully migrated to v2 from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 @@ -393,7 +422,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) + spec.num_speculative_blocks max_num_blocks_per_group.append(max_num_blocks) - self.attn_groups, attn_cg_support, kernel_block_sizes = init_attn_backend( + self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) self.block_tables = BlockTables( @@ -402,7 +431,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_num_batched_tokens=self.max_num_tokens, max_num_blocks_per_group=max_num_blocks_per_group, device=self.device, - kernel_block_sizes=kernel_block_sizes, + kernel_block_sizes=self.kernel_block_sizes, cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, @@ -442,11 +471,22 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.attn_groups, self.device, self.cache_config.cache_dtype, - kernel_block_sizes, + self.kernel_block_sizes, self.vllm_config, ) self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) + def _init_kv_zero_meta(self) -> None: + """Build KV-block zeroing metadata; invoked from gpu_worker.""" + self.kv_block_zeroer = KVBlockZeroer( + self.device, + is_pin_memory_available(), + attn_groups_iter=(g for groups in self.attn_groups for g in groups), + kernel_block_sizes=self.kernel_block_sizes, + cache_dtype=self.cache_config.cache_dtype, + static_forward_context=self.compilation_config.static_forward_context, + ) + @torch.inference_mode() @step_eplb_after(is_dummy=True) def _dummy_run( @@ -665,8 +705,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): return cuda_graph_size def _remove_request(self, req_id: str) -> bool: - if not self.req_states.remove_request(req_id): + req_idx = self.req_states.remove_request(req_id) + if req_idx is None: return False + if self.pp_handler is not None: + self.pp_handler.on_req_idx_freed(req_idx) if self.encoder_cache is not None: self.encoder_cache.remove_request(req_id) if self.prompt_logprobs_worker is not None: @@ -687,6 +730,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): for mm_hash in scheduler_output.free_encoder_mm_hashes: self.encoder_cache.free_encoder_cache(mm_hash) + def update_pp_decode_requests(self): + # For non-last PP ranks, update decode requests with sampler output from + # the prior step in which they were scheduled (pp_size steps ago). + if self.pp_handler is not None: + outputs = self.pp_handler.get_prev_sampled_outputs() + if outputs is not None: + self.postprocess_sampled(**outputs) + def add_requests(self, scheduler_output: SchedulerOutput) -> None: for new_req_data in scheduler_output.scheduled_new_reqs: assert new_req_data.prompt_token_ids is not None @@ -699,11 +750,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): self._remove_request(req_id) prompt_len = len(new_req_data.prompt_token_ids) + sampling_params = new_req_data.sampling_params self.req_states.add_request( req_id=req_id, prompt_len=prompt_len, all_token_ids=new_req_data.prefill_token_ids, num_computed_tokens=new_req_data.num_computed_tokens, + max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type] ) req_index = self.req_states.req_id_to_index[req_id] @@ -746,13 +799,19 @@ class GPUModelRunner(LoRAModelRunnerMixin): req_index, req_new_block_ids, overwrite=False ) - # Update num_computed_prefill_tokens. + # Update CPU num_computed_prefill_tokens. np.minimum( self.req_states.num_computed_tokens_np, self.req_states.prefill_len.np, out=self.req_states.num_computed_prefill_tokens, ) + # Zero GPU memory for freshly allocated cache blocks to prevent + # stale NaN/data from corrupting attention or SSM computation. + if scheduler_output.new_block_ids_to_zero: + assert self.kv_block_zeroer is not None + self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + def prepare_inputs( self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor ) -> InputBatch: @@ -819,7 +878,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc) query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1] query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1] - is_prefilling_np = self.req_states.is_prefilling(idx_mapping_np) + prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np] + computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens + num_computed_prefill_tokens_np = computed_prefill_tokens_np[idx_mapping_np] + is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np # Get prefill tokens if any. if np.any(is_prefilling_np): @@ -871,13 +933,19 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) # CPU upper bound on seq_lens; padded entries left at zero. + num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np] seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32) np.add( - self.req_states.num_computed_tokens_np[idx_mapping_np], + num_computed_tokens_np, num_scheduled_tokens, out=seq_lens_cpu_upper_bound_np[:num_reqs], ) seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np) + + max_seq_len_np = None + if self.use_pp: + # max_seq_len is only consumed by the PP `compute_need_sampled_mask` + max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np] return InputBatch( req_ids=req_ids, num_reqs=num_reqs, @@ -896,7 +964,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=dcp_local_seq_lens, + num_computed_tokens_np=num_computed_tokens_np, + prefill_len_np=prefill_len_np, + num_computed_prefill_tokens_np=num_computed_prefill_tokens_np, is_prefilling_np=is_prefilling_np, + max_seq_len_np=max_seq_len_np, input_ids=self.input_buffers.input_ids[:num_tokens_after_padding], positions=self.input_buffers.positions[:num_tokens_after_padding], logits_indices=logits_indices, @@ -976,12 +1048,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) return sampler_output, num_sampled, num_rejected - def postprocess( + def postprocess_sampled( self, - input_batch: InputBatch, + idx_mapping: torch.Tensor, # May include -1 for masked entries sampled_tokens: torch.Tensor, num_sampled: torch.Tensor, num_rejected: torch.Tensor, + query_start_loc: torch.Tensor | None = None, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -990,19 +1063,19 @@ class GPUModelRunner(LoRAModelRunnerMixin): else: output_bin_counts = None post_update( - input_batch.idx_mapping, + idx_mapping, self.req_states.num_computed_tokens.gpu, self.req_states.last_sampled_tokens, output_bin_counts, sampled_tokens, num_sampled, num_rejected, - input_batch.query_start_loc, + query_start_loc, self.req_states.all_token_ids.gpu, self.req_states.total_len.gpu, ) - self.model_state.postprocess_state(input_batch, num_sampled) + self.model_state.postprocess_state(idx_mapping, num_sampled) @torch.inference_mode() def execute_model( @@ -1015,6 +1088,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: # Update the request states. + self.update_pp_decode_requests() self.finish_requests(scheduler_output) self.free_states(scheduler_output) self.add_requests(scheduler_output) @@ -1133,9 +1207,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs # to obtain inputs_embeds, because the compiled model expects this input. inputs_embeds = self.model_state.get_mm_embeddings( - scheduler_output.scheduled_encoder_inputs, - input_batch, - self.req_states, + scheduler_output.scheduled_encoder_inputs, input_batch ) model_inputs = { @@ -1155,12 +1227,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): assert intermediate_tensors is not None assert self.intermediate_tensors is not None n = input_batch.num_tokens_after_padding - model_inputs["intermediate_tensors"] = IntermediateTensors( - { - k: v[:n].copy_(intermediate_tensors.tensors[k][:n]) - for k, v in self.intermediate_tensors.tensors.items() - } - ) + new_tensors = { + k: v[:n] + if dummy_run + else v[:n].copy_(intermediate_tensors.tensors[k][:n]) + for k, v in self.intermediate_tensors.tensors.items() + } + model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) del intermediate_tensors # Run model. @@ -1252,10 +1325,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Non-last PP rank: hidden_states is None because this rank produced # IntermediateTensors instead of final hidden states. Receive the # sampled tokens broadcast from the last rank and update local state. - sampled, num_sampled, num_rejected = pp_receive( - input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1 - ) - self.postprocess(input_batch, sampled, num_sampled, num_rejected) + assert self.pp_handler is not None + all_decode_next = self.pp_handler.receive(input_batch) + # Optimistically update num_computed_tokens for entire batch here. + # Will be adjusted for rejections if necessary in update_requests. + self.postprocess_num_computed_tokens(input_batch) + if not all_decode_next: + # Might contain non-final prefill chunks, which will be scheduled + # in the immediate next step (rather than in pp_size steps). + self.model_state.postprocess_state(input_batch.idx_mapping, 0) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) @@ -1266,9 +1344,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): hidden_states, input_batch, grammar_output ) - if self.use_pp: + if self.pp_handler is not None: # Broadcast to non-last PP ranks (handles spec decode multi-token). - pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected) + self.pp_handler.broadcast( + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch, + ) assert self.prompt_logprobs_worker is not None prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs( @@ -1278,8 +1361,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.all_token_ids.gpu, self.req_states.num_computed_tokens.gpu, self.req_states.prompt_len.np, - self.req_states.prefill_len.np, - self.req_states.num_computed_prefill_tokens, ) # Prepare the model runner output. @@ -1305,17 +1386,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Get cached multimodal embeddings for draft forward. # NOTE: This is done here because postprocess updates # num_computed_prefill_tokens. - prefill_lens = self.req_states.prefill_len.np[input_batch.idx_mapping_np] - computed_prefill_lens = self.req_states.num_computed_prefill_tokens[ - input_batch.idx_mapping_np - ] mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings( input_batch.req_ids, input_batch.num_tokens, input_batch.num_scheduled_tokens, input_batch.query_start_loc_np, - prefill_lens, - computed_prefill_lens + 1, # +1 to consider the skew in eagle + input_batch.prefill_len_np, + # +1 to consider the skew in eagle + input_batch.num_computed_prefill_tokens_np + 1, ) # Postprocess results and update request states. @@ -1323,8 +1401,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): # ensuring that `copy_event` is recorded before calling postprocess. # This sequencing may slightly reduce latency as async D2H copy does not # need to wait for the postprocess to finish. - self.postprocess( - input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected + self.postprocess_sampled( + input_batch.idx_mapping, + sampler_output.sampled_token_ids, + num_sampled, + num_rejected, + input_batch.query_start_loc, ) if self.speculator is not None: @@ -1381,7 +1463,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): kv_connector_output = self.kv_connector.post_forward(finished_req_ids) if not self.is_last_pp_rank: - self.postprocess_pool(input_batch) + self.postprocess_num_computed_tokens(input_batch) return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) assert self.pooling_runner is not None @@ -1403,14 +1485,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): copy_stream=self.output_copy_stream, ) - self.postprocess_pool(input_batch) + self.postprocess_num_computed_tokens(input_batch) if self.use_async_scheduling: return async_output return async_output.get_output() - def postprocess_pool(self, input_batch: InputBatch) -> None: + def postprocess_num_computed_tokens(self, input_batch: InputBatch) -> None: # Update the number of computed tokens. - post_update_pool( + post_update_num_computed_tokens( input_batch.idx_mapping, self.req_states.num_computed_tokens.gpu, input_batch.query_start_loc, diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 7f7955a58ab..ee5d9384fa3 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -102,7 +102,6 @@ class DefaultModelState(ModelState): self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch, - req_states: RequestState, ) -> torch.Tensor: mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs( scheduled_encoder_inputs @@ -118,8 +117,8 @@ class DefaultModelState(ModelState): input_batch.num_tokens, input_batch.num_scheduled_tokens, input_batch.query_start_loc_np, - req_states.prefill_len.np[input_batch.idx_mapping_np], - req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np], + input_batch.prefill_len_np, + input_batch.num_computed_prefill_tokens_np, ) # Use unpadded input_ids to match is_mm_embed size (num_tokens). # input_batch.input_ids may be padded for CUDA graphs. @@ -178,7 +177,7 @@ class DefaultModelState(ModelState): # Capture with worst-case max_seq_len so the graph is valid at any replay. max_seq_len = self.max_model_len else: - max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 721e5c2013d..55bf8d473cc 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -57,18 +57,13 @@ class ModelState(ABC): return None def postprocess_state( - self, - input_batch: InputBatch, - num_sampled: torch.Tensor, + self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor ) -> None: return None @abstractmethod def get_mm_embeddings( - self, - scheduled_encoder_inputs: dict[str, list[int]], - input_batch: InputBatch, - req_states: RequestState, + self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch ) -> torch.Tensor | None: raise NotImplementedError diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 93115fdf64d..ced97c4f277 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -9,6 +9,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder from vllm.v1.kv_cache_interface import KVCacheConfig @@ -86,6 +87,12 @@ class MambaHybridModelState(DefaultModelState): num_tokens = input_batch.num_tokens query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) max_query_len = input_batch.num_scheduled_tokens.max().item() + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound + if for_capture: + # Capture with worst-case max_seq_len so the graph is valid at any replay. + max_seq_len = self.max_model_len + else: + max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") is_prefilling[: input_batch.num_reqs] = torch.from_numpy( @@ -106,13 +113,10 @@ class MambaHybridModelState(DefaultModelState): # 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: - spec_decode_mask = ( - input_batch.num_draft_tokens_per_req > 0 - ) & ~input_batch.is_prefilling_np + has_draft_tokens = input_batch.num_draft_tokens_per_req > 0 + spec_decode_mask = has_draft_tokens & ~input_batch.is_prefilling_np 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, input_batch.num_draft_tokens_per_req, -1 ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) @@ -129,7 +133,7 @@ class MambaHybridModelState(DefaultModelState): query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, seq_lens=input_batch.seq_lens, - max_seq_len=self.max_model_len, + max_seq_len=max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=kv_cache_config, @@ -139,12 +143,33 @@ class MambaHybridModelState(DefaultModelState): ) def postprocess_state( - self, - input_batch: InputBatch, - num_sampled: torch.Tensor, + self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int ) -> None: # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. - self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp( - num_sampled, min=1 - ) + if not isinstance(num_sampled, int): + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the + # kernel skips them rather than scattering with a host-side gather. + num_reqs = idx_mapping.shape[0] + if num_reqs: + _scatter_num_accepted_kernel[(num_reqs,)]( + idx_mapping, num_sampled, self.num_accepted_tokens_gpu + ) + return + + # Fill with single value. + self.num_accepted_tokens_gpu.index_fill_(0, idx_mapping, max(num_sampled, 1)) + + +@triton.jit +def _scatter_num_accepted_kernel( + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) + num_sampled_ptr, # [num_reqs] + num_accepted_ptr, # [max_num_reqs] +): + row = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + row) + if req_state_idx < 0: + return + num_sampled = tl.load(num_sampled_ptr + row) + tl.store(num_accepted_ptr + req_state_idx, tl.maximum(num_sampled, 1)) diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/whisper.py index 0ef3cadc87a..b38cdae9033 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/whisper.py @@ -84,10 +84,7 @@ class WhisperModelState(ModelState): return ("transcription",) def get_mm_embeddings( - self, - scheduled_encoder_inputs: dict[str, list[int]], - input_batch: InputBatch, - req_states: RequestState, + self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch ) -> None: # Ensure encoder inputs are ordered consistently with input_batch.req_ids. encoder_inputs: dict[str, list[int]] = {} diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index bf379b5fb5a..9f5d4c2d807 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -2,40 +2,193 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Pipeline Parallelism utils for V2 Model Runner.""" +from collections import deque +from dataclasses import dataclass + +import numpy as np import torch from vllm.distributed.parallel_state import get_pp_group +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.v1.worker.gpu.input_batch import InputBatch -def pp_broadcast( - sampled_token_ids: torch.Tensor, - num_sampled: torch.Tensor, - num_rejected: torch.Tensor, -) -> None: - pp = get_pp_group() - assert pp.is_last_rank +@dataclass +class PendingRecv: + """Per-step slot data for a deferred postprocess on the main stream.""" - assert sampled_token_ids.dtype == torch.int64 - torch.distributed.broadcast( - sampled_token_ids.contiguous(), src=pp.last_rank, group=pp.device_group - ) + event: torch.cuda.Event - combined = torch.stack((num_sampled, num_rejected), dim=0) - torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group) + sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] + num_sampled: torch.Tensor # [num_reqs] + num_rejected: torch.Tensor # [num_reqs] + idx_mapping: torch.Tensor # [num_reqs] + idx_mapping_np: np.ndarray # [num_reqs] + # Records which rows need a deferred postprocess (bool). + need_sampled_mask: np.ndarray # [num_reqs] + # Snapshot of slot generation counters at receive time, used to + # detect requests aborted since then. + gen_at_receive_np: np.ndarray # [num_reqs] -def pp_receive( - num_reqs: int, max_sample_len: int = 1 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - pp = get_pp_group() - assert not pp.is_last_rank +def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: + """Return a bool array of shape `[input_batch.num_reqs]` marking requests + with outputs that might be needed in a subsequent (decode) step. + Returns None if no sampled outputs are needed in the requests' next step.""" - sampled_tokens = torch.empty( - num_reqs, max_sample_len, dtype=torch.int64, device=pp.device - ) - torch.distributed.broadcast(sampled_tokens, src=pp.last_rank, group=pp.device_group) + old_computed = input_batch.num_computed_tokens_np + prefill_len = input_batch.prefill_len_np + max_seq_len = input_batch.max_seq_len_np + assert max_seq_len is not None # always populated under PP + # Exclude non-final prefill chunks (they don't produce a sample). + produces_sample = old_computed + input_batch.num_scheduled_tokens >= prefill_len + # Exclude requests that we know are finished. + not_finishing = np.maximum(old_computed, prefill_len) + 1 < max_seq_len + need_sampled_mask = produces_sample & not_finishing + return need_sampled_mask if need_sampled_mask.any() else None - combined = torch.empty(2, num_reqs, dtype=torch.int32, device=pp.device) - torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group) - num_sampled, num_rejected = combined.unbind(dim=0) - return sampled_tokens, num_sampled, num_rejected + +class PPHandler: + """Runs the PP sampled-token broadcast/recv on a side stream so the + default stream isn't gated by the matching peer call. Step T's recv is + consumed at step T+pp_size via `get_prev_sampled_outputs`. + + Uses a dedicated NCCL communicator (sibling of the PP `device_group`) + for the broadcast so it does not serialize on the wire with the + inter-stage hidden-state p2p send/recv ops. + """ + + def __init__( + self, max_num_reqs: int, num_speculative_steps: int, device: torch.device + ): + self.is_last_rank = get_pp_group().is_last_rank + self.last_rank = get_pp_group().last_rank + self.max_sample_len = num_speculative_steps + 1 + self.device = device + self.main_stream = torch.cuda.current_stream(device) + self.broadcast_stream = torch.cuda.Stream(device) + + # On non-last ranks, a FIFO with one entry per in-flight step: the entry + # pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded + # with pp_size None placeholders so the first pp_size consumes are no-ops. + # None means no postprocess is pending for that step (broadcast skipped). + self.queue: deque[PendingRecv | None] = ( + deque() if self.is_last_rank else deque([None] * get_pp_group().world_size) + ) + + # Per req-index generation counter, incremented every time a request + # index is freed in RequestStats. Used for invalidating freed req data + # between PP decodes. + self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32) + + # Dedicated subgroup for the sampled-token broadcast. + self.broadcast_group = get_pp_group().make_sibling_device_group( + group_desc="pp_broadcast" + ) + + def on_req_idx_freed(self, req_idx: int) -> None: + self.req_idx_gen_np[req_idx] += 1 + + def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: + """Consume the entry from pp_size steps ago and wait for its recv event, + then filter out entries whose request was freed since `receive`. + """ + if not self.queue: + return None + slot = self.queue.popleft() + # Reserve this step's slot; `receive` overwrites it if applicable. + self.queue.append(None) + if slot is None: + return None + + # Skip requests which did not need sampled output and/or those already + # finished. The post_update kernel skips the -1 entries. + freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np + exclude_mask = freed | ~slot.need_sampled_mask + idx_mapping = slot.idx_mapping + if exclude_mask.any(): + if exclude_mask.all(): + # No states require update anymore. + return None + # Filter excluded request indices. + idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) + idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + + self.main_stream.wait_event(slot.event) + return dict( + sampled_tokens=slot.sampled_tokens, + num_sampled=slot.num_sampled, + num_rejected=slot.num_rejected, + idx_mapping=idx_mapping, + ) + + def receive(self, input_batch: InputBatch) -> bool: + """Returns True iff sampled tokens need to be gathered from *all* + requests in the batch.""" + assert not self.is_last_rank + need_sampled_mask = compute_need_sampled_mask(input_batch) + if need_sampled_mask is None: + # Leave this step's reserved slot as None. + return False + + # Snapshot the per-slot generation counter so a later free of any of + # these RequestStates request indices is detectable at consume time. + gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np] + + num_reqs = input_batch.num_reqs + with torch.cuda.stream(self.broadcast_stream): + self.broadcast_stream.wait_stream(self.main_stream) + sampled_tokens = torch.empty( + num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device + ) + combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) + torch.distributed.broadcast( + sampled_tokens, src=self.last_rank, group=self.broadcast_group + ) + torch.distributed.broadcast( + combined, src=self.last_rank, group=self.broadcast_group + ) + event = self.broadcast_stream.record_event() + num_sampled, num_rejected = combined.unbind(dim=0) + # Must record_stream since these were allocated on broadcast stream but + # later used on the main stream. + sampled_tokens.record_stream(self.main_stream) + combined.record_stream(self.main_stream) + self.queue[-1] = PendingRecv( + event, + sampled_tokens, + num_sampled, + num_rejected, + input_batch.idx_mapping, + input_batch.idx_mapping_np, + need_sampled_mask, + gen_at_receive_np, + ) + return bool(need_sampled_mask.all()) + + def broadcast( + self, + sampled_token_ids: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + input_batch: InputBatch, + ) -> None: + assert self.is_last_rank + if compute_need_sampled_mask(input_batch) is None: + # No request needs sampled outputs for a subsequent decode step. + return + + assert sampled_token_ids.dtype == torch.int64 + with torch.cuda.stream(self.broadcast_stream): + self.broadcast_stream.wait_stream(self.main_stream) + torch.distributed.broadcast( + sampled_token_ids.contiguous(), + src=self.last_rank, + group=self.broadcast_group, + ) + combined = torch.stack((num_sampled, num_rejected), dim=0) + torch.distributed.broadcast( + combined, src=self.last_rank, group=self.broadcast_group + ) + for tensor in (sampled_token_ids, num_sampled, num_rejected): + tensor.record_stream(self.broadcast_stream) diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 71feb7cf0e9..b89ebac35d9 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -42,10 +42,6 @@ class PromptLogprobsWorker: num_computed_tokens: torch.Tensor, # [max_num_reqs] prompt_lens: np.ndarray, - # [max_num_reqs] - prefill_lens: np.ndarray, - # [max_num_reqs] - num_computed_prefill_tokens: np.ndarray, ) -> dict[str, LogprobsTensors]: idx_mapping_np = input_batch.idx_mapping_np needs_prompt_logprobs = self.uses_prompt_logprobs[idx_mapping_np] @@ -55,11 +51,11 @@ class PromptLogprobsWorker: num_prompt_logprobs = self.num_prompt_logprobs[idx_mapping_np] prompt_lens = prompt_lens[idx_mapping_np] - computed_prefill = num_computed_prefill_tokens[idx_mapping_np] + computed_prefill = input_batch.num_computed_prefill_tokens_np includes_prompt = computed_prefill < prompt_lens # NOTE(woosuk): If the request was resumed after preemption, its prompt # logprobs must have been computed before preemption. Skip. - resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np] + resumed_after_prompt = prompt_lens < input_batch.prefill_len_np needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt if not np.any(needs_prompt_logprobs): return {} diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 8bf884fd9b3..6b545aef3a2 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -7,6 +7,11 @@ import torch import vllm.envs as envs from vllm.config.model import LogprobsMode from vllm.sampling_params import SamplingParams +from vllm.v1.sample.ops.topk_topp_sampler import ( + apply_top_k_top_p, + flashinfer_sample, + flashinfer_sampler_supported, +) from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.metrics.logits import get_num_nans from vllm.v1.worker.gpu.sample.bad_words import BadWordsState @@ -45,6 +50,7 @@ class Sampler: self.bad_words_state = BadWordsState(req_states) self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.num_speculative_tokens = num_speculative_tokens + self.use_flashinfer = flashinfer_sampler_supported() def add_request( self, req_idx: int, prompt_len: int, sampling_params: SamplingParams @@ -77,6 +83,13 @@ class Sampler: # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear # that num_nans is computed before applying penalties and temperature. num_nans = get_num_nans(logits) if self.compute_nans else None + + max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) + max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids( + idx_mapping_np + ) + return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0 + sampled, processed_logits = self.sample( logits, expanded_idx_mapping, @@ -84,13 +97,10 @@ class Sampler: pos, input_ids, expanded_local_pos, + return_logprobs=return_logprobs, ) - max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) - max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids( - idx_mapping_np - ) - if max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0: + if return_logprobs: if self.logprobs_mode == "processed_logprobs": logits = processed_logits expanded_logits = logits.shape[0] != idx_mapping_np.shape[0] @@ -128,6 +138,7 @@ class Sampler: pos: torch.Tensor, input_ids: torch.Tensor, expanded_local_pos: torch.Tensor, + skip_top_k_top_p: bool = False, ) -> torch.Tensor: # Copy logits to a new FP32 tensor. logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) @@ -163,6 +174,9 @@ class Sampler: # Apply min_p in place. self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np) + if skip_top_k_top_p: + return logits + # Apply top_k and/or top_p. This might or might not return a new tensor. return self.sampling_states.apply_top_k_top_p( logits, expanded_idx_mapping, idx_mapping_np @@ -176,6 +190,7 @@ class Sampler: pos: torch.Tensor, input_ids: torch.Tensor, expanded_local_pos: torch.Tensor, + return_logprobs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: processed_logits = self.apply_sampling_params( logits, @@ -184,16 +199,33 @@ class Sampler: pos, input_ids, expanded_local_pos, + skip_top_k_top_p=True, + ) + top_k, top_p = self.sampling_states.get_top_k_top_p( + expanded_idx_mapping, idx_mapping_np + ) + use_flashinfer = self.use_flashinfer and not ( + # Don't use FI sampler if no requests use top_k/top_p, if there are + # any greedy requests or per-request seeds, or if post-processed + # logprobs need to be returned for any requests. + (top_k is None and top_p is None) + or (return_logprobs and self.logprobs_mode == "processed_logprobs") + or self.sampling_states.any_greedy(idx_mapping_np) + or self.sampling_states.any_explicit_seed(idx_mapping_np) ) # Sample the next token. - sampled = gumbel_sample( - processed_logits, - expanded_idx_mapping, - self.sampling_states.temperature.gpu, - self.sampling_states.seeds.gpu, - pos, - apply_temperature=False, - use_fp64=self.use_fp64_gumbel, - ) + if use_flashinfer: + sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64) + else: + processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) + sampled = gumbel_sample( + processed_logits, + expanded_idx_mapping, + self.sampling_states.temperature.gpu, + self.sampling_states.seeds.gpu, + pos, + apply_temperature=False, + use_fp64=self.use_fp64_gumbel, + ) return sampled, processed_logits diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index f247acba07c..bf2f1ce78fe 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -24,6 +24,9 @@ class SamplingStates: self.top_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.min_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.seeds = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + # Tracks whether `seed` was set explicitly by the user, so callers + # can fall back from RNG paths that don't honor per-request seeds. + self.seeds_set = np.zeros(max_num_reqs, dtype=bool) # Initialize top_k and top_p manually because 0 is an invalid value for them. self.top_k.np.fill(self.vocab_size) @@ -45,6 +48,7 @@ class SamplingStates: self.min_p.np[req_idx] = sampling_params.min_p seed = sampling_params.seed + self.seeds_set[req_idx] = seed is not None if seed is None: seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX) self.seeds.np[req_idx] = seed @@ -85,20 +89,31 @@ class SamplingStates: return apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu) + def get_top_k_top_p( + self, expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size) + do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0) + top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None + top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None + return top_k, top_p + def apply_top_k_top_p( self, logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, ) -> torch.Tensor: - do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size) - do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0) - if not (do_top_k or do_top_p): + top_k, top_p = self.get_top_k_top_p(expanded_idx_mapping, idx_mapping_np) + if top_k is None and top_p is None: return logits - - top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None - top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None return apply_top_k_top_p(logits, top_k, top_p) + def any_greedy(self, idx_mapping_np: np.ndarray) -> bool: + return bool(np.any(self.temperature.np[idx_mapping_np] == 0.0)) + + def any_explicit_seed(self, idx_mapping_np: np.ndarray) -> bool: + return bool(np.any(self.seeds_set[idx_mapping_np])) + def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int: return int(np.max(self.num_logprobs[idx_mapping_np])) diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8ca2882e3a9..1a1ae1f63e9 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -52,6 +52,7 @@ class EagleSpeculator: self.max_num_reqs = self.scheduler_config.max_num_seqs self.max_num_tokens = self.scheduler_config.max_num_batched_tokens self.max_model_len = vllm_config.model_config.max_model_len + self.draft_max_seq_len = self.max_model_len # We need to get the hidden size from the draft model config because # the draft model's hidden size can be different from the target model's # hidden size (e.g., Llama 3.3 70B). @@ -416,7 +417,7 @@ class EagleSpeculator: query_start_loc_cpu=query_start_loc_cpu, max_query_len=1, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], - max_seq_len=self.max_model_len, + max_seq_len=self.draft_max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, kv_cache_config=self.kv_cache_config, @@ -494,6 +495,10 @@ class EagleSpeculator: num_tokens = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() + self.draft_max_seq_len = min( + max_seq_len + self.num_speculative_steps, self.max_model_len + ) # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of eagle's input_ids and diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index cdd7286fa56..be7bae7f17e 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -65,6 +65,9 @@ class RequestState: self.max_num_reqs, 1, dtype=torch.int64, device=device ) + # Max total seq length (prompt_len + max_tokens). + self.max_seq_len = np.zeros(self.max_num_reqs, dtype=np.int32) + # Draft tokens. self.draft_tokens = torch.zeros( self.max_num_reqs, @@ -87,12 +90,14 @@ class RequestState: prompt_len: int, all_token_ids: list[int], num_computed_tokens: int, + max_tokens: int, ) -> None: assert len(self.free_indices) > 0, "No free indices" req_idx = self.free_indices.pop() self.req_id_to_index[req_id] = req_idx self.index_to_req_id[req_idx] = req_id + self.max_seq_len[req_idx] = prompt_len + max_tokens self.prompt_len.np[req_idx] = prompt_len prefill_len = len(all_token_ids) assert prefill_len >= prompt_len, ( @@ -124,17 +129,11 @@ class RequestState: self.all_token_ids.apply_write() self.num_computed_tokens.apply_write() - def remove_request(self, req_id: str) -> bool: + def remove_request(self, req_id: str) -> int | None: + """Return the freed slot index, or None if the request was not found.""" req_idx = self.req_id_to_index.pop(req_id, None) if req_idx is None: - # Request not found. - return False + return None self.index_to_req_id.pop(req_idx, None) self.free_indices.append(req_idx) - return True - - def is_prefilling(self, idx_mapping_np: np.ndarray) -> np.ndarray: - return ( - self.num_computed_prefill_tokens[idx_mapping_np] - < self.prefill_len.np[idx_mapping_np] - ) + return req_idx diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 9a05c765894..5265c3a43a2 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -104,7 +104,7 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItem, PlaceholderRange, ) -from vllm.multimodal.utils import group_and_batch_mm_kwargs +from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingType @@ -152,6 +152,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.outputs import ( EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, @@ -1084,11 +1085,11 @@ class GPUModelRunner( def _init_kv_zero_meta(self) -> None: """One-time precomputation for _zero_block_ids. - Delegates to KVBlockZeroer.init_meta with the runner's state. Called from gpu_worker.py outside the CuMem pool context. """ - self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory) - self._kv_block_zeroer.init_meta( + self._kv_block_zeroer = KVBlockZeroer( + self.device, + self.pin_memory, attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), kernel_block_sizes=self._kernel_block_sizes, cache_dtype=self.cache_config.cache_dtype, @@ -3100,23 +3101,18 @@ class GPUModelRunner( req_state = self.requests[req_id] num_computed_tokens = req_state.num_computed_tokens + shift_computed_tokens - for mm_feature in req_state.mm_features: + mm_features = req_state.mm_features + lo, hi = get_mm_features_in_window( + mm_features, + start=num_computed_tokens, + end=num_computed_tokens + num_scheduled_tokens, + ) + for i in range(lo, hi): + mm_feature = mm_features[i] pos_info = mm_feature.mm_position start_pos = pos_info.offset num_encoder_tokens = pos_info.length - # The encoder output is needed if the two ranges overlap: - # [num_computed_tokens, - # num_computed_tokens + num_scheduled_tokens) and - # [start_pos, start_pos + num_encoder_tokens) - if start_pos >= num_computed_tokens + num_scheduled_tokens: - # The encoder output is not needed in this step. - break - if start_pos + num_encoder_tokens <= num_computed_tokens: - # The encoder output is already processed and stored - # in the decoder's KV cache. - continue - start_idx = max(num_computed_tokens - start_pos, 0) end_idx = min( num_computed_tokens - start_pos + num_scheduled_tokens, @@ -4362,6 +4358,21 @@ class GPUModelRunner( return None + def _input_fits_in_drafter( + self, common_attn_metadata: CommonAttentionMetadata | None + ) -> bool: + if common_attn_metadata is None: + return False + assert self.speculative_config is not None + # DFlash queries one extra token (the bonus token) beyond num_spec_tokens + num_drafter_query_tokens = self.num_spec_tokens + ( + 1 if self.speculative_config.use_dflash() else 0 + ) + return ( + common_attn_metadata.max_seq_len + num_drafter_query_tokens + <= self.effective_drafter_max_model_len + ) + @torch.inference_mode def sample_tokens( self, grammar_output: "GrammarOutput | None" @@ -4441,9 +4452,8 @@ class GPUModelRunner( propose_drafts_after_bookkeeping = False if spec_config is not None: # Decide whether to run the drafter or zero out draft tokens. - input_fits_in_drafter = spec_decode_common_attn_metadata is not None and ( - spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens - <= self.effective_drafter_max_model_len + input_fits_in_drafter = self._input_fits_in_drafter( + spec_decode_common_attn_metadata ) use_gpu_toks = ( spec_config.use_eagle() @@ -6226,6 +6236,7 @@ class GPUModelRunner( ) kv_cache_spec = self.get_kv_cache_spec() + KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec) kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec) min_blocks = self.compilation_config.max_cudagraph_capture_size or 1 @@ -6264,11 +6275,21 @@ class GPUModelRunner( # Calls torch.accelerator.synchronize() self._cleanup_profiling_kv_cache() + if current_platform.is_rocm(): + # Drop captured graphs before distributed teardown. On ROCm, delayed + # graph destruction can surface HSA faults in the next engine startup. + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + self.encoder_cudagraph_manager = None self.compilation_config.static_forward_context.clear() self.model = None # type: ignore[assignment] _ROPE_DICT.clear() reset_workspace_manager() + if current_platform.is_rocm(): + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() def _cleanup_profiling_kv_cache(self) -> None: torch.accelerator.synchronize() @@ -6304,6 +6325,42 @@ class GPUModelRunner( logger.debug("Cleaned up profiling KV cache and CUDA graphs") + @torch.inference_mode() + def _create_encoder_cudagraph_manager(self) -> "EncoderCudaGraphManager | None": + if not ( + self.compilation_config.cudagraph_mm_encoder and self.supports_mm_inputs + ): + return None + + # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper, because + # @runtime_checkable Protocol isinstance() checks do not work through + # __getattr__ forwarding. + from vllm.model_executor.models.interfaces import ( + SupportsEncoderCudaGraph, + supports_encoder_cudagraph, + ) + from vllm.v1.worker.encoder_cudagraph import ( + EncoderCudaGraphManager, + ) + + raw_model = self.get_model() + if not supports_encoder_cudagraph(raw_model): + return None + + return EncoderCudaGraphManager( + vllm_config=self.vllm_config, + device=self.device, + dtype=self.dtype, + model=cast(SupportsEncoderCudaGraph, raw_model), + ) + + @torch.inference_mode() + def _maybe_init_encoder_cudagraph_manager(self) -> None: + if self.encoder_cudagraph_manager is None: + self.encoder_cudagraph_manager = self._create_encoder_cudagraph_manager() + if self.encoder_cudagraph_manager is not None: + logger.info("Initialized EncoderCudaGraphManager for vision encoder") + @torch.inference_mode() def profile_cudagraph_memory(self) -> int: with set_current_vllm_config(self.vllm_config): @@ -6312,24 +6369,40 @@ class GPUModelRunner( saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured capture_descs = self.cudagraph_dispatcher.get_capture_descs() + # Use a temporary manager for memory profiling. The persistent manager + # is initialized later so it does not keep profiling-only graph state. + encoder_cudagraph_manager = self._create_encoder_cudagraph_manager() - total_graphs = sum(len(descs) for _, descs in capture_descs) + decoder_graphs = sum(len(descs) for _, descs in capture_descs) + encoder_graphs = ( + encoder_cudagraph_manager.get_num_graphs_to_capture() + if encoder_cudagraph_manager is not None + else 0 + ) + total_graphs = decoder_graphs + encoder_graphs if total_graphs == 0: logger.debug("No CUDA graphs will be captured, skipping profiling") self._cleanup_profiling_kv_cache() return 0 - logger.info( - "Profiling CUDA graph memory: %s", - ", ".join( + graph_groups = [ + *( f"{mode.name}={len(descs)} (largest={descs[0].num_tokens})" for mode, descs in capture_descs if descs ), - ) + ] + if encoder_graphs > 0: + graph_groups.append( + f"ENCODER={encoder_graphs} " + f"(largest={encoder_cudagraph_manager.token_budgets[-1]})" + ) + + logger.info("Profiling CUDA graph memory: %s", ", ".join(graph_groups)) # Use a temporary pool for profiling to avoid fragmentation in the main pool. profiling_pool = current_platform.graph_pool_handle() + encoder_profiling_pool = current_platform.graph_pool_handle() original_pools: dict[int, Any] = {} all_wrappers = list(CUDAGraphWrapper._all_instances) + list( BreakableCUDAGraphWrapper._all_instances @@ -6338,73 +6411,98 @@ class GPUModelRunner( original_pools[id(instance)] = instance.graph_pool instance.graph_pool = profiling_pool - set_cudagraph_capturing_enabled(True) - with self._freeze_gc(), graph_capture(device=self.device): - shared_memory_estimate = {} - per_graph_estimate = {} - torch.accelerator.synchronize() - torch.accelerator.empty_cache() + shared_memory_estimate = {} + per_graph_estimate = {} + encoder_memory_estimate = 0 - for mode, descs in capture_descs: - profile_descs = descs[:2] - mem_samples: list[int] = [] + # Cleanup-only guard: CUDA graph capture errors should still propagate + # because encoder graph capture is opt-in. + try: + set_cudagraph_capturing_enabled(True) + with self._freeze_gc(), graph_capture(device=self.device): + torch.accelerator.synchronize() + torch.accelerator.empty_cache() - for i, desc in enumerate(profile_descs): - mem_before = torch.cuda.mem_get_info()[0] - self._warmup_and_capture( - desc, - cudagraph_runtime_mode=mode, - profile_seq_lens=( - min( - self.max_model_len, - self.max_num_tokens // desc.num_tokens, - ) - if mode == CUDAGraphMode.FULL and i == 0 - else None - ), + for mode, descs in capture_descs: + profile_descs = descs[:2] + mem_samples: list[int] = [] + + for i, desc in enumerate(profile_descs): + mem_before = torch.cuda.mem_get_info()[0] + self._warmup_and_capture( + desc, + cudagraph_runtime_mode=mode, + profile_seq_lens=( + min( + self.max_model_len, + self.max_num_tokens // desc.num_tokens, + ) + if mode == CUDAGraphMode.FULL and i == 0 + else None + ), + ) + torch.accelerator.synchronize() + free_after = torch.cuda.mem_get_info()[0] + mem_samples.append(mem_before - free_after) + + first_capture = mem_samples[0] + # Use at least 1 MiB per graph for driver overhead + per_graph = max( + mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20 ) + + shared_memory_estimate[mode] = first_capture + per_graph_estimate[mode] = per_graph * (len(descs) - 1) + + logger.debug( + "Estimated %s CUDA graph memory: " + "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph", + mode.name, + first_capture / (1 << 20), + len(descs), + per_graph / (1 << 20), + ) + + if encoder_cudagraph_manager is not None: + mem_before = torch.cuda.mem_get_info()[0] + encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool) torch.accelerator.synchronize() free_after = torch.cuda.mem_get_info()[0] - mem_samples.append(mem_before - free_after) + encoder_memory_estimate = max(mem_before - free_after, 0) - first_capture = mem_samples[0] - # Use at least 1 MiB per graph for driver overhead - per_graph = max(mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20) - - shared_memory_estimate[mode] = first_capture - per_graph_estimate[mode] = per_graph * (len(descs) - 1) - - logger.debug( - "Estimated %s CUDA graph memory: " - "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph", - mode.name, - first_capture / (1 << 20), - len(descs), - per_graph / (1 << 20), - ) - - set_cudagraph_capturing_enabled(False) - CUDAGraphWrapper.clear_all_graphs() - BreakableCUDAGraphWrapper.clear_all_graphs() - all_wrappers = list(CUDAGraphWrapper._all_instances) + list( - BreakableCUDAGraphWrapper._all_instances - ) - for instance in all_wrappers: - if id(instance) in original_pools: - instance.graph_pool = original_pools[id(instance)] - for key_set in self.cudagraph_dispatcher.cudagraph_keys.values(): - key_set.clear() - self.cudagraph_dispatcher.keys_initialized = False - self.maybe_remove_all_loras(self.lora_config) - self._cleanup_profiling_kv_cache() - compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured + logger.debug( + "Estimated encoder CUDA graph memory: %.2f MiB for %d graphs", + encoder_memory_estimate / (1 << 20), + encoder_graphs, + ) + finally: + set_cudagraph_capturing_enabled(False) + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + if encoder_cudagraph_manager is not None: + encoder_cudagraph_manager.clear() + all_wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for instance in all_wrappers: + if id(instance) in original_pools: + instance.graph_pool = original_pools[id(instance)] + for key_set in self.cudagraph_dispatcher.cudagraph_keys.values(): + key_set.clear() + self.cudagraph_dispatcher.keys_initialized = False + self.maybe_remove_all_loras(self.lora_config) + self._cleanup_profiling_kv_cache() + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured # FULL and PIECEWISE graphs share the global pool at runtime and are # never replayed concurrently, so the pool overlays their memory. # Take the max to avoid double-counting the overlap. - total_estimate = max(shared_memory_estimate.values()) + sum( + decoder_estimate = max(shared_memory_estimate.values(), default=0) + sum( per_graph_estimate.values() ) + # Encoder graphs use a manager-local pool at runtime, separate from the + # decoder pool, so add their estimate instead of overlaying it. + total_estimate = decoder_estimate + encoder_memory_estimate logger.info( "Estimated CUDA graph memory: %.2f GiB total", total_estimate / (1 << 30), @@ -6422,31 +6520,7 @@ class GPUModelRunner( return 0 # Initialize encoder CUDA graph manager if enabled. - # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper, - # because @runtime_checkable Protocol isinstance() checks do not - # work through __getattr__ forwarding. - if ( - self.compilation_config.cudagraph_mm_encoder - and self.supports_mm_inputs - and self.encoder_cudagraph_manager is None - ): - from vllm.model_executor.models.interfaces import ( - SupportsEncoderCudaGraph, - supports_encoder_cudagraph, - ) - from vllm.v1.worker.encoder_cudagraph import ( - EncoderCudaGraphManager, - ) - - raw_model = self.get_model() - if supports_encoder_cudagraph(raw_model): - self.encoder_cudagraph_manager = EncoderCudaGraphManager( - vllm_config=self.vllm_config, - device=self.device, - dtype=self.dtype, - model=cast(SupportsEncoderCudaGraph, raw_model), - ) - logger.info("Initialized EncoderCudaGraphManager for vision encoder") + self._maybe_init_encoder_cudagraph_manager() compilation_counter.num_gpu_runner_capture_triggers += 1 @@ -6473,7 +6547,8 @@ class GPUModelRunner( # Capture encoder CUDA graphs if enabled if self.encoder_cudagraph_manager is not None: - self.encoder_cudagraph_manager.capture() + encoder_graph_pool = current_platform.graph_pool_handle() + self.encoder_cudagraph_manager.capture(graph_pool=encoder_graph_pool) torch.accelerator.synchronize() end_free_gpu_memory = torch.cuda.mem_get_info()[0] diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 121fc69f532..259cd05554c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1147,7 +1147,10 @@ def init_worker_distributed_environment( from vllm.model_executor.layers.batch_invariant import init_batch_invariance init_batch_invariance() - override_envs_for_eplb(parallel_config) + override_envs_for_eplb( + parallel_config, + moe_backend=getattr(vllm_config.kernel_config, "moe_backend", None), + ) set_custom_all_reduce(not parallel_config.disable_custom_all_reduce) init_method = distributed_init_method or "env://" diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 7cb1620c95e..c0f44b6db0c 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -80,30 +80,23 @@ def _zero_kv_blocks_kernel( class KVBlockZeroer: """Manages efficient zeroing of KV cache blocks via a Triton kernel. - Call :meth:`init_meta` once after KV caches are allocated to precompute - segment addresses, then call :meth:`zero_block_ids` each step to zero + Construct once after KV caches are allocated to precompute segment + addresses, then call :meth:`zero_block_ids` each step to zero newly-allocated blocks. """ - def __init__(self, device: torch.device, pin_memory: bool): - self.device = device - self.pin_memory = pin_memory - self._meta: tuple[torch.Tensor, int, int, int] | None = None - self._id_cap: int = 0 - self._ids_pinned: torch.Tensor | None = None - self._ids_gpu: torch.Tensor | None = None - - def init_meta( + def __init__( self, + device: torch.device, + pin_memory: bool, attn_groups_iter: Iterable["AttentionGroup"], kernel_block_sizes: list[int], cache_dtype: str, - runner_only_attn_layers: set[str], static_forward_context: dict[str, Any], + runner_only_attn_layers: set[str] | None = None, ) -> None: - """One-time precomputation for zero_block_ids. + """Precompute the absolute-address table for the Triton zeroing kernel. - Builds absolute-address table for the Triton zeroing kernel. Each entry is the absolute byte address of a segment start on the GPU, so segments in different CUDA allocations work correctly. @@ -114,6 +107,15 @@ class KVBlockZeroer: Only AttentionSpec layers are processed; Mamba layers are skipped. """ + self.device = device + self.pin_memory = pin_memory + self._meta: tuple[torch.Tensor, int, int, int] | None = None + self._id_cap: int = 0 + self._ids_pinned: torch.Tensor | None = None + self._ids_gpu: torch.Tensor | None = None + + if runner_only_attn_layers is None: + runner_only_attn_layers = set() seen_ptrs: set[int] = set() seg_addrs: list[int] = [] page_size_el: int | None = None