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/ci-fetch-log.sh b/.buildkite/scripts/ci-fetch-log.sh index 02798b56f4a..3f99bc50a57 100755 --- a/.buildkite/scripts/ci-fetch-log.sh +++ b/.buildkite/scripts/ci-fetch-log.sh @@ -9,6 +9,13 @@ # Find and via: # gh pr checks --repo vllm-project/vllm # Each failing row's URL is .../builds/#. +# +# Default output path: ci--.log (e.g. +# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's +# first 8 chars, so the second segment is needed for uniqueness when +# fetching multiple jobs in parallel. The script refuses to overwrite an +# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1 +# to override. set -euo pipefail @@ -26,12 +33,12 @@ if [ $# -lt 1 ]; then usage; fi if [[ "$1" == https://* ]]; then BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p') JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1) - OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}" + OUT="${2:-}" else if [ $# -lt 2 ]; then usage; fi BUILD="$1" JOB="$2" - OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}" + OUT="${3:-}" fi if [ -z "$BUILD" ] || [ -z "$JOB" ]; then @@ -39,6 +46,18 @@ if [ -z "$BUILD" ] || [ -z "$JOB" ]; then usage fi +# Jobs in the same build share the UUID's first segment, so include the +# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames +# unique when fetching multiple jobs from one build in parallel. +if [ -z "$OUT" ]; then + OUT="ci-${BUILD}-${JOB:0:13}.log" +fi + +if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then + echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2 + exit 1 +fi + COOKIES=$(mktemp) trap 'rm -f "$COOKIES"' EXIT 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/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh index e26273bba39..2ff0b2e1251 100644 --- a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh @@ -49,6 +49,7 @@ for BACK in "${BACKENDS[@]}"; do --data-parallel-size 2 \ --enable-expert-parallel \ --enable-eplb \ + --eplb-config '{"use_async": false}' \ --trust-remote-code \ --max-model-len 2048 \ --all2all-backend "$BACK" \ diff --git a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh index 729a0fb7f68..87168553190 100644 --- a/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh +++ b/.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh @@ -48,7 +48,7 @@ for BACK in "${BACKENDS[@]}"; do --enforce-eager \ --enable-eplb \ --all2all-backend "$BACK" \ - --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \ + --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true, "use_async":false}' \ --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \ --data-parallel-size "${DATA_PARALLEL_SIZE}" \ --enable-expert-parallel \ diff --git a/.buildkite/scripts/tool_call/run-bfcl-eval.sh b/.buildkite/scripts/tool_call/run-bfcl-eval.sh index 3748cab62c7..d50767ef0f2 100755 --- a/.buildkite/scripts/tool_call/run-bfcl-eval.sh +++ b/.buildkite/scripts/tool_call/run-bfcl-eval.sh @@ -70,7 +70,7 @@ echo "============================================" # ---- Install bfcl-eval if missing ---- if ! python3 -c "import bfcl_eval" 2>/dev/null; then echo "Installing bfcl-eval..." - pip install "bfcl-eval>=2025.10.20.1,<2026" + uv pip install "bfcl-eval>=2025.10.20.1,<2026" fi # ---- Cleanup handler ---- @@ -100,7 +100,7 @@ SERVE_ARGS=( --tensor-parallel-size "$TP_SIZE" --max-model-len "$MAX_MODEL_LEN" --enforce-eager - --no-enable-prefix-caching + --enable-prefix-caching ) # Append reasoning parser if specified diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c7ef61b719a..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 @@ -1261,7 +1258,7 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - label: Entrypoints Integration (API Server openai - Part 2) # TBD timeout_in_minutes: 180 @@ -1275,10 +1272,14 @@ steps: - vllm/ - 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 @@ -1368,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 + - 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 @@ -1484,7 +1485,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD +- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_4 @@ -1526,7 +1527,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 -- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_4 @@ -2745,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 @@ -2768,7 +2766,7 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py - label: Entrypoints Integration (API Server openai - Part 2) # TBD timeout_in_minutes: 180 @@ -2782,10 +2780,14 @@ steps: - vllm/ - 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 @@ -2895,7 +2897,7 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt -- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 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/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index bb8aa14eac1..88039a33960 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -2,8 +2,8 @@ group: E2E Integration depends_on: - image-build steps: -- label: DeepSeek V2-Lite Accuracy - key: deepseek-v2-lite-accuracy +- label: DeepSeek V2-Lite Sync EPLB Accuracy + key: deepseek-v2-lite-sync-eplb-accuracy timeout_in_minutes: 60 device: h100 optional: true @@ -12,8 +12,8 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010 -- label: Qwen3-30B-A3B-FP8-block Accuracy - key: qwen3-30b-a3b-fp8-block-accuracy +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy timeout_in_minutes: 60 device: h100 optional: true @@ -22,8 +22,8 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 -- label: Qwen3-30B-A3B-FP8-block Accuracy (B200) - key: qwen3-30b-a3b-fp8-block-accuracy-b200 +- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200) + key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200 timeout_in_minutes: 60 device: b200-k8s optional: true diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 57bde22194a..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 + - 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 @@ -43,7 +43,7 @@ steps: - tests/entrypoints/test_chat_utils commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py + - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py mirror: amd: device: mi325_1 @@ -60,9 +60,13 @@ steps: - vllm/ - 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 @@ -98,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 @@ -153,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_executor.yaml b/.buildkite/test_areas/model_executor.yaml index c41ef8a7110..e34b7eadfac 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -14,5 +14,12 @@ steps: commands: - apt-get update && apt-get install -y curl libsodium23 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s model_executor -m '(not slow_test)' - - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py + # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA + # init surfaces a stack instead of silently stalling. + - export PYTHONFAULTHANDLER=1 + # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA + # init) fails fast with a traceback instead of running until the global + # build timeout. The `thread` method also handles hangs inside C/CUDA + # calls that the signal method cannot interrupt. + - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread + - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread 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 f750d58be58..16d69f77345 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -16,7 +16,7 @@ steps: - tests/benchmarks/test_serve_cli.py - tests/entrypoints/openai/chat_completion/test_chat_completion.py # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py - # - tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py + # - tests/entrypoints/openai/completion/test_prompt_validation.py - tests/entrypoints/openai/completion/test_shutdown.py # - tests/entrypoints/openai/test_return_token_ids.py @@ -28,7 +28,7 @@ steps: - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid" - # - pytest -v -s entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py + # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds" - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly" # - pytest -v -s entrypoints/openai/test_return_token_ids.py @@ -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/CODEOWNERS b/.github/CODEOWNERS index 540540e5132..beaaa5d8642 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -40,6 +40,12 @@ /vllm/entrypoints/chat_utils.py @DarkLight1337 /vllm/entrypoints/llm.py @DarkLight1337 +# Rust Frontend +/rust/ @BugenZhao @njhill +/build_rust.sh @BugenZhao @njhill +/rust-toolchain.toml @BugenZhao @njhill +/.buildkite/test_areas/rust* @BugenZhao @njhill + # Input/Output Processing /vllm/sampling_params.py @njhill @NickLucche /vllm/pooling_params.py @noooop @DarkLight1337 @@ -72,11 +78,13 @@ /vllm/v1/worker/gpu/kv_connector.py @orozery # CI & building -/.buildkite @Harry-Chen -/docker/Dockerfile @Harry-Chen +/.buildkite @Harry-Chen @khluu +/docker/Dockerfile @Harry-Chen @khluu +/pyproject.toml @khluu +/setup.py @khluu # Test ownership -/.buildkite/lm-eval-harness @mgoin +/.buildkite/lm-eval-harness @mgoin /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao 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/.readthedocs.yaml b/.readthedocs.yaml index 1dabec70ba5..d2a400f46b9 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -9,8 +9,8 @@ build: python: "3.12" jobs: post_checkout: - - bash docs/pre_run_check.sh - git fetch origin main --unshallow --no-tags --filter=blob:none || true + - bash docs/pre_run_check.sh pre_create_environment: - pip install uv create_environment: diff --git a/CMakeLists.txt b/CMakeLists.txt index 706f390948f..9e8ccfa8718 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -144,14 +144,14 @@ endif() # Set up GPU language and check the torch version and warn if it isn't # what is expected. # -if (NOT HIP_FOUND AND CUDA_FOUND) +if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND) set(VLLM_GPU_LANG "CUDA") if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA}) message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} " "expected for CUDA build, saw ${Torch_VERSION} instead.") endif() -elseif(HIP_FOUND) +elseif(HIP_FOUND OR PYTORCH_FOUND_HIP) set(VLLM_GPU_LANG "HIP") # Importing torch recognizes and sets up some HIP/ROCm configuration but does @@ -305,10 +305,6 @@ endif() # set(VLLM_EXT_SRC - "csrc/cache_kernels.cu" - "csrc/cache_kernels_fused.cu" - "csrc/attention/paged_attention_v1.cu" - "csrc/attention/paged_attention_v2.cu" "csrc/cuda_view.cu" "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" "csrc/quantization/activation_kernels.cu" @@ -637,6 +633,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 +645,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/attention/merge_attn_states.cu" "csrc/libtorch_stable/sampler.cu" "csrc/libtorch_stable/topk.cu" - "csrc/libtorch_stable/mamba/selective_scan_fwd.cu") + "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" + "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") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC @@ -657,8 +659,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu" "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") set_gencode_flags_for_srcs( @@ -683,6 +683,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "in CUDA target architectures.") endif() + # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0. + cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS) + set(SRCS + "csrc/libtorch_stable/fp32_router_gemm_entry.cu" + "csrc/libtorch_stable/fp32_router_gemm.cu") + set_gencode_flags_for_srcs( + SRCS "${SRCS}" + CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}") + list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") + message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}") + else() + message(STATUS "Not building fp32_router_gemm as no compatible archs found " + "(requires SM90+ and CUDA >= 12.0).") + endif() + # Only build AllSpark kernels if we are building for at least some compatible archs. cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}") if (ALLSPARK_ARCHS) @@ -924,13 +940,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly - # from cache_kernels.cu, so it belongs in _C rather than _C_stable. - set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${NVFP4_KV_SRC}" CUDA_ARCHS "${FP4_ARCHS}") - target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1") @@ -960,11 +974,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${SRCS}" CUDA_ARCHS "${FP4_ARCHS}") list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}") - set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu") + set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") set_gencode_flags_for_srcs( SRCS "${NVFP4_KV_SRC}" CUDA_ARCHS "${FP4_ARCHS}") - target_sources(_C PRIVATE ${NVFP4_KV_SRC}) + list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}") target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1) list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1") list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1") @@ -1242,24 +1256,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") " in CUDA target architectures") endif() - # DeepSeek V3 router GEMM kernel - requires SM90+ - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}") - endif() - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS) + # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0. + # (fp32_router_gemm has been migrated to _C_stable_libtorch above.) + cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS) set(DSV3_ROUTER_GEMM_SRC "csrc/moe/dsv3_router_gemm_entry.cu" "csrc/moe/dsv3_router_gemm_float_out.cu" "csrc/moe/dsv3_router_gemm_bf16_out.cu") set_gencode_flags_for_srcs( SRCS "${DSV3_ROUTER_GEMM_SRC}" - CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}") + CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}") list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}") - message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}") + + message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}") else() - message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found" + message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found" " (requires SM90+ and CUDA >= 12.0)") endif() endif() @@ -1286,6 +1298,14 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_HAS_GFX1100 OFF) + if(VLLM_GPU_ARCHES MATCHES "gfx1100") + set(VLLM_ROCM_HAS_GFX1100 ON) + list(APPEND VLLM_ROCM_EXT_SRC + "csrc/rocm/q_gemm_rdna3.cu" + "csrc/rocm/q_gemm_rdna3_wmma.cu") + endif() + define_extension_target( _rocm_C DESTINATION vllm @@ -1295,6 +1315,10 @@ if(VLLM_GPU_LANG STREQUAL "HIP") ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + + if(VLLM_ROCM_HAS_GFX1100) + target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) + endif() endif() # Must run after the last HIP `define_extension_target` so every extension diff --git a/benchmarks/kernels/benchmark_router_gemm.py b/benchmarks/kernels/benchmark_router_gemm.py new file mode 100644 index 00000000000..ba46a7fd78d --- /dev/null +++ b/benchmarks/kernels/benchmark_router_gemm.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform +from vllm.transformers_utils.config import get_config +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser + +# Dimensions supported by the DSV3 specialized kernel +DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] +DSV3_SUPPORTED_HIDDEN_SIZES = [7168] + +# Dimensions supported by the gpt-oss specialized kernel +GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128] +GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880] + +# Dimensions supported by the fp32 specialized kernel (MiniMax-M2) +FP32_SUPPORTED_NUM_EXPERTS = [256] +FP32_SUPPORTED_HIDDEN_SIZES = [3072] +FP32_MAX_TOKENS = 32 + + +def get_batch_size_range(max_batch_size): + return [2**x for x in range(14) if 2**x <= max_batch_size] + + +def get_model_params(config): + if config.architectures[0] in ( + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + ): + num_experts = config.n_routed_experts + hidden_size = config.hidden_size + elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[ + 0 + ] in ("MiniMaxM2ForCausalLM",): + num_experts = config.num_local_experts + hidden_size = config.hidden_size + else: + raise ValueError(f"Unsupported architecture: {config.architectures}") + return num_experts, hidden_size + + +def get_benchmark(model, max_batch_size, trust_remote_code): + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=get_batch_size_range(max_batch_size), + x_log=False, + line_arg="provider", + line_vals=[ + "torch", + "vllm", + ], + line_names=["PyTorch", "vLLM"], + styles=([("blue", "-"), ("red", "-")]), + ylabel="TFLOPs", + plot_name=f"{model} router gemm throughput", + args={}, + ) + ) + def benchmark(batch_size, provider): + config = get_config(model=model, trust_remote_code=trust_remote_code) + num_experts, hidden_size = get_model_params(config) + + is_hopper_or_blackwell = current_platform.is_device_capability( + 90 + ) or current_platform.is_device_capability_family(100) + allow_dsv3_router_gemm = ( + is_hopper_or_blackwell + and num_experts in DSV3_SUPPORTED_NUM_EXPERTS + and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES + ) + allow_gpt_oss_router_gemm = ( + is_hopper_or_blackwell + and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS + and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES + ) + is_fp32_router_model = ( + is_hopper_or_blackwell + and num_experts in FP32_SUPPORTED_NUM_EXPERTS + and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES + ) + allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS + + # Weight dtype: fp32 kernel requires fp32 weights; others use bf16. + weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16 + mat_a = torch.randn( + (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda" + ).contiguous() + mat_b = torch.randn( + (num_experts, hidden_size), dtype=weight_dtype, device="cuda" + ).contiguous() + bias = torch.randn( + num_experts, dtype=torch.bfloat16, device="cuda" + ).contiguous() + + has_bias = allow_gpt_oss_router_gemm + + quantiles = [0.5, 0.2, 0.8] + + if provider == "torch": + + def runner(): + if allow_fp32_router_gemm: + F.linear(mat_a.float(), mat_b) + elif has_bias: + F.linear(mat_a, mat_b, bias) + else: + F.linear(mat_a, mat_b) + elif provider == "vllm": + + def runner(): + if allow_dsv3_router_gemm: + ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16) + elif allow_fp32_router_gemm: + ops.fp32_router_gemm(mat_a, mat_b) + elif allow_gpt_oss_router_gemm: + ops.gpt_oss_router_gemm(mat_a, mat_b, bias) + elif is_fp32_router_model: + # batch_size > FP32_MAX_TOKENS: fall back to F.linear + F.linear(mat_a.float(), mat_b) + else: + F.linear(mat_a, mat_b) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + runner, quantiles=quantiles + ) + + def tflops(t_ms): + flops = 2 * batch_size * hidden_size * num_experts + return flops / (t_ms * 1e-3) / 1e12 + + return tflops(ms), tflops(max_ms), tflops(min_ms) + + return benchmark + + +if __name__ == "__main__": + parser = FlexibleArgumentParser() + parser.add_argument("--model", type=str, default="openai/gpt-oss-20b") + parser.add_argument("--max-batch-size", default=16, type=int) + parser.add_argument("--trust-remote-code", action="store_true") + args = parser.parse_args() + + # Get the benchmark function + benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code) + # Run performance benchmark + benchmark.run(print_data=True) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index ffab4015f49..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" @@ -396,6 +409,13 @@ set(VLLM_EXT_SRC "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp") +if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND + VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND)) + set(VLLM_EXT_SRC + "csrc/cpu/cpu_wna16.cpp" + ${VLLM_EXT_SRC}) +endif() + if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) set(VLLM_EXT_SRC "csrc/cpu/shm.cpp" @@ -403,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" @@ -411,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" @@ -423,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" @@ -439,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" @@ -512,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 f81882ccbc2..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) @@ -476,6 +484,16 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE) endfunction() + +function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}") + endif() + set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE) +endfunction() + # # Override the GPU architectures detected by cmake/torch and filter them by # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in 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_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index d6cae76c45c..06a38c780a2 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec { : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)( static_cast(ptr), VEC_ELEM_NUM)) {}; + explicit FP16Vec16(const c10::Half v) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)( + RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {}; + explicit FP16Vec16(const FP32Vec16& vec); void save(void* ptr) const { @@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec { reinterpret_cast(ptr), VEC_ELEM_NUM))) {}; explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {}; + explicit BF16Vec16(const c10::BFloat16 v) + : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)( + RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {}; explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { @@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec { } reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16); } + explicit BF16Vec16(const c10::BFloat16 v) + : reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast(v), + VEC_ELEM_NUM)) {} explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { float tmp[16]; @@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec { : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)( data.reg, data.reg)) {}; explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; + explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { + const uint64_t q_values = static_cast(value); + auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM); + auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM); + auto shifts = + RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM); + auto shifted = + RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM); + auto idx64 = + RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM); + auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM); + reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM); + } explicit FP32Vec16(const FP16Vec16& v); #ifdef __riscv_zvfbfmin @@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {}; #endif + // FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for + // load_b_pair_vec template to compile on all platforms. + explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {} + FP32Vec16 operator+(const FP32Vec16& b) const { return FP32Vec16( RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM)); @@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) { acc = acc.fma(a, b); } +template +static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) { + alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM]; + alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM]; + vec0.save(values0); + vec1.save(values1); + + auto* packed = reinterpret_cast(ptr); + for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) { + packed[i] = static_cast(values0[i]) | + (static_cast(values1[i]) << 16); + } +} + +static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1, + void* ptr) { + interleave_save_16b(vec0, vec1, ptr); +} + +static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1, + void* ptr) { + interleave_save_16b(vec0, vec1, ptr); +} + #ifdef __riscv_zvfbfmin template <> inline void storeFP32(float v, c10::BFloat16* ptr) { 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 35350cf247e..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, " @@ -518,7 +519,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op); // WNA16 -#if defined(__AVX512F__) +#if defined(__AVX512F__) || defined(__riscv_v) ops.def( "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, " "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt " diff --git a/csrc/libtorch_stable/activation_kernels.cu b/csrc/libtorch_stable/activation_kernels.cu index 28fdce5c305..cdab456348e 100644 --- a/csrc/libtorch_stable/activation_kernels.cu +++ b/csrc/libtorch_stable/activation_kernels.cu @@ -4,7 +4,7 @@ #include #include "../cuda_compat.h" -#include "../cuda_vec_utils.cuh" +#include "cuda_vec_utils.cuh" #include "dispatch_utils.h" #include "torch_utils.h" 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/attention/attention_kernels.cuh b/csrc/libtorch_stable/attention/attention_kernels.cuh similarity index 99% rename from csrc/attention/attention_kernels.cuh rename to csrc/libtorch_stable/attention/attention_kernels.cuh index 052ff168cec..c5f9a9876c3 100644 --- a/csrc/attention/attention_kernels.cuh +++ b/csrc/libtorch_stable/attention/attention_kernels.cuh @@ -17,21 +17,18 @@ * limitations under the License. */ -#include -#include -#include #include -#include "attention_dtypes.h" +#include "../../attention/attention_dtypes.h" #include "attention_utils.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #ifdef USE_ROCM #include - #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh" typedef __hip_bfloat16 __nv_bfloat16; #else - #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #define MAX(a, b) ((a) > (b) ? (a) : (b)) diff --git a/csrc/attention/attention_utils.cuh b/csrc/libtorch_stable/attention/attention_utils.cuh similarity index 95% rename from csrc/attention/attention_utils.cuh rename to csrc/libtorch_stable/attention/attention_utils.cuh index 826b0edffae..29783cb0e5f 100644 --- a/csrc/attention/attention_utils.cuh +++ b/csrc/libtorch_stable/attention/attention_utils.cuh @@ -18,8 +18,8 @@ */ #pragma once -#include "../cuda_compat.h" -#include "attention_dtypes.h" +#include "../../cuda_compat.h" +#include "../../attention/attention_dtypes.h" #include #include diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index 88e7723ad2f..b132e82e253 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -7,7 +7,7 @@ #include #include "../../attention/attention_dtypes.h" -#include "../../attention/attention_utils.cuh" +#include "attention_utils.cuh" #include "../../quantization/w8a8/fp8/common.cuh" namespace vllm { diff --git a/csrc/attention/paged_attention_v1.cu b/csrc/libtorch_stable/attention/paged_attention_v1.cu similarity index 80% rename from csrc/attention/paged_attention_v1.cu rename to csrc/libtorch_stable/attention/paged_attention_v1.cu index 307300e5566..8fa41791593 100644 --- a/csrc/attention/paged_attention_v1.cu +++ b/csrc/libtorch_stable/attention/paged_attention_v1.cu @@ -16,8 +16,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "../torch_utils.h" #include "attention_kernels.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #define MAX(a, b) ((a) > (b) ? (a) : (b)) #define MIN(a, b) ((a) < (b) ? (a) : (b)) @@ -44,13 +45,15 @@ template void paged_attention_v1_launcher( - torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int num_kv_heads, float scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int num_kv_heads, float scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int max_seq_len, + const std::optional& alibi_slopes, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, + const int tp_rank, const int blocksparse_local_blocks, + const int blocksparse_vert_stride, const int blocksparse_block_size, + const int blocksparse_head_sliding_step) { int num_seqs = query.size(0); int num_heads = query.size(1); int head_size = query.size(2); @@ -69,8 +72,8 @@ void paged_attention_v1_launcher( T* query_ptr = reinterpret_cast(query.data_ptr()); CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); + int* block_tables_ptr = block_tables.mutable_data_ptr(); + int* seq_lens_ptr = seq_lens.mutable_data_ptr(); const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); @@ -85,8 +88,9 @@ void paged_attention_v1_launcher( dim3 grid(num_heads, num_seqs, 1); dim3 block(NUM_THREADS); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + query.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); switch (head_size) { // NOTE(woosuk): To reduce the compilation time, we only compile for the // head sizes that we use in the model. However, we can easily extend this @@ -119,7 +123,7 @@ void paged_attention_v1_launcher( LAUNCH_PAGED_ATTENTION_V1(256); break; default: - TORCH_CHECK(false, "Unsupported head size: ", head_size); + STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); break; } } @@ -141,43 +145,43 @@ void paged_attention_v1_launcher( // NOTE(woosuk): To reduce the compilation time, we omitted block sizes // 1, 2, 4, 64, 128, 256. -#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ +#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ + switch (block_size) { \ + case 8: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ + break; \ + case 16: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ + break; \ + case 32: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ } void paged_attention_v1( - torch::Tensor& out, // [num_seqs, num_heads, head_size] - torch::Tensor& query, // [num_seqs, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] int64_t num_kv_heads, // [num_heads] double scale, - torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::Tensor& seq_lens, // [num_seqs] + torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::stable::Tensor& seq_lens, // [num_seqs] int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, const int64_t blocksparse_local_blocks, const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, const int64_t blocksparse_head_sliding_step) { const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, CALL_V1_LAUNCHER_BLOCK_SIZE) } diff --git a/csrc/attention/paged_attention_v2.cu b/csrc/libtorch_stable/attention/paged_attention_v2.cu similarity index 78% rename from csrc/attention/paged_attention_v2.cu rename to csrc/libtorch_stable/attention/paged_attention_v2.cu index eb9b4feb4a8..4e8e56ae05c 100644 --- a/csrc/attention/paged_attention_v2.cu +++ b/csrc/libtorch_stable/attention/paged_attention_v2.cu @@ -16,8 +16,9 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include "../torch_utils.h" #include "attention_kernels.cuh" -#include "../cuda_compat.h" +#include "../../cuda_compat.h" #define MAX(a, b) ((a) > (b) ? (a) : (b)) #define MIN(a, b) ((a) < (b) ? (a) : (b)) @@ -44,14 +45,16 @@ template void paged_attention_v2_launcher( - torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, - torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int num_kv_heads, float scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { + torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, + torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, + torch::stable::Tensor& query, torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, int num_kv_heads, float scale, + torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, + int max_seq_len, const std::optional& alibi_slopes, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, + const int tp_rank, const int blocksparse_local_blocks, + const int blocksparse_vert_stride, const int blocksparse_block_size, + const int blocksparse_head_sliding_step) { int num_seqs = query.size(0); int num_heads = query.size(1); int head_size = query.size(2); @@ -73,8 +76,8 @@ void paged_attention_v2_launcher( T* query_ptr = reinterpret_cast(query.data_ptr()); CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); + int* block_tables_ptr = block_tables.mutable_data_ptr(); + int* seq_lens_ptr = seq_lens.mutable_data_ptr(); const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); @@ -91,8 +94,9 @@ void paged_attention_v2_launcher( int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float); dim3 block(NUM_THREADS); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + query.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); switch (head_size) { // NOTE(woosuk): To reduce the compilation time, we only compile for the // head sizes that we use in the model. However, we can easily extend this @@ -125,7 +129,7 @@ void paged_attention_v2_launcher( LAUNCH_PAGED_ATTENTION_V2(256); break; default: - TORCH_CHECK(false, "Unsupported head size: ", head_size); + STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); break; } } @@ -148,46 +152,48 @@ void paged_attention_v2_launcher( // NOTE(woosuk): To reduce the compilation time, we omitted block sizes // 1, 2, 4, 64, 128, 256. -#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ +#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ + switch (block_size) { \ + case 8: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ + break; \ + case 16: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ + break; \ + case 32: \ + CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ } void paged_attention_v2( - torch::Tensor& out, // [num_seqs, num_heads, head_size] - torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions] - torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions] - torch::Tensor& + torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& + exp_sums, // [num_seqs, num_heads, max_num_partitions] + torch::stable::Tensor& + max_logits, // [num_seqs, num_heads, max_num_partitions] + torch::stable::Tensor& tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] - torch::Tensor& query, // [num_seqs, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] int64_t num_kv_heads, // [num_heads] double scale, - torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::Tensor& seq_lens, // [num_seqs] + torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::stable::Tensor& seq_lens, // [num_seqs] int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, const int64_t blocksparse_local_blocks, const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, const int64_t blocksparse_head_sliding_step) { const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, CALL_V2_LAUNCHER_BLOCK_SIZE) } diff --git a/csrc/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu similarity index 73% rename from csrc/cache_kernels.cu rename to csrc/libtorch_stable/cache_kernels.cu index 9130dd2ccae..eac93ac9a9f 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1,20 +1,16 @@ -#include -#include -#include -#include -#include - -#include "cuda_utils.h" -#include "cuda_compat.h" +#include "torch_utils.h" #include "dispatch_utils.h" -#include "libtorch_stable/quantization/vectorization_utils.cuh" +#include "../cuda_utils.h" +#include "../cuda_compat.h" + +#include "quantization/vectorization_utils.cuh" #include "concat_mla_q.cuh" #ifdef USE_ROCM - #include "quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #include @@ -34,40 +30,45 @@ constexpr float kFp8ScaleDivisor = 224.f; constexpr float kFp8ScaleDivisor = 448.f; #endif -void swap_blocks(torch::Tensor& src, torch::Tensor& dst, +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, int64_t block_size_in_bytes, - const torch::Tensor& block_mapping) { - torch::Device src_device = src.device(); - torch::Device dst_device = dst.device(); + const torch::stable::Tensor& block_mapping) { + torch::stable::Device src_device = src.device(); + torch::stable::Device dst_device = dst.device(); cudaMemcpyKind memcpy_type; if (src_device.is_cuda() && dst_device.is_cuda()) { - TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); memcpy_type = cudaMemcpyDeviceToDevice; } else if (src_device.is_cuda() && dst_device.is_cpu()) { memcpy_type = cudaMemcpyDeviceToHost; } else if (src_device.is_cpu() && dst_device.is_cuda()) { memcpy_type = cudaMemcpyHostToDevice; } else { - TORCH_CHECK(false, "Invalid device combination"); + STD_TORCH_CHECK(false, "Invalid device combination"); } // NOTE(youkaichao): keep in mind that `block_mapping` should be // a cpu tensor, otherwise every `item` call will require a gpu-cpu // synchronization. - TORCH_CHECK(block_mapping.device().is_cpu(), "block_mapping must be on CPU"); + STD_TORCH_CHECK(block_mapping.device().is_cpu(), + "block_mapping must be on CPU"); char* src_ptr = static_cast(src.data_ptr()); char* dst_ptr = static_cast(dst.data_ptr()); - const at::cuda::OptionalCUDAGuard device_guard( - src_device.is_cuda() ? src_device : dst_device); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + auto guard_device = src_device.is_cuda() ? src_device : dst_device; + const torch::stable::accelerator::DeviceGuard device_guard( + guard_device.index()); + const cudaStream_t stream = get_current_cuda_stream(); // NOTE(woosuk): This can be slow if the number of blocks is large. const int64_t num_blocks = block_mapping.size(0); + const int64_t* bm_ptr = block_mapping.const_data_ptr(); + const int64_t bm_stride0 = block_mapping.stride(0); + const int64_t bm_stride1 = block_mapping.stride(1); for (size_t i = 0; i < num_blocks; i++) { - int64_t src_block_number = block_mapping[i][0].item(); - int64_t dst_block_number = block_mapping[i][1].item(); + int64_t src_block_number = bm_ptr[i * bm_stride0]; + int64_t dst_block_number = bm_ptr[i * bm_stride0 + bm_stride1]; int64_t src_offset = src_block_number * block_size_in_bytes; int64_t dst_offset = dst_block_number * block_size_in_bytes; cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset, @@ -75,20 +76,23 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst, } } -void swap_blocks_batch(const torch::Tensor& src_ptrs, - const torch::Tensor& dst_ptrs, - const torch::Tensor& sizes, +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, bool is_src_access_order_any) { - TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); - TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); - TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); - TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64"); - TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64"); - TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64"); + STD_TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + STD_TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + STD_TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + STD_TORCH_CHECK(src_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "src_ptrs must be int64"); + STD_TORCH_CHECK(dst_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "dst_ptrs must be int64"); + STD_TORCH_CHECK(sizes.scalar_type() == torch::headeronly::ScalarType::Long, + "sizes must be int64"); const int64_t n = src_ptrs.size(0); - TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); - TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + STD_TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + STD_TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); if (n == 0) return; @@ -96,7 +100,7 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, int64_t* dst_data = dst_ptrs.mutable_data_ptr(); int64_t* size_data = sizes.mutable_data_ptr(); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a // single driver call, amortizing per-copy submission overhead. int64_t @@ -138,8 +142,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, reinterpret_cast(size_data), static_cast(n), &attr, &attrs_idx, 1, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); + STD_TORCH_CHECK(result == CUDA_SUCCESS, + "cuMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); return; } #elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000 @@ -155,8 +160,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, reinterpret_cast(dst_data), reinterpret_cast(src_data), reinterpret_cast(size_data), static_cast(n), &attr, &attrs_idx, 0, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == hipSuccess, "hipMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); + STD_TORCH_CHECK(result == hipSuccess, + "hipMemcpyBatchAsync failed at index ", fail_idx, + " with error ", result); return; } #endif @@ -675,28 +681,28 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel( // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - vllm::reshape_and_cache_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.data_ptr(), key_stride, value_stride, \ - num_heads, head_size, block_size, x, \ - reinterpret_cast(k_scale.data_ptr()), \ +#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), key_stride, value_stride, \ + num_heads, head_size, block_size, x, \ + reinterpret_cast(k_scale.data_ptr()), \ reinterpret_cast(v_scale.data_ptr())); void reshape_and_cache( - torch::Tensor& key, // [num_tokens, num_heads, head_size] - torch::Tensor& value, // [num_tokens, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::Tensor& + torch::stable::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] - torch::Tensor& slot_mapping, // [num_tokens] - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale) { + torch::stable::Tensor& slot_mapping, // [num_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); @@ -709,39 +715,41 @@ void reshape_and_cache( dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_div_x, 512)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE); } // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ - vllm::reshape_and_cache_flash_kernel \ - <<>>( \ - reinterpret_cast(key.data_ptr()), \ - reinterpret_cast(value.data_ptr()), \ - reinterpret_cast(key_cache.data_ptr()), \ - reinterpret_cast(value_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, page_stride, \ - head_stride, key_stride, value_stride, num_heads, head_size, \ - block_size, reinterpret_cast(k_scale.data_ptr()), \ - reinterpret_cast(v_scale.data_ptr()), \ +#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_flash_kernel \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, page_stride, \ + head_stride, key_stride, value_stride, num_heads, head_size, \ + block_size, reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr()), \ kv_scale_stride); void reshape_and_cache_flash( - torch::Tensor& key, // [num_tokens, num_heads, head_size] - torch::Tensor& value, // [num_tokens, num_heads, head_size] - torch::Tensor& key_cache, // [num_blocks, block_size, num_heads, head_size] - torch::Tensor& + torch::stable::Tensor& key, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& value, // [num_tokens, num_heads, head_size] + torch::stable::Tensor& + key_cache, // [num_blocks, block_size, num_heads, head_size] + torch::stable::Tensor& value_cache, // [num_blocks, block_size, num_heads, head_size] - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] const std::string& kv_cache_dtype, - torch::Tensor& k_scale, // [1] or [num_heads] - torch::Tensor& v_scale) { // [1] or [num_heads] + torch::stable::Tensor& k_scale, // [1] or [num_heads] + torch::stable::Tensor& v_scale) { // [1] or [num_heads] // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -756,23 +764,26 @@ void reshape_and_cache_flash( int num_heads = key.size(1); int head_size = key.size(2); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "nvfp4") { #if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) // NVFP4 dispatch is compiled separately for SM100+. extern void reshape_and_cache_nvfp4_dispatch( - torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache, - torch::Tensor & value_cache, torch::Tensor & slot_mapping, - torch::Tensor & k_scale, torch::Tensor & v_scale); + torch::stable::Tensor & key, torch::stable::Tensor & value, + torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache, + torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale, + torch::stable::Tensor & v_scale); reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, slot_mapping, k_scale, v_scale); return; #else - TORCH_CHECK(false, - "NVFP4 KV cache requires SM100+ (Blackwell). " - "Please rebuild vllm with a Blackwell-compatible CUDA target."); + STD_TORCH_CHECK( + false, + "NVFP4 KV cache requires SM100+ (Blackwell). " + "Please rebuild vllm with a Blackwell-compatible CUDA target."); #endif } @@ -784,53 +795,53 @@ void reshape_and_cache_flash( int64_t block_stride = key_cache.stride(0); int64_t page_stride = key_cache.stride(1); int64_t head_stride = key_cache.stride(2); - TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); + STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); - TORCH_CHECK(k_scale.sizes() == v_scale.sizes(), - "k_scale and v_scale must have the same shape"); - TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, - "k_scale and v_scale must be of shape [1] or [num_heads]"); + STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), + "k_scale and v_scale must have the same shape"); + STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, + "k_scale and v_scale must be of shape [1] or [num_heads]"); int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; dim3 grid(num_tokens); dim3 block(std::min(num_heads * head_size, 512)); - DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, CALL_RESHAPE_AND_CACHE_FLASH); } // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ +#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ + vllm::concat_and_cache_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_ds_mla_kernel \ - <<>>( \ - reinterpret_cast(kv_c.data_ptr()), \ - reinterpret_cast(k_pe.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, entry_stride, \ - kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ +#define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ + vllm::concat_and_cache_ds_mla_kernel \ + <<>>( \ + reinterpret_cast(kv_c.data_ptr()), \ + reinterpret_cast(k_pe.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, entry_stride, \ + kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); void concat_and_cache_mla( - torch::Tensor& kv_c, // [num_tokens, kv_lora_rank] - torch::Tensor& k_pe, // [num_tokens, pe_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + - // pe_dim)] - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, torch::Tensor& scale) { + torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] + torch::stable::Tensor& k_pe, // [num_tokens, pe_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + // + pe_dim)] + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -847,16 +858,17 @@ void concat_and_cache_mla( int block_size = kv_cache.size(1); if (kv_cache_dtype == "fp8_ds_mla") { - TORCH_CHECK(kv_lora_rank == 512, "kv_lora_rank must be 512 for fp8_ds_mla"); - TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); - TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.itemsize(), - "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); - TORCH_CHECK(kv_c.itemsize() == 2, - "kv_c.itemsize() must be 2 for fp8_ds_mla"); - TORCH_CHECK(k_pe.itemsize() == 2, - "k_pe.itemsize() must be 2 for fp8_ds_mla"); + STD_TORCH_CHECK(kv_lora_rank == 512, + "kv_lora_rank must be 512 for fp8_ds_mla"); + STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla"); + STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), + "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); + STD_TORCH_CHECK(kv_c.element_size() == 2, + "kv_c.element_size() must be 2 for fp8_ds_mla"); + STD_TORCH_CHECK(k_pe.element_size() == 2, + "k_pe.element_size() must be 2 for fp8_ds_mla"); } else { - TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); + STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim); } int kv_c_stride = kv_c.stride(0); @@ -864,8 +876,9 @@ void concat_and_cache_mla( int block_stride = kv_cache.stride(0); int entry_stride = kv_cache.stride(1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + kv_c.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "fp8_ds_mla") { dim3 grid(num_tokens); @@ -875,12 +888,12 @@ void concat_and_cache_mla( // The RoPE part (last 64 elements) is handled by another 1 warp (32 // threads). So in total, we use 3 warps (96 threads) per block. dim3 block(96); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_DS_MLA); } else { dim3 grid(num_tokens); dim3 block(std::min(kv_lora_rank, 512)); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_MLA); } } @@ -908,55 +921,62 @@ __global__ void convert_fp8_kernel(const Tin* __restrict__ src_cache, reinterpret_cast(dst_cache.data_ptr()), scale, block_stride); // Only for testing. -void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache, - const double scale, const std::string& kv_cache_dtype) { - torch::Device src_device = src_cache.device(); - torch::Device dst_device = dst_cache.device(); - TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") - TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") - TORCH_CHECK(src_device.index() == dst_device.index(), - "src and dst must be on the same GPU"); - at::cuda::OptionalCUDAGuard device_guard(src_device); +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype) { + torch::stable::Device src_device = src_cache.device(); + torch::stable::Device dst_device = dst_cache.device(); + STD_TORCH_CHECK(src_device.is_cuda(), "src must be on a GPU") + STD_TORCH_CHECK(dst_device.is_cuda(), "dst must be on a GPU") + STD_TORCH_CHECK(src_device.index() == dst_device.index(), + "src and dst must be on the same GPU"); + torch::stable::accelerator::DeviceGuard device_guard(src_device.index()); int64_t num_blocks = src_cache.size(0); int64_t block_stride = src_cache.stride(0); dim3 grid(num_blocks); dim3 block(std::min(block_stride, int64_t(512))); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); if (kv_cache_dtype == "auto") { - if (src_cache.dtype() == at::ScalarType::Float) { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kAuto); - } else if (src_cache.dtype() == at::ScalarType::Half) { + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (src_cache.dtype() == at::ScalarType::BFloat16) { + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::Float) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::Half) { + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kAuto); - } else if (dst_cache.dtype() == at::ScalarType::BFloat16) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kAuto); } } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - if (src_cache.dtype() == at::ScalarType::Float) { + if (src_cache.scalar_type() == torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(uint8_t, float, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.dtype() == at::ScalarType::Half) { + } else if (src_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint8_t, uint16_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (src_cache.dtype() == at::ScalarType::BFloat16) { + } else if (src_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(uint8_t, __nv_bfloat16, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::Float) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::Float) { CALL_CONVERT_FP8(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::Half) { + } else if (dst_cache.scalar_type() == torch::headeronly::ScalarType::Half) { CALL_CONVERT_FP8(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); - } else if (dst_cache.dtype() == at::ScalarType::BFloat16) { + } else if (dst_cache.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { CALL_CONVERT_FP8(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); } } else { - TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); + STD_TORCH_CHECK(false, "Unsupported data type: ", kv_cache_dtype); } } @@ -1053,8 +1073,9 @@ __global__ void gather_and_maybe_dequant_cache( <<>>( \ reinterpret_cast(src_cache.data_ptr()), \ reinterpret_cast(dst.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - token_to_seq.data_ptr(), num_tokens, block_size, \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), \ + token_to_seq.const_data_ptr(), num_tokens, block_size, \ block_table_stride, cache_block_stride, cache_entry_stride, \ dst_entry_stride, reinterpret_cast(scale.data_ptr()), \ seq_starts_ptr); @@ -1072,42 +1093,47 @@ __global__ void gather_and_maybe_dequant_cache( // - Optionally, seq_starts (if provided) offsets the starting block index by // (seq_starts[bid] / page_size) void gather_and_maybe_dequant_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] - torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] int64_t num_tokens, const std::string& kv_cache_dtype, - torch::Tensor const& scale, - std::optional seq_starts = std::nullopt) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); int32_t head_dim = dst.size(-1); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32, - "cu_seq_lens must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); if (seq_starts.has_value()) { - TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32, - "seq_starts must be int32"); + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); } - TORCH_CHECK( + STD_TORCH_CHECK( head_dim == 320 || head_dim == 576, "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 " "for better performance") - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); if (seq_starts.has_value()) { - TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); } int64_t block_table_stride = block_table.stride(0); @@ -1120,13 +1146,14 @@ void gather_and_maybe_dequant_cache( dim3 block(thread_block_size); const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().data_ptr() : nullptr; + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; if (head_dim == 576) { - DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, CALL_GATHER_CACHE_576); } else { - DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype, CALL_GATHER_CACHE_320); } } @@ -1267,13 +1294,14 @@ __global__ void cp_gather_cache( } // namespace vllm // Macro to dispatch the kernel based on the data type. -#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ - vllm::cp_gather_cache<<>>( \ - reinterpret_cast(src_cache.data_ptr()), \ - reinterpret_cast(dst.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - block_size, entry_size, block_table_stride, cache_block_stride, \ - cache_entry_stride, dst_entry_stride, seq_starts_ptr); +#define CALL_CP_GATHER_CACHE(CPY_DTYPE) \ + vllm::cp_gather_cache<<>>( \ + reinterpret_cast(src_cache.data_ptr()), \ + reinterpret_cast(dst.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), block_size, entry_size, \ + block_table_stride, cache_block_stride, cache_entry_stride, \ + dst_entry_stride, seq_starts_ptr); // Gather sequences from the cache into the destination tensor. // - cu_seq_lens contains the cumulative sequence lengths for each batch @@ -1281,36 +1309,41 @@ __global__ void cp_gather_cache( // - Optionally, seq_starts (if provided) offsets the starting slot index by // seq_starts[bid] void cp_gather_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] - torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& + src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] int64_t batch_size, - std::optional seq_starts = std::nullopt) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + std::optional seq_starts = std::nullopt) { + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); - int32_t entry_size = src_cache.flatten(2, -1).size(2); + int32_t entry_size = torch::stable::flatten(src_cache, 2, -1).size(2); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32, - "cu_seq_lens must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK( + cu_seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "cu_seq_lens must be int32"); if (seq_starts.has_value()) { - TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32, - "seq_starts must be int32"); + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); } - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), - "src_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == cu_seq_lens.device(), + "src_cache and cu_seq_lens must be on the same device"); if (seq_starts.has_value()) { - TORCH_CHECK(src_cache.device() == seq_starts.value().device(), - "src_cache and seq_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); } int64_t block_table_stride = block_table.stride(0); @@ -1323,12 +1356,13 @@ void cp_gather_cache( dim3 grid(batch_size, num_splits); dim3 block(1024); - TORCH_CHECK(src_cache.dtype() == dst.dtype(), - "src_cache and dst must have the same dtype"); + STD_TORCH_CHECK(src_cache.scalar_type() == dst.scalar_type(), + "src_cache and dst must have the same dtype"); const int dtype_bits = src_cache.element_size() * 8; const int32_t* seq_starts_ptr = - seq_starts.has_value() ? seq_starts.value().data_ptr() : nullptr; + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; if (dtype_bits == 32) { CALL_CP_GATHER_CACHE(uint32_t); @@ -1337,46 +1371,51 @@ void cp_gather_cache( } else if (dtype_bits == 8) { CALL_CP_GATHER_CACHE(uint8_t); } else { - TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); + STD_TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits); } } void cp_gather_and_upconvert_fp8_kv_cache( - torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] - torch::Tensor const& dst, // [TOT_TOKENS, 576] - torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& seq_lens, // [BATCH] - torch::Tensor const& workspace_starts, // [BATCH] + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& seq_lens, // [BATCH] + torch::stable::Tensor const& workspace_starts, // [BATCH] int64_t batch_size) { - at::cuda::OptionalCUDAGuard device_guard(src_cache.device()); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::stable::accelerator::DeviceGuard device_guard( + src_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); int32_t block_size = src_cache.size(1); int32_t head_dim = dst.size(1); - TORCH_CHECK(block_table.dtype() == torch::kInt32, - "block_table must be int32"); - TORCH_CHECK(seq_lens.dtype() == torch::kInt32, "seq_lens must be int32"); - TORCH_CHECK(workspace_starts.dtype() == torch::kInt32, - "workspace_starts must be int32"); + STD_TORCH_CHECK( + block_table.scalar_type() == torch::headeronly::ScalarType::Int, + "block_table must be int32"); + STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, + "seq_lens must be int32"); + STD_TORCH_CHECK( + workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, + "workspace_starts must be int32"); - TORCH_CHECK(src_cache.device() == dst.device(), - "src_cache and dst must be on the same device"); - TORCH_CHECK(src_cache.device() == block_table.device(), - "src_cache and block_table must be on the same device"); - TORCH_CHECK(src_cache.device() == seq_lens.device(), - "src_cache and seq_lens must be on the same device"); - TORCH_CHECK(src_cache.device() == workspace_starts.device(), - "src_cache and workspace_starts must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == dst.device(), + "src_cache and dst must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == block_table.device(), + "src_cache and block_table must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == seq_lens.device(), + "src_cache and seq_lens must be on the same device"); + STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(), + "src_cache and workspace_starts must be on the same device"); auto dtype = src_cache.scalar_type(); - TORCH_CHECK( - dtype == at::ScalarType::Byte || // uint8 - dtype == at::ScalarType::Float8_e4m3fn || // fp8 e4m3 - dtype == at::ScalarType::Float8_e5m2, // fp8 e5m2 + STD_TORCH_CHECK( + dtype == torch::headeronly::ScalarType::Byte || // uint8 + dtype == torch::headeronly::ScalarType::Float8_e4m3fn || // fp8 e4m3 + dtype == torch::headeronly::ScalarType::Float8_e5m2, // fp8 e5m2 "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ", - src_cache.dtype()); - TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16"); - TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); + src_cache.scalar_type()); + STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "dst must be bfloat16"); + STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); int64_t block_table_stride = block_table.stride(0); int64_t cache_block_stride = src_cache.stride(0); @@ -1384,8 +1423,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( int64_t dst_entry_stride = dst.stride(0); const uint8_t* src_ptr = nullptr; - if (dtype == at::ScalarType::Byte) { - src_ptr = src_cache.data_ptr(); + if (dtype == torch::headeronly::ScalarType::Byte) { + src_ptr = src_cache.const_data_ptr(); } else { // float8_e4m3fn or float8_e5m2 src_ptr = reinterpret_cast(src_cache.data_ptr()); @@ -1399,26 +1438,27 @@ void cp_gather_and_upconvert_fp8_kv_cache( vllm::cp_gather_and_upconvert_fp8_kv_cache<<>>( src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()), - block_table.data_ptr(), workspace_starts.data_ptr(), + block_table.const_data_ptr(), + workspace_starts.const_data_ptr(), static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, dst_entry_stride); } // Macro to dispatch the kernel based on the data type. -#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ - vllm::indexer_k_quant_and_cache_kernel \ - <<>>( \ - reinterpret_cast(k.data_ptr()), \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), head_dim, quant_block_size, \ +#define CALL_INDEXER_K_QUANT_AND_CACHE(KV_T, CACHE_T, KV_DTYPE) \ + vllm::indexer_k_quant_and_cache_kernel \ + <<>>( \ + reinterpret_cast(k.data_ptr()), \ + reinterpret_cast(kv_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ cache_block_size, cache_stride, use_ue8m0); void indexer_k_quant_and_cache( - torch::Tensor& k, // [num_tokens, head_dim] - torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& slot_mapping, // [num_tokens] - int64_t quant_block_size, // quantization block size + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size const std::string& scale_fmt) { int num_tokens = k.size(0); int head_dim = k.size(1); @@ -1426,65 +1466,70 @@ void indexer_k_quant_and_cache( int cache_stride = kv_cache.size(2); bool use_ue8m0 = scale_fmt == "ue8m0"; - TORCH_CHECK(k.device() == kv_cache.device(), - "k and kv_cache must be on the same device"); - TORCH_CHECK(k.device() == slot_mapping.device(), - "k and slot_mapping must be on the same device"); - TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); + STD_TORCH_CHECK(k.device() == kv_cache.device(), + "k and kv_cache must be on the same device"); + STD_TORCH_CHECK(k.device() == slot_mapping.device(), + "k and slot_mapping must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); constexpr int vec_size = 4; dim3 grid(num_tokens, (head_dim + quant_block_size * vec_size - 1) / (quant_block_size * vec_size)); dim3 block(32, vec_size); - const at::cuda::OptionalCUDAGuard device_guard(device_of(k)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + k.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); static const std::string kv_cache_dtype = "fp8_e4m3"; - DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype, CALL_INDEXER_K_QUANT_AND_CACHE); } // Macro to dispatch the kernel based on the data amount. -#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ - vllm::cp_gather_indexer_k_quant_cache_kernel \ - <<>>( \ - reinterpret_cast(kv_cache.data_ptr()), \ - reinterpret_cast(dst_k.data_ptr()), \ - reinterpret_cast(dst_scale.data_ptr()), \ - block_table.data_ptr(), cu_seq_lens.data_ptr(), \ - batch_size, dst_k.stride(0), dst_k.size(1), kv_cache.stride(0), \ - kv_cache.stride(1), kv_cache.size(1), block_table.size(1), \ - num_tokens, quant_block_size); +#define CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(BLOCK_Y_SIZE) \ + vllm::cp_gather_indexer_k_quant_cache_kernel \ + <<>>( \ + reinterpret_cast(kv_cache.data_ptr()), \ + reinterpret_cast(dst_k.data_ptr()), \ + reinterpret_cast(dst_scale.data_ptr()), \ + block_table.const_data_ptr(), \ + cu_seq_lens.const_data_ptr(), batch_size, dst_k.stride(0), \ + dst_k.size(1), kv_cache.stride(0), kv_cache.stride(1), \ + kv_cache.size(1), block_table.size(1), num_tokens, \ + quant_block_size); void cp_gather_indexer_k_quant_cache( - const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride] - torch::Tensor& dst_k, // [num_tokens, head_dim] - torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4] - const torch::Tensor& block_table, // [batch_size, num_blocks] - const torch::Tensor& cu_seq_lens // [batch_size + 1] + const torch::stable::Tensor& + kv_cache, // [num_blocks, block_size, cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& + dst_scale, // [num_tokens, head_dim / quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens // [batch_size + 1] ) { int batch_size = block_table.size(0); int num_tokens = dst_k.size(0); int head_dim = dst_k.size(1); int quant_block_size = head_dim * 4 / dst_scale.size(1); - TORCH_CHECK(kv_cache.device() == dst_k.device(), - "kv_cache and dst_k must be on the same device"); - TORCH_CHECK(kv_cache.device() == dst_scale.device(), - "kv_cache and dst_scale must be on the same device"); - TORCH_CHECK(kv_cache.device() == block_table.device(), - "kv_cache and block_table must be on the same device"); - TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), - "kv_cache and cu_seq_lens must be on the same device"); - TORCH_CHECK(head_dim % quant_block_size == 0, - "head_dim must be divisible by quant_block_size"); + STD_TORCH_CHECK(kv_cache.device() == dst_k.device(), + "kv_cache and dst_k must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == dst_scale.device(), + "kv_cache and dst_scale must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == block_table.device(), + "kv_cache and block_table must be on the same device"); + STD_TORCH_CHECK(kv_cache.device() == cu_seq_lens.device(), + "kv_cache and cu_seq_lens must be on the same device"); + STD_TORCH_CHECK(head_dim % quant_block_size == 0, + "head_dim must be divisible by quant_block_size"); constexpr int vec_size = 16; - const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_cache)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + kv_cache.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); if (num_tokens < 32) { CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(1); @@ -1503,27 +1548,30 @@ void cp_gather_indexer_k_quant_cache( // Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA. // Replaces torch.cat((ql_nope, q_pe), dim=-1). -void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] - torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] - torch::Tensor& q_out // [num_tokens, num_heads, nope_dim + - // rope_dim] +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out // [num_tokens, num_heads, nope_dim + + // rope_dim] ) { const int num_tokens = ql_nope.size(0); const int num_heads = ql_nope.size(1); const int nope_dim = ql_nope.size(2); const int rope_dim = q_pe.size(2); - TORCH_CHECK(nope_dim % 512 == 0, "nope_dim must be a multiple of 512, got ", - nope_dim); - TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); - TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); + STD_TORCH_CHECK(nope_dim % 512 == 0, + "nope_dim must be a multiple of 512, got ", nope_dim); + STD_TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim); + STD_TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim); - TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2"); - TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); - TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); - TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half || - ql_nope.scalar_type() == at::ScalarType::BFloat16, - "ql_nope must be float16 or bfloat16 dtype"); + STD_TORCH_CHECK(ql_nope.stride(2) == 1, + "ql_nope must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2"); + STD_TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2"); + STD_TORCH_CHECK( + ql_nope.scalar_type() == torch::headeronly::ScalarType::Half || + ql_nope.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "ql_nope must be float16 or bfloat16 dtype"); if (num_tokens == 0) return; @@ -1532,13 +1580,14 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block; const int block_size = warps_per_block * 32; - const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + ql_nope.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { + VLLM_STABLE_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] { vllm::ConcatMLAQKernel<<>>( - q_out.data_ptr(), ql_nope.data_ptr(), - q_pe.data_ptr(), num_tokens, num_heads, q_out.stride(0), + q_out.mutable_data_ptr(), ql_nope.const_data_ptr(), + q_pe.const_data_ptr(), num_tokens, num_heads, q_out.stride(0), q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0), q_pe.stride(1)); }); diff --git a/csrc/cache_kernels_fused.cu b/csrc/libtorch_stable/cache_kernels_fused.cu similarity index 53% rename from csrc/cache_kernels_fused.cu rename to csrc/libtorch_stable/cache_kernels_fused.cu index 8687ebe1f14..744e3a2dcdc 100644 --- a/csrc/cache_kernels_fused.cu +++ b/csrc/libtorch_stable/cache_kernels_fused.cu @@ -1,15 +1,13 @@ -#include -#include -#include - -#include "cuda_compat.h" +#include "torch_utils.h" #include "dispatch_utils.h" -#include "quantization/w8a8/fp8/common.cuh" +#include "../cuda_compat.h" + +#include "../quantization/w8a8/fp8/common.cuh" #ifdef USE_ROCM - #include "quantization/w8a8/fp8/amd/quant_utils.cuh" + #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" #else - #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh" + #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif #ifdef USE_ROCM @@ -164,43 +162,52 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( } // namespace vllm -#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ - do { \ - VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \ - using qk_t = scalar_t; \ - VLLM_DISPATCH_FLOATING_TYPES( \ - rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \ - [&] { \ - using cos_sin_t = scalar_t; \ - if (rope_is_neox) { \ - vllm::concat_and_cache_mla_rope_fused_kernel< \ - qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ - kv_c_stride, num_q_heads, \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } else { \ - vllm::concat_and_cache_mla_rope_fused_kernel< \ - qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ - <<>>( \ - positions.data_ptr(), q_pe.data_ptr(), \ - k_pe.data_ptr(), kv_c.data_ptr(), \ - rope_cos_sin_cache.data_ptr(), rot_dim, \ - q_pe_stride_token, q_pe_stride_head, k_pe_stride, \ - kv_c_stride, num_q_heads, \ - reinterpret_cast(kv_cache.data_ptr()), \ - slot_mapping.data_ptr(), block_stride, \ - entry_stride, kv_lora_rank, block_size, \ - kv_cache_quant_scale.data_ptr()); \ - } \ - }); \ - }); \ +#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ + do { \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + q_pe.scalar_type(), "qk_scalar_type", [&] { \ + using qk_t = scalar_t; \ + VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ + rope_cos_sin_cache.scalar_type(), \ + "rope_cos_sin_cache_scalar_type", [&] { \ + using cos_sin_t = scalar_t; \ + if (rope_is_neox) { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.const_data_ptr(), \ + q_pe.mutable_data_ptr(), \ + k_pe.mutable_data_ptr(), \ + kv_c.const_data_ptr(), \ + rope_cos_sin_cache.const_data_ptr(), \ + rot_dim, q_pe_stride_token, q_pe_stride_head, \ + k_pe_stride, kv_c_stride, num_q_heads, \ + reinterpret_cast( \ + kv_cache.mutable_data_ptr()), \ + slot_mapping.const_data_ptr(), \ + block_stride, entry_stride, kv_lora_rank, \ + block_size, \ + kv_cache_quant_scale.const_data_ptr()); \ + } else { \ + vllm::concat_and_cache_mla_rope_fused_kernel< \ + qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ + <<>>( \ + positions.const_data_ptr(), \ + q_pe.mutable_data_ptr(), \ + k_pe.mutable_data_ptr(), \ + kv_c.const_data_ptr(), \ + rope_cos_sin_cache.const_data_ptr(), \ + rot_dim, q_pe_stride_token, q_pe_stride_head, \ + k_pe_stride, kv_c_stride, num_q_heads, \ + reinterpret_cast( \ + kv_cache.mutable_data_ptr()), \ + slot_mapping.const_data_ptr(), \ + block_stride, entry_stride, kv_lora_rank, \ + block_size, \ + kv_cache_quant_scale.const_data_ptr()); \ + } \ + }); \ + }); \ } while (false) // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache. @@ -208,64 +215,69 @@ __global__ void concat_and_cache_mla_rope_fused_kernel( // Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and // concat_and_cache_mla. void concat_and_cache_mla_rope_fused( - torch::Tensor& positions, // [num_tokens] - torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim] - torch::Tensor& k_pe, // [num_tokens, rot_dim] - torch::Tensor& kv_c, // [num_tokens, kv_lora_rank] - torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] + torch::stable::Tensor& positions, // [num_tokens] + torch::stable::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim] + torch::stable::Tensor& k_pe, // [num_tokens, rot_dim] + torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] + torch::stable::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] bool rope_is_neox, - torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - torch::Tensor& + torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)] - const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) { + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale) { // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. - // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because - // both include padding. - // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0) - // since key includes padding for CUDA graphs, while slot_mapping does not. - // In this case, slot_mapping.size(0) represents the actual number of tokens + // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) + // because both include padding. + // In vLLM V1, however, key.size(0) can be larger than + // slot_mapping.size(0) since key includes padding for CUDA graphs, + // while slot_mapping does not. In this case, + // slot_mapping.size(0) represents the actual number of tokens // before padding. - // For compatibility with both cases, we use slot_mapping.size(0) as the - // number of tokens. - int num_tokens = slot_mapping.size(0); - int num_padded_tokens = q_pe.size(0); - TORCH_CHECK_GE(num_padded_tokens, num_tokens); + // For compatibility with both cases, we use slot_mapping.size(0) as + // the number of tokens. + const int64_t num_tokens = slot_mapping.size(0); + const int64_t num_padded_tokens = q_pe.size(0); + STD_TORCH_CHECK(num_padded_tokens >= num_tokens); const int num_q_heads = q_pe.size(1); const int rot_dim = q_pe.size(2); const int kv_lora_rank = kv_c.size(1); - TORCH_CHECK_EQ(positions.size(0), num_padded_tokens); - TORCH_CHECK_EQ(positions.dim(), 1); - TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long); + STD_TORCH_CHECK(positions.size(0) == num_padded_tokens); + STD_TORCH_CHECK(positions.dim() == 1); + STD_TORCH_CHECK(positions.scalar_type() == + torch::headeronly::ScalarType::Long); - TORCH_CHECK_EQ(q_pe.dim(), 3); - TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens); - TORCH_CHECK_EQ(q_pe.size(1), num_q_heads); - TORCH_CHECK_EQ(q_pe.size(2), rot_dim); + STD_TORCH_CHECK(q_pe.dim() == 3); + STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens); + STD_TORCH_CHECK(q_pe.size(1) == num_q_heads); + STD_TORCH_CHECK(q_pe.size(2) == rot_dim); - TORCH_CHECK_EQ(k_pe.dim(), 2); - TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens); - TORCH_CHECK_EQ(k_pe.size(1), rot_dim); - TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type()); + STD_TORCH_CHECK(k_pe.dim() == 2); + STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens); + STD_TORCH_CHECK(k_pe.size(1) == rot_dim); + STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_c.dim(), 2); - TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens); - TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank); - TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type()); - TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype()); + STD_TORCH_CHECK(kv_c.dim() == 2); + STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens); + STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank); + STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim); + STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim); + STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type()); - TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens); - TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long); + STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens); + STD_TORCH_CHECK(slot_mapping.scalar_type() == + torch::headeronly::ScalarType::Long); - TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim); - TORCH_CHECK_EQ(kv_cache.dim(), 3); + STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim); + STD_TORCH_CHECK(kv_cache.dim() == 3); - TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1); - TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float); + STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1); + STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() == + torch::headeronly::ScalarType::Float); int64_t q_pe_stride_token = q_pe.stride(0); int64_t q_pe_stride_head = q_pe.stride(1); @@ -286,9 +298,10 @@ void concat_and_cache_mla_rope_fused( dim3 grid(num_tokens, 1, 1); dim3 block(thread_block_size, 1, 1); - const at::cuda::OptionalCUDAGuard device_guard(device_of(positions)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + positions.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype, + DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED); } diff --git a/csrc/concat_mla_q.cuh b/csrc/libtorch_stable/concat_mla_q.cuh similarity index 100% rename from csrc/concat_mla_q.cuh rename to csrc/libtorch_stable/concat_mla_q.cuh diff --git a/csrc/cuda_vec_utils.cuh b/csrc/libtorch_stable/cuda_vec_utils.cuh similarity index 100% rename from csrc/cuda_vec_utils.cuh rename to csrc/libtorch_stable/cuda_vec_utils.cuh 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/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu new file mode 100644 index 00000000000..04397e0893c --- /dev/null +++ b/csrc/libtorch_stable/fp32_router_gemm.cu @@ -0,0 +1,223 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32. +// Supports bf16 or fp32 activation; weight is always fp32. +// Adapted from dsv3_router_gemm_float_out.cu. + +#include +#include + +// --------------------------------------------------------------------------- +// Load helpers +// --------------------------------------------------------------------------- + +// Load VPT fp32 values from the weight matrix (always fp32). +// VPT=4 when activation is fp32 (one float4 load) +// VPT=8 when activation is bf16 (two float4 loads) +template +__device__ __forceinline__ void load_weight(float const* ptr, float* dst); + +template <> +__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) { + float4 v = *reinterpret_cast(ptr); + dst[0] = v.x; + dst[1] = v.y; + dst[2] = v.z; + dst[3] = v.w; +} + +template <> +__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) { + float4 v0 = *reinterpret_cast(ptr); + float4 v1 = *reinterpret_cast(ptr + 4); + dst[0] = v0.x; + dst[1] = v0.y; + dst[2] = v0.z; + dst[3] = v0.w; + dst[4] = v1.x; + dst[5] = v1.y; + dst[6] = v1.z; + dst[7] = v1.w; +} + +// Load VPT activation values and convert to fp32. +template +__device__ __forceinline__ void load_activation(T const* ptr, float* dst); + +// fp32 activation: one float4 load, no conversion needed. +template <> +__device__ __forceinline__ void load_activation(float const* ptr, + float* dst) { + float4 v = *reinterpret_cast(ptr); + dst[0] = v.x; + dst[1] = v.y; + dst[2] = v.z; + dst[3] = v.w; +} + +// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion. +template <> +__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>( + __nv_bfloat16 const* ptr, float* dst) { + uint4 v = *reinterpret_cast(ptr); + __nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v); +#pragma unroll + for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]); +} + +// --------------------------------------------------------------------------- +// Kernel +// --------------------------------------------------------------------------- + +// InputT : type of activation (float or __nv_bfloat16) +// Weight is always fp32; output is always fp32. +// VPT = 16 / sizeof(InputT): 4 for fp32, 8 for bf16 +template +__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel( + float* out, InputT const* mat_a, float const* mat_b) { + constexpr int VPT = 16 / sizeof(InputT); + constexpr int k_elems_per_k_iteration = VPT * kBlockSize; + constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; + constexpr int kWarpSize = 32; + constexpr int kNumWarps = kBlockSize / kWarpSize; + + int const n_idx = blockIdx.x; + int const tid = threadIdx.x; + int const warpId = tid / kWarpSize; + int const laneId = tid % kWarpSize; + + float acc[kNumTokens] = {}; + __shared__ float sm_reduction[kNumTokens][kNumWarps]; + + float const* b_col = mat_b + n_idx * kHiddenDim; + + int k_bases[k_iterations]; +#pragma unroll + for (int ki = 0; ki < k_iterations; ki++) { + k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;"); +#endif + + for (int ki = 0; ki < k_iterations; ki++) { + int const k_base = k_bases[ki]; + + float b_float[VPT]; + load_weight(b_col + k_base, b_float); + +#pragma unroll + for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { + float a_float[VPT]; + load_activation(mat_a + m_idx * kHiddenDim + k_base, + a_float); +#pragma unroll + for (int k = 0; k < VPT; k++) { + acc[m_idx] += a_float[k] * b_float[k]; + } + } + } + + // Warp-level butterfly reduction +#pragma unroll + for (int m = 0; m < kNumTokens; m++) { + float sum = acc[m]; + sum += __shfl_xor_sync(0xffffffff, sum, 16); + sum += __shfl_xor_sync(0xffffffff, sum, 8); + sum += __shfl_xor_sync(0xffffffff, sum, 4); + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 1); + if (laneId == 0) sm_reduction[m][warpId] = sum; + } + + __syncthreads(); + + if (tid == 0) { +#pragma unroll + for (int m = 0; m < kNumTokens; m++) { + float final_sum = 0.0f; +#pragma unroll + for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w]; + out[m * kNumExperts + n_idx] = final_sum; + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +// --------------------------------------------------------------------------- +// Launcher +// --------------------------------------------------------------------------- + +template +void invokeFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + constexpr int kBlockSize = 128; + cudaLaunchConfig_t config; + config.gridDim = kNumExperts; + config.blockDim = kBlockSize; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.numAttrs = 1; + config.attrs = attrs; + cudaLaunchKernelEx(&config, + fp32_router_gemm_kernel, + output, mat_a, mat_b); +} + +// --------------------------------------------------------------------------- +// Explicit instantiations: M=1..32, E=256, H=3072, for both input types +// --------------------------------------------------------------------------- + +#define INSTANTIATE(T, M) \ + template void invokeFp32RouterGemm( \ + float*, T const*, float const*, cudaStream_t); + +#define INSTANTIATE_ALL(T) \ + INSTANTIATE(T, 1) \ + INSTANTIATE(T, 2) \ + INSTANTIATE(T, 3) \ + INSTANTIATE(T, 4) \ + INSTANTIATE(T, 5) \ + INSTANTIATE(T, 6) \ + INSTANTIATE(T, 7) \ + INSTANTIATE(T, 8) \ + INSTANTIATE(T, 9) \ + INSTANTIATE(T, 10) \ + INSTANTIATE(T, 11) \ + INSTANTIATE(T, 12) \ + INSTANTIATE(T, 13) \ + INSTANTIATE(T, 14) \ + INSTANTIATE(T, 15) \ + INSTANTIATE(T, 16) \ + INSTANTIATE(T, 17) \ + INSTANTIATE(T, 18) \ + INSTANTIATE(T, 19) \ + INSTANTIATE(T, 20) \ + INSTANTIATE(T, 21) \ + INSTANTIATE(T, 22) \ + INSTANTIATE(T, 23) \ + INSTANTIATE(T, 24) \ + INSTANTIATE(T, 25) \ + INSTANTIATE(T, 26) \ + INSTANTIATE(T, 27) \ + INSTANTIATE(T, 28) \ + INSTANTIATE(T, 29) \ + INSTANTIATE(T, 30) \ + INSTANTIATE(T, 31) \ + INSTANTIATE(T, 32) + +INSTANTIATE_ALL(float) +INSTANTIATE_ALL(__nv_bfloat16) + +#undef INSTANTIATE_ALL +#undef INSTANTIATE diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu new file mode 100644 index 00000000000..4baa740de93 --- /dev/null +++ b/csrc/libtorch_stable/fp32_router_gemm_entry.cu @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include + +#include "core/registration.h" +#include "libtorch_stable/torch_utils.h" + +#include +#include + +#include + +namespace { + +inline int getSMVersion() { + auto* props = get_device_prop(); + return props->major * 10 + props->minor; +} + +} // namespace + +static constexpr int FP32_NUM_EXPERTS = 256; +static constexpr int FP32_HIDDEN_DIM = 3072; +static constexpr int FP32_MAX_TOKENS = 32; + +// Forward declarations — 4 template params must match fp32_router_gemm.cu +template +void invokeFp32RouterGemm(float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream); + +// LoopUnroller templated on InputT +template +struct Fp32LoopUnroller { + static void unroll(int num_tokens, float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + if (num_tokens == kBegin) { + invokeFp32RouterGemm( + output, mat_a, mat_b, stream); + } else { + Fp32LoopUnroller::unroll(num_tokens, output, + mat_a, mat_b, stream); + } + } +}; + +template +struct Fp32LoopUnroller { + static void unroll(int num_tokens, float* output, InputT const* mat_a, + float const* mat_b, cudaStream_t stream) { + if (num_tokens == kEnd) { + invokeFp32RouterGemm( + output, mat_a, mat_b, stream); + } else { + throw std::invalid_argument( + "fp32_router_gemm: num_tokens must be in [1, 32]"); + } + } +}; + +void fp32_router_gemm( + torch::stable::Tensor& output, // [num_tokens, num_experts] + torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim] + torch::stable::Tensor const& mat_b // [num_experts, hidden_dim] +) { + STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); + STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(), + "fp32_router_gemm: all tensors must be CUDA tensors"); + STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() && + output.get_device_index() == mat_b.get_device_index(), + "fp32_router_gemm: all tensors must be on the same device"); + STD_TORCH_CHECK( + output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(), + "fp32_router_gemm: all tensors must be contiguous"); + + const int num_tokens = mat_a.size(0); + const int num_experts = mat_b.size(0); + const int hidden_dim = mat_a.size(1); + + STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts, + "fp32_router_gemm: output must have shape [num_tokens, " + "num_experts]"); + STD_TORCH_CHECK( + mat_a.size(1) == mat_b.size(1), + "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim"); + STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM, + "fp32_router_gemm: expected hidden_dim=3072"); + STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS, + "fp32_router_gemm: expected num_experts=256"); + STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, + "fp32_router_gemm: num_tokens must be in [0, 32]"); + STD_TORCH_CHECK( + mat_a.scalar_type() == torch::headeronly::ScalarType::Float || + mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "fp32_router_gemm: mat_a must be float32 or bfloat16"); + STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float, + "fp32_router_gemm: mat_b (weight) must be float32"); + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float, + "fp32_router_gemm: output must be float32"); + + if (num_tokens == 0) { + return; + } + + STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+"); + + auto stream = get_current_cuda_stream(mat_a.get_device_index()); + float* out_ptr = reinterpret_cast(output.mutable_data_ptr()); + float const* mat_b_ptr = reinterpret_cast(mat_b.data_ptr()); + + if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + auto const* mat_a_ptr = + reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()); + Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll( + num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream); + } else { + auto const* mat_a_ptr = reinterpret_cast(mat_a.data_ptr()); + Fp32LoopUnroller::unroll( + num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream); + } +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm)); +} 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/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index fb714b1b1e0..37df6be329f 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -78,8 +78,7 @@ __global__ void rms_norm_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - float w = static_cast(src2.val[j]); - dst.val[j] = static_cast(x * s_variance * w); + dst.val[j] = static_cast(x * s_variance) * src2.val[j]; } v_out[i] = dst; } @@ -143,8 +142,7 @@ fused_add_rms_norm_kernel( #pragma unroll for (int j = 0; j < width; ++j) { float x = Converter::convert(res.data[j]); - float wf = Converter::convert(w.data[j]); - out.data[j] = Converter::convert(x * s_variance * wf); + out.data[j] = Converter::convert(x * s_variance) * w.data[j]; } input_v[strided_id] = out; } @@ -183,8 +181,8 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - float w = (float)weight[idx]; - input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); + input[blockIdx.x * input_stride + idx] = + (scalar_t)(x * s_variance) * weight[idx]; } } diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index 26ffa76d6e1..32f3495f4e9 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -66,13 +66,8 @@ __global__ void rms_norm_static_fp8_quant_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - float w = static_cast(src2.val[j]); - // Round normalized result through scalar_t to match the precision of the - // unfused composite (rms_norm writes scalar_t, then - // static_scaled_fp8_quant re-loads it as float before FP8 conversion). - // Without this round, the fused path is strictly more accurate and - // disagrees with the composite at exact E4M3 quantization tie boundaries. - scalar_t out_norm = static_cast(x * s_variance * w); + // Multiply in weight's native dtype to match rms_norm_kernel. + scalar_t out_norm = static_cast(x * s_variance) * src2.val[j]; out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] = scaled_fp8_conversion(static_cast(out_norm), scale_inv); @@ -142,12 +137,8 @@ fused_add_rms_norm_static_fp8_quant_kernel( #pragma unroll for (int i = 0; i < width; ++i) { float x = Converter::convert(res.data[i]); - float wf = Converter::convert(w.data[i]); - // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t - // to match the unfused composite path at FP8 boundaries. We use the - // backend's hip_type for the intermediate since c10::Half/BFloat16 has - // ambiguous conversions on CUDA and no implicit conversion on ROCm. - HipT out_norm_h = Converter::convert(x * s_variance * wf); + // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. + HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i]; out[id * width + i] = scaled_fp8_conversion( Converter::convert(out_norm_h), scale_inv); } @@ -192,10 +183,8 @@ fused_add_rms_norm_static_fp8_quant_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - float w = (float)weight[idx]; - // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t - // to match the unfused composite path at FP8 boundaries. - scalar_t out_norm = static_cast(x * s_variance * w); + // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. + scalar_t out_norm = static_cast(x * s_variance) * weight[idx]; out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion( static_cast(out_norm), scale_inv); } diff --git a/csrc/nvfp4_kv_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu similarity index 81% rename from csrc/nvfp4_kv_cache_kernels.cu rename to csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu index d6aa715c203..693d55f231e 100644 --- a/csrc/nvfp4_kv_cache_kernels.cu +++ b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu @@ -17,11 +17,8 @@ #define NVFP4_ENABLE_ELTS16 1 #include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh" -#include -#include -#include - -#include "dispatch_utils.h" +#include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/torch_utils.h" namespace vllm { @@ -184,12 +181,13 @@ __global__ void reshape_and_cache_nvfp4_kernel( // Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1]. // Each KV side contains both data and scale: // page = [K_data | K_scale | V_data | V_scale] -void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, - torch::Tensor& key_cache, - torch::Tensor& value_cache, - torch::Tensor& slot_mapping, - torch::Tensor& k_scale, - torch::Tensor& v_scale) { +void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key, + torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); @@ -200,17 +198,18 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, // key_cache is kv_cache[:, 0] with shape // [num_blocks, block_size, num_heads, full_dim] in logical order. // Strides encode the physical layout (HND or NHD). - TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); - TORCH_CHECK(key_cache.size(3) == full_dim, - "key_cache last dim must be data_dim + scale_dim, got ", - key_cache.size(3), " expected ", full_dim); + STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D"); + STD_TORCH_CHECK(key_cache.size(3) == full_dim, + "key_cache last dim must be data_dim + scale_dim, got ", + key_cache.size(3), " expected ", full_dim); int block_size = key_cache.size(1); - TORCH_CHECK(head_size % 16 == 0, - "head_size must be divisible by 16 for NVFP4 KV cache"); - TORCH_CHECK(block_size % 4 == 0, - "block_size must be divisible by 4 for NVFP4 KV cache swizzle"); + STD_TORCH_CHECK(head_size % 16 == 0, + "head_size must be divisible by 16 for NVFP4 KV cache"); + STD_TORCH_CHECK(block_size % 4 == 0, + "block_size must be divisible by 4 for NVFP4 KV cache " + "swizzle"); // Detect physical layout from strides (based on full_dim). // HND: head stride > block_offset stride. @@ -230,8 +229,9 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, // Scale follows data within each KV side. int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim; - uint8_t* key_scale_ptr = key_cache.data_ptr() + data_per_kv; - uint8_t* value_scale_ptr = value_cache.data_ptr() + data_per_kv; + uint8_t* key_scale_ptr = key_cache.mutable_data_ptr() + data_per_kv; + uint8_t* value_scale_ptr = + value_cache.mutable_data_ptr() + data_per_kv; // Scale strides: same page stride, inner strides from layout. int64_t scale_block_stride = data_block_stride; @@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, scale_block_offset_stride = (int64_t)num_heads * scale_dim; } - const float* k_scale_ptr = k_scale.data_ptr(); - const float* v_scale_ptr = v_scale.data_ptr(); + const float* k_scale_ptr = k_scale.const_data_ptr(); + const float* v_scale_ptr = v_scale.const_data_ptr(); int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE; int total_groups = num_heads * groups_per_head; @@ -256,20 +256,22 @@ void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value, dim3 grid(num_tokens); dim3 block(num_threads); - const at::cuda::OptionalCUDAGuard device_guard(device_of(key)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + key.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); - AT_DISPATCH_REDUCED_FLOATING_TYPES( + VLLM_STABLE_DISPATCH_HALF_TYPES( key.scalar_type(), "reshape_and_cache_nvfp4", [&] { vllm::reshape_and_cache_nvfp4_kernel <<>>( - key.data_ptr(), value.data_ptr(), - key_cache.data_ptr(), value_cache.data_ptr(), - key_scale_ptr, value_scale_ptr, - slot_mapping.data_ptr(), k_scale_ptr, v_scale_ptr, - key.stride(0), value.stride(0), num_heads, head_size, - block_size, data_block_stride, data_head_stride, - data_block_offset_stride, scale_block_stride, scale_head_stride, - scale_block_offset_stride); + key.const_data_ptr(), + value.const_data_ptr(), + key_cache.mutable_data_ptr(), + value_cache.mutable_data_ptr(), key_scale_ptr, + value_scale_ptr, slot_mapping.const_data_ptr(), + k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), + num_heads, head_size, block_size, data_block_stride, + data_head_stride, data_block_offset_stride, scale_block_stride, + scale_head_stride, scale_block_offset_stride); }); } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 90082ac3d29..0363ec7cdfc 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); @@ -355,3 +355,132 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X, int64_t tokens); int64_t ggml_moe_get_block_size(int64_t type); + +void paged_attention_v1( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +void paged_attention_v2( + torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, + torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, + torch::stable::Tensor& query, torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale, + torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, + int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +// Cache ops (shared CUDA/ROCm) +void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, + int64_t block_size_in_bytes, + const torch::stable::Tensor& block_mapping); + +// Batch swap: submit all block copies in a single driver call. +void swap_blocks_batch(const torch::stable::Tensor& src_ptrs, + const torch::stable::Tensor& dst_ptrs, + const torch::stable::Tensor& sizes, + bool is_src_access_order_any); + +void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, + torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale); + +void reshape_and_cache_flash( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, + torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale); + +void concat_and_cache_mla(torch::stable::Tensor& kv_c, + torch::stable::Tensor& k_pe, + torch::stable::Tensor& kv_cache, + torch::stable::Tensor& slot_mapping, + const std::string& kv_cache_dtype, + torch::stable::Tensor& scale); + +// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla +void concat_and_cache_mla_rope_fused( + torch::stable::Tensor& positions, torch::stable::Tensor& q_pe, + torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c, + torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache, + const std::string& kv_cache_dtype, + torch::stable::Tensor& kv_cache_quant_scale); + +// Just for unittest +void convert_fp8(torch::stable::Tensor& dst_cache, + torch::stable::Tensor& src_cache, const double scale, + const std::string& kv_cache_dtype); + +void gather_and_maybe_dequant_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS] + int64_t num_tokens, const std::string& kv_cache_dtype, + torch::stable::Tensor const& scale, + std::optional seq_starts = std::nullopt); + +// TODO(hc): cp_gather_cache need support scaled kvcahe in the future. +void cp_gather_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // ENTRIES...] + torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& cu_seq_lens, // [BATCH+1] + int64_t batch_size, + std::optional seq_starts = std::nullopt); + +// Gather and upconvert FP8 KV cache to BF16 workspace +void cp_gather_and_upconvert_fp8_kv_cache( + torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, + // 656] + torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] + torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] + torch::stable::Tensor const& seq_lens, // [BATCH] + torch::stable::Tensor const& workspace_starts, // [BATCH] + int64_t batch_size); + +// Indexer K quantization and cache function +void indexer_k_quant_and_cache( + torch::stable::Tensor& k, // [num_tokens, head_dim] + torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& slot_mapping, // [num_tokens] + int64_t quant_block_size, // quantization block size + const std::string& scale_fmt); + +// Concatenate query nope and rope for MLA/DSA attention +void concat_mla_q( + torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim] + torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim] + torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim + + // rope_dim] + +// Extract function to gather quantized K cache +void cp_gather_indexer_k_quant_cache( + const torch::stable::Tensor& kv_cache, // [num_blocks, block_size, + // cache_stride] + torch::stable::Tensor& dst_k, // [num_tokens, head_dim] + torch::stable::Tensor& dst_scale, // [num_tokens, head_dim / + // quant_block_size * 4] + const torch::stable::Tensor& block_table, // [batch_size, num_blocks] + const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1] 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 49f2944f3fd..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 @@ -17,7 +17,7 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include #include @@ -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 78e4eda0c01..062f6018653 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -27,14 +27,14 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include "cuda_utils.h" #include "nvfp4_utils.cuh" 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 744ae4f7311..92b139b7e40 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu @@ -17,7 +17,7 @@ #include #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #include #include @@ -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 e1d101f9d86..f7c965dbc1b 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu @@ -23,10 +23,10 @@ #include "libtorch_stable/torch_utils.h" #include "libtorch_stable/dispatch_utils.h" -#include "cuda_vec_utils.cuh" +#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/libtorch_stable/quantization/fp4/nvfp4_utils.cuh b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh index 590e4c06b62..0c04f010888 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh @@ -20,7 +20,7 @@ #include #include -#include "cuda_vec_utils.cuh" +#include "../../cuda_vec_utils.cuh" #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \ CUDA_VERSION >= 12090 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 37a612b4394..316a7d37522 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; } @@ -237,10 +252,12 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, VLLM_STABLE_DISPATCH_FLOATING_TYPES( input.scalar_type(), "per_token_group_quant_8bit", ([&] { - if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) { - LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3); - } else if (dst_type == torch::headeronly::ScalarType::Char) { + if (dst_type == torch::headeronly::ScalarType::Char) { LAUNCH_KERNEL(scalar_t, int8_t); + } else { + VLLM_STABLE_DISPATCH_FP8_TYPES( + dst_type, "per_token_group_quant_8bit_fp8", + ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); })); } })); @@ -317,10 +334,18 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the // 8 lanes within the warp that share a group. +#ifdef USE_ROCM + const int lane_in_wave = threadIdx.x % warpSize; + const unsigned long long mask = 0xFFull << (lane_in_wave & ~7); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8)); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8)); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8)); +#else unsigned mask = 0xffu << (threadIdx.x & 24u); local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4)); local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2)); local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1)); +#endif float y_s = local_absmax / max_8bit; y_s = fmaxf(y_s, 1e-10f); @@ -503,15 +528,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 23934193230..98cd31df13b 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( @@ -247,6 +247,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { ops.def( "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). + // conditionally compiled so impl registration is in source file + ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); + // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel ops.def( "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, " @@ -474,14 +478,36 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? initial_state_idx," "Tensor? cu_chunk_seqlen," "Tensor? last_chunk_indices) -> ()"); + + // Attention ops + // Compute the attention between an input query and the cached + // keys/values using PagedAttention. + ops.def( + "paged_attention_v1(" + " Tensor! out, Tensor query, Tensor key_cache," + " Tensor value_cache, int num_kv_heads, float scale," + " Tensor block_tables, Tensor seq_lens, int block_size," + " int max_seq_len, Tensor? alibi_slopes," + " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," + " int tp_rank, int blocksparse_local_blocks," + " int blocksparse_vert_stride, int blocksparse_block_size," + " int blocksparse_head_sliding_step) -> ()"); + + // PagedAttention V2. + ops.def( + "paged_attention_v2(" + " Tensor! out, Tensor! exp_sums, Tensor! max_logits," + " Tensor! tmp_out, Tensor query, Tensor key_cache," + " Tensor value_cache, int num_kv_heads, float scale," + " Tensor block_tables, Tensor seq_lens, int block_size," + " int max_seq_len, Tensor? alibi_slopes," + " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," + " int tp_rank, int blocksparse_local_blocks," + " int blocksparse_vert_stride, int blocksparse_block_size," + " int blocksparse_head_sliding_step) -> ()"); } 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", @@ -489,6 +515,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)); @@ -581,6 +612,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); + + ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1)); + ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2)); } // These capability-check functions take only primitive args (no tensors), so @@ -603,4 +637,115 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size)); } +// Cache ops +STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { + // Swap in (out) the cache blocks from src to dst. + ops.def( + "swap_blocks(Tensor src, Tensor! dst," + " int block_size_in_bytes, Tensor block_mapping) -> ()"); + + // Batch swap: submit all block copies in a single driver call. + ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes," + " bool is_src_access_order_any=False) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache(Tensor key, Tensor value," + " Tensor! key_cache, Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Reshape the key and value tensors and cache them. + ops.def( + "reshape_and_cache_flash(Tensor key, Tensor value," + " Tensor! key_cache," + " Tensor! value_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor k_scale, Tensor v_scale) -> ()"); + + // Concat kv_c and k_pe and cache them. + ops.def( + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," + " Tensor! kv_cache," + " Tensor slot_mapping," + " str kv_cache_dtype," + " Tensor scale) -> ()"); + + // Rotate Q and K, then write to kv cache for MLA + ops.def( + "concat_and_cache_mla_rope_fused(" + " Tensor positions," + " Tensor! q_pe," + " Tensor! k_pe," + " Tensor kv_c," + " Tensor cos_sin_cache," + " bool is_neox," + " Tensor slot_mapping," + " Tensor! kv_cache," + " str kv_cache_dtype," + " Tensor kv_cache_scale) -> ()"); + + // Convert the key and value cache to fp8 data type. + ops.def( + "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " + "str kv_cache_dtype) -> ()"); + + // Gather cache blocks from src_cache to dst, dequantizing from + // src_cache's dtype to dst's dtype if necessary. + ops.def( + "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " + " Tensor block_table, Tensor cu_seq_lens, " + " Tensor token_to_seq, " + " int num_tokens, " + " str kv_cache_dtype, " + " Tensor scale, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " + "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); + + ops.def( + "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " + "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int " + "batch_size) -> ()"); + + ops.def( + "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " + "slot_mapping, " + "int quant_block_size, str kv_cache_dtype) -> ()"); + + ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); + + ops.def( + "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " + "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) { + ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch)); +} + +STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) { + ops.impl("swap_blocks", TORCH_BOX(&swap_blocks)); + ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache)); + ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash)); + ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla)); + ops.impl("concat_and_cache_mla_rope_fused", + TORCH_BOX(&concat_and_cache_mla_rope_fused)); + ops.impl("convert_fp8", TORCH_BOX(&convert_fp8)); + ops.impl("gather_and_maybe_dequant_cache", + TORCH_BOX(&gather_and_maybe_dequant_cache)); + ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache)); + ops.impl("cp_gather_and_upconvert_fp8_kv_cache", + TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache)); + ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache)); + ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q)); + ops.impl("cp_gather_indexer_k_quant_cache", + TORCH_BOX(&cp_gather_indexer_k_quant_cache)); +} + REGISTER_EXTENSION(_C_stable_libtorch) 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/ops.h b/csrc/ops.h index 34da56fec4b..f458f79d6f4 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -31,29 +31,6 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) { return new_tensor; } -void paged_attention_v1( - torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size, - int64_t max_seq_len, const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - -void paged_attention_v2( - torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, - torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size, - int64_t max_seq_len, const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::Tensor& k_scale, - torch::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 7ae644d81d4..58d19de7349 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -6,6 +6,7 @@ #include #include "../../../../attention/attention_dtypes.h" +#include namespace vllm { #ifdef USE_ROCM @@ -642,27 +643,29 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 3b7e25dc56b..6e95d72ed5e 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -1,6 +1,7 @@ #pragma once #include "../../../../attention/attention_dtypes.h" +#include #include #include #include @@ -546,37 +547,40 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ + if (SRC_DTYPE == torch::headeronly::ScalarType::Float) { \ FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) { \ FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) { \ FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ } else { \ - TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + STD_TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index dbc466f036e..73197d8a5e2 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -18,6 +18,15 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, const at::Tensor& scale_a, const at::Tensor& scale_b, const int64_t CuCount); +torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, diff --git a/csrc/rocm/q_gemm_rdna3.cu b/csrc/rocm/q_gemm_rdna3.cu new file mode 100644 index 00000000000..fb178f6d3b1 --- /dev/null +++ b/csrc/rocm/q_gemm_rdna3.cu @@ -0,0 +1,780 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the +// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit +// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes: +// +// 1. Direct write to the T-typed output via packed CAS-loop on a 64-bit +// word (atomic_add_pk4_{f16,bf16}). gfx11 has no native +// v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with +// global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch +// buffer + memset + cast-pass that an fp32-accumulator design would +// need; the caller passes a zero-initialised T-typed output tensor +// and every block atomically adds its partial sum into it. +// +// 2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only +// "upper nibble * 16" trick, which would overflow the 7-bit bf16 +// mantissa. See qdq_4_rdna3.cuh for details. +// +// 3. Wave32 geometry sized for high CU saturation: THREADS_X=256 +// (8 waves per block) and BLOCK_KN_SIZE=256, with each thread +// computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE +// splits K and the output is atomically accumulated. fp16 uses +// v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot; +// bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates +// with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch +// based on size_m. +// +// 4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in +// q_gemm_rdna3_wmma.cu (separate translation unit) where +// v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays +// scalar (the bit-trick dequant beats WMMA below M=64). + +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace gptq_rdna3 { + +// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K +// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this +// halves gridDim.z (32 → 16) and therefore halves the atomic count per +// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3 +// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak. +// +// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved +// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to +// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16 +// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs, +// breaking memory-latency hiding for the fp16 path which is already +// memory-bound. Reverted to 256; bf16 keeps most of its gains from the +// fp32 dequant rewrite alone. +#define BLOCK_KN_SIZE 256 +#define THREADS_X 256 + +// Device code below is RDNA3-only; non-RDNA3 device passes fall through to +// the empty __global__ stub at the #else below for symbol parity. +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// --------------------------------------------------------------------------- +// Per-dtype helpers. We avoid heavy template metaprogramming and just provide +// overloaded inline functions; the kernel below selects via `if constexpr`. +// --------------------------------------------------------------------------- + +// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting +// constructor from float, but going through __float2half_rn / __float2bfloat16 +// is the unambiguously correct path on every ROCm version. +template +__forceinline__ __device__ T tzero(); + +template <> +__forceinline__ __device__ half tzero() { + return __float2half_rn(0.0f); +} + +template <> +__forceinline__ __device__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) { + // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes + // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator + // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious + // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA + // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in + // the M_COUNT=8 kernel before this change), so we issue the builtin + // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops) + // per dot22_8_f call vs the half2-accumulator form. With 128 calls per + // K=32 step that's ~384 ops/K-step less issue pressure on the VALU. + // + // Numerical bonus: accumulator stays fp32 throughout the dot. The old + // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting, + // which could lose ~3 bits of precision on borderline magnitudes. + float result = 0.0f; + const half2* a2_ptr = (const half2*)a_ptr; + #pragma unroll + for (int i = 0; i < 4; i++) { + result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false); + } + return result; +} + +__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4], + const bf16_t* a_ptr) { + // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in + // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc + // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised + // fallback (single-element FMAs or fp32 round-trips), which empirically + // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16 + // decode path was paying that tax in full, scaling linearly with M (the + // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate + // and the kernel becomes memory-bound). + // + // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs) + // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA + // count, but each FMA is fast. Bonus: the accumulator is now fp32 + // throughout instead of bf16, which is also numerically more accurate + // (no compounding bf16-rounding inside the dot loop). + float result = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t aw, dw; + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); + __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t)); + // bf16 in low 16 bits → fp32 by left-shifting into the upper half. + // bf16 in high 16 bits → already aligned with fp32's upper half. + float a_x = __uint_as_float((aw & 0xFFFFu) << 16); + float a_y = __uint_as_float(aw & 0xFFFF0000u); + float d_x = __uint_as_float((dw & 0xFFFFu) << 16); + float d_y = __uint_as_float(dw & 0xFFFF0000u); + result = __fmaf_rn(d_x, a_x, result); + result = __fmaf_rn(d_y, a_y, result); + } + return result; +} + +// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already +// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t +// overload above does for dq (still need to widen A from bf16). Wins more +// at high N: the bf162_t version's per-call widening cost scales with the +// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only +// for A widening (M_COUNT × 4 × 4 widens, half as many). +__forceinline__ __device__ float dot22_8_f(float (&dq)[8], + const bf16_t* a_ptr) { + float result = 0.0f; + #pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t aw; + __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t)); + float a_x = __uint_as_float((aw & 0xFFFFu) << 16); + float a_y = __uint_as_float(aw & 0xFFFF0000u); + result = __fmaf_rn(dq[2 * i + 0], a_x, result); + result = __fmaf_rn(dq[2 * i + 1], a_y, result); + } + return result; +} + +// --------------------------------------------------------------------------- +// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS). +// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those +// landed on gfx940 / gfx1250 respectively), so this lowers to +// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to +// write 4 output columns per row in a single atomic operation — half the +// atomic instruction count and half the contention vs two 32-bit CAS calls. +// +// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer + +// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue +// cast pass that an fp32-accumulator design would need. +// +// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4 +// (always multiple of 4), and partition_weight_shape[1] is required to be a +// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte +// aligned. Required by global_atomic_cmpswap_b64. +// --------------------------------------------------------------------------- + +__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01, + half2 v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + half2 h2[2]; + } cur, sum; + cur.u = old; + sum.h2[0] = __hadd2(cur.h2[0], v01); + sum.h2[1] = __hadd2(cur.h2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01, + bf162_t v23) { + unsigned long long* addr_u = reinterpret_cast(addr); + unsigned long long old = *addr_u; + while (true) { + union { + unsigned long long u; + bf162_t b2[2]; + } cur, sum; + cur.u = old; + sum.b2[0] = __hadd2(cur.b2[0], v01); + sum.b2[1] = __hadd2(cur.b2[1], v23); + unsigned long long prev = atomicCAS(addr_u, old, sum.u); + if (prev == old) break; + old = prev; + } +} + +// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8] +// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with +// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two +// uint32 words; in practice within one because n & 7 is 0 or 4. +__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n, + int (&zeros)[4]) { + int qcol = n / 8; + int shift = (n & 0x07) * 4; + uint32_t d = qzeros_row[qcol] >> shift; + zeros[0] = (int)(d & 0xF); + zeros[1] = (int)((d >> 4) & 0xF); + zeros[2] = (int)((d >> 8) & 0xF); + zeros[3] = (int)((d >> 12) & 0xF); +} + +template +__forceinline__ __device__ void load4_scales(const T* scales_row, int n, + T (&scales)[4]) { + scales[0] = scales_row[n + 0]; + scales[1] = scales_row[n + 1]; + scales[2] = scales_row[n + 2]; + scales[3] = scales_row[n + 3]; +} + +// --------------------------------------------------------------------------- +// Main kernel. +// --------------------------------------------------------------------------- + +template +__global__ void gemm_q4_kernel_rdna3( + const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + const int t = threadIdx.x; + const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4; + const int offset_m = blockIdx.y * M_COUNT; + const int offset_k = blockIdx.z * BLOCK_KN_SIZE; + const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k); + const int n = offset_n + t * 4; + + // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M + // row break the natural 256-element/512-byte alignment that would otherwise + // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k] + // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132 + // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for + // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant. + constexpr int LDS_PAD = 8; + __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD]; + + // Stage A: each thread loads 1 K element per M row into LDS (with optional + // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map. + // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m + // are zero-padded so the dot product contribution is 0 (we then skip the + // atomic write for those rows below). + // + // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads + // read the SAME 8-element A window per inner step (a_off is uniform across + // the block), so the cache-line broadcast through L1 makes global reads as + // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128. + static_assert(BLOCK_KN_SIZE == THREADS_X, + "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)"); + // The M=1 fast path (skip LDS) only has a global-read code path for bf16 + // (the v_dot2_f32_bf16 branch). The fp16 inner loop still indexes + // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through + // LDS even at M=1 to avoid reading uninitialized shared memory. + constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same::value; + if constexpr (USE_LDS_A) { + if (offset_k + t < end_k) { + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + T av; + if (offset_m + m < size_m) { + const T* a_row = a + (offset_m + m) * size_k; + if (b_q_perm) + av = a_row[b_q_perm[offset_k + t]]; + else + av = a_row[offset_k + t]; + } else { + av = tzero(); // zero-pad invalid M rows + } + block_a[m][t] = av; + } + } + + // Threads beyond the right edge of N have nothing to do. Note: we must NOT + // return before __syncthreads() if any thread in the block participates in + // the LDS load above — but here all THREADS_X (=256) threads always do, + // regardless of whether their `n` is in bounds. + __syncthreads(); + } else if (b_q_perm) { + // bf16 M=1 fast path skips LDS, but its global read below is sequential + // and cannot apply act-order. When a permutation is present, stage the + // single A row through LDS (as fp16 / M>1 do) so the read picks it up. + // b_q_perm is block-uniform, so the __syncthreads is non-divergent. + if (offset_k + t < end_k) + block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]]; + __syncthreads(); + } + if (n >= size_n) return; + + // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K). + const int groupsize = size_k / groups; + int group = offset_k / groupsize; + int nextgroup = (group + 1) * groupsize; + + // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0. + int qk = offset_k / 8; + const uint32_t* b_ptr = b_q_weight + qk * size_n + n; + + // Per-column dequant constants. We hold one set of (z, y) pairs per column. + // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper- + // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant + // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA + // bypass for the missing v_pk_fma_bf16 on gfx11. + half2 z1z16_h[4][2], y1y16_h[4][2]; + float z_b_f[4], y_b_f[4]; + + auto refresh_group = [&](int g) { + const uint32_t* qz_row = b_qzeros + g * (size_n / 8); + const T* sc_row = b_scales + g * size_n; + int zeros[4]; + T scales[4]; + load4_zeros(qz_row, n, zeros); + load4_scales(sc_row, n, scales); + if constexpr (std::is_same::value) { + #pragma unroll + for (int i = 0; i < 4; ++i) { + prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i], + z1z16_h[i], y1y16_h[i]); + } + } else { + #pragma unroll + for (int i = 0; i < 4; ++i) { + prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i], + z_b_f[i], y_b_f[i]); + } + } + }; + + refresh_group(group); + + float block_c[M_COUNT][4]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + #pragma unroll + for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f; + } + + // Note on group-transition granularity: we check `k == nextgroup` at the + // start of each outer iteration (which advances K by 32). This is correct + // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one + // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the + // inner loop would cross a group boundary between j-iterations; we require + // group_size >= 32 here, mirroring exllama's assumption. + // + // Software pipelining: we issue all 4 vectorized weight loads up front + // before any dequant/FMA depends on them. This gives the AMDGPU backend + // freedom to schedule the global_loads early and overlap their latency + // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16 + // VGPRs in flight per thread, plenty of headroom on RDNA3. + int k = offset_k; + while (k < end_k) { + if (k == nextgroup) { + group++; + nextgroup += groupsize; + refresh_group(group); + } + + // Prefetch all four j-iterations' weight words. The compiler emits 4 + // global_load_b128 instructions back-to-back; the dependent dequant + + // FMA work below hides their latency. + int4 b_w[4]; + #pragma unroll + for (int j = 0; j < 4; ++j) { + b_w[j] = *(const int4*)(b_ptr + j * size_n); + } + b_ptr += 4 * size_n; + + #pragma unroll + for (int j = 0; j < 4; ++j) { + const int a_off = (k - offset_k) + 8 * j; + + if constexpr (std::is_same::value) { + half2 dq[4][4]; + dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]); + dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]); + dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]); + dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]); + + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + const half* a_ptr = reinterpret_cast(&block_a[m][a_off]); + block_c[m][0] += dot22_8_f(dq[0], a_ptr); + block_c[m][1] += dot22_8_f(dq[1], a_ptr); + block_c[m][2] += dot22_8_f(dq[2], a_ptr); + block_c[m][3] += dot22_8_f(dq[3], a_ptr); + } + } else if constexpr (M_COUNT == 1) { + // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of + // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's + // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000): + // * activations and magic-value weights share a fp32-aliased + // union (bytes written as uint32, read as bf16x2_t for the + // dot — pointer-cast opacity defeats the fold) + // * sum_a computed via a *second* v_dot2 with bf162(1,1) as the + // second operand, avoiding any explicit bf16→fp32 widen of A + // * bias correction y_b_f * partial + z_b_f * sum_a, identical + // to the previous fp32-FMA-chain path + // + // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the + // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so + // the substitution is ~2× cheaper for the inner accumulator. + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) + constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a + // fp32-aliased union. Storing as uint32 keeps the IR-level type + // opaque so the inner v_dot2 cannot be folded to fp32 widening. + // + // A is read direct from global (no LDS staging — see USE_LDS_A above), + // except under act-order, where it comes from the permuted LDS copy. + pack4 a_pack; + { + const uint32_t* a_words = + b_q_perm + ? reinterpret_cast(&block_a[0][a_off]) + : reinterpret_cast(a + offset_k + a_off); + a_pack.u[0] = a_words[0]; + a_pack.u[1] = a_words[1]; + a_pack.u[2] = a_words[2]; + a_pack.u[3] = a_words[3]; + } + + // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as + // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi. + // No fp32 widening of activations: the bytes go straight from LDS + // through v_dot2 into the fp32 accumulator. + float sum_a = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + sum_a = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES), + sum_a, /*clamp=*/false); + } + + // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs + // recycled across cols), avoiding straight-line expansion that + // would inflate live-range to 32 VGPRs. + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + // Build dequant magic values bf16(128 + nibble) directly into a + // fp32-aliased union via uint32 stores. No fp32 in the data flow + // until v_dot2 consumes the bytes. + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16. + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + + // block_c += y_b_f * partial + z_b_f * sum_a + // y_b_f = scale, z_b_f = -(128+zero)*scale + // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a + // and scaling yields scale · (nibble - zero) · a as required. + block_c[0][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a, block_c[0][col])); + } + } else { + // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as + // the M=1 branch: activations + magic-value weights stored in + // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with + // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2 + // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a) + // applied after the dot. Magic values built once per col and reused + // across all M rows — amortizes dequant cost across M_COUNT. + typedef short __attribute__((ext_vector_type(2))) bf16x2_t; + constexpr uint32_t BF16_MAGIC = 0x43004300u; // bf162(128, 128) + constexpr uint32_t BF16_ONES = 0x3F803F80u; // bf162(1.0, 1.0) + union pack4 { + float f[4]; + uint32_t u[4]; + }; + + uint32_t w[4]; + __builtin_memcpy(w, &b_w[j], sizeof(int4)); + + // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4 + // unions. Stored as uint32 to keep IR-level types opaque (defeats + // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3 + // budget. + pack4 a_pack[M_COUNT]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + const uint32_t* a_words = + reinterpret_cast(&block_a[m][a_off]); + a_pack[m].u[0] = a_words[0]; + a_pack[m].u[1] = a_words[1]; + a_pack[m].u[2] = a_words[2]; + a_pack[m].u[3] = a_words[3]; + } + + // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen. + float sum_a[M_COUNT]; + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + float s = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])), + *((const bf16x2_t*)&BF16_ONES), + s, /*clamp=*/false); + } + sum_a[m] = s; + } + + // Per col: build magic-value pack, dot against all M activations. + // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled) + // — same register-pressure trick as the previous fp32 path. + #pragma unroll 1 + for (int col = 0; col < 4; ++col) { + pack4 q_pack; + const uint32_t qa = w[col]; + q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC; + q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC; + + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + float partial = 0.0f; + #pragma unroll + for (int b = 0; b < 4; ++b) { + partial = __builtin_amdgcn_fdot2_f32_bf16( + *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])), + partial, /*clamp=*/false); + } + // block_c += y_b_f * partial + z_b_f * sum_a (same correction as + // M=1) + block_c[m][col] = + __fmaf_rn(y_b_f[col], partial, + __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col])); + } + } + } + } + k += 32; // 4 weight words * 8 nibbles = 32 K elements + } + + // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all + // four lanes in a single 64-bit CAS write directly to the T-typed output + // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with + // a single b64 op we halve the atomic instruction count vs two b32 CAS + // calls, AND save the FP32 buffer + memset + cast pass entirely. + #pragma unroll + for (int m = 0; m < M_COUNT; ++m) { + if (offset_m + m >= size_m) continue; // skip padding rows past size_m + T* out = c + (offset_m + m) * size_n + n; + if constexpr (std::is_same::value) { + half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]), + __float2half_rn(block_c[m][1])); + half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]), + __float2half_rn(block_c[m][3])); + atomic_add_pk4_f16(out, r01, r23); + } else { + bf162_t r01; + r01.x = __float2bfloat16(block_c[m][0]); + r01.y = __float2bfloat16(block_c[m][1]); + bf162_t r23; + r23.x = __float2bfloat16(block_c[m][2]); + r23.y = __float2bfloat16(block_c[m][3]); + atomic_add_pk4_bf16(out, r01, r23); + } + } +} + +#else // non-RDNA3 device pass: empty __global__ for symbol parity. + +template +__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*, + const T*, T*, const int, const int, + const int, const int, const int, + const int*) {} + +#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__ + +// --------------------------------------------------------------------------- +// Launcher. +// --------------------------------------------------------------------------- + +template +void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + dim3 block(THREADS_X); + dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4), + (size_m + M_COUNT - 1) / M_COUNT, + (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE); + + gemm_q4_kernel_rdna3<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +// Dispatch to the largest M_COUNT template that doesn't waste more than +// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is +// the right tool, not bigger M_COUNT in the scalar dot-product path. +// +// Tile-waste table: +// M=1 -> M_COUNT=1 (no waste) +// M=2,3 -> M_COUNT=2 (M=3 wastes 1/2 of last tile) +// M=4-7 -> M_COUNT=4 (worst case M=5: wastes 3/4 of last tile) +// M=8-15-> M_COUNT=8 (worst case M=9: wastes 7/8 of last tile) +// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they +// only burn instructions on the last block, never affect correctness. +template +void launch_gemm_q4(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, int size_n, + int size_k, int groups, bool use_v2_format, + cudaStream_t stream) { + const int zero_offset = use_v2_format ? 0 : 1; + + if (size_m == 1) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else if (size_m <= 3) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else if (size_m <= 7) { + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } else { + // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the + // WMMA path, but if it falls through we still produce correct output — + // just leaving 3-5× of throughput on the table for prefill workloads. + launch_gemm_q4_for_mcount(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + } +} + +} // namespace gptq_rdna3 +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] half or bfloat16 +// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) +// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) +// b_scales [groups, N] half or bfloat16 +// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) +// use_v2_format bool (true = GPTQv2, no +1 zero offset) +// +// Output: +// c [M, N] same dtype as a + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format); + +torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format) { + if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 && + b_q_weight.size(1) % 16 == 0 && + ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) || + (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) { + return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx, + use_v2_format); + } + + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), + "b_scales dtype must match a"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto stream = at::cuda::getCurrentCUDAStream(); + + int size_m = (int)a.size(0); + int size_k = (int)a.size(1); + int size_n = (int)b_q_weight.size(1); + int groups = (int)b_qzeros.size(0); + + TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, + "b_q_weight first dim must be K/8"); + TORCH_CHECK(b_scales.size(0) == groups, + "b_scales must have same group count as qzeros"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); + TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)"); + + auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + at::Tensor c = torch::zeros({size_m, size_n}, opts); + + const int* g_idx_ptr = nullptr; + if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { + TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, + "b_g_idx must be int32"); + g_idx_ptr = (const int*)b_g_idx.data_ptr(); + } + + if (a.scalar_type() == torch::kHalf) { + vllm::gptq_rdna3::launch_gemm_q4( + (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), + g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, + use_v2_format, stream); + } else { + vllm::gptq_rdna3::launch_gemm_q4( + (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(), + (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), + (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups, + use_v2_format, stream); + } + + return c; +} diff --git a/csrc/rocm/q_gemm_rdna3_wmma.cu b/csrc/rocm/q_gemm_rdna3_wmma.cu new file mode 100644 index 00000000000..966d5df403f --- /dev/null +++ b/csrc/rocm/q_gemm_rdna3_wmma.cu @@ -0,0 +1,2165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 GPTQ WMMA prefill kernel for AMD RDNA3 (gfx1100). This is the +// matrix-instruction path for M >= 16; small-M decode lives in the sibling +// file q_gemm_rdna3.cu and is exposed via a different op +// (`gptq_gemm_rdna3`). Keeping the two paths in separate translation units +// is intentional: an earlier attempt at putting WMMA in the same TU as the +// scalar dot-product kernel introduced a compile-time interaction that +// silently miscompiled the M=1 path even though the WMMA template was never +// instantiated for M=1. Hipcc's optimizer appears to scope some decisions +// at the TU level (likely register file / SGPR pressure heuristics across +// all kernels in the TU), so we isolate. +// +// Hardware notes (RDNA3 / gfx1100 / RX 7900 XTX): +// * v_wmma_f32_16x16x16_{f16,bf16}_w32 — 16×16×16 GEMM in one instruction +// (~16 cycles per WMMA on wave32). Accumulator dtype is FP32; inputs +// are fp16 or bf16. There is NO 16x16x16 with fp16/bf16 accumulator on +// gfx11 that we'd want here (we always need fp32 accum to avoid loss +// across many K iterations). +// * Wave32 input fragment storage is "doubled" — lanes 16..31 hold a +// copy of lanes 0..15 for the A and B fragments. The output C +// fragment uses a different mapping: lane t holds COLUMN n=lane_lo +// of the 16x16 output, with 8 elements alternating M rows by lane_hi +// (lanes 0..15 = even rows, lanes 16..31 = odd rows). See the layout +// diagram on `gemm_q4_wmma_kernel_16x16_1w` below for the full mapping. +// * No native v_global_atomic_pk_add_{f16,bf16} on gfx11; the K-split +// epilogue (gridDim.z > 1) emulates packed atomic add via a CAS-32 +// retry loop on a uint32 word covering 2 fp16/bf16 lanes. Within a +// block we shuffle adjacent lanes via shfl_xor first so each pair of +// output cols goes through a single atomic — no intra-block +// contention. K_SPLIT == 1 keeps the original direct-write path with +// each block owning its 16M × 16N output tile. + +#include + +#include +#include +#include + +#include +#include +#include + +#include "qdq_4_rdna3.cuh" + +#if defined(__HIPCC__) && defined(__gfx1100__) + #define __HIP__RDNA3__ +#endif + +namespace vllm { +namespace gptq_rdna3_wmma { + +// Pull dequant types from the sibling namespace. +using vllm::gptq_rdna3::bf162_t; +using vllm::gptq_rdna3::bf16_t; + +// Device code below uses RDNA3-only __builtin_amdgcn_wmma_* intrinsics; +// non-RDNA3 device passes fall through to empty __global__ stubs at the +// #else block at the end of this TU. +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// PRECISE dequant variants live HERE, not in the shared qdq_4_rdna3.cuh +// header. Reason: hipcc takes different register/scheduling decisions when +// the shared header grows, which caused a measurable decode-tk/s regression +// in the scalar kernel even when these new functions were never called from +// it. Keeping them in the WMMA TU only restores scalar's binary identity +// to its tuned baseline. +// +// Numerics: the classic fp16 bit-trick (FMA form: q*scale + +// (-(1024+zero)*scale)) loses up to ~0.025 per cell at scale=0.1 because +// fp16(scale) ≠ scale and the FMA amplifies that by 1024× without cancelling +// against the precomputed (1024+zero)*scale (which is rounded to fp16 BEFORE +// the FMA). +// +// Fix: subtract (1024+zero) as an integer FIRST — exact in fp16 because +// integers in [1024, 2047] are exactly representable — then multiply by +// scale, incurring at most one half-ULP rounding. Costs one extra +// instruction per dequant pair (sub+mul vs single FMA), worth it for the +// WMMA path because K can be small (16) and errors don't average. +__forceinline__ __device__ void prep_zero_scale_fp16_precise(uint32_t zero, + half scale, + half2& z_prep, + half2& y_prep) { + union { + uint16_t u; + half h; + } zu; + zu.u = (uint16_t)(0x6400 | zero); + z_prep = __half2half2(zu.h); + y_prep = __half2half2(scale); +} + +__forceinline__ __device__ void dequant_4bit_8_fp16_precise(uint32_t qa, + half2 (&dq)[4], + half2 z_prep, + half2 y_prep) { + const uint32_t c0 = 0x64006400; + union { + uint32_t u; + half2 h2; + } q0, q1, q2, q3; + q0.u = ((qa >> 0) & 0x000F000F) | c0; + q1.u = ((qa >> 4) & 0x000F000F) | c0; + q2.u = ((qa >> 8) & 0x000F000F) | c0; + q3.u = ((qa >> 12) & 0x000F000F) | c0; + dq[0] = __hmul2(__hsub2(q0.h2, z_prep), y_prep); + dq[1] = __hmul2(__hsub2(q1.h2, z_prep), y_prep); + dq[2] = __hmul2(__hsub2(q2.h2, z_prep), y_prep); + dq[3] = __hmul2(__hsub2(q3.h2, z_prep), y_prep); +} + +// fp32 prep for the bf16→bf16 fp32-internal dequant. Returns: +// z_prep = -(128 + zero) * scale (folded bias for FMA) +// y_prep = scale +// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero). +__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, + bf16_t scale, + float& z_prep, + float& y_prep) { + float scale_f = __bfloat162float(scale); + z_prep = -(128.0f + (float)zero) * scale_f; + y_prep = scale_f; +} + +// fp32 → bf16 narrow that skips the defensive NaN-canonicalisation hipcc +// emits for __float2bfloat16. Round-half-to-even via add (0x7FFF + lsb) +// + truncate; no NaN check, since dequant outputs are bounded products +// of (nibble - zero) ∈ [-15, 15] and bf16 scales — never NaN/Inf. +__forceinline__ __device__ bf16_t f32_to_bf16_no_canon(float f) { + uint32_t fu = __float_as_uint(f); + uint32_t lsb = (fu >> 16) & 1u; + uint16_t out_u = (uint16_t)((fu + 0x7FFFu + lsb) >> 16); + bf16_t out; + __builtin_memcpy(&out, &out_u, sizeof(out)); + return out; +} + +// 4-bit GPTQ dequant for the bf16 WMMA path: 8 nibbles per int32 in qa, +// outputs bf162_t dq[4]. Implementation rationale: +// +// gfx11 has no v_pk_fma_bf16, so __hmul2/__hsub2 on bf16 lower to a +// widen-fp32-op-narrow chain that hipcc decorates with NaN canonicalisation +// (28 extra VALU ops per call observed in the v5 ISA dump: v_cmp_u_f32 + +// v_cndmask_b32 around every bf16 sub/mul). Doing the math in fp32 directly +// — bit-cast widening (`__uint_as_float((bf16_bits) << 16)` is just a shift, +// no NaN canon), one fused FMA per element, single `__float2bfloat16` narrow +// at the end — eliminates the canon and is also strictly more precise (one +// rounding step instead of two). +// +// Inputs match prep_zero_scale_bf16_f32: +// z_prep = -(128 + zero) * scale (folded bias for FMA) +// y_prep = scale +// Per-element FMA `q*y_prep + z_prep` then yields scale * (nibble - zero). +__forceinline__ __device__ void dequant_4bit_8_bf16_to_bf16(uint32_t qa, + bf162_t (&dq)[4], + float z_prep, + float y_prep) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + // bf16(128+nibble) bits → fp32 via left-shift by 16 (zero-extends mantissa). + const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16); + const float q0y = __uint_as_float(q0 & 0xFFFF0000u); + const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16); + const float q1y = __uint_as_float(q1 & 0xFFFF0000u); + const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16); + const float q2y = __uint_as_float(q2 & 0xFFFF0000u); + const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16); + const float q3y = __uint_as_float(q3 & 0xFFFF0000u); + // r = q*scale + (-(128+zero)*scale) = (nibble - zero)*scale, then narrow. + dq[0].x = f32_to_bf16_no_canon(__fmaf_rn(q0x, y_prep, z_prep)); + dq[0].y = f32_to_bf16_no_canon(__fmaf_rn(q0y, y_prep, z_prep)); + dq[1].x = f32_to_bf16_no_canon(__fmaf_rn(q1x, y_prep, z_prep)); + dq[1].y = f32_to_bf16_no_canon(__fmaf_rn(q1y, y_prep, z_prep)); + dq[2].x = f32_to_bf16_no_canon(__fmaf_rn(q2x, y_prep, z_prep)); + dq[2].y = f32_to_bf16_no_canon(__fmaf_rn(q2y, y_prep, z_prep)); + dq[3].x = f32_to_bf16_no_canon(__fmaf_rn(q3x, y_prep, z_prep)); + dq[3].y = f32_to_bf16_no_canon(__fmaf_rn(q3y, y_prep, z_prep)); +} + +// --------------------------------------------------------------------------- +// Packed atomic-add helpers used by the K-split epilogue. +// +// When the kernel is launched with gridDim.z > 1, multiple K-segments +// accumulate into the same 16x16 output tile and need atomic write-back. +// gfx11 has no native v_global_atomic_pk_add_{f16,bf16}, so we issue a +// CAS-loop on a 32-bit word covering 2 packed fp16/bf16 lanes. Within a +// block the kernel pairs adjacent lanes via shfl_xor first, so each pair +// of cols (n=lane_lo even, lane_lo+1) goes through a SINGLE atomic — no +// intra-block contention on the same uint32 target. Inter-block +// contention from gridDim.z (4-way at K_SPLIT=4) is the residual cost. +// --------------------------------------------------------------------------- + +__forceinline__ __device__ void atomic_add_pk_f16(half2* addr, half2 val) { + uint32_t* addr_u = reinterpret_cast(addr); + uint32_t old = *addr_u; + while (true) { + half2 cur; + __builtin_memcpy(&cur, &old, sizeof(cur)); + half2 sum = __hadd2(cur, val); + uint32_t sum_u; + __builtin_memcpy(&sum_u, &sum, sizeof(sum_u)); + uint32_t prev = atomicCAS(addr_u, old, sum_u); + if (prev == old) break; + old = prev; + } +} + +__forceinline__ __device__ void atomic_add_pk_bf16(bf162_t* addr, bf162_t val) { + uint32_t* addr_u = reinterpret_cast(addr); + uint32_t old = *addr_u; + while (true) { + bf162_t cur; + __builtin_memcpy(&cur, &old, sizeof(cur)); + bf162_t sum = __hadd2(cur, val); + uint32_t sum_u; + __builtin_memcpy(&sum_u, &sum, sizeof(sum_u)); + uint32_t prev = atomicCAS(addr_u, old, sum_u); + if (prev == old) break; + old = prev; + } +} + +#endif // helpers guard; K-split heuristics below are pure host/device + // arithmetic, called from launch_* on non-RDNA3 device passes too. + +// K-split factor heuristic. Returns the gridDim.z to use for a given K. +// Aim: each block does at least ~16 K-tiles (= K=256) so the per-block +// constant overhead (LDS init, kernel prologue) is amortised. Upper +// bound K_SPLIT=4 to cap inter-block atomic contention to 4-way. +// +// For typical Qwen-class shapes K ∈ {4096, 5120, 11008}, all return 4. +// Smaller K (e.g., embedding lookups) fall back to 1 (no split, no +// atomic). K must be divisible by (K_SPLIT × 16) for the split to be +// valid; the heuristic checks divisibility before raising the factor. +__host__ __device__ static inline int compute_wmma_k_split(int size_k) { + if (size_k >= 1024 && size_k % 64 == 0) return 4; + if (size_k >= 512 && size_k % 32 == 0) return 2; + return 1; +} + +// M-and-N-aware K-split heuristic for the v3/v4/v5 launchers. +// +// The original `compute_wmma_k_split` was K-only and always returns 4 for +// Qwen-class K, which over-subscribes wave slots and pays the atomic CAS +// epilogue once-per-K-segment per output cell. With v3/v4/v5's larger +// tiles (64M × 16/32/64N) and 4 resident waves per block, the no-split +// grid is often already well-saturated on gfx1100's 96 CUs / 3072 wave +// slots — adding gridDim.z just adds atomic overhead. +// +// Heuristic: compute the no-split block count gridDim.x × gridDim.y, then +// pick the smallest K_SPLIT that brings total waves to at least +// ~2× over-subscription (~6000 waves for our 3072 slots, i.e. 1500 blocks +// at 4 waves/block). Above that threshold, K_SPLIT=1 — direct write, no +// atomic. +// +// Args: +// size_m, size_n, size_k — GEMM dims +// m_tile, n_tile — block-level M and N tile (64×16 for v3, +// 64×32 for v4, 64×64 for v5) +// +// Returns: gridDim.z divisor (1, 2, or 4), respecting K-divisibility. +__host__ __device__ static inline int compute_wmma_k_split_mn( + int size_m, int size_n, int size_k, int m_tile, int n_tile) { + const int blocks_xy = + ((size_n + n_tile - 1) / n_tile) * ((size_m + m_tile - 1) / m_tile); + // Target: enough blocks to keep ~2× oversubscription on 96 CUs / 3072 + // wave slots at 4 waves/block ⇒ ~1500 blocks no-split. + constexpr int kTargetBlocksXY = 1500; + if (blocks_xy >= kTargetBlocksXY) return 1; + if (blocks_xy * 2 >= kTargetBlocksXY && size_k >= 512 && size_k % 32 == 0) + return 2; + if (blocks_xy * 4 >= kTargetBlocksXY && size_k >= 1024 && size_k % 64 == 0) + return 4; + // Fall back to the K-only heuristic when blocks are very few (small + // models with small N): K-split is the only way to add parallelism. + return compute_wmma_k_split(size_k); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// Native AMDGPU vector types expected by the WMMA built-ins. +using v16fp16 = _Float16 __attribute__((ext_vector_type(16))); +using v16bf16 = __bf16 __attribute__((ext_vector_type(16))); +using v8fp32 = float __attribute__((ext_vector_type(8))); + +__device__ __forceinline__ v8fp32 wmma_mma(v16fp16 a, v16fp16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c); +} +__device__ __forceinline__ v8fp32 wmma_mma(v16bf16 a, v16bf16 b, v8fp32 c) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c); +} + +// Map HIP wrapper types (half, __hip_bfloat16) to native compiler types +// (_Float16, __bf16) used by the WMMA built-ins. Bitcast is a register +// reinterpret in practice. +template +struct WmmaNative; +template <> +struct WmmaNative { + using elem = _Float16; + using v16 = v16fp16; +}; +template <> +struct WmmaNative { + using elem = __bf16; + using v16 = v16bf16; +}; + +template +__device__ __forceinline__ TO bitcast_elem(FROM x) { + static_assert(sizeof(FROM) == sizeof(TO), + "bitcast_elem requires equal-sized types"); + TO r; + __builtin_memcpy(&r, &x, sizeof(TO)); + return r; +} + +// Per-T tzero (matches the helper in the scalar TU). +template +__device__ __forceinline__ T tzero(); +template <> +__device__ __forceinline__ half tzero() { + return __float2half_rn(0.0f); +} +template <> +__device__ __forceinline__ bf16_t tzero() { + return __float2bfloat16(0.0f); +} + +#endif // helpers guard (each __global__ below has its own guard so launch_* + // host code remains visible to the parser on non-RDNA3 device passes) + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// WMMA kernel: 16M × 16N tile per block, 1 wave, full K traversal. +// +// Wave32 fragment layout (verified empirically with all-modes diagnostic +// against a random A and B and the eight candidate output mappings): +// +// * A frag (row-major in M, K in slot): +// lane t, slot i → A[lane_lo][k = i] +// Lane axis encodes M (A's row), slot encodes K. +// * B frag (col-major in N, K in slot): +// lane t, slot i → B[k = i][lane_lo] +// Lane axis encodes N (B's column), slot encodes K. K-axis aligns +// with A's K-axis (same slot index). +// * C frag (output, lane = N, slot = M with hi-bit interleave): +// lane t, slot i → C[m = 2*i + lane_hi][n = lane_lo] +// Lane axis encodes N (C's column). Each lane holds 8 elements of +// its output column, alternating rows: lanes 0..15 (hi=0) hold even +// rows m=0,2,4,...,14; lanes 16..31 (hi=1) hold odd rows +// m=1,3,5,...,15. +// * Both halves of the wave (lanes 0..15 and 16..31) hold IDENTICAL input +// fragments (AMD's "doubled" wave32 input layout). Output is split +// between halves via lane_hi. +// +// History note: an earlier version of this kernel loaded B row-major in K +// and assumed the output was C[lane_lo][2*i+lane_hi]. That layout passes +// all-A=identity tests because A=I makes the K-axis sum collapse, but +// implements C = A @ B^T for non-trivial A — the bug only shows up against +// random A. A layout probe iterating all four +// {row,col} × {row,col} loadings identified mode 1 (A row, B col) with +// output [m=2*i+hi][n=lane_lo] as the unique mapping that yields A @ B. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_16x16_1w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 16; + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int lane = threadIdx.x; // 0..31 + const int lane_lo = lane & 15; // row index within fragment + const int lane_hi = lane >> 4; // 0 or 1 + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block in the gridDim.z dimension processes a contiguous + // K-segment [k_start, k_end). With gridDim.z > 1, multiple blocks + // accumulate into the same output tile and need atomic write-back at + // the end. With gridDim.z == 1 the kernel falls back to the original + // behaviour: full K range, single writer per cell, direct write. + // + // The split multiplies the wave count by gridDim.z and proportionally + // raises CU saturation — this is the dominant lever for closing the + // throughput gap to the fp16 scalar kernel at M >= 16, which already + // uses K-split natively (gridDim.z = K/256). At gridDim.z = 4 with + // K=4096, WMMA jumps from 17% to ~67% wave-slot saturation. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // LDS tile of dequantized B. 16 K rows × 16 N cols. + __shared__ T b_lds[16][16]; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + // ---- Dequant 16x16 B tile into LDS ---- + // 32 lanes split 16 N cols × 2 K-octets per col = 32 dequant tasks. + const int my_n = lane_lo; + const int my_k_octet = lane_hi; // 0 → K[0..7], 1 → K[8..15] + const int actual_n = n_tile + my_n; + + if (actual_n < size_n) { + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[k_base + 0][my_n] = __low2half(dq[0]); + b_lds[k_base + 1][my_n] = __high2half(dq[0]); + b_lds[k_base + 2][my_n] = __low2half(dq[1]); + b_lds[k_base + 3][my_n] = __high2half(dq[1]); + b_lds[k_base + 4][my_n] = __low2half(dq[2]); + b_lds[k_base + 5][my_n] = __high2half(dq[2]); + b_lds[k_base + 6][my_n] = __low2half(dq[3]); + b_lds[k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[k_base + 0][my_n] = dq[0].x; + b_lds[k_base + 1][my_n] = dq[0].y; + b_lds[k_base + 2][my_n] = dq[1].x; + b_lds[k_base + 3][my_n] = dq[1].y; + b_lds[k_base + 4][my_n] = dq[2].x; + b_lds[k_base + 5][my_n] = dq[2].y; + b_lds[k_base + 6][my_n] = dq[3].x; + b_lds[k_base + 7][my_n] = dq[3].y; + } + } + + // No __syncthreads() needed: the launch is `dim3 block(32)` = exactly one + // wave32, so there is no inter-wave concurrency. Within a wave the + // compiler emits `s_waitcnt lgkmcnt(0)` between dependent ds_write/ds_read + // pairs automatically, so cross-lane LDS reads (lane 0 reading what + // lane 16 wrote into b_lds[8..15][0]) still observe the writes. Keeping + // the explicit `__syncthreads()` would emit a wave-level `s_barrier` that + // costs ~10-20 cycles every iteration but provides no semantic guarantee + // we don't already have for free in single-wave mode. + + // ---- Build A and B fragments, run WMMA ---- + V16 a_frag, b_frag; + const int m_row = m_tile + lane_lo; + + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + // Permuted (act-order): scattered global reads, no vectorization. + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + // Sequential A reads: replace 16 single-element global_load_b16 with + // a bulk 32-byte copy. The AMDGPU backend lowers a memcpy of this + // size + alignment to two `global_load_b128` instructions. size_k is + // a multiple of 16 (TORCH_CHECK above) and k_tile increments by 16, + // so k_tile + 16 is always within bounds — no tail handling needed. + // Note: we memcpy into the whole vector (`&a_frag`) rather than + // `&a_frag[0]`; ext_vector_type element addresses aren't reliably + // valid C pointers across compiler versions. + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes (16 × 2)"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // B fragment: lane t holds COLUMN n=lane_lo of the B tile (K-axis in + // slot, N-axis in lane). This is the AMD WMMA convention for the right + // operand of a matrix multiply — K-axis aligns with A's K-axis (also + // in slot), enabling per-lane inner products. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[i][lane_lo]); + } + + #ifdef VLLM_WMMA_LAYOUT_DEBUG + // Diagnostic: skip WMMA, force c_acc to encode (lane, slot) so the + // store pattern reveals the C-output lane→matrix mapping. Compile with + // -DVLLM_WMMA_LAYOUT_DEBUG to enable. Output: c[m][n] = lane + slot/16. + (void)a_frag; + (void)b_frag; + #pragma unroll + for (int i = 0; i < 8; i++) { + c_acc[i] = (float)lane + (float)i / 16.0f; + } + // Run only one K iteration in debug mode so c_acc isn't overwritten. + if (k_tile == 0) { + k_tile = size_k; // exit loop on next check + } + #else + c_acc = wmma_mma(a_frag, b_frag, c_acc); + #endif + + // No __syncthreads() needed before the next iter overwrites b_lds: + // single-wave block, and the next iter's ds_write to b_lds is preceded + // by a `s_waitcnt lgkmcnt(0)` from the compiler that ensures the WMMA's + // ds_read of b_frag has completed before the new ds_write issues. + } + + // ---- Store C ---- + // Lane t holds column n=lane_lo of the output tile, 8 rows determined by + // slot i and lane_hi: + // lane_hi == 0 → rows 0, 2, 4, ..., 14 (even rows) + // lane_hi == 1 → rows 1, 3, 5, ..., 15 (odd rows) + // c_acc[i] corresponds to actual row m = 2*i + lane_hi at column lane_lo. + if (gridDim.z > 1) { + // K-split path: 4 (or whatever the split factor is) K-segments per + // output cell contend → atomic accumulation. Caller has zero-init'd c. + // + // Pair-shuffle to avoid intra-block CAS contention: lanes lane_lo and + // lane_lo+1 share the same uint32 atomic target (4 bytes = 2 fp16), + // so without pairing they'd hammer the same word. Instead, swap the + // c_acc[i] value with the lane_lo+1 neighbour via shfl_xor and have + // ONLY the even lane issue a single packed CAS. Inter-block contention + // (gridDim.z-way per cell) remains and is the residual atomic cost. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; // valid only on even lane + #pragma unroll + for (int i = 0; i < 8; i++) { + // Wave-wide shuffle: every lane participates so the side-effect is + // visible. Only even lanes use the result. shfl_xor with mask 1 + // swaps with the lane_lo XOR 1 neighbour (same lane_hi → same row). + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + // Pack: .x = mine (col=lane_lo even), .y = neighbour (col=lane_lo+1) + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // gridDim.z == 1: single writer per cell, direct non-atomic write. + // Caller can leave c uninitialised (torch::empty) since every cell is + // assigned exactly once. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_16x16_1w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_16x16_1w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // 1 wave per block (32 lanes), 16x16 C tile per block. gridDim.z splits + // K so that more blocks (and therefore more waves) are in flight; with + // K_SPLIT > 1 the kernel switches to atomic write-back at the epilogue. + const int k_split = compute_wmma_k_split(size_k); + dim3 block(32); + dim3 grid((size_n + 15) / 16, (size_m + 15) / 16, k_split); + gemm_q4_wmma_kernel_16x16_1w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 32x16_2w kernel: 2 waves per block, 32M × 16N tile, double-buffered LDS. +// +// Targets the bf16-WMMA prefill regime (M >= 128) where the v1 single-wave +// kernel saturates only ~24% of WMMA peak because each wave does roughly +// one v_wmma every ~30-40 cycles (16-cycle wmma latency + dequant + LDS + +// global A load all serial inside the single resident wave). +// +// Two structural changes vs v1: +// +// * 2 waves per block (64 threads). Both waves cooperate on a 32M×16N +// output tile: wave 0 produces rows [0..15], wave 1 rows [16..31]. +// The B-tile in LDS is shared (only wave 0 dequants); each wave loads +// its own A slice from global. With two resident waves the SIMD +// scheduler can keep the WMMA pipeline full by interleaving wmmas +// from the two waves while the other does dequant / LDS / load work. +// +// * Double-buffered LDS B-tile (b_lds[2][16][16]). Wave 0 dequants +// the K-tile for iter k+1 while both waves consume the K-tile for +// iter k. Pulls dequant out of the WMMA-critical path. Costs ~512 B +// extra LDS per block (irrelevant — gfx1100 has 64 KB LDS/CU). +// +// One __syncthreads() per K-iter remains: it ensures wave 0's dequant of +// the next K-tile has committed AND both waves have finished reading the +// current K-tile before wave 0 wraps around and overwrites it. +// +// The v2 launcher (`launch_gemm_q4_wmma_32x16_2w`) is the production entry on +// the WMMA path and falls back to v1 internally for size_m < 32 — see the +// comment at the top of `launch_gemm_q4_wmma_32x16_2w` for the M=16 regression +// rationale that justifies the fallback. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_32x16_2w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 32; // 32-row stride per block + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..63 + const int wave_id = tid >> 5; // 0 or 1 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block handles a contiguous K-segment when gridDim.z > 1. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Double-buffered LDS B-tile. 2 × 16K × 16N × sizeof(T) = 1024 B for + // fp16/bf16. + __shared__ T b_lds[2][16][16]; + + // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0 participates + // (32 lanes do 32 dequant tasks: 16 N-cols × 2 K-octets per col). + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id != 0) return; + + const int my_n = lane_lo; + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n] = dq[0].x; + b_lds[buf][k_base + 1][my_n] = dq[0].y; + b_lds[buf][k_base + 2][my_n] = dq[1].x; + b_lds[buf][k_base + 3][my_n] = dq[1].y; + b_lds[buf][k_base + 4][my_n] = dq[2].x; + b_lds[buf][k_base + 5][my_n] = dq[2].y; + b_lds[buf][k_base + 6][my_n] = dq[3].x; + b_lds[buf][k_base + 7][my_n] = dq[3].y; + } + }; + + // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume. + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + // Issue dequant of next K-tile (wave 0 only) — overlaps with current + // iter's WMMA work below. The LDS write is non-blocking; the sync at + // end of iter ensures wave 1 sees it before iter k+1. + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its own M slice in parallel. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B from current buffer (both waves read identical data). + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + } + + // Each wave issues its own WMMA against its own a_frag + shared b_frag. + c_acc = wmma_mma(a_frag, b_frag, c_acc); + + // Sync ensures: (a) wave 0's dequant into next_buf has committed, + // (b) both waves are done reading cur_buf — so next iter's overwrite + // (cur_buf becomes the previous next_buf and gets reused two iters + // later) is race-free. + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- + // Each wave owns rows [m_tile + wave_id*16 .. + 16) of the output tile. + const int m_tile_wave = m_tile + wave_id * 16; + + if (gridDim.z > 1) { + // K-split atomic path. Pair-shuffle within wave to halve atomic count. + // shfl_xor here is wave-local (wave32 semantics) so each wave does its + // own pairing — the two waves don't interact during the store. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // Single writer per cell, direct non-atomic write. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_32x16_2w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_32x16_2w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fallback to v1 for size_m < 32. With M-tile=32 the v2 block has 2 waves + // working on rows [0..15] and [16..31]; at M < 32 the second wave processes + // out-of-range M rows (zero-padded a_frag → wmma produces nothing useful) + // and just wastes SIMD cycles. Bench measured a +47 % regression at M=16 + // vs v1 for this reason. The M < 32 case is rare in serving (decode at + // max-num-seqs=32 lands at M≈32 steady-state; the M=16 sliver is edge), + // but the fallback costs nothing and is the right shape. + if (size_m < 32) { + launch_gemm_q4_wmma_16x16_1w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 2 waves per block (64 threads), 32M × 16N C tile per block. + // K-split heuristic shared with v1. With M-tile=32, the natural + // grid blocks are halved on Y vs v1 — but each block does 2× the work, + // so total wave count is unchanged at the same M when K_SPLIT is equal. + const int k_split = compute_wmma_k_split(size_k); + dim3 block(64); + dim3 grid((size_n + 15) / 16, (size_m + 31) / 32, k_split); + gemm_q4_wmma_kernel_32x16_2w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x16_4w kernel: 4 waves per block, 64M × 16N tile, double-buffered LDS. +// +// Targets the prefill plateau observed at M >= 128 in v2 (~144 K tk/s bf16, +// ~28 % of WMMA peak). The bottleneck is wmma issue rate per resident wave: +// each wave issues at most 1 wmma per ~30-40 cycles. With only 2 waves per +// block, two wmmas overlap; the wmma pipeline (16-cycle latency) is mostly +// idle. +// +// Doubling the resident wave count (2 → 4) targets ~2× wmma throughput by +// keeping the pipeline closer to full. 64M tile keeps N-tile at 16 (so the +// b_lds layout, dequant pattern, and store mapping carry over from v2) — the +// only structural changes are: +// +// * 4 waves cooperate on the 64M × 16N output tile. Wave w produces rows +// [16w .. 16w+15] of the M tile. +// * Dequant remains on wave 0 only (32 lanes do 32 dequant slots, identical +// to v2). Waves 1-3 idle through the dequant phase but their wmmas can +// issue concurrently with wave 0's dequant of the *next* K-tile thanks +// to the double buffer — net wave occupancy is dominated by the wmma +// phase, not the dequant phase. +// * One __syncthreads() per K-iter still required (same race as v2). +// +// Costs: same LDS as v2 (1024 B for the b_lds double buffer). Block has 128 +// threads vs 64 in v2; gfx1100 supports up to 1024 threads/block so this is +// well within budget. Doubles VGPR pressure slightly because the four waves +// each hold their own a_frag + c_acc — but each wave's working set is +// independent so per-thread VGPR is unchanged. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x16_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = + blockIdx.y * 64; // 64-row stride per block (4 waves × 16M) + const int n_tile = blockIdx.x * 16; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + // K-split: each block handles a contiguous K-segment when gridDim.z > 1. + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Double-buffered LDS B-tile. Same layout as v2. + __shared__ T b_lds[2][16][16]; + + // Dequant a 16K × 16N B-tile into b_lds[buf]. Only wave 0's 32 lanes + // participate (16 N-cols × 2 K-octets = 32 dequant slots). Waves 1-3 + // skip — their wmma can run concurrently with wave 0's next-iter dequant + // through the double buffer. + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id != 0) return; + + const int my_n = lane_lo; + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n] = dq[0].x; + b_lds[buf][k_base + 1][my_n] = dq[0].y; + b_lds[buf][k_base + 2][my_n] = dq[1].x; + b_lds[buf][k_base + 3][my_n] = dq[1].y; + b_lds[buf][k_base + 4][my_n] = dq[2].x; + b_lds[buf][k_base + 5][my_n] = dq[2].y; + b_lds[buf][k_base + 6][my_n] = dq[3].x; + b_lds[buf][k_base + 7][my_n] = dq[3].y; + } + }; + + // Pre-fill buffer 0 with the first K-tile so iter 0 has data to consume. + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + // Issue dequant of next K-tile (wave 0 only) — overlaps with the wmma + // work below across all 4 waves. + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Each wave loads its own 16M slice of A (wave w handles M-rows + // [m_tile + 16w .. m_tile + 16w + 15]). + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B from current buffer (all 4 waves read identical data). + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + } + + // Each wave issues its own WMMA against its own a_frag + shared b_frag. + // 4 wmmas in flight per block per K-iter. + c_acc = wmma_mma(a_frag, b_frag, c_acc); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16 M-rows of the output tile. + const int m_tile_wave = m_tile + wave_id * 16; + + if (gridDim.z > 1) { + // K-split atomic path. Pair-shuffle within wave to halve atomic count. + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_tile + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(c_acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(c_acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(c_acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + // Single writer per cell, direct non-atomic write. + const int out_n = n_tile + lane_lo; + if (out_n < size_n) { + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(c_acc[i]); + } else { + *dst = __float2bfloat16(c_acc[i]); + } + } + } + } + } +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_64x16_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_64x16_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v2 for M < 64 (would waste 1+ waves on out-of-range rows). + if (size_m < 64) { + launch_gemm_q4_wmma_32x16_2w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 4 waves per block (128 threads), 64M × 16N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 16); + dim3 block(128); + dim3 grid((size_n + 15) / 16, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x16_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x32_4w kernel: 4 waves per block, 64M × 32N tile, double-buffered LDS. +// +// Builds on v3 by doubling the N-tile from 16 → 32. Each wave now issues +// 2 wmmas per K-iter (one for cols 0-15, one for cols 16-31, sharing the +// same a_frag). With 4 waves × 2 wmmas = 8 wmmas in flight per K-iter, +// the wmma pipeline gets ~2× more in-flight work than v3 — targeting the +// remaining wmma issue gap on Qwen-class shapes (gate/up, down) where v3 +// plateaus at ~22 TFLOPS effective. +// +// Costs: +// * LDS B-tile doubles (2 × 16K × 32N × sizeof(T) = 2048 B for fp16/bf16). +// * Dequant doubles (32 N-cols × 2 K-octets = 64 slots/K-tile). Distributed +// across waves 0-1: wave 0 dequants n=[0..15], wave 1 dequants n=[16..31]. +// Waves 2-3 idle on dequant but do wmma work. +// * Per-wave registers: 2 × v8fp32 accumulator (16 fp32 = 32 VGPRs) plus +// b_frag0 + b_frag1 (32 VGPRs total). Within budget. +// +// Mapping invariant (wave-id → output tile slice): +// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15] +// * c_acc0 holds N cols [n_tile + 0 .. n_tile + 15] +// * c_acc1 holds N cols [n_tile + 16 .. n_tile + 31] +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x32_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 64; + const int n_tile = blockIdx.x * 32; // 32-col stride per block (was 16 in v3) + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + // Two accumulators per wave: c_acc0 covers cols [n_tile..n_tile+15], + // c_acc1 covers cols [n_tile+16..n_tile+31]. + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Doubled LDS B-tile: [buf][k][n_in_tile=0..31]. + __shared__ T b_lds[2][16][32]; + + // Dequant: 64 slots per K-tile (32 N-cols × 2 K-octets). Distributed + // across waves 0,1: wave w handles n_in_tile in [16w..16w+15], k_oct in + // {0,1}. Each dequanting wave has 32 lanes for 32 dequant slots — perfect + // mapping, identical layout to v2/v3 dequant per wave. Waves 2,3 stay idle + // during dequant; they catch up via the double buffer overlapping the + // wmma of iter k with dequant of iter k+1. + // + // Tried distributing dequant across all 4 waves (16 slots/wave): regressed + // 3% on gate/up M=2048 (357K → 347K tk/s). The dequant is not on the + // critical path; spreading it just adds LDS bank pressure with no gain. + auto dequant_into = [&](int buf, int k_tile) { + if (wave_id >= 2) return; + + const int my_n_local = lane_lo; // 0..15 (within wave's N-half) + const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..31 (in 32N tile) + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n_in_tile; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x; + b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y; + b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x; + b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y; + b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x; + b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y; + b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x; + b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its 16M slice, shared across both wmmas. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag0, b_frag1; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B for cols [0..15] and cols [16..31]. Both halves of the 32N tile. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag0[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo]); + b_frag1[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 16]); + } + + // Two wmmas per wave per K-iter, sharing a_frag. + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- + // Each wave owns 16M rows × 32N cols. c_acc0 → cols [n_tile..n_tile+15], + // c_acc1 → cols [n_tile+16..n_tile+31]. + const int m_tile_wave = m_tile + wave_id * 16; + + // Helper: store one v8fp32 accumulator's 8 outputs (covers 16 N-cols at + // n_base via lane_lo + interleaved M rows m_tile_wave + 2i + lane_hi). + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile); + store_acc(c_acc1, n_tile + 16); +} + +#else // non-RDNA3 device pass: empty kernel for symbol parity. +template +__global__ void gemm_q4_wmma_kernel_64x32_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +#endif + +template +void launch_gemm_q4_wmma_64x32_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v3 when M < 64 (small-M decode/prefill stays on the + // narrower 64M × 16N path) or when N < 32 (tile would waste a wave on + // out-of-range cols). + if (size_m < 64 || size_n < 32) { + launch_gemm_q4_wmma_64x16_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // 4 waves per block (128 threads), 64M × 32N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 32); + dim3 block(128); + dim3 grid((size_n + 31) / 32, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x32_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__) + +// =========================================================================== +// 64x64_4w kernel: 4 waves per block, 64M × 64N tile, 4 wmmas per wave per +// K-iter. +// +// Doubles the N-tile from 32 → 64. Each wave issues 4 wmmas per K-iter +// (cols 0-15, 16-31, 32-47, 48-63), all sharing the same a_frag. With +// 4 waves × 4 wmmas = 16 wmmas in flight per K-iter, the wmma pipeline +// is fully saturated (16-cycle latency × 1 issue/cycle = 16 wmmas in +// flight at peak). +// +// Costs: +// * LDS B-tile: 2 × 16K × 64N × sizeof(T) = 4 KB. Within budget. +// * Dequant: 64 N × 2 K-oct = 128 slots/K-tile. Distributed across all +// 4 waves: 32 slots/wave (16 N-cols × 2 K-octets per wave) — perfect +// 32-lane fit, full lane utilization on dequant. +// * Per-wave registers: 4 × v8fp32 acc (64 VGPRs) + 4 b_frag (64 VGPRs) +// + a_frag (16 VGPRs) ≈ 144 VGPRs/thread + locals ≈ ~170 total. +// Under the 192-VGPR gfx1100 cap. +// +// Mapping invariant: +// * Wave w produces M rows [m_tile + 16w .. m_tile + 16w + 15] +// * c_acc[i] holds N cols [n_tile + 16i .. n_tile + 16i + 15] for i=0..3 +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_64x64_4w( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 64; + const int n_tile = blockIdx.x * 64; // 64-col stride per block + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..127 + const int wave_id = tid >> 5; // 0..3 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + // Four accumulators per wave, each covering 16 N-cols. + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Larger LDS B-tile: [buf][k][n_in_tile=0..63]. + __shared__ T b_lds[2][16][64]; + + // Dequant: 128 slots per K-tile (64 N-cols × 2 K-octets). All 4 waves + // participate; wave w covers n_in_tile in [16w..16w+15], k_oct in {0,1} + // — 32 slots per wave (16 N-cols × 2 K-octets), perfect 32-lane fit. + auto dequant_into = [&](int buf, int k_tile) { + const int my_n_local = lane_lo; // 0..15 within wave + const int my_n_in_tile = wave_id * 16 + my_n_local; // 0..63 + const int my_k_octet = lane_hi; + const int actual_n = n_tile + my_n_in_tile; + + if (actual_n >= size_n) return; + + const int qk_row = (k_tile / 8) + my_k_octet; + const uint32_t qa = b_q[qk_row * size_n + actual_n]; + + const int g = k_tile / groupsize; + const int qz_idx = g * (size_n / 8) + actual_n / 8; + const int qz_shift = (actual_n & 7) * 4; + const uint32_t zero_v = + ((b_qzeros[qz_idx] >> qz_shift) & 0xF) + (uint32_t)zero_offset; + const T scale_t = b_scales[g * size_n + actual_n]; + + const int k_base = my_k_octet * 8; + + if constexpr (std::is_same::value) { + half2 z_prep, y_prep; + prep_zero_scale_fp16_precise(zero_v, scale_t, z_prep, y_prep); + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, z_prep, y_prep); + b_lds[buf][k_base + 0][my_n_in_tile] = __low2half(dq[0]); + b_lds[buf][k_base + 1][my_n_in_tile] = __high2half(dq[0]); + b_lds[buf][k_base + 2][my_n_in_tile] = __low2half(dq[1]); + b_lds[buf][k_base + 3][my_n_in_tile] = __high2half(dq[1]); + b_lds[buf][k_base + 4][my_n_in_tile] = __low2half(dq[2]); + b_lds[buf][k_base + 5][my_n_in_tile] = __high2half(dq[2]); + b_lds[buf][k_base + 6][my_n_in_tile] = __low2half(dq[3]); + b_lds[buf][k_base + 7][my_n_in_tile] = __high2half(dq[3]); + } else { + float z_f, y_f; + prep_zero_scale_bf16_f32(zero_v, scale_t, z_f, y_f); + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, z_f, y_f); + b_lds[buf][k_base + 0][my_n_in_tile] = dq[0].x; + b_lds[buf][k_base + 1][my_n_in_tile] = dq[0].y; + b_lds[buf][k_base + 2][my_n_in_tile] = dq[1].x; + b_lds[buf][k_base + 3][my_n_in_tile] = dq[1].y; + b_lds[buf][k_base + 4][my_n_in_tile] = dq[2].x; + b_lds[buf][k_base + 5][my_n_in_tile] = dq[2].y; + b_lds[buf][k_base + 6][my_n_in_tile] = dq[3].x; + b_lds[buf][k_base + 7][my_n_in_tile] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // Load A: each wave loads its 16M slice, shared across all 4 wmmas. + const int m_row = m_tile + wave_id * 16 + lane_lo; + V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; + if (m_row < size_m) { + const T* a_row = a + m_row * size_k; + if (b_q_perm) { + #pragma unroll + for (int i = 0; i < 16; i++) { + T v = a_row[b_q_perm[k_tile + i]]; + a_frag[i] = bitcast_elem(v); + } + } else { + static_assert(sizeof(a_frag) == 32, "V16 must be 32 bytes"); + __builtin_memcpy(&a_frag, a_row + k_tile, sizeof(a_frag)); + } + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Load B for all four 16-col halves. + #pragma unroll + for (int i = 0; i < 16; i++) { + b_frag0[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 0]); + b_frag1[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 16]); + b_frag2[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 32]); + b_frag3[i] = bitcast_elem(b_lds[cur_buf][i][lane_lo + 48]); + } + + // Four wmmas per wave per K-iter, sharing a_frag. + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16M × 64N. Helper writes one acc slice. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +// =========================================================================== +// 128x64_k16 kernel: 8 waves per block, 128M × 64N tile, K=16 per iteration. +// +// Doubles M-tile from 64 → 128. Each B-tile in LDS is reused by 8 waves +// (8 independent A-row slices) instead of 4, halving the effective B-load +// cost per output element. This matches Hybrid Triton's BLOCK_M=128. +// +// Dequant: same 128 slots as V5 (64N × 2 K-octets). Only waves 0-3 +// participate in dequant (same mapping). Waves 4-7 are pure compute. +// LDS: [2][16][64] × sizeof(T) — unchanged from V5. +// =========================================================================== + +template +__global__ void gemm_q4_wmma_kernel_128x64_k16( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 128; // 128-row M tile + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..255 + const int wave_id = tid >> 5; // 0..7 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // LDS with K-contiguous layout [buf][N][K] — enables vectorized ds_load_b128. + // Triton uses this layout to read 8 bf16 per instruction instead of 1. + __shared__ T b_lds[2][64][16]; + + // Dequant with scale/zero caching: 1 global_load per iter (7/8 of the time). + const bool dq_ok = + (wave_id < 4) && (n_tile + wave_id * 16 + lane_lo < size_n); + const int dq_n = wave_id * 16 + lane_lo; + const int dq_an = n_tile + dq_n; + const int dq_oct = lane_hi; + const int dq_kb = dq_oct * 8; + half2 ch_z = {}, ch_y = {}; + float cf_z = 0, cf_y = 0; + int cached_g = -1; + + auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) { + if (!dq_ok) return; + + // Reload scale/zero only on group boundary (every groupsize/16 iters). + const int g = k_tile / groupsize; + if (g != cached_g) { + cached_g = g; + const int qz_idx = g * (size_n / 8) + dq_an / 8; + const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) + + (uint32_t)zero_offset; + const T sc = b_scales[g * size_n + dq_an]; + if constexpr (std::is_same::value) + prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y); + else + prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y); + } + + const int qk_row = (k_tile / 8) + dq_oct; + const uint32_t qa = b_q[qk_row * size_n + dq_an]; + const int k_base = dq_kb; + + if constexpr (std::is_same::value) { + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y); + b_lds[buf][dq_n][k_base + 0] = __low2half(dq[0]); + b_lds[buf][dq_n][k_base + 1] = __high2half(dq[0]); + b_lds[buf][dq_n][k_base + 2] = __low2half(dq[1]); + b_lds[buf][dq_n][k_base + 3] = __high2half(dq[1]); + b_lds[buf][dq_n][k_base + 4] = __low2half(dq[2]); + b_lds[buf][dq_n][k_base + 5] = __high2half(dq[2]); + b_lds[buf][dq_n][k_base + 6] = __low2half(dq[3]); + b_lds[buf][dq_n][k_base + 7] = __high2half(dq[3]); + } else { + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y); + b_lds[buf][dq_n][k_base + 0] = dq[0].x; + b_lds[buf][dq_n][k_base + 1] = dq[0].y; + b_lds[buf][dq_n][k_base + 2] = dq[1].x; + b_lds[buf][dq_n][k_base + 3] = dq[1].y; + b_lds[buf][dq_n][k_base + 4] = dq[2].x; + b_lds[buf][dq_n][k_base + 5] = dq[2].y; + b_lds[buf][dq_n][k_base + 6] = dq[3].x; + b_lds[buf][dq_n][k_base + 7] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + const int m_row = m_tile + wave_id * 16 + lane_lo; + const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 16) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 16; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // A-load: vectorized 256-bit. b_q_perm branch removed from V7 hot path + // to eliminate ~450 ISA instructions of dead code (6× icache bloat). + V16 a_frag, b_frag0, b_frag1, b_frag2, b_frag3; + if (a_row_ptr) { + __builtin_memcpy(&a_frag, a_row_ptr + k_tile, sizeof(a_frag)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag[i] = (E)0; + } + + // Vectorized LDS reads: 32 bytes (16 bf16) per b_frag → ds_load_b128 × 2. + static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy"); + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32); + + c_acc0 = wmma_mma(a_frag, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Each wave owns 16M × 64N. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +// =========================================================================== +// 128x64_k32 kernel: K=32 per iteration, all 8 waves dequant. +// +// Same 128M × 64N tile as V7, but processes 32 K-elements per iteration +// instead of 16. Halves iteration count and __syncthreads() calls. +// +// Dequant mapping: +// Waves 0-3 dequant N[wave*16 +: 16] × K[0:15] (same as V7) +// Waves 4-7 dequant N[(wave-4)*16 +: 16] × K[16:31] (new) +// +// LDS layout: b_lds[2][64][34] — 2 extra padding elements per row to +// avoid 8-way bank conflicts (row stride 68 bytes gives 16 unique banks +// across the 16 lanes). +// +// Per iteration: 8 WMMAs (2 A-frags × 4 N-groups) vs V7's 4. +// Requires K divisible by 32 and groupsize ≥ 32. +// =========================================================================== +template +__global__ void gemm_q4_wmma_kernel_128x64_k32( + const T* __restrict__ a, const uint32_t* __restrict__ b_q, + const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales, + T* __restrict__ c, const int size_m, const int size_n, const int size_k, + const int groups, const int zero_offset, const int* __restrict__ b_q_perm) { + using E = typename WmmaNative::elem; + using V16 = typename WmmaNative::v16; + + const int m_tile = blockIdx.y * 128; + const int n_tile = blockIdx.x * 64; + if (m_tile >= size_m || n_tile >= size_n) return; + + const int tid = threadIdx.x; // 0..255 + const int wave_id = tid >> 5; // 0..7 + const int lane = tid & 31; + const int lane_lo = lane & 15; + const int lane_hi = lane >> 4; + + v8fp32 c_acc0 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc1 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc2 = {0, 0, 0, 0, 0, 0, 0, 0}; + v8fp32 c_acc3 = {0, 0, 0, 0, 0, 0, 0, 0}; + + const int groupsize = size_k / groups; + + const int k_per_split = size_k / gridDim.z; + const int k_start = blockIdx.z * k_per_split; + const int k_end = k_start + k_per_split; + + // Padded LDS: +2 elements per row to break bank conflicts. + // Row stride = 34 elements × 2 bytes = 68 bytes → 16 unique banks. + __shared__ T b_lds[2][64][34]; + + // All 8 waves dequant. Waves 0-3 fill K[0:15], waves 4-7 fill K[16:31]. + const int dq_wave4 = wave_id & 3; // 0-3 for both halves + const int dq_k_half = (wave_id >= 4) ? 1 : 0; + const int dq_n = dq_wave4 * 16 + lane_lo; // N position 0..63 + const int dq_an = n_tile + dq_n; + const bool dq_ok = (dq_an < size_n); + const int dq_oct = lane_hi + dq_k_half * 2; // K octet 0-3 + const int dq_kb = dq_oct * 8; // K base 0,8,16,24 + half2 ch_z = {}, ch_y = {}; + float cf_z = 0, cf_y = 0; + int cached_g = -1; + + auto dequant_into = [&](int buf, int k_tile) __attribute__((always_inline)) { + if (!dq_ok) return; + + const int g = k_tile / groupsize; + if (g != cached_g) { + cached_g = g; + const int qz_idx = g * (size_n / 8) + dq_an / 8; + const uint32_t zero_v = ((b_qzeros[qz_idx] >> ((dq_an & 7) * 4)) & 0xF) + + (uint32_t)zero_offset; + const T sc = b_scales[g * size_n + dq_an]; + if constexpr (std::is_same::value) + prep_zero_scale_fp16_precise(zero_v, sc, ch_z, ch_y); + else + prep_zero_scale_bf16_f32(zero_v, sc, cf_z, cf_y); + } + + const int qk_row = (k_tile / 8) + dq_oct; + const uint32_t qa = b_q[qk_row * size_n + dq_an]; + + if constexpr (std::is_same::value) { + half2 dq[4]; + dequant_4bit_8_fp16_precise(qa, dq, ch_z, ch_y); + b_lds[buf][dq_n][dq_kb + 0] = __low2half(dq[0]); + b_lds[buf][dq_n][dq_kb + 1] = __high2half(dq[0]); + b_lds[buf][dq_n][dq_kb + 2] = __low2half(dq[1]); + b_lds[buf][dq_n][dq_kb + 3] = __high2half(dq[1]); + b_lds[buf][dq_n][dq_kb + 4] = __low2half(dq[2]); + b_lds[buf][dq_n][dq_kb + 5] = __high2half(dq[2]); + b_lds[buf][dq_n][dq_kb + 6] = __low2half(dq[3]); + b_lds[buf][dq_n][dq_kb + 7] = __high2half(dq[3]); + } else { + bf162_t dq[4]; + dequant_4bit_8_bf16_to_bf16(qa, dq, cf_z, cf_y); + b_lds[buf][dq_n][dq_kb + 0] = dq[0].x; + b_lds[buf][dq_n][dq_kb + 1] = dq[0].y; + b_lds[buf][dq_n][dq_kb + 2] = dq[1].x; + b_lds[buf][dq_n][dq_kb + 3] = dq[1].y; + b_lds[buf][dq_n][dq_kb + 4] = dq[2].x; + b_lds[buf][dq_n][dq_kb + 5] = dq[2].y; + b_lds[buf][dq_n][dq_kb + 6] = dq[3].x; + b_lds[buf][dq_n][dq_kb + 7] = dq[3].y; + } + }; + + dequant_into(0, k_start); + __syncthreads(); + + int cur_buf = 0; + const int m_row = m_tile + wave_id * 16 + lane_lo; + const T* a_row_ptr = (m_row < size_m) ? (a + m_row * size_k) : nullptr; + + for (int k_tile = k_start; k_tile < k_end; k_tile += 32) { + const int next_buf = 1 - cur_buf; + const int k_next = k_tile + 32; + + if (k_next < k_end) { + dequant_into(next_buf, k_next); + } + + // A-load: two 16-element fragments for K[0:15] and K[16:31]. + V16 a_frag_lo, a_frag_hi; + V16 b_frag0, b_frag1, b_frag2, b_frag3; + if (a_row_ptr) { + __builtin_memcpy(&a_frag_lo, a_row_ptr + k_tile, sizeof(V16)); + __builtin_memcpy(&a_frag_hi, a_row_ptr + k_tile + 16, sizeof(V16)); + } else { + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_lo[i] = (E)0; + #pragma unroll + for (int i = 0; i < 16; i++) a_frag_hi[i] = (E)0; + } + + // --- Lower K half [0:15]: 4 WMMAs --- + static_assert(sizeof(V16) == 32, "V16 must be 32 bytes for memcpy"); + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][0], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][0], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][0], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][0], 32); + + c_acc0 = wmma_mma(a_frag_lo, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_lo, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_lo, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_lo, b_frag3, c_acc3); + + // --- Upper K half [16:31]: 4 WMMAs --- + __builtin_memcpy(&b_frag0, &b_lds[cur_buf][lane_lo + 0][16], 32); + __builtin_memcpy(&b_frag1, &b_lds[cur_buf][lane_lo + 16][16], 32); + __builtin_memcpy(&b_frag2, &b_lds[cur_buf][lane_lo + 32][16], 32); + __builtin_memcpy(&b_frag3, &b_lds[cur_buf][lane_lo + 48][16], 32); + + c_acc0 = wmma_mma(a_frag_hi, b_frag0, c_acc0); + c_acc1 = wmma_mma(a_frag_hi, b_frag1, c_acc1); + c_acc2 = wmma_mma(a_frag_hi, b_frag2, c_acc2); + c_acc3 = wmma_mma(a_frag_hi, b_frag3, c_acc3); + + __syncthreads(); + cur_buf = next_buf; + } + + // ---- Store C ---- Same as V7. + const int m_tile_wave = m_tile + wave_id * 16; + auto store_acc = [&](const v8fp32& acc, int n_base) { + if (gridDim.z > 1) { + const bool is_even_lane = (lane_lo & 1) == 0; + const int out_n_pair = n_base + lane_lo; + #pragma unroll + for (int i = 0; i < 8; i++) { + float other_f = __shfl_xor(acc[i], 1); + if (!is_even_lane) continue; + + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m >= size_m || out_n_pair >= size_n) continue; + + T* dst = c + out_m * size_n + out_n_pair; + if constexpr (std::is_same::value) { + half2 packed = + __halves2half2(__float2half_rn(acc[i]), __float2half_rn(other_f)); + atomic_add_pk_f16(reinterpret_cast(dst), packed); + } else { + bf162_t packed; + packed.x = __float2bfloat16(acc[i]); + packed.y = __float2bfloat16(other_f); + atomic_add_pk_bf16(reinterpret_cast(dst), packed); + } + } + } else { + const int out_n = n_base + lane_lo; + if (out_n >= size_n) return; + #pragma unroll + for (int i = 0; i < 8; i++) { + const int out_m = m_tile_wave + 2 * i + lane_hi; + if (out_m < size_m) { + T* dst = c + out_m * size_n + out_n; + if constexpr (std::is_same::value) { + *dst = __float2half_rn(acc[i]); + } else { + *dst = __float2bfloat16(acc[i]); + } + } + } + } + }; + + store_acc(c_acc0, n_tile + 0); + store_acc(c_acc1, n_tile + 16); + store_acc(c_acc2, n_tile + 32); + store_acc(c_acc3, n_tile + 48); +} + +#else // non-RDNA3 device pass: empty kernels for symbol parity (covers the + // three kernels that share this launcher). +template +__global__ void gemm_q4_wmma_kernel_64x64_4w(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, const int*) { +} +template +__global__ void gemm_q4_wmma_kernel_128x64_k16(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, + const int*) {} +template +__global__ void gemm_q4_wmma_kernel_128x64_k32(const T*, const uint32_t*, + const uint32_t*, const T*, T*, + const int, const int, const int, + const int, const int, + const int*) {} +#endif + +template +void launch_gemm_q4_wmma_64x64_4w(const T* a, const uint32_t* b_q_weight, + const uint32_t* b_qzeros, const T* b_scales, + const int* b_q_perm, T* c, int size_m, + int size_n, int size_k, int groups, + int zero_offset, cudaStream_t stream) { + // Fall back to v4 when N < 64 (would waste 1+ waves on out-of-range cols). + if (size_m < 64 || size_n < 64) { + launch_gemm_q4_wmma_64x32_4w(a, b_q_weight, b_qzeros, b_scales, b_q_perm, + c, size_m, size_n, size_k, groups, + zero_offset, stream); + return; + } + + // V8 (128M × 64N, K=32/iter, 8-wave dequant) when K%32==0 and gs≥32. + // Falls back to V7 otherwise. V7/V8 read A sequentially, so act-order + // (b_q_perm != null) must skip them and use v5, which honors the perm. + if (size_m >= 128 && b_q_perm == nullptr) { + const int k_split = + compute_wmma_k_split_mn(size_m, size_n, size_k, 128, 64); + const int groupsize = size_k / groups; + dim3 block(256); + dim3 grid((size_n + 63) / 64, (size_m + 127) / 128, k_split); + if (size_k % 32 == 0 && groupsize >= 32 && (size_k / k_split) % 32 == 0) { + gemm_q4_wmma_kernel_128x64_k32<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); + } else { + gemm_q4_wmma_kernel_128x64_k16<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); + } + return; + } + + // 4 waves per block (128 threads), 64M × 64N tile per block. + const int k_split = compute_wmma_k_split_mn(size_m, size_n, size_k, 64, 64); + dim3 block(128); + dim3 grid((size_n + 63) / 64, (size_m + 63) / 64, k_split); + gemm_q4_wmma_kernel_64x64_4w<<>>( + a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups, + zero_offset, b_q_perm); +} + +} // namespace gptq_rdna3_wmma +} // namespace vllm + +// --------------------------------------------------------------------------- +// Public entry point. +// --------------------------------------------------------------------------- +// +// Inputs: +// a [M, K] half or bfloat16 +// b_q_weight[K/8, N] uint32 (already shuffled via gptq_shuffle) +// b_qzeros [groups, N/8] uint32 (packed 4-bit zeros) +// b_scales [groups, N] half or bfloat16 +// b_g_idx [K] or empty int32 (act-order permutation; empty=identity) +// use_v2_format bool (true = GPTQv2, no +1 zero offset) +// +// Output: +// c [M, N] same dtype as a +// +// Requirements: +// * size_m >= 16 (otherwise prefer the scalar gptq_gemm_rdna3 op) +// * size_n % 16 == 0 (WMMA tile size) +// * size_k % 16 == 0 (WMMA tile size) + +torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight, + torch::Tensor b_qzeros, + torch::Tensor b_scales, + torch::Tensor b_g_idx, bool use_v2_format) { + TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor"); + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor"); + TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor"); + TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor"); + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]"); + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "a must be half or bfloat16"); + TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(), + "b_scales dtype must match a"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto stream = at::cuda::getCurrentCUDAStream(); + + int size_m = (int)a.size(0); + int size_k = (int)a.size(1); + int size_n = (int)b_q_weight.size(1); + int groups = (int)b_qzeros.size(0); + + TORCH_CHECK(b_q_weight.size(0) * 8 == size_k, + "b_q_weight first dim must be K/8"); + TORCH_CHECK(b_scales.size(0) == groups, + "b_scales must have same group count as qzeros"); + TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N"); + TORCH_CHECK(size_n % 16 == 0, "WMMA path requires N % 16 == 0"); + TORCH_CHECK(size_k % 16 == 0, "WMMA path requires K % 16 == 0"); + + auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device()); + // Always zero-init the output: some V3-V8 boundary threads may exit + // without writing their output cell (e.g. out_m >= size_m), leaving + // uninitialized garbage when torch::empty is used. The cost is + // negligible (< 1.5% of prefill time on gfx1100). + at::Tensor c = torch::zeros({size_m, size_n}, opts); + + const int* g_idx_ptr = nullptr; + if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) { + TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32, + "b_g_idx must be int32"); + g_idx_ptr = (const int*)b_g_idx.data_ptr(); + } + + const int zero_offset = use_v2_format ? 0 : 1; + + // launch_gemm_q4_wmma_64x64_4w dispatches: + // M >= 128 → 128x64_k32 / 128x64_k16 (8 waves, K=32/16 per iter) + // 64 <= M < 128 & N >= 64 → 64x64_4w (4 waves, 4 wmma/wave/K-iter) + // M >= 64 && 32 <= N < 64 → 64x32_4w (4 waves, 2 wmma/wave/K-iter) + // M >= 64 && N < 32 → 64x16_4w (4 waves) + // 32 <= M < 64 → 32x16_2w (2 waves) + // M < 32 → 16x16_1w (1 wave) + if (a.scalar_type() == torch::kHalf) { + vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w( + (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(), + g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups, + zero_offset, stream); + } else { + vllm::gptq_rdna3_wmma::launch_gemm_q4_wmma_64x64_4w< + vllm::gptq_rdna3_wmma::bf16_t>( + (const vllm::gptq_rdna3_wmma::bf16_t*)a.data_ptr(), + (const uint32_t*)b_q_weight.data_ptr(), + (const uint32_t*)b_qzeros.data_ptr(), + (const vllm::gptq_rdna3_wmma::bf16_t*)b_scales.data_ptr(), g_idx_ptr, + (vllm::gptq_rdna3_wmma::bf16_t*)c.data_ptr(), size_m, size_n, size_k, + groups, zero_offset, stream); + } + + return c; +} diff --git a/csrc/rocm/qdq_4_rdna3.cuh b/csrc/rocm/qdq_4_rdna3.cuh new file mode 100644 index 00000000000..f4be668c1f0 --- /dev/null +++ b/csrc/rocm/qdq_4_rdna3.cuh @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on +// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses +// the classic exllamav2 bit-trick: +// +// (qa & 0x000F000F) | 0x64006400 -> half2(1024+q_lo, 1024+q_hi) +// (qa & 0x00F000F0) | 0x64006400 -> half2(1024+q_lo*16, 1024+q_hi*16) +// +// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs +// works in fp16 because the mantissa (10 bits) is wide enough to hold a value +// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper +// nibble into bits [7:4] would spill into the exponent. To avoid that, the +// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a +// single right-shift before the OR with 0x43004300 (= bf162(128, 128)). + +#ifndef _qdq_4_rdna3_cuh +#define _qdq_4_rdna3_cuh + +#include + +#include +#include + +namespace vllm { +namespace gptq_rdna3 { + +using bf16_t = __hip_bfloat16; +using bf162_t = __hip_bfloat162; + +// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]: +// in: q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0] (LSB first) +// out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0] (even/odd interleaved) +// +// After shuffle, q[2k] sits at bits [4k : 4k+3] (lower 16) +// q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16) +// so a single mask 0x000F000F selects the matching even/odd pair, ready to +// bitcast to half2 / bfloat162 after OR-ing with the magic constant. +__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) { + uint32_t qa = q[0]; + uint32_t qb = 0; +#pragma unroll + for (int i = 0; i < 4; i++) { + uint32_t qa0 = qa & 0x0F; + uint32_t qa1 = (qa & 0xF0) >> 4; + qa >>= 8; + qb |= (qa1 << (i * 4 + 16)); + qb |= (qa0 << (i * 4)); + } + q[0] = qb; +} + +// --------------------------------------------------------------------------- +// fp16 path +// --------------------------------------------------------------------------- + +// Precompute scale-baked constants for a single zero/scale pair. +// z1z16[0] = scale * (-1024 - zero) (used for "low" pairs) +// z1z16[1] = scale * (-64 - zero) (used for "high" pairs) +// y1y16[0] = scale * 1 (low pairs are q + 1024) +// y1y16[1] = scale * (1/16) (high pairs are q*16 + 1024) +__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale, + half2 (&z1z16)[2], + half2 (&y1y16)[2]) { + // half(-1024 - zero) via the exllamav2 bit-trick: + // half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it. + union { + uint16_t u; + half h; + } z1u; + z1u.u = (uint16_t)(0xE400 | zero); + half z1 = z1u.h; + half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero)); + + half2 scale2 = __half2half2(scale); + z1z16[0] = __hmul2(scale2, __half2half2(z1)); + z1z16[1] = __hmul2(scale2, __half2half2(z16)); + + half y1 = __float2half_rn(1.0f); + half y16 = __float2half_rn(1.0f / 16.0f); + y1y16[0] = __hmul2(scale2, __half2half2(y1)); + y1y16[1] = __hmul2(scale2, __half2half2(y16)); +} + +// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs: +// dq[0] = (q[0], q[1]) * scale - zero*scale +// dq[1] = (q[2], q[3]) * scale - zero*scale +// dq[2] = (q[4], q[5]) * scale - zero*scale +// dq[3] = (q[6], q[7]) * scale - zero*scale +__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4], + half2 (&z1z16)[2], + half2 (&y1y16)[2]) { + const uint32_t c0 = 0x64006400; + + union { + uint32_t u; + half2 h2; + } q0, q1, q2, q3; + q0.u = (qa & 0x000F000F) | c0; // half2(q[0]+1024, q[1]+1024) + q1.u = (qa & 0x00F000F0) | c0; // half2(q[2]*16+1024, q[3]*16+1024) + uint32_t qa_hi = qa >> 8; + q2.u = (qa_hi & 0x000F000F) | c0; // half2(q[4]+1024, q[5]+1024) + q3.u = (qa_hi & 0x00F000F0) | c0; // half2(q[6]*16+1024, q[7]*16+1024) + + dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]); + dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]); + dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]); + dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]); +} + +// --------------------------------------------------------------------------- +// bf16 path +// --------------------------------------------------------------------------- + +// Bit-trick magic for bf16: +// bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0). +// For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because +// bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So +// ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi). +// +// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble +// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16] +// via a single 4/8/12-bit right-shift before the OR. That costs one extra +// shift per pair vs fp16, but keeps the FMA structure identical. +__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero, + bf16_t scale, + bf162_t& z_prep, + bf162_t& y_prep) { + // z = scale * -(128 + zero); y = scale. + float scale_f = __bfloat162float(scale); + float zf = -(128.0f + (float)zero) * scale_f; + bf16_t zb = __float2bfloat16(zf); + z_prep = __bfloat162bfloat162(zb); + y_prep = __bfloat162bfloat162(scale); +} + +__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa, + bf162_t (&dq)[4], + bf162_t z_prep, + bf162_t y_prep) { + const uint32_t c0 = 0x43004300; + + union { + uint32_t u; + bf162_t b2; + } q0, q1, q2, q3; + q0.u = ((qa >> 0) & 0x000F000F) | c0; // bf162(128+q[0], 128+q[1]) + q1.u = ((qa >> 4) & 0x000F000F) | c0; // bf162(128+q[2], 128+q[3]) + q2.u = ((qa >> 8) & 0x000F000F) | c0; // bf162(128+q[4], 128+q[5]) + q3.u = ((qa >> 12) & 0x000F000F) | c0; // bf162(128+q[6], 128+q[7]) + + // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale + dq[0] = __hfma2(q0.b2, y_prep, z_prep); + dq[1] = __hfma2(q1.b2, y_prep, z_prep); + dq[2] = __hfma2(q2.b2, y_prep, z_prep); + dq[3] = __hfma2(q3.b2, y_prep, z_prep); +} + +// --------------------------------------------------------------------------- +// bf16-input → fp32-output dequant (RDNA3 scalar path). +// +// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow +// fallback. Rather than computing dq in bf16 and widening at FMA time in +// the dot product, we widen to fp32 here once (a free left-shift by 16) and +// emit the (q - zero) * scale FMA directly in fp32. This: +// * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32. +// * Eliminates 4× bf16→fp32 widens that the dot product would do. +// * Keeps the dot product accumulator in fp32 without a roundtrip. +// +// Output: fp32 dq[8], one element per K position (consumed by the +// fp32-overload of dot22_8_f in q_gemm_rdna3.cu). +__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero, + bf16_t scale, + float& z_prep, + float& y_prep) { + float scale_f = __bfloat162float(scale); + z_prep = -(128.0f + (float)zero) * scale_f; + y_prep = scale_f; +} + +// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32 +// values 128+nibble, without folding scale/zero. The caller folds scale/zb +// into the accumulator outside the inner loop using a precomputed sum_a, +// which saves ~27% of the FMA count vs the per-col-dequant approach above +// (only beneficial at M_COUNT=1; break-even at M_COUNT=2). +// +// Cost: 0 FMAs (pure bit-trick + as_float reinterprets). +__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa, + float (&q_f32)[8]) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16); + q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u); + q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16); + q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u); + q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16); + q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u); + q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16); + q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u); +} + +__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa, + float (&dq)[8], + float z_prep, + float y_prep) { + const uint32_t c0 = 0x43004300; + const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0; + const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0; + const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0; + const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0; + // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16 + // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved). + const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16); + const float q0y = __uint_as_float(q0 & 0xFFFF0000u); + const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16); + const float q1y = __uint_as_float(q1 & 0xFFFF0000u); + const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16); + const float q2y = __uint_as_float(q2 & 0xFFFF0000u); + const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16); + const float q3y = __uint_as_float(q3 & 0xFFFF0000u); + // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale + dq[0] = __fmaf_rn(q0x, y_prep, z_prep); + dq[1] = __fmaf_rn(q0y, y_prep, z_prep); + dq[2] = __fmaf_rn(q1x, y_prep, z_prep); + dq[3] = __fmaf_rn(q1y, y_prep, z_prep); + dq[4] = __fmaf_rn(q2x, y_prep, z_prep); + dq[5] = __fmaf_rn(q2y, y_prep, z_prep); + dq[6] = __fmaf_rn(q3x, y_prep, z_prep); + dq[7] = __fmaf_rn(q3y, y_prep, z_prep); +} + +} // namespace gptq_rdna3 +} // namespace vllm + +#endif // _qdq_4_rdna3_cuh diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index b0b44964c24..1e589598c74 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -39,6 +39,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#ifdef VLLM_ROCM_GFX1100 + // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). + rocm_ops.def( + "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " + "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3); + + rocm_ops.def( + "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " + "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor"); + rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma); +#endif + // Custom attention op // Compute the attention between an input query and the cached // keys/values using PagedAttention. diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a2a5d809745..01869474e0f 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1,4 +1,7 @@ -#include "cache.h" +// Provides torch::Tensor for ops.h (previously included transitively via +// cache.h, which is no longer included here after cache ops moved to +// _C_stable_libtorch). +#include #include "cuda_utils.h" #include "ops.h" #include "core/registration.h" @@ -33,35 +36,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU, &get_cuda_view_from_cpu_tensor); - // Attention ops - // Compute the attention between an input query and the cached - // keys/values using PagedAttention. - ops.def( - "paged_attention_v1(" - " Tensor! out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1); - - // PagedAttention V2. - ops.def( - "paged_attention_v2(" - " Tensor! out, Tensor! exp_sums, Tensor! max_logits," - " Tensor! tmp_out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2); - // Activation ops (quantized only — basic ops moved to _C_stable_libtorch) ops.def( "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); @@ -217,114 +191,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif } -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) { - // Cache ops - // Swap in (out) the cache blocks from src to dst. - cache_ops.def( - "swap_blocks(Tensor src, Tensor! dst," - " int block_size_in_bytes, Tensor block_mapping) -> ()"); - cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks); - - // Batch swap: submit all block copies in a single driver call. - cache_ops.def( - "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," - " Tensor sizes," - " bool is_src_access_order_any=False) -> ()"); - cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch); - - // Reshape the key and value tensors and cache them. - cache_ops.def( - "reshape_and_cache(Tensor key, Tensor value," - " Tensor! key_cache, Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache); - - // Reshape the key and value tensors and cache them. - cache_ops.def( - "reshape_and_cache_flash(Tensor key, Tensor value," - " Tensor! key_cache," - " Tensor! value_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor k_scale, Tensor v_scale) -> ()"); - cache_ops.impl("reshape_and_cache_flash", torch::kCUDA, - &reshape_and_cache_flash); - - // Concat kv_c and k_pe and cache them. - cache_ops.def( - "concat_and_cache_mla(Tensor kv_c, Tensor k_pe," - " Tensor! kv_cache," - " Tensor slot_mapping," - " str kv_cache_dtype," - " Tensor scale) -> ()"); - cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla); - - // Rotate Q and K, then write to kv cache for MLA - cache_ops.def( - "concat_and_cache_mla_rope_fused(" - " Tensor positions," - " Tensor! q_pe," - " Tensor! k_pe," - " Tensor kv_c," - " Tensor cos_sin_cache," - " bool is_neox," - " Tensor slot_mapping," - " Tensor! kv_cache," - " str kv_cache_dtype," - " Tensor kv_cache_scale) -> ()"); - cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA, - &concat_and_cache_mla_rope_fused); - - // Convert the key and value cache to fp8 data type. - cache_ops.def( - "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, " - "str kv_cache_dtype) -> ()"); - cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8); - - // Gather cache blocks from src_cache to dst, dequantizing from - // src_cache's dtype to dst's dtype if necessary. - cache_ops.def( - "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, " - " Tensor block_table, Tensor cu_seq_lens, " - " Tensor token_to_seq, " - " int num_tokens, " - " str kv_cache_dtype, " - " Tensor scale, Tensor? seq_starts) -> ()"); - cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA, - &gather_and_maybe_dequant_cache); - - cache_ops.def( - "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, " - "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()"); - cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache); - - cache_ops.def( - "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int " - "batch_size) -> ()"); - cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA, - &cp_gather_and_upconvert_fp8_kv_cache); - - cache_ops.def( - "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " - "slot_mapping, " - "int quant_block_size, str kv_cache_dtype) -> ()"); - cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA, - &indexer_k_quant_and_cache); - - cache_ops.def( - "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()"); - cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q); - - cache_ops.def( - "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! " - "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()"); - cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA, - &cp_gather_indexer_k_quant_cache); -} - TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { // Cuda utils 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.cpu b/docker/Dockerfile.cpu index c2ce0e88f29..e185c00cb2f 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -165,6 +165,23 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \ VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 +######################### TRITON-CPU BUILD IMAGE ######################### +FROM base AS vllm-triton-cpu-build + +WORKDIR /vllm-workspace + +RUN mkdir dist + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/ccache \ + --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \ + if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \ + git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \ + cd triton-cpu; \ + git checkout "270e696d"; \ + uv build --wheel --out-dir=../dist; \ + fi + ######################### TEST DEPS ######################### FROM base AS vllm-test-deps @@ -257,7 +274,14 @@ WORKDIR /vllm-workspace RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/ccache \ --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \ - uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]" + uv pip install "$(realpath dist/*.whl)[audio]" + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/ccache \ + --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \ + if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \ + uv pip install "$(realpath dist/*.whl)"; \ + fi # Add labels to document build configuration LABEL org.opencontainers.image.title="vLLM CPU" 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/configuration/optimization.md b/docs/configuration/optimization.md index eb6bdce37b9..5bf789a0919 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -46,14 +46,14 @@ In V1, **chunked prefill is enabled by default whenever possible**. With chunked This policy has two benefits: -- It improves ITL and generation decode because decode requests are prioritized. +- It improves inter-token latency (ITL) and generation decode because decode requests are prioritized. - It helps achieve better GPU utilization by locating compute-bound (prefill) and memory-bound (decode) requests to the same batch. ### Performance Tuning with Chunked Prefill You can tune the performance by adjusting `max_num_batched_tokens`: -- Smaller values (e.g., 2048) achieve better inter-token latency (ITL) because there are fewer prefills slowing down decodes. +- Smaller values (e.g., 2048) achieve better ITL because there are fewer prefills slowing down decodes. - Higher values achieve better time to first token (TTFT) as you can process more prefill tokens in a batch. - For optimal throughput, we recommend setting `max_num_batched_tokens > 8192` especially for smaller models on large GPUs. - If `max_num_batched_tokens` is the same as `max_model_len`, that's almost the equivalent to the V0 default scheduling policy (except that it still prioritizes decodes). 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/attention_backends.md b/docs/design/attention_backends.md index ed9f0d30162..329a4aacfb6 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first). | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A | | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any | 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/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index f6d4f3f86d8..847743dfff1 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows: base_url=openai_api_base, ) - video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4" + video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" ## Use video url in the payload chat_completion_from_url = client.chat.completions.create( diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index e8b74a06f07..f6cd88b97fc 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -41,7 +41,7 @@ pip install -v -r requirements/xpu.txt ```bash pip uninstall -y triton triton-xpu - pip install triton-xpu==3.6.0 --extra-index-url https://download.pytorch.org/whl/xpu + pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu ``` !!! note diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 386e4f2a4bb..738c59df445 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -17,6 +17,7 @@ Sorted alphabetically by GitHub handle: - [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser - [@benchislett](https://github.com/benchislett): Engine core and spec decode - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration +- [@BugenZhao](https://github.com/BugenZhao): Rust frontend - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser - [@DarkLight1337](https://github.com/DarkLight1337): Multimodality, API server - [@esmeetu](https://github.com/esmeetu): developer marketing, community @@ -130,6 +131,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r - @DarkLight1337 - API Server: The OpenAI-compatible API server - @DarkLight1337, @njhill, @aarnphm, @simon-mo, @heheda12345 (Responses API) +- Rust Frontend: The experimental API server in Rust + - @BugenZhao, @njhill - Batch Runner: The OpenAI-compatible batch runner - @simon-mo diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1f38200a786..4612b4c423f 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -428,7 +428,6 @@ th { | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | | | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | | -| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ | | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ | | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ | @@ -438,6 +437,7 @@ th { | `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ | | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | +| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ | | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ | | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ | @@ -634,6 +634,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | | | `Step3VLForConditionalGeneration` | Step3-VL | T + I+ | `stepfun-ai/step3` | | ✅︎ | | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I+ | `stepfun-ai/Step3-VL-10B` | | ✅︎ | +| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I+ | `stepfun-ai/Step-3.7-Flash` | | ✅︎ | | `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ | | `Tarsier2ForConditionalGeneration`^ | Tarsier2 | T + IE+ + VE+ | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ | | `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ | diff --git a/docs/pre_run_check.sh b/docs/pre_run_check.sh index de93f82faf1..464766c42ec 100644 --- a/docs/pre_run_check.sh +++ b/docs/pre_run_check.sh @@ -1,41 +1,60 @@ -if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then - MAX_WAIT=300 - INTERVAL=60 - ELAPSED=0 - while :; do - RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest") - HTTP_CODE=$(printf %s "$RAW" | tail -n1) - BODY=$(printf %s "$RAW" | sed '$d') - if [ "$HTTP_CODE" != "200" ]; then - echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate." - break - fi - STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")") - CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")") - CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")") - if [ "$STATUS" = "none" ]; then - echo "no pre-run-check found for this commit; skipping gate." - break - fi - if [ -n "$CONCLUSION" ]; then - echo "pre-run-check conclusion: $CONCLUSION" - if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then - echo "pre-run-check did not pass; skipping docs build." - if [ -n "$CHECK_URL" ]; then - echo "pre-run-check failure reason: $CHECK_URL" - fi - exit 1 - fi - break - fi - if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then - echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate." - break - fi - echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..." - sleep "$INTERVAL" - ELAPSED=$((ELAPSED + INTERVAL)) - done -else +if [ "$READTHEDOCS_VERSION_TYPE" != "external" ]; then echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate." -fi \ No newline at end of file + exit 0 +fi + +echo "Checking for changes to docs-affecting files vs origin/main..." +DOCS_PATHS=( + docs/ # Actual docs content + examples/ # Examples are rendered in docs + vllm/ # API & CLI reference + requirements/test/cuda.txt # CLI reference (see docs/mkdocs/hooks/generate_argparse.py) + mkdocs.yaml # Affects build process + .readthedocs.yaml # Affects build process + requirements/docs.txt # Affects build process + requirements/docs.in # Affects build process +) +if git diff --quiet origin/main -- "${DOCS_PATHS[@]}"; then + echo "No docs-affecting files changed vs origin/main; cancelling build." + # See https://docs.readthedocs.com/platform/latest/guides/build/skip-build.html for info on exit code + exit 183 +fi +echo "Docs-affecting files changed; continuing pre-run-check." +echo "Checking pre-commit/pre-run-check status..." +MAX_WAIT=300 +INTERVAL=60 +ELAPSED=0 +while :; do + RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest") + HTTP_CODE=$(printf %s "$RAW" | tail -n1) + BODY=$(printf %s "$RAW" | sed '$d') + if [ "$HTTP_CODE" != "200" ]; then + echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-commit/pre-run-check gate." + break + fi + STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")") + CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")") + CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")") + if [ "$STATUS" = "none" ]; then + echo "no pre-commit/pre-run-check found for this commit; skipping gate." + break + fi + if [ -n "$CONCLUSION" ]; then + echo "pre-commit/pre-run-check conclusion: $CONCLUSION" + if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then + echo "pre-commit/pre-run-check did not pass; skipping docs build." + if [ -n "$CHECK_URL" ]; then + echo "pre-commit/pre-run-check failure reason: $CHECK_URL" + fi + exit 1 + fi + break + fi + if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then + echo "pre-commit/pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate." + break + fi + echo "pre-commit/pre-run-check status=$STATUS; waiting ${INTERVAL}s..." + sleep "$INTERVAL" + ELAPSED=$((ELAPSED + INTERVAL)) +done diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index fef4df770fa..b7c2ee87375 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -151,7 +151,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `step_interval` | Frequency of rebalancing (every N engine steps) | 3000 | | `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` | | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | -| `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | +| `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` | | `policy` | The policy type for expert parallel load balancing | `"default"` | | `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` | 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/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index bfde1ee2ae3..7b531218568 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/docs/training/weight_transfer/nccl.md @@ -84,7 +84,10 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT ## Receiving Weights (Inference Side) -The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are: +The inference side triggers weight reception using the four-phase protocol: +`init_weight_transfer_engine`, `start_weight_update`, `update_weights`, +`finish_weight_update`. The init phase is shown [above](#initialization). The +remaining three steps are: ```python from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest @@ -108,12 +111,24 @@ llm.update_weights( llm.finish_weight_update() ``` -The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters. +The `names`, `dtype_names`, and `shapes` lists describe each parameter. These +must match the order in which the trainer iterates over its parameters. -`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights). +`start_weight_update` must be called before `update_weights`, and +`finish_weight_update` must be called after all weight chunks have been +transferred. The `is_checkpoint_format` flag controls whether layerwise reload +processing is applied (`True` for checkpoint-format weights, `False` for +pre-processed kernel-format weights). + +Sparse NCCL patches still use `update_kind="sparse_flat"` inside +`update_info`, but they should be wrapped in +`start_weight_update(is_checkpoint_format=False)` because sparse patches apply +directly to runtime/kernel-format parameters. The current sparse MVP requires +`TP=1` and `PP=1`. ## Examples - [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast +- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1` - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model - [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane diff --git a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py index 3a007731c74..e1c4fd76d7d 100644 --- a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py +++ b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py @@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None: # Video input inference def run_video(model: str, max_completion_tokens: int) -> None: - video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4" + video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" video_base64 = encode_base64_content_from_url(video_url) ## Use video url in the payload diff --git a/examples/rl/rlhf_sparse_nccl.py b/examples/rl/rlhf_sparse_nccl.py new file mode 100644 index 00000000000..bddd28b6485 --- /dev/null +++ b/examples/rl/rlhf_sparse_nccl.py @@ -0,0 +1,526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Demonstrates dense-vs-sparse NCCL weight syncing with a real model. + +This example mirrors the validation story used for the sparse NCCL MVP: +both the dense update path and the sparse patch path start from the same real +checkpoint and apply the same deterministic trainer-side patch. The script then +checks that greedy 1-token outputs match between the dense and sparse vLLM +engines after the update. + +The example performs the following steps: +* Load a training model on one GPU via a Ray actor. +* Launch a vLLM engine with the same real model on a second GPU. +* Verify trainer vs vLLM baseline agreement before any update. +* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer. +* Run a dense NCCL update into a fresh vLLM engine and collect post-update + outputs. +* Reset the trainer back to the baseline checkpoint. +* Apply the same deterministic patch again. +* Run a sparse NCCL update into another fresh vLLM engine and collect + post-update outputs. +* Compare dense vs sparse baseline outputs, dense vs sparse post-update + outputs, estimated payload sizes, and trainer-side send times. + +Current sparse weight transfer MVP limitations: +* ``TP=1`` and ``PP=1`` only +* sparse updates use runtime/kernel-format parameter names +* sparse updates are not composable with checkpoint-format or packed updates + +This example assumes a single-node cluster with two GPUs. +""" + +import hashlib +import os +import time +from collections.abc import Sequence + +import ray +import torch +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from transformers import AutoModelForCausalLM, AutoTokenizer + +from vllm import LLM, SamplingParams +from vllm.config import WeightTransferConfig +from vllm.distributed.weight_transfer.base import SparseWeightPatch +from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerSendWeightsArgs, + NCCLWeightTransferEngine, +) +from vllm.utils.network_utils import get_ip, get_open_port + +MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" +PATCHED_PARAM_NAME = "model.embed_tokens.weight" +MAX_PATCH_ROWS = 32 +PROMPTS = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", +] +SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1) + + +class MyLLM(LLM): + """Configure the vLLM worker for Ray placement group execution.""" + + def __init__(self, *args, **kwargs): + os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0" + super().__init__(*args, **kwargs) + + +@ray.remote(num_gpus=1) +class TrainModel: + """Ray actor that owns the trainer-side model and deterministic patch state.""" + + def __init__(self, model_name: str): + self.model_name = model_name + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + if self.tokenizer.pad_token_id is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + self.model = None + self.patched_param = None + self.pending_sparse_patches: list[SparseWeightPatch] | None = None + self.model_update_group = None + self.master_address = get_ip() + self.port = get_open_port() + self.reset_model() + + def reset_model(self) -> None: + self.model = AutoModelForCausalLM.from_pretrained( + self.model_name, + torch_dtype=torch.bfloat16, + ).to("cuda:0") + self.model.eval() + + try: + self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME) + except AttributeError as exc: + raise RuntimeError( + f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`" + ) from exc + + self.pending_sparse_patches = None + + def create_rendezvous(self) -> tuple[str, int]: + self.port = get_open_port() + return self.master_address, self.port + + def init_weight_transfer_group(self, world_size: int) -> None: + self.model_update_group = NCCLWeightTransferEngine.trainer_init( + dict( + master_address=self.master_address, + master_port=self.port, + world_size=world_size, + ) + ) + + def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]: + names = [] + dtype_names = [] + shapes = [] + payload_bytes = 0 + for name, param in self.model.named_parameters(): + names.append(name) + dtype_names.append(str(param.dtype).split(".")[-1]) + shapes.append(list(param.shape)) + payload_bytes += param.numel() * param.element_size() + + return ( + dict( + names=names, + dtype_names=dtype_names, + shapes=shapes, + packed=packed, + ), + payload_bytes, + ) + + @torch.inference_mode() + def generate( + self, + prompts: Sequence[str], + max_new_tokens: int = 1, + ) -> list[dict[str, object]]: + generations = [] + for prompt in prompts: + model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0") + output = self.model.generate( + **model_inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + pad_token_id=self.tokenizer.pad_token_id, + ) + new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist() + generations.append( + { + "token_ids": new_token_ids, + "text": self.tokenizer.decode( + new_token_ids, + skip_special_tokens=False, + ), + } + ) + return generations + + def prepare_sparse_patch( + self, + prompts: Sequence[str], + max_patch_rows: int = MAX_PATCH_ROWS, + ) -> tuple[dict[str, object], list[int], str, int]: + selected_token_ids: list[int] = [] + special_ids = set(self.tokenizer.all_special_ids) + for prompt in prompts: + token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + for token_id in token_ids: + if token_id in special_ids or token_id in selected_token_ids: + continue + selected_token_ids.append(token_id) + if len(selected_token_ids) == max_patch_rows: + break + if len(selected_token_ids) == max_patch_rows: + break + + if not selected_token_ids: + raise ValueError("Could not derive any non-special token IDs to patch") + + vocab_size = self.patched_param.shape[0] + next_token_id = selected_token_ids[-1] + while len(selected_token_ids) < max_patch_rows: + next_token_id = (next_token_id + 1) % vocab_size + if next_token_id in special_ids or next_token_id in selected_token_ids: + continue + selected_token_ids.append(next_token_id) + + row_ids = torch.tensor( + selected_token_ids, + device=self.patched_param.device, + dtype=torch.long, + ) + hidden_size = self.patched_param.shape[1] + column_offsets = torch.arange( + hidden_size, + device=self.patched_param.device, + dtype=torch.long, + ) + + with torch.no_grad(): + # Rotate the selected embedding rows instead of zeroing them so the + # patch remains deterministic while avoiding a degenerate collapse + # to the same special token after the update. + replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0) + self.patched_param[row_ids] = replacement_rows + + flat_indices = ( + row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1) + ) + flat_values = self.patched_param[row_ids].reshape(-1).contiguous() + self.pending_sparse_patches = [ + SparseWeightPatch( + name=PATCHED_PARAM_NAME, + indices=flat_indices.to(torch.int32), + values=flat_values, + ) + ] + patch_digest = hashlib.sha256( + self.pending_sparse_patches[0].indices.cpu().numpy().tobytes() + + self.pending_sparse_patches[0] + .values.detach() + .float() + .cpu() + .numpy() + .tobytes() + ).hexdigest() + + sparse_payload_bytes = ( + flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size() + + flat_values.numel() * flat_values.element_size() + ) + update_info = dict( + names=[PATCHED_PARAM_NAME], + dtype_names=[str(self.patched_param.dtype).split(".")[-1]], + shapes=[list(self.patched_param.shape)], + num_updates_list=[flat_indices.numel()], + update_kind="sparse_flat", + ) + return update_info, selected_token_ids, patch_digest, sparse_payload_bytes + + def broadcast_weights(self, packed: bool = False) -> float: + if self.model_update_group is None: + raise RuntimeError("Weight transfer group is not initialized") + + trainer_args = NCCLTrainerSendWeightsArgs( + group=self.model_update_group, + packed=packed, + ) + start = time.perf_counter() + NCCLWeightTransferEngine.trainer_send_weights( + iterator=self.model.named_parameters(), + trainer_args=trainer_args, + ) + torch.accelerator.synchronize() + return (time.perf_counter() - start) * 1000.0 + + def broadcast_pending_sparse_patch(self) -> float: + if self.model_update_group is None: + raise RuntimeError("Weight transfer group is not initialized") + if self.pending_sparse_patches is None: + raise RuntimeError("Sparse patch has not been prepared") + + start = time.perf_counter() + NCCLWeightTransferEngine.trainer_send_sparse_weights( + iter(self.pending_sparse_patches), + NCCLTrainerSendWeightsArgs(group=self.model_update_group), + ) + torch.accelerator.synchronize() + self.pending_sparse_patches = None + return (time.perf_counter() - start) * 1000.0 + + +def launch_llm( + scheduling_inference: PlacementGroupSchedulingStrategy, +): + return ray.remote( + num_cpus=0, + num_gpus=0, + scheduling_strategy=scheduling_inference, + )(MyLLM).remote( + model=MODEL_NAME, + enforce_eager=True, + tensor_parallel_size=1, + distributed_executor_backend="ray", + gpu_memory_utilization=0.7, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + +def collect_vllm_generations(llm_handle) -> list[dict[str, object]]: + outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS)) + generations = [] + for output in outputs: + generations.append( + { + "token_ids": output.outputs[0].token_ids, + "text": output.outputs[0].text, + } + ) + return generations + + +def token_sequences_match( + left: Sequence[dict[str, object]], + right: Sequence[dict[str, object]], +) -> bool: + return [item["token_ids"] for item in left] == [item["token_ids"] for item in right] + + +def print_generations(label: str, prompts: Sequence[str], generations) -> None: + print(f"\n{label}") + print("-" * 50) + for prompt, generation in zip(prompts, generations): + print(f"Prompt: {prompt!r}") + print(f"Token IDs: {generation['token_ids']}") + print(f"Text: {generation['text']!r}") + print("-" * 50) + + +def run_dense_phase( + train_model, + scheduling_inference: PlacementGroupSchedulingStrategy, +) -> dict[str, object]: + ray.get(train_model.reset_model.remote()) + llm = launch_llm(scheduling_inference) + try: + dense_before = collect_vllm_generations(llm) + + ray.get(llm.sleep.remote(level=0)) + master_address, master_port = ray.get(train_model.create_rendezvous.remote()) + world_size = ray.get(llm.get_world_size.remote()) + 1 + inference_init = llm.init_weight_transfer_engine.remote( + dict( + init_info=dict( + master_address=master_address, + master_port=master_port, + rank_offset=1, + world_size=world_size, + ) + ) + ) + trainer_init = train_model.init_weight_transfer_group.remote(world_size) + ray.get([trainer_init, inference_init]) + ray.get(llm.start_weight_update.remote(is_checkpoint_format=True)) + + dense_update_info, dense_payload_bytes = ray.get( + train_model.get_dense_update_info.remote() + ) + _, selected_token_ids, patch_digest, _ = ray.get( + train_model.prepare_sparse_patch.remote(PROMPTS) + ) + + inference_update = llm.update_weights.remote( + dict(update_info=dense_update_info) + ) + dense_send_ms, _ = ray.get( + [ + train_model.broadcast_weights.remote(packed=False), + inference_update, + ] + ) + ray.get(llm.finish_weight_update.remote()) + ray.get(llm.wake_up.remote(tags=["scheduling"])) + + dense_after = collect_vllm_generations(llm) + + return { + "dense_before": dense_before, + "dense_after": dense_after, + "selected_token_ids": selected_token_ids, + "patch_digest": patch_digest, + "dense_payload_bytes": dense_payload_bytes, + "dense_send_ms": dense_send_ms, + } + finally: + ray.kill(llm) + + +def run_sparse_phase( + train_model, + scheduling_inference: PlacementGroupSchedulingStrategy, +) -> dict[str, object]: + ray.get(train_model.reset_model.remote()) + llm = launch_llm(scheduling_inference) + try: + sparse_before = collect_vllm_generations(llm) + + ray.get(llm.sleep.remote(level=0)) + master_address, master_port = ray.get(train_model.create_rendezvous.remote()) + world_size = ray.get(llm.get_world_size.remote()) + 1 + inference_init = llm.init_weight_transfer_engine.remote( + dict( + init_info=dict( + master_address=master_address, + master_port=master_port, + rank_offset=1, + world_size=world_size, + ) + ) + ) + trainer_init = train_model.init_weight_transfer_group.remote(world_size) + ray.get([trainer_init, inference_init]) + ray.get(llm.start_weight_update.remote(is_checkpoint_format=False)) + + sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = ( + ray.get(train_model.prepare_sparse_patch.remote(PROMPTS)) + ) + + inference_update = llm.update_weights.remote( + dict(update_info=sparse_update_info) + ) + sparse_send_ms, _ = ray.get( + [ + train_model.broadcast_pending_sparse_patch.remote(), + inference_update, + ] + ) + ray.get(llm.finish_weight_update.remote()) + ray.get(llm.wake_up.remote(tags=["scheduling"])) + + sparse_after = collect_vllm_generations(llm) + + return { + "sparse_before": sparse_before, + "sparse_after": sparse_after, + "selected_token_ids": selected_token_ids, + "patch_digest": patch_digest, + "sparse_payload_bytes": sparse_payload_bytes, + "sparse_send_ms": sparse_send_ms, + } + finally: + ray.kill(llm) + + +ray.init() + +try: + train_model = TrainModel.remote(MODEL_NAME) + + pg_inference = placement_group([{"GPU": 1, "CPU": 0}]) + ray.get(pg_inference.ready()) + scheduling_inference = PlacementGroupSchedulingStrategy( + placement_group=pg_inference, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=0, + ) + + dense_results = run_dense_phase(train_model, scheduling_inference) + sparse_results = run_sparse_phase(train_model, scheduling_inference) + + baseline_equal = token_sequences_match( + dense_results["dense_before"], + sparse_results["sparse_before"], + ) + patch_selection_equal = ( + dense_results["selected_token_ids"] == sparse_results["selected_token_ids"] + ) + patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"] + after_equal = token_sequences_match( + dense_results["dense_after"], + sparse_results["sparse_after"], + ) + any_output_changed = any( + before["token_ids"] != after["token_ids"] + for before, after in zip( + dense_results["dense_before"], + dense_results["dense_after"], + ) + ) + dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024) + sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024) + + print_generations( + "Dense baseline outputs", + PROMPTS, + dense_results["dense_before"], + ) + print_generations( + "Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"] + ) + print_generations( + "Dense outputs after update", PROMPTS, dense_results["dense_after"] + ) + print_generations( + "Sparse outputs after update", + PROMPTS, + sparse_results["sparse_after"], + ) + + print(f"patched_token_ids = {dense_results['selected_token_ids']}") + print(f"patch_selection_equal = {patch_selection_equal}") + print(f"dense_patch_digest = {dense_results['patch_digest']}") + print(f"sparse_patch_digest = {sparse_results['patch_digest']}") + print(f"patch_digest_equal = {patch_digest_equal}") + print(f"baseline_equal = {baseline_equal}") + print(f"after_equal = {after_equal}") + print(f"any_output_changed = {any_output_changed}") + print(f"dense_payload_mb = {dense_payload_mb:.2f}") + print(f"sparse_payload_mb = {sparse_payload_mb:.2f}") + print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}") + print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}") + + if not baseline_equal: + raise RuntimeError( + "Dense and sparse phases did not start from the same baseline" + ) + if not patch_selection_equal: + raise RuntimeError("Dense and sparse phases used different sparse patches") + if not patch_digest_equal: + raise RuntimeError("Dense and sparse phases produced different patch values") + if not after_equal: + raise RuntimeError("Dense and sparse updates produced different outputs") + if not any_output_changed: + raise RuntimeError("Patch did not change the observed outputs") +finally: + ray.shutdown() diff --git a/examples/tool_chat_template_gemma4.jinja b/examples/tool_chat_template_gemma4.jinja index d61dd795b58..ef765823106 100644 --- a/examples/tool_chat_template_gemma4.jinja +++ b/examples/tool_chat_template_gemma4.jinja @@ -295,6 +295,15 @@ {%- endif -%} {%- endfor -%} {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') == 'image' -%} + {{- '<|image|>' -}} + {%- elif part.get('type') == 'audio' -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} {%- else -%} {{- format_tool_response_block(ns_tname.name, tool_body) -}} {%- endif -%} 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/requirements/xpu.txt b/requirements/xpu.txt index 1bbf777b520..897e2080daf 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -17,4 +17,4 @@ torchaudio torchvision auto_round_lib>=0.13.0 -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9/vllm_xpu_kernels-0.1.9-cp38-abi3-manylinux_2_28_x86_64.whl 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 1cf213ca604..5b6f66cf417 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -93,7 +93,7 @@ pub struct ChatLlm { text: TextLlm, backend: DynChatBackend, /// Effective model dtype reported by the engine. - model_dtype: Option, + model_dtype: ModelDtype, /// Tool-call parser selection. tool_call_parser: ParserSelection, /// Reasoning parser selection. @@ -135,7 +135,7 @@ impl ChatLlm { } /// Override the effective model dtype used for multimodal tensor encoding. - pub fn with_model_dtype(mut self, model_dtype: Option) -> Self { + pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self { self.model_dtype = model_dtype; self } @@ -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(); @@ -233,7 +234,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string()); } #[test] diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 8526d4fa0ef..fcfee0ccb33 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -12,7 +12,7 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::path::Path; -use std::sync::{Arc, LazyLock, Once}; +use std::sync::{Arc, LazyLock}; use itertools::izip; use llm_multimodal::{ @@ -239,7 +239,7 @@ pub(crate) async fn finalize_rendered_prompt( request: &ChatRequest, rendered: RenderedPrompt, info: Option<&MultimodalModelInfo>, - model_dtype: Option, + model_dtype: ModelDtype, ) -> Result<(Prompt, Option)> { if !request.has_multimodal() { return Ok((rendered.prompt, None)); @@ -249,16 +249,6 @@ pub(crate) async fn finalize_rendered_prompt( bail_multimodal!("multimodal chat renderer must return a text prompt before expansion"); }; let media_parts = extract_media_parts(request)?; - let model_dtype = model_dtype.unwrap_or_else(|| { - static WARN_ONCE: Once = Once::new(); - WARN_ONCE.call_once(|| { - warn!( - "engine handshake did not report model dtype; \ - falling back to float32 for multimodal tensor encoding" - ); - }); - ModelDtype::Float32 - }); let mut prompt_token_ids = info .context 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/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 7dc72672299..ad220b5a787 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -5,9 +5,9 @@ use std::sync::LazyLock; pub use vllm_tool_parser::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser, - KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, - ToolParserOutput, + Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, + MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, + ToolParserError, ToolParserOutput, }; use crate::parser::ParserFactory; @@ -24,6 +24,9 @@ pub mod names { pub const GEMMA4: &str = "gemma4"; pub const HERMES: &str = "hermes"; pub const HY_V3: &str = "hy_v3"; + // Matches the Python CLI name `--tool-call-parser internlm`, which Python + // also routes to `Internlm2ToolParser` despite the version-agnostic name. + pub const INTERNLM: &str = "internlm"; pub const KIMI_K2: &str = "kimi_k2"; pub const LLAMA3_JSON: &str = "llama3_json"; pub const LLAMA4_JSON: &str = "llama4_json"; @@ -62,6 +65,7 @@ impl ToolParserFactory { .register_parser::(names::GEMMA4) .register_parser::(names::HERMES) .register_parser::(names::HY_V3) + .register_parser::(names::INTERNLM) .register_parser::(names::KIMI_K2) .register_parser::(names::LLAMA3_JSON) .register_parser::(names::LLAMA4_JSON) @@ -80,6 +84,12 @@ impl ToolParserFactory { .register_pattern("hermes", names::HERMES) .register_pattern("hy3", names::HY_V3) .register_pattern("hy_v3", names::HY_V3) + // Narrow to `internlm2` substring so it matches `internlm2-chat-7b` + // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1, + // routes to Llama), `internlm3-*` (also Llama-architecture per + // vllm/model_executor/models/registry.py:146), or `Intern-S1` / + // `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115). + .register_pattern("internlm2", names::INTERNLM) .register_pattern("llama-4", names::LLAMA4_JSON) .register_pattern("llama-3.2", names::LLAMA3_JSON) .register_pattern("llama-3.1", names::LLAMA3_JSON) diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index f18aa69f950..65e9f4e075b 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -161,4 +161,33 @@ fn factory_new_resolves_default_patterns() { factory.resolve_name_for_model("org/mm-m2-base"), Some(names::MINIMAX_M2) ); + + // InternLM2 positive: both dashed and underscored versioned names route. + assert_eq!( + factory.resolve_name_for_model("internlm/internlm2-chat-7b"), + Some(names::INTERNLM) + ); + assert_eq!( + factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"), + Some(names::INTERNLM) + ); + + // Negative: other internlm-org models do NOT route to the InternLM2 parser, + // since they use unrelated prompt formats. + // - InternLM v1 (`internlm-chat-7b`) routes to Llama + // - InternLM3 (`internlm3-8b-instruct`) routes to Llama + // - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115) + assert_eq!( + factory.resolve_name_for_model("internlm/internlm-chat-7b"), + None + ); + assert_eq!( + factory.resolve_name_for_model("internlm/internlm3-8b-instruct"), + None + ); + assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None); + assert_eq!( + factory.resolve_name_for_model("internlm/Intern-S1-Pro"), + None + ); } 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 dd3000f3c6a..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 @@ -290,10 +292,8 @@ impl EngineCoreClient { // If any engine reported a dp_stats_address in its ready response, use it // as the external coordinator address. - let dp_stats_address: Option = engines - .iter() - .filter_map(|e| e.ready_response.as_ref()) - .find_map(|r| r.dp_stats_address.clone()); + let dp_stats_address: Option = + engines.iter().find_map(|engine| engine.ready_response.dp_stats_address.clone()); let (coordinator, coordinator_output_task, coordinator_task) = if let Some(coordinator_transport) = connected.coordinator { @@ -368,40 +368,44 @@ impl EngineCoreClient { /// Return the ready responses received from all engines on the input /// socket. pub fn ready_responses(&self) -> Vec<&EngineCoreReadyResponse> { - self.engines - .iter() - .filter_map(|engine| engine.ready_response.as_ref()) - .collect() + self.engines.iter().map(|engine| &engine.ready_response).collect() } - /// Return the engine-reported effective model dtype, when available. - pub fn model_dtype(&self) -> Option { + /// Return the engine-reported effective model dtype. + pub fn model_dtype(&self) -> ModelDtype { self.engines - .iter() - .filter_map(|engine| engine.ready_response.as_ref()) - .find_map(|response| response.dtype) + .first() + .expect("engine core client requires at least one engine") + .ready_response + .dtype + } + + /// Return the engine-reported Python vLLM version. + pub fn vllm_version(&self) -> &str { + self.engines + .first() + .expect("engine core client requires at least one engine") + .ready_response + .vllm_version + .as_str() } /// Return the total number of GPU blocks summed across all connected /// engines. pub fn total_num_gpu_blocks(&self) -> u64 { - self.engines - .iter() - .filter_map(|engine| engine.ready_response.as_ref()) - .map(|r| r.num_gpu_blocks) - .sum() + self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum() } /// Return the minimum engine-reported `max_model_len` across all engines. /// /// This is the auto-fitted value after KV cache profiling and may differ /// from the originally configured value. - pub fn max_model_len(&self) -> Option { + pub fn max_model_len(&self) -> u32 { self.engines .iter() - .filter_map(|e| e.ready_response.as_ref()) - .map(|r| r.max_model_len as u32) + .map(|engine| engine.ready_response.max_model_len as u32) .min() + .expect("engine core client requires at least one engine") } /// Get the model name associated with this client used for metrics @@ -657,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/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index e432638f350..9a66ad84cc1 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -382,6 +382,7 @@ mod tests { use zeromq::{RouterSocket, Socket}; use super::*; + use crate::mock_engine::default_ready_response; async fn test_inner() -> ClientInner { let mut socket = RouterSocket::new(); @@ -392,7 +393,7 @@ mod tests { "test-model".to_string(), &[ConnectedEngine { engine_id: EngineId::from(b"engine-0"), - ready_response: None, + ready_response: default_ready_response(), }], ) } diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index d47c5a80719..99302e4f8cc 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -339,15 +339,20 @@ mod tests { use super::{EngineRoutingState, RequestRegistry, UtilityRegistry}; use crate::EngineId; use crate::client::state::EngineLoadSnapshot; + use crate::mock_engine::default_ready_response; use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput}; use crate::transport::ConnectedEngine; + fn connected_engine(engine_id: EngineId) -> ConnectedEngine { + ConnectedEngine { + engine_id, + ready_response: default_ready_response(), + } + } + #[test] fn registry_rejects_duplicate_request_ids() { - let mut registry = RequestRegistry::new(&[ConnectedEngine { - engine_id: EngineId::from(b"engine-0"), - ready_response: None, - }]); + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry.register("req-1".to_string(), None).unwrap(); let error = registry.register("req-1".to_string(), None).unwrap_err(); assert!(matches!( @@ -358,10 +363,7 @@ mod tests { #[test] fn registry_removes_finished_request_on_output() { - let mut registry = RequestRegistry::new(&[ConnectedEngine { - engine_id: EngineId::from(b"engine-0"), - ready_response: None, - }]); + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry.register("req-1".to_string(), None).unwrap(); let sender = registry.sender_for_output(&EngineCoreOutput { @@ -376,10 +378,7 @@ mod tests { #[test] fn registry_closes_all_requests_on_failure() { - let mut registry = RequestRegistry::new(&[ConnectedEngine { - engine_id: EngineId::from(b"engine-0"), - ready_response: None, - }]); + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry.register("req-1".to_string(), None).unwrap(); registry.register("req-2".to_string(), None).unwrap(); @@ -394,14 +393,8 @@ mod tests { let engine_0 = EngineId::from_engine_index(0); let engine_1 = EngineId::from_engine_index(1); let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_1.clone(), - ready_response: None, - }, + connected_engine(engine_0.clone()), + connected_engine(engine_1.clone()), ]); let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap(); let (chosen_1, _) = registry.register("req-2".to_string(), None).unwrap(); @@ -428,14 +421,8 @@ mod tests { let engine_0 = EngineId::from_engine_index(0); let engine_1 = EngineId::from_engine_index(1); let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_1.clone(), - ready_response: None, - }, + connected_engine(engine_0.clone()), + connected_engine(engine_1.clone()), ]); let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap(); @@ -488,14 +475,8 @@ mod tests { let engine_0 = EngineId::from_engine_index(0); let engine_1 = EngineId::from_engine_index(1); let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_1.clone(), - ready_response: None, - }, + connected_engine(engine_0.clone()), + connected_engine(engine_1.clone()), ]); assert!(registry.apply_scheduler_counts( @@ -523,18 +504,9 @@ mod tests { let engine_1 = EngineId::from_engine_index(1); let engine_2 = EngineId::from_engine_index(2); let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_1.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_2.clone(), - ready_response: None, - }, + connected_engine(engine_0.clone()), + connected_engine(engine_1.clone()), + connected_engine(engine_2.clone()), ]); // Explicitly target rank 2 (third engine). @@ -555,14 +527,8 @@ mod tests { let engine_0 = EngineId::from_engine_index(0); let engine_1 = EngineId::from_engine_index(1); let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }, - ConnectedEngine { - engine_id: engine_1.clone(), - ready_response: None, - }, + connected_engine(engine_0.clone()), + connected_engine(engine_1.clone()), ]); // Load-balance: first two go to engine_0 and engine_1. @@ -577,14 +543,8 @@ mod tests { #[test] fn register_with_out_of_range_rank_returns_error() { let mut registry = RequestRegistry::new(&[ - ConnectedEngine { - engine_id: EngineId::from_engine_index(0), - ready_response: None, - }, - ConnectedEngine { - engine_id: EngineId::from_engine_index(1), - ready_response: None, - }, + connected_engine(EngineId::from_engine_index(0)), + connected_engine(EngineId::from_engine_index(1)), ]); let error = registry.register("req-1".to_string(), Some(2)).unwrap_err(); @@ -600,10 +560,7 @@ mod tests { #[test] fn register_with_rank_on_single_engine_only_accepts_zero() { let engine_0 = EngineId::from_engine_index(0); - let mut registry = RequestRegistry::new(&[ConnectedEngine { - engine_id: engine_0.clone(), - ready_response: None, - }]); + let mut registry = RequestRegistry::new(&[connected_engine(engine_0.clone())]); let (chosen, _) = registry.register("req-ok".to_string(), Some(0)).unwrap(); assert_eq!(chosen, engine_0); diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index aa1cfecaee4..32cd48c396f 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -47,7 +47,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN, num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS, dp_stats_address: None, - dtype: Some(ModelDtype::Float32), + dtype: ModelDtype::Float32, + vllm_version: "test-vllm-version".to_string(), } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index 622a032772e..d659dc8a244 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -39,10 +39,9 @@ pub struct EngineCoreReadyResponse { /// DP coordinator stats publish address, if applicable. pub dp_stats_address: Option, /// Effective model dtype after Python vLLM resolves `--dtype`. - // TODO: This is currently not wired up on the engine side. After it's added, remove `Option` - // and `serde(default)`. - #[serde(default)] - pub dtype: Option, + pub dtype: ModelDtype, + /// Python vLLM version reported by the engine process. + pub vllm_version: String, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup @@ -69,22 +68,3 @@ pub struct HandshakeInitMessage { pub addresses: HandshakeAddresses, pub parallel_config: BTreeMap, } - -#[cfg(test)] -mod tests { - use super::EngineCoreReadyResponse; - use crate::protocol::ModelDtype; - - #[test] - fn ready_response_accepts_effective_dtype() { - let response: EngineCoreReadyResponse = serde_json::from_value(serde_json::json!({ - "max_model_len": 4096, - "num_gpu_blocks": 2, - "dp_stats_address": null, - "dtype": "bfloat16" - })) - .unwrap(); - - assert_eq!(response.dtype, Some(ModelDtype::BFloat16)); - } -} 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/engine-core-client/src/test_utils.rs b/rust/src/engine-core-client/src/test_utils.rs index f1c5c65503f..06f56380ab1 100644 --- a/rust/src/engine-core-client/src/test_utils.rs +++ b/rust/src/engine-core-client/src/test_utils.rs @@ -10,9 +10,9 @@ use crate::EngineId; pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets}; use crate::mock_engine::{ MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend, + default_ready_response, }; -use crate::protocol::ModelDtype; -use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage}; +use crate::protocol::handshake::HandshakeInitMessage; /// Per-test IPC endpoint namespace backed by a unique temporary directory. /// @@ -57,12 +57,7 @@ fn test_mock_engine_config() -> MockEngineConfig { MockEngineConfig { local: true, headless: true, - ready_response: EngineCoreReadyResponse { - max_model_len: 4096, - num_gpu_blocks: 0, - dp_stats_address: None, - dtype: Some(ModelDtype::Float32), - }, + ready_response: default_ready_response(), ..Default::default() } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index af6cf3ea9a5..9a92ffe447e 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -925,6 +925,7 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() { .await; assert_eq!(client.engine_identities()[0], b"engine-0"); assert!(client.ready_responses()[0].max_model_len > 0); + assert_eq!(client.vllm_version(), "test-vllm-version"); let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index 0d6c49340af..360f94eda12 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -104,8 +104,8 @@ pub struct ConnectedEngine { /// The identity of the connected engine. pub engine_id: EngineId, /// Post-initialization configuration received from the engine on the input - /// socket registration message. `None` until the registration is received. - pub ready_response: Option, + /// socket registration message. + pub ready_response: EngineCoreReadyResponse, } /// Represents the connected shared transport plus all registered engines after @@ -295,18 +295,9 @@ pub async fn connect_handshake( } } - // 4. Wait for every engine to connect to the shared input socket and register itself. The - // `ready_response` is a placeholder; it is populated for each engine by - // `wait_for_input_registrations` below. - let mut engines: Vec<_> = engines - .into_keys() - .map(|engine_id| ConnectedEngine { - engine_id, - ready_response: None, - }) - .collect(); - - wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?; + // 6. Wait for every engine to connect to the shared input socket and register itself. + let engines = + wait_for_input_registrations(&mut input_socket, engines.into_keys(), ready_timeout).await?; debug!( engine_count = engines.len(), "all engines registered on shared input socket" @@ -349,15 +340,13 @@ pub async fn connect_bootstrapped( let mut output_socket = PullSocket::new(); let output_address = output_socket.bind(output_address).await?.to_string(); - // TODO: follow start rank - let mut engines = (0..engine_count) - .map(|index| ConnectedEngine { - engine_id: EngineId::from((index as u16).to_le_bytes().to_vec()), - ready_response: None, - }) - .collect::>(); - - wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?; + let engines = wait_for_input_registrations( + &mut input_socket, + // TODO: follow start rank + (0..engine_count).map(|index| EngineId::from((index as u16).to_le_bytes().to_vec())), + ready_timeout, + ) + .await?; info!( engine_count = engines.len(), "bootstrapped engines connected" @@ -455,17 +444,14 @@ async fn send_init_message( /// Simplify API server handshake"), the payload is a msgpack-encoded /// [`EngineCoreReadyResponse`] carrying post-initialization values such as /// `max_model_len`. -/// -/// Older engines sent an empty second frame here just to establish the -/// ROUTER/DEALER backchannel, with no structured payload on the input socket. -/// We continue to tolerate that legacy shape so the frontend can still connect -/// to slightly older local engine checkouts. async fn wait_for_input_registrations( input_socket: &mut RouterSocket, - engines: &mut [ConnectedEngine], + expected_engines: impl IntoIterator, ready_timeout: Duration, -) -> Result<()> { - let mut pending = engines.iter().map(|e| e.engine_id.clone()).collect::>(); +) -> Result> { + let expected_engines = expected_engines.into_iter().collect::>(); + let mut pending = expected_engines.iter().cloned().collect::>(); + let mut ready_responses = BTreeMap::new(); while !pending.is_empty() { let registration = timeout(ready_timeout, input_socket.recv()).await.map_err(|_| { @@ -489,29 +475,33 @@ async fn wait_for_input_registrations( ); } - let ready_response = if frames[1].is_empty() { - debug!( - ?actual_id, - "received legacy empty input registration from engine" + if frames[1].is_empty() { + bail_unexpected_handshake_message!( + "expected msgpack EngineCoreReadyResponse for engine input registration, got empty payload from engine id {actual_id:?}" ); - None - } else { - let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?; - debug!( - ?actual_id, - ?ready_response, - "received input registration from engine" - ); - Some(ready_response) - }; - - // Store the ready response in the corresponding engine entry. - if let Some(engine) = engines.iter_mut().find(|e| e.engine_id == actual_id) { - engine.ready_response = ready_response; } + + let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?; + debug!( + ?actual_id, + ?ready_response, + "received input registration from engine" + ); + ready_responses.insert(actual_id, ready_response); } - Ok(()) + Ok(expected_engines + .into_iter() + .map(|engine_id| { + let ready_response = ready_responses + .remove(&engine_id) + .expect("every expected engine id has a decoded ready response"); + ConnectedEngine { + engine_id, + ready_response, + } + }) + .collect()) } /// Send an encoded message to the engine through the input socket. 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/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index fd04761090a..a80aef40300 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -98,7 +98,8 @@ async fn mock_engine_connects_over_tcp() { let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; assert_eq!(client.engine_count(), 1); assert_eq!(client.engine_identities()[0], &[0, 0]); - assert_eq!(client.max_model_len(), Some(1024 * 1024)); + assert_eq!(client.max_model_len(), 1024 * 1024); + assert_eq!(client.vllm_version(), "test-vllm-version"); shutdown_mock(client, shutdown, task).await; } 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 ccf90db9aa8..a0473c783a0 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -3,9 +3,12 @@ mod collective_rpc; mod health; mod inference; mod load; +mod lora; mod metrics; pub(crate) mod openai; +mod server_info; mod sleep; +mod version; use std::sync::Arc; @@ -24,17 +27,46 @@ 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)) .route("/metrics", get(metrics::scrape)) .route("/load", get(load::load)) + .route("/version", get(version::version)) // OpenAI-compatible endpoints .route("/v1/models", get(openai::list_models)) .route("/v1/completions", post(openai::completions)) @@ -42,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 @@ -52,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 ff7ea3c6302..f15f757c09a 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -3,23 +3,33 @@ mod types; mod validate; use std::collections::HashMap; +use std::convert::Infallible; +use std::result::Result; use std::sync::Arc; +use asynk_strim_attr::{TryYielder, try_stream}; use axum::Json; use axum::extract::State; use axum::http::HeaderMap; +use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response}; +use futures::{Stream, StreamExt as _, pin_mut}; use thiserror_ext::AsReport as _; -use tracing::info; +use tracing::{error, info, trace}; use tracing_futures::Instrument as _; use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs}; -use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _}; +use vllm_llm::{ + CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _, +}; use self::convert::prepare_generate_request; -use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice}; -use crate::error::{ApiError, server_error}; +use self::types::{ + GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice, + GenerateResponseStreamChoice, GenerateStreamResponse, +}; +use crate::error::{ApiError, bail_server_error, server_error}; use crate::routes::openai::utils::logprobs::clamp_logprob; -use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb}; +use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage}; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::resolve_request_context; @@ -32,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(), }; @@ -46,6 +56,7 @@ pub async fn generate( let log_request = state.enable_log_requests; let include_logprobs = prepared.include_logprobs; let include_prompt_logprobs = prepared.include_prompt_logprobs; + let stream = prepared.stream; let raw_stream = match state .chat @@ -64,6 +75,20 @@ pub async fn generate( } }; + if stream { + let chunk_stream = generate_chunk_stream( + raw_stream, + prepared.request_id, + log_request, + prepared.include_usage, + prepared.include_continuous_usage, + include_logprobs, + ); + let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span); + + return Sse::new(sse_stream).into_response(); + } + let collected = match raw_stream.collect_output().instrument(request_span.clone()).await { Ok(collected) => collected, Err(error) => { @@ -98,6 +123,102 @@ pub async fn generate( Json(response).into_response() } +#[try_stream] +async fn generate_chunk_stream( + stream: impl Stream>, + request_id: String, + log_request: bool, + include_usage: bool, + include_continuous_usage: bool, + include_logprobs: bool, + mut y: TryYielder, +) -> Result<(), ApiError> { + pin_mut!(stream); + let mut prompt_tokens: Option = None; + let mut output_tokens = 0_u32; + + while let Some(next) = stream.next().await { + match next { + Ok(output) => { + if prompt_tokens.is_none() { + prompt_tokens = + output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32); + } + let usage_prompt_tokens = prompt_tokens.unwrap_or_default(); + + let token_ids = output.token_ids; + output_tokens = output_tokens.saturating_add(token_ids.len() as u32); + let finish_reason = output.finish_reason; + + if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) { + bail_server_error!("Internal server error"); + } + + if let Some(finish_reason) = finish_reason.as_ref() + && log_request + { + info!( + stream = true, + prompt_tokens = usage_prompt_tokens, + output_tokens, + finish_reason = finish_reason.as_str(), + "generate finished" + ); + } + + if token_ids.is_empty() && finish_reason.is_none() { + continue; + } + + let logprobs = if include_logprobs && !token_ids.is_empty() { + let logprobs = output.logprobs.as_ref().ok_or_else(|| { + server_error!( + "raw generate stream requested logprobs but generation returned none" + ) + })?; + Some(raw_logprobs_to_openai_chat(logprobs)?) + } else { + None + }; + + y.yield_ok(GenerateStreamResponse { + request_id: request_id.clone(), + choices: vec![GenerateResponseStreamChoice { + index: 0, + logprobs, + finish_reason: finish_reason.map(|reason| reason.as_str().to_string()), + token_ids, + }], + usage: include_continuous_usage + .then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)), + }) + .await; + } + Err(error) => { + error!( + error = %error.as_report(), + "raw generate stream failed" + ); + bail_server_error!("{}", error.to_report_string()); + } + } + } + + if include_usage { + y.yield_ok(GenerateStreamResponse { + request_id, + choices: Vec::new(), + usage: Some(Usage::from_counts( + prompt_tokens.unwrap_or_default(), + output_tokens, + )), + }) + .await; + } + + Ok(()) +} + fn collect_generate( collected: CollectedGenerateOutput, request_id: String, @@ -213,3 +334,94 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap String { format!("token_id:{token_id}") } + +/// Convert one raw-generate chunk stream into SSE events. +#[try_stream] +async fn generate_sse_stream( + stream: impl Stream>, + mut y: TryYielder, +) -> Result<(), Infallible> { + pin_mut!(stream); + + while let Some(next) = stream.next().await { + match next { + Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await, + Err(error) => { + y.yield_ok(to_error_sse_event(&error)).await; + break; + } + } + } + + y.yield_ok(done_sse_event()).await; + Ok(()) +} + +fn to_sse_event(chunk: &GenerateStreamResponse) -> Event { + let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON"); + trace!(payload, "generate emitting chunk"); + Event::default().data(payload) +} + +fn to_error_sse_event(error: &ApiError) -> Event { + let payload = serde_json::to_string(&error.to_error_response()) + .expect("ErrorResponse must serialize to JSON"); + trace!(payload, "generate emitting error"); + Event::default().data(payload) +} + +fn done_sse_event() -> Event { + trace!("generate emitting done"); + Event::default().data("[DONE]") +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use futures::{TryStreamExt as _, stream}; + use vllm_llm::GeneratePromptInfo; + + use super::*; + + #[tokio::test] + async fn generate_chunk_stream_captures_late_prompt_info() { + let stream = stream::iter(vec![ + Ok(GenerateOutput { + request_id: String::new(), + prompt_info: None, + token_ids: Vec::new(), + logprobs: None, + finish_reason: None, + kv_transfer_params: None, + }), + Ok(GenerateOutput { + request_id: String::new(), + prompt_info: Some(GeneratePromptInfo { + prompt_token_ids: Arc::from([11_u32, 22_u32]), + prompt_logprobs: None, + }), + token_ids: vec![33], + logprobs: None, + finish_reason: Some(FinishReason::stop_eos()), + kv_transfer_params: None, + }), + ]); + + let chunks: Vec<_> = + generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false) + .try_collect() + .await + .expect("collect chunks"); + + assert_eq!(chunks.len(), 2); + assert_eq!( + chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens, + 2 + ); + assert_eq!( + chunks[1].usage.as_ref().expect("final usage").prompt_tokens, + 2 + ); + } +} diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 70374c2f1eb..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. @@ -10,6 +11,9 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params}; pub struct PreparedRequest { pub request_id: String, pub text_request: TextRequest, + pub stream: bool, + pub include_usage: bool, + pub include_continuous_usage: bool, pub include_logprobs: bool, pub include_prompt_logprobs: bool, } @@ -18,11 +22,23 @@ 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 + .stream_options + .as_ref() + .and_then(|options| options.include_usage) + .unwrap_or(false); + let include_continuous_usage = include_usage + && request + .stream_options + .as_ref() + .and_then(|options| options.continuous_usage_stats) + .unwrap_or(false); let include_logprobs = request.sampling_params.logprobs.is_some(); let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some(); let mut sampling_params = request.sampling_params; @@ -42,11 +58,15 @@ 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 { request_id: ctx.request_id, text_request, + stream, + include_usage, + include_continuous_usage, include_logprobs, include_prompt_logprobs, }) @@ -58,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!({ @@ -82,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"); @@ -109,4 +137,28 @@ mod tests { Some(json!({"connector": "x"})) ); } + + #[test] + fn prepare_generate_request_gates_continuous_usage_on_include_usage() { + let request: GenerateRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22], + "stream": true, + "stream_options": { + "continuous_usage_stats": true + }, + "sampling_params": {} + })) + .expect("parse request"); + + let prepared = prepare_generate_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("prepare"); + + assert!(!prepared.include_usage); + assert!(!prepared.include_continuous_usage); + } } diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index de7a196c3c6..d4567c44aa6 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -5,7 +5,7 @@ use serde_json::{Map, Value}; use validator::Validate; use vllm_text::SamplingParams; -use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable}; +use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage}; /// vLLM-compatible request type for the token-in/token-out generate API. #[serde_with::skip_serializing_none] @@ -17,6 +17,7 @@ pub struct GenerateRequest { pub sampling_params: SamplingParams, #[serde(default)] pub stream: bool, + pub stream_options: Option, pub cache_salt: Option, #[serde(default)] pub priority: i32, @@ -37,6 +38,25 @@ pub(super) struct GenerateResponseChoice { pub token_ids: Vec, } +/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, Serialize)] +pub(super) struct GenerateResponseStreamChoice { + pub index: u32, + pub logprobs: Option, + pub finish_reason: Option, + pub token_ids: Vec, +} + +/// Mirrors the Python vLLM `GenerateStreamResponse` class. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, Serialize)] +pub(super) struct GenerateStreamResponse { + pub request_id: String, + pub choices: Vec, + pub usage: Option, +} + /// Mirrors the Python vLLM `GenerateResponse` class. #[serde_with::skip_serializing_none] #[derive(Debug, Clone, Serialize)] diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs index 74a5bbb690a..43347c60b57 100644 --- a/rust/src/server/src/routes/inference/generate/validate.rs +++ b/rust/src/server/src/routes/inference/generate/validate.rs @@ -13,8 +13,11 @@ pub(super) fn validate_request_compat( return Err(ApiError::model_not_found(model.clone())); } - if request.stream { - bail_invalid_request!(param = "stream", "stream=true is not supported."); + if request.stream_options.is_some() && !request.stream { + bail_invalid_request!( + param = "stream_options", + "stream_options are only supported when stream=true." + ); } if request.token_ids.is_empty() { @@ -65,11 +68,24 @@ mod tests { } #[test] - fn validate_request_compat_rejects_streaming() { + fn validate_request_compat_accepts_streaming() { let request = GenerateRequest { stream: true, ..base_request() }; + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok()); + } + + #[test] + fn validate_request_compat_rejects_stream_options_without_streaming() { + let request: GenerateRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "token_ids": [11, 22], + "stream": false, + "stream_options": {"include_usage": true}, + "sampling_params": {} + })) + .expect("parse request"); assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err()); } 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 e1a16abcda5..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,459 @@ 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() { + let mut app = test_app().await; + let response = app + .call(Request::builder().uri("/version").body(Body::empty()).expect("build request")) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!( + json, + json!({ + "version": "test-vllm-version", + "rust_frontend_version": env!("CARGO_PKG_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() { @@ -2404,8 +2899,72 @@ async fn non_stream_raw_generate_returns_token_output_envelope() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] -async fn raw_generate_rejects_streaming() { - let mut app = test_app().await; +async fn stream_raw_generate_returns_sse_chunks_and_usage() { + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-raw-generate-stream".to_vec(); + + let engine_task = MockEngineTask::new(spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + |dealer, push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..])); + assert_eq!(request.external_req_id.as_deref(), Some("raw-stream")); + + send_outputs( + push, + EngineCoreOutputs { + engine_index: 0, + outputs: vec![ + request_output_with_logprobs( + &request.request_id, + vec![33], + None, + None, + Some(sample_logprobs_for_token(33, 34)), + None, + ), + request_output_with_logprobs( + &request.request_id, + vec![44], + Some(EngineCoreFinishReason::Stop), + None, + Some(sample_logprobs_for_token(44, 45)), + None, + ), + ], + scheduler_stats: None, + timestamp: 0.0, + utility_output: None, + finished_requests: None, + wave_complete: None, + start_wave: None, + }, + ) + .await; + }) + }, + )); + + let client = EngineCoreClient::connect( + EngineCoreClientConfig::new_single(handshake_address) + .with_model_name("test-model") + .with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ), + ) + .await + .expect("connect client"); + let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new())); + let mut app = build_router(Arc::new(AppState::new( + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], + chat, + ))); let response = app .call( @@ -2416,9 +2975,17 @@ async fn raw_generate_rejects_streaming() { .body(Body::from( json!({ "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "raw-stream", "token_ids": [11, 22], "stream": true, - "sampling_params": {} + "stream_options": { + "include_usage": true, + "continuous_usage_stats": true + }, + "sampling_params": { + "max_tokens": 2, + "logprobs": 1 + } }) .to_string(), )) @@ -2427,10 +2994,196 @@ async fn raw_generate_rejects_streaming() { .await .expect("call app"); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers().get("content-type").and_then(|value| value.to_str().ok()), + Some("text/event-stream") + ); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); - assert_eq!(json["error"]["param"], "stream"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_data_payloads(&text); + assert_eq!(payloads.len(), 4, "{text}"); + + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json"); + assert_eq!(first["request_id"], "raw-stream"); + assert_eq!(first["choices"][0]["index"], 0); + assert_eq!(first["choices"][0]["token_ids"], json!([33])); + assert_eq!( + first["choices"][0]["logprobs"]["content"][0]["token"], + "token_id:33" + ); + assert_eq!(first["usage"]["prompt_tokens"], 2); + assert_eq!(first["usage"]["completion_tokens"], 1); + + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json"); + assert_eq!(second["choices"][0]["token_ids"], json!([44])); + assert_eq!(second["choices"][0]["finish_reason"], "stop"); + assert_eq!(second["usage"]["completion_tokens"], 2); + + let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json"); + assert_eq!(usage["choices"], json!([])); + assert_eq!(usage["usage"]["prompt_tokens"], 2); + assert_eq!(usage["usage"]["completion_tokens"], 2); + assert_eq!(usage["usage"]["total_tokens"], 4); + assert_eq!(payloads[3], "[DONE]"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_raw_generate_emits_final_usage_without_continuous_usage() { + let (mut app, engine_task) = test_app_with_stream_output_specs(vec![ + (vec![33], None), + (vec![44], Some(EngineCoreFinishReason::Stop)), + ]) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "raw-stream-final-usage", + "token_ids": [11, 22], + "stream": true, + "stream_options": { + "include_usage": true + }, + "sampling_params": { + "max_tokens": 2 + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_data_payloads(&text); + assert_eq!(payloads.len(), 4, "{text}"); + + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json"); + assert_eq!(first["choices"][0]["token_ids"], json!([33])); + assert!(first.get("usage").is_none()); + + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json"); + assert_eq!(second["choices"][0]["token_ids"], json!([44])); + assert_eq!(second["choices"][0]["finish_reason"], "stop"); + assert!(second.get("usage").is_none()); + + let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json"); + assert_eq!(usage["choices"], json!([])); + assert_eq!(usage["usage"]["prompt_tokens"], 2); + assert_eq!(usage["usage"]["completion_tokens"], 2); + assert_eq!(usage["usage"]["total_tokens"], 4); + assert_eq!(payloads[3], "[DONE]"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_raw_generate_emits_empty_finish_chunk() { + let (mut app, engine_task) = test_app_with_stream_output_specs(vec![ + (vec![33], None), + (vec![], Some(EngineCoreFinishReason::Stop)), + ]) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "raw-stream-empty-finish", + "token_ids": [11, 22], + "stream": true, + "sampling_params": { + "max_tokens": 2 + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_data_payloads(&text); + assert_eq!(payloads.len(), 3, "{text}"); + + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json"); + assert_eq!(first["choices"][0]["token_ids"], json!([33])); + assert!(first["choices"][0].get("finish_reason").is_none()); + + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json"); + assert_eq!(second["choices"][0]["token_ids"], json!([])); + assert_eq!(second["choices"][0]["finish_reason"], "stop"); + assert_eq!(payloads[2], "[DONE]"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_raw_generate_error_finish_returns_sse_error() { + let (mut app, engine_task) = + test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))]) + .await; + + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "request_id": "raw-stream-error", + "token_ids": [11, 22], + "stream": true, + "stream_options": { + "include_usage": true + }, + "sampling_params": { + "max_tokens": 2 + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + + assert!(text.contains("\"type\":\"server_error\""), "{text}"); + assert!(text.contains("Internal server error"), "{text}"); + assert!(!text.contains("\"finish_reason\":\"error\""), "{text}"); + assert!(!text.contains("\"usage\":"), "{text}"); + assert!(text.trim_end().ends_with("data: [DONE]"), "{text}"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -2658,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/routes/version.rs b/rust/src/server/src/routes/version.rs new file mode 100644 index 00000000000..07a92e6d88b --- /dev/null +++ b/rust/src/server/src/routes/version.rs @@ -0,0 +1,23 @@ +use std::sync::Arc; + +use axum::Json; +use axum::extract::State; +use serde::Serialize; + +use crate::state::AppState; + +#[derive(Serialize)] +pub(crate) struct VersionResponse { + version: String, + rust_frontend_version: &'static str, +} + +/// Get engine and Rust frontend version metadata. +pub async fn version(State(state): State>) -> Json { + let version = state.engine_core_client().vllm_version().to_string(); + + Json(VersionResponse { + version, + rust_frontend_version: env!("CARGO_PKG_VERSION"), + }) +} 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/lib.rs b/rust/src/text/src/lib.rs index ef5615ec6d4..48828045a2d 100644 --- a/rust/src/text/src/lib.rs +++ b/rust/src/text/src/lib.rs @@ -45,9 +45,9 @@ pub struct TextLlm { /// Tokenizer/model metadata backend responsible for prompt encode/decode /// and sampling hints. backend: DynTextBackend, - /// Context window size derived by the backend or from engine startup - /// handshake, with optional override from config. - max_model_len: Option, + /// Context window size reported by the engine startup handshake, with + /// optional override from config. + max_model_len: u32, } impl TextLlm { @@ -71,7 +71,7 @@ impl TextLlm { /// This takes priority over both the engine-reported default and any /// tokenizer/model metadata exposed by the backend. pub fn with_max_model_len(mut self, max_model_len: u32) -> Self { - self.max_model_len = Some(max_model_len); + self.max_model_len = max_model_len; self } @@ -129,9 +129,7 @@ impl TextLlm { }; let mut sampling_hints = self.backend.sampling_hints()?; - if let Some(max_model_len) = self.max_model_len { - sampling_hints.max_model_len = Some(max_model_len); - } + sampling_hints.max_model_len = Some(self.max_model_len); let PreparedTextRequest { text_request, generate_request, 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/hermes.rs b/rust/src/tool-parser/src/json/hermes.rs index c8e6a49a7c5..f6b130ec472 100644 --- a/rust/src/tool-parser/src/json/hermes.rs +++ b/rust/src/tool-parser/src/json/hermes.rs @@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig { marker_whitespace: JsonToolCallWhitespace::Optional, delimiter: None, name_key: "name", - arguments_key: "arguments", + arguments_key: &["arguments"], }; /// Tool parser for Hermes XML-wrapped JSON tool calls. diff --git a/rust/src/tool-parser/src/json/internlm2.rs b/rust/src/tool-parser/src/json/internlm2.rs new file mode 100644 index 00000000000..8284a4d0e1d --- /dev/null +++ b/rust/src/tool-parser/src/json/internlm2.rs @@ -0,0 +1,345 @@ +use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; +use crate::{Result, Tool, ToolParser, ToolParserOutput}; + +const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig { + parser_name: "InternLM2", + start_marker: "<|action_start|><|plugin|>", + end_marker: "<|action_end|>", + marker_whitespace: JsonToolCallWhitespace::Optional, + delimiter: None, + name_key: "name", + // The Python parser's `get_arguments()` accepts either `parameters` or + // `arguments` and prefers `parameters` when both are present. This Rust + // parser uses first-encountered semantics because the header parser only + // permits one args key per tool-call object; if a future model emits + // both keys in the same object, the Rust port will accept the first one + // and reject the trailing one as a syntax error rather than silently + // shadowing it. + arguments_key: &["parameters", "arguments"], +}; + +/// Tool parser for InternLM2 special-token wrapped JSON tool calls. +/// +/// Example tool call content: +/// +/// ```text +/// <|action_start|><|plugin|>{"name": "get_weather", "parameters": {"location":"Tokyo"}}<|action_end|> +/// ``` +/// +/// Arguments are already OpenAI-style JSON text, so they are streamed as raw +/// argument deltas without schema conversion or JSON normalization. +/// +/// # Divergences from the Python reference +/// +/// 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. +/// - **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. +/// +/// # Known unaddressed divergences (TODO) +/// +/// The following Python behaviors are NOT yet matched. They are deferred to +/// follow-up work because they require non-local changes to the shared +/// `JsonToolCallParser` core that would affect Hermes / Llama / Mistral / +/// 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- +/// 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. +pub struct Internlm2ToolParser { + inner: JsonToolCallParser, +} + +impl Internlm2ToolParser { + /// Create an InternLM2 tool parser. + fn new(_tools: &[Tool]) -> Self { + Self { + inner: JsonToolCallParser::new(INTERNLM2_CONFIG), + } + } +} + +impl ToolParser for Internlm2ToolParser { + /// Create a boxed InternLM2 tool parser. + fn create(tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self::new(tools))) + } + + /// Preserve special-token markers while decoding, since + /// `<|action_start|>`, `<|plugin|>`, and `<|action_end|>` are tokenizer + /// special tokens in InternLM2 models. + fn preserve_special_tokens(&self) -> bool { + true + } + + /// Feed one decoded text chunk through the InternLM2 parser. + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + self.inner.parse_into(chunk, output) + } + + /// Flush any buffered partial state at end of stream. + fn finish(&mut self) -> Result { + self.inner.finish() + } + + /// Clear parser state and return currently uncommitted buffered text. + fn reset(&mut self) -> String { + self.inner.reset() + } +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use thiserror_ext::AsReport; + + use super::Internlm2ToolParser; + use crate::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + + const ACTION_START: &str = "<|action_start|><|plugin|>"; + const ACTION_END: &str = "<|action_end|>"; + + fn build_tool_call(function_name: &str, args_key: &str, arguments: &str) -> String { + format!( + r#"{ACTION_START}{{"name":"{function_name}","{args_key}":{arguments}}}{ACTION_END}"# + ) + } + + #[test] + fn internlm2_parse_complete_without_tool_call_keeps_text() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let result = parser.parse_complete("Hello, world!").unwrap(); + + assert_eq!(result.normal_text, "Hello, world!"); + assert!(result.calls.is_empty()); + } + + #[test] + fn internlm2_parse_complete_extracts_parameters_key() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let arguments = r#"{"location":"Tokyo","days":"3"}"#; + let result = parser + .parse_complete(&format!( + "Let me check.\n{}", + build_tool_call("get_weather", "parameters", arguments) + )) + .unwrap(); + + assert_eq!(result.normal_text, "Let me check.\n"); + assert_eq!(result.calls.len(), 1); + assert_eq!(result.calls[0].tool_index, 0); + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls[0].arguments, arguments); + } + + #[test] + fn internlm2_parse_complete_extracts_arguments_key_fallback() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let arguments = r#"{"location":"Tokyo"}"#; + let result = parser + .parse_complete(&build_tool_call("get_weather", "arguments", arguments)) + .unwrap(); + + assert_eq!(result.calls.len(), 1); + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls[0].arguments, arguments); + } + + #[test] + fn internlm2_accepts_whitespace_after_plugin_marker() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let result = parser + .parse_complete(&format!( + r#"{ACTION_START} +{{"name":"get_weather","parameters":{{}}}}{ACTION_END}"# + )) + .unwrap(); + + assert_eq!(result.calls.len(), 1); + assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); + } + + #[test] + fn internlm2_does_not_validate_or_normalize_arguments() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let arguments = r#"{"location":"Tokyo",}"#; + let result = parser + .parse_complete(&build_tool_call("get_weather", "parameters", arguments)) + .unwrap(); + + assert_eq!(result.calls[0].arguments, arguments); + } + + #[test] + fn internlm2_streaming_emits_argument_deltas() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let chunks = [ + "preface <|action", + "_start|><|plugin|>", + r#"{"name":"get_weather","parameters":"#, + r#"{"location":"#, + r#""Beijing""#, + r#"}"#, + r#"}<|action_end|> suffix"#, + ]; + + let mut result = ToolParserOutput::default(); + let mut observed_arguments = Vec::new(); + for chunk in chunks { + let next = parser.parse_chunk(chunk).unwrap(); + observed_arguments.extend( + next.calls + .iter() + .filter(|call| call.name.is_none()) + .map(|call| call.arguments.clone()), + ); + result.append(next); + } + result.append(parser.finish().unwrap()); + + assert_eq!( + observed_arguments, + [r#"{"location":"#, r#""Beijing""#, r#"}"#] + ); + assert_eq!(result.normal_text, "preface suffix"); + assert_eq!( + result.coalesce_calls().calls[0].arguments, + r#"{"location":"Beijing"}"# + ); + } + + #[test] + fn internlm2_streaming_handles_split_markers() { + let input = format!( + "hello {}", + build_tool_call("get_weather", "parameters", r#"{"location":"Tokyo"}"#) + ); + let chunks = split_by_chars(&input, 5); + let mut parser = Internlm2ToolParser::new(&test_tools()); + + let result = collect_stream(&mut parser, &chunks); + + assert_eq!(result.normal_text, "hello "); + assert_eq!(result.calls.len(), 1); + assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#); + } + + #[test] + fn internlm2_streaming_extracts_multiple_blocks() { + let input = format!( + "{}{}", + build_tool_call("get_weather", "parameters", r#"{"location":"Shanghai"}"#), + build_tool_call("add", "arguments", r#"{"x":1,"y":2}"#), + ); + let chunks = split_by_chars(&input, 7); + let mut parser = Internlm2ToolParser::new(&test_tools()); + + let result = collect_stream(&mut parser, &chunks); + + expect![[r#" + ToolParserOutput { + normal_text: "", + calls: [ + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ], + } + "#]] + .assert_debug_eq(&result); + } + + #[test] + fn internlm2_keeps_end_marker_literal_inside_json_string() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let arguments = format!(r#"{{"text":"literal {ACTION_END} inside"}}"#); + let input = build_tool_call("echo", "parameters", &arguments); + + let result = parser.parse_complete(&input).unwrap(); + + assert_eq!(result.calls.len(), 1); + assert_eq!(result.calls[0].arguments, arguments); + } + + #[test] + fn internlm2_finish_errors_on_truncated_tool_call() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let pre_finish = parser + .parse_chunk(&format!( + r#"{ACTION_START}{{"name":"get_weather","parameters":{{"location""# + )) + .unwrap(); + let error = parser.finish().unwrap_err(); + + assert_eq!( + pre_finish.calls[0].name.as_deref(), + Some("get_weather"), + "name delta is still emitted from parse_chunk() before truncation", + ); + assert!( + error.to_report_string().contains("incomplete InternLM2 tool call"), + "finish() reports the truncated tool call as incomplete: {}", + error.to_report_string(), + ); + } + + #[test] + fn internlm2_unknown_arguments_key_fails() { + let mut parser = Internlm2ToolParser::new(&test_tools()); + let input = build_tool_call("get_weather", "params", r#"{"location":"Tokyo"}"#); + + let error = parser.parse_chunk(&input).unwrap_err(); + + expect![[r#" + tool parser parsing failed: invalid InternLM2 + expected `parameters`, `arguments`"#]] + .assert_eq(&error.to_report_string()); + } + + #[test] + fn internlm2_preserve_special_tokens_is_true() { + let parser = Internlm2ToolParser::new(&test_tools()); + assert!(parser.preserve_special_tokens()); + } +} diff --git a/rust/src/tool-parser/src/json/llama.rs b/rust/src/tool-parser/src/json/llama.rs index 5ad6918b617..36bc8a8347d 100644 --- a/rust/src/tool-parser/src/json/llama.rs +++ b/rust/src/tool-parser/src/json/llama.rs @@ -203,7 +203,7 @@ fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult, name_key: &'static str, - arguments_key: &'static str, + /// Candidate JSON keys naming the arguments payload, tried in order. + /// Most parsers use a single key like `["arguments"]`, but some accept + /// multiple (e.g. InternLM2 accepts `parameters` or `arguments`). + arguments_key: &'static [&'static str], } #[derive(Debug, Clone, Copy)] @@ -224,7 +229,7 @@ fn tool_call_header_event( _: ws0, _: literal(","), _: ws0, - _: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key), + _: |input: &mut JsonToolInput<'_>| json_arguments_key(input, config.arguments_key), _: ws0, _: literal(":"), _: ws0, @@ -246,6 +251,39 @@ fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()> .parse_next(input) } +/// Parse a JSON object key accepting any of `candidates`. +/// +/// The full quoted key is consumed and compared against the candidate list, +/// so this works correctly under partial input regardless of key lengths. +/// +/// On mismatch, each candidate is attached as its own `Expected` context so the +/// error enumerates every valid key ("expected `a`, expected `b`"). Because +/// `StrContextValue::StringLiteral` carries a single `&'static str`, the +/// contexts are added in a loop over `candidates` rather than through chained +/// `.context(...)` calls, which keeps the diagnostics complete for any number +/// of candidates. +fn json_arguments_key( + input: &mut JsonToolInput<'_>, + candidates: &'static [&'static str], +) -> ModalResult<()> { + let start = input.checkpoint(); + json_str + .verify(|key: &String| candidates.contains(&key.as_str())) + .void() + .parse_next(input) + .map_err(|err| { + err.map(|context_error| { + candidates.iter().fold(context_error, |context_error, candidate| { + context_error.add_context( + &*input, + &start, + StrContext::Expected(StrContextValue::StringLiteral(candidate)), + ) + }) + }) + }) +} + /// Parse one event inside a marker-wrapped JSON tool-call arguments payload. fn parse_arguments_event( input: &mut JsonToolInput<'_>, @@ -341,7 +379,7 @@ mod tests { marker_whitespace: JsonToolCallWhitespace::Optional, delimiter: Some("<"), name_key: "function", - arguments_key: "parameters", + arguments_key: &["parameters"], }; fn build_tool_call(function_name: &str, arguments: &str) -> String { diff --git a/rust/src/tool-parser/src/json/qwen.rs b/rust/src/tool-parser/src/json/qwen.rs index f58ca6e0fa6..b8caff0fefd 100644 --- a/rust/src/tool-parser/src/json/qwen.rs +++ b/rust/src/tool-parser/src/json/qwen.rs @@ -8,7 +8,7 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { marker_whitespace: JsonToolCallWhitespace::Exact("\n"), delimiter: None, name_key: "name", - arguments_key: "arguments", + arguments_key: &["arguments"], }; /// Tool parser for Qwen XML-wrapped JSON tool calls. diff --git a/rust/src/tool-parser/src/lib.rs b/rust/src/tool-parser/src/lib.rs index d8f648b81d5..f1dc0455843 100644 --- a/rust/src/tool-parser/src/lib.rs +++ b/rust/src/tool-parser/src/lib.rs @@ -24,7 +24,10 @@ pub use error::{Result, ToolParserError}; pub use gemma4::Gemma4ToolParser; pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser}; pub use hy_v3::HyV3ToolParser; -pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser}; +pub use json::{ + HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser, + Qwen3XmlToolParser, +}; pub use kimi_k2::KimiK2ToolParser; pub use minimax_m2::MinimaxM2ToolParser; pub use qwen_coder::Qwen3CoderToolParser; 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/setup.py b/setup.py index a95ee3451b5..07374807bee 100644 --- a/setup.py +++ b/setup.py @@ -1165,9 +1165,7 @@ setup( install_requires=get_requirements(), extras_require={ # AMD Zen CPU optimizations via zentorch - "zen": [ - "zentorch-weekly==5.2.1.dev20260408" - ], # Zentorch has weekly releases. This pulls the known-good version. + "zen": ["zentorch==2.11.0.0"], "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"], "tensorizer": ["tensorizer==2.10.1"], "fastsafetensors": ["fastsafetensors >= 0.2.2"], @@ -1195,11 +1193,6 @@ setup( "opentelemetry-exporter-otlp>=1.26.0", "opentelemetry-semantic-conventions-ai>=0.4.1", ], - "triton-cpu": [ - "triton @ " - "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; " - "platform_machine == 'x86_64'", - ], # Remove after stable release }, cmdclass=cmdclass, package_data=package_data, 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/fusions_e2e/common.py b/tests/compile/fusions_e2e/common.py index 2c6dc2b3ebb..c239e93785b 100644 --- a/tests/compile/fusions_e2e/common.py +++ b/tests/compile/fusions_e2e/common.py @@ -20,6 +20,7 @@ class Matches(NamedTuple): attn_quant_fusion: int = 0 # distributed ar_rms_fusion: int = 0 + aiter_ar_rms_fusion: int = 0 sequence_parallel: int = 0 async_tp: int = 0 @@ -97,6 +98,9 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = { "ar_rms_fusion": re.compile( r"allreduce_rms_fusion.py:\d+] Replaced (\d+) patterns" ), + "aiter_ar_rms_fusion": re.compile( + r"RocmAiterAllReduceFusionPass Replaced (\d+) patterns" + ), "sequence_parallel": re.compile( r"sequence_parallelism.py:\d+] Replaced (\d+) patterns" ), diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 3a060874720..9f34d25c46d 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -97,6 +97,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): f"attention backend '{attn_backend.backend.name}'" ) + if backend_name == "rocm_attn" and model_name == "openai/gpt-oss-20b": + pytest.skip( + "ROCM_ATTN does not support attention sinks (required by gpt-oss-20b)" + ) + if attn_backend.backend.name == "FLASHINFER": from vllm.utils.flashinfer import supports_trtllm_attention diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 32f1ea35063..847783d7033 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -75,6 +75,7 @@ llama3_8b = ModelFusionInfo( model_name="meta-llama/Llama-3.1-8B-Instruct", matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, + aiter_ar_rms_fusion=n_layers * 2, sequence_parallel=n_layers * 2 + 1, async_tp=n_layers * 4, ), @@ -136,6 +137,7 @@ qwen3_a3b = ModelFusionInfo( matches=lambda n_layers: Matches( norm_rope_fusion=n_layers, ar_rms_fusion=n_layers * 2 + 1, + aiter_ar_rms_fusion=n_layers * 2, sequence_parallel=n_layers * 2 + 1, async_tp=n_layers * 2, ), @@ -211,6 +213,7 @@ gpt_oss_20b = ModelFusionInfo( model_name="openai/gpt-oss-20b", matches=lambda n_layers: Matches( ar_rms_fusion=n_layers * 2 + 1, + aiter_ar_rms_fusion=n_layers + 1, sequence_parallel=n_layers * 2 + 1, async_tp=n_layers * 2, ), diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index b5e2b2dc07e..b18c41658fd 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -84,6 +84,7 @@ def test_tp2_ar_rms_fp8_fusions( model_kwargs["load_format"] = "dummy" model_kwargs["max_model_len"] = 1024 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + model_kwargs["disable_custom_all_reduce"] = False compilation_config = dict( use_inductor_graph_partition=inductor_graph_partition, @@ -149,6 +150,7 @@ def test_tp2_ar_rms_fp4_fusions( model_kwargs["load_format"] = "dummy" model_kwargs["max_model_len"] = 1024 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + model_kwargs["disable_custom_all_reduce"] = False compilation_config = dict( use_inductor_graph_partition=inductor_graph_partition, @@ -213,6 +215,7 @@ def test_tp2_ar_rms_fusions( model_kwargs["load_format"] = "dummy" model_kwargs["max_model_len"] = 1024 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + model_kwargs["disable_custom_all_reduce"] = False compilation_config = dict( use_inductor_graph_partition=inductor_graph_partition, @@ -225,9 +228,13 @@ def test_tp2_ar_rms_fusions( matches_check = [ "norm_rope_fusion", - "ar_rms_fusion", ] + if current_platform.is_rocm(): + matches_check.append("aiter_ar_rms_fusion") + else: + matches_check.append("ar_rms_fusion") + run_e2e_fusion_test( model_name, matches, diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 33e050d776e..7d1320bccaa 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import ( initialize_model_parallel, ) from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed @@ -234,8 +235,20 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel): TestAGMMModel, TestScaledMMRSModel, TestAGScaledMMModel, - TestCutlassScaledMMRSModel, - TestAGCutlassScaledMMModel, + pytest.param( + TestCutlassScaledMMRSModel, + marks=pytest.mark.skipif( + not hasattr(torch.ops._C, "cutlass_scaled_mm"), + reason="Requires cutlass_scaled_mm", + ), + ), + pytest.param( + TestAGCutlassScaledMMModel, + marks=pytest.mark.skipif( + not hasattr(torch.ops._C, "cutlass_scaled_mm"), + reason="Requires cutlass_scaled_mm", + ), + ), ], ) @pytest.mark.parametrize("batch_size", [8]) @@ -268,6 +281,7 @@ def test_async_tp_pass_replace( ) num_processes = 2 + master_port = str(get_open_port()) def run_torch_spawn(fn, nprocs): # need to use torch.mp.spawn otherwise will have problems with @@ -282,6 +296,7 @@ def test_async_tp_pass_replace( hidden_size, dtype, dynamic, + master_port, ), nprocs=nprocs, ) @@ -314,6 +329,7 @@ def async_tp_pass_on_test_model( hidden_size: int, dtype: torch.dtype, dynamic: bool, + master_port: str = "0", ): set_random_seed(0) @@ -328,7 +344,7 @@ def async_tp_pass_on_test_model( "LOCAL_RANK": str(local_rank), "WORLD_SIZE": str(world_size), "MASTER_ADDR": "localhost", - "MASTER_PORT": "12345", + "MASTER_PORT": master_port, } ) 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_comm_ops.py b/tests/distributed/test_comm_ops.py index 2804c95d32a..48b007da664 100644 --- a/tests/distributed/test_comm_ops.py +++ b/tests/distributed/test_comm_ops.py @@ -292,7 +292,7 @@ def test_async_intermediate_tensors_lazy_wait() -> None: ) # accessing non-tensor attributes should not trigger wait. - assert it.kv_connector_output is None + assert it._comm_handles is not None assert work.wait_calls == 0 assert post_calls["n"] == 0 diff --git a/tests/distributed/test_eplb_spec_decode.py b/tests/distributed/test_eplb_spec_decode.py index 22977ce9440..b4211c74c17 100644 --- a/tests/distributed/test_eplb_spec_decode.py +++ b/tests/distributed/test_eplb_spec_decode.py @@ -15,7 +15,7 @@ def get_model_args( spec_method: str, tp_size: int, model_max_len: int, - use_async: bool = False, + use_async: bool = True, ) -> dict: speculative_config = { "method": spec_method, @@ -28,9 +28,8 @@ def get_model_args( "window_size": 128, "step_interval": 1024, "log_balancedness": False, + "use_async": use_async, } - if use_async: - eplb_config["use_async"] = True model_args = { "pretrained": model_name, "dtype": "auto", 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_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 2742663093f..c2dda1b51cf 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -127,7 +127,6 @@ TEXT_GENERATION_MODELS = { # Uses Llama # "internlm/internlm-chat-7b": PPTestSettings.fast(), "internlm/internlm2-chat-7b": PPTestSettings.fast(), - "inceptionai/jais-13b-chat": PPTestSettings.fast(), "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(), "pfnet/plamo-2-1b": PPTestSettings.fast(), "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(), 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 295e812a124..2df0d9e71c3 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -18,6 +18,7 @@ from torch.multiprocessing.reductions import reduce_tensor from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer import WeightTransferEngineFactory +from vllm.distributed.weight_transfer.base import SparseWeightPatch from vllm.distributed.weight_transfer.ipc_engine import ( IPCWeightTransferEngine, IPCWeightTransferInitInfo, @@ -89,6 +90,67 @@ class TestNCCLWeightTransferUpdateInfoValidation: ) assert len(info.names) == 0 + def test_valid_sparse_update_info(self): + """Test creating valid sparse NCCL update info.""" + info = NCCLWeightTransferUpdateInfo( + names=["layer.weight", "layer.bias"], + dtype_names=["float32", "bfloat16"], + shapes=[[10, 10], [10]], + num_updates_list=[4, 2], + update_kind="sparse_flat", + ) + assert info.update_kind == "sparse_flat" + assert info.num_updates_list == [4, 2] + + def test_sparse_update_requires_num_updates_list(self): + with pytest.raises(ValueError, match="`num_updates_list` is required"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + update_kind="sparse_flat", + ) + + def test_sparse_update_rejects_empty_num_updates_list(self): + with pytest.raises(ValueError, match="cannot be empty"): + NCCLWeightTransferUpdateInfo( + names=[], + dtype_names=[], + shapes=[], + num_updates_list=[], + update_kind="sparse_flat", + ) + + def test_sparse_update_rejects_packed(self): + with pytest.raises(ValueError, match="cannot be combined with `packed=True`"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + num_updates_list=[3], + update_kind="sparse_flat", + packed=True, + ) + + def test_sparse_update_rejects_mismatched_num_updates(self): + with pytest.raises(ValueError, match="`num_updates_list`"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight", "layer.bias"], + dtype_names=["float32", "float32"], + shapes=[[10, 10], [10]], + num_updates_list=[3], + update_kind="sparse_flat", + ) + + def test_dense_update_rejects_sparse_metadata(self): + with pytest.raises(ValueError, match="Sparse metadata"): + NCCLWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + num_updates_list=[3], + ) + # --- Unit Tests: Engine Parsing --- @@ -222,6 +284,29 @@ def test_nccl_receive_weights_without_init_raises(): engine.receive_weights(update_info, lambda x: None) +def test_nccl_receive_sparse_weights_without_init_raises(): + """Test that sparse receive raises if init_transfer_engine wasn't called.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + config = WeightTransferConfig(backend="nccl") + parallel_config = create_mock_parallel_config() + engine = NCCLWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) + + update_info = NCCLWeightTransferUpdateInfo( + names=["w"], + dtype_names=["float32"], + shapes=[[10]], + num_updates_list=[2], + update_kind="sparse_flat", + ) + + with pytest.raises(RuntimeError, match="not initialized"): + engine.receive_sparse_weights(update_info, lambda x: None) + + # --- Integration Test: NCCL Weight Transfer Between Ray Tasks --- @@ -379,6 +464,138 @@ def test_nccl_weight_transfer_between_processes(): ) +@ray.remote(num_gpus=1) +def trainer_broadcast_sparse_tensor( + master_address: str, + master_port: int, + world_size: int, +) -> bool: + """Trainer task that broadcasts sparse patches via NCCL.""" + import torch + + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + from vllm.distributed.weight_transfer.base import SparseWeightPatch + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerSendWeightsArgs, + NCCLWeightTransferEngine, + ) + + pg = StatelessProcessGroup.create( + host=master_address, + port=master_port, + rank=0, + world_size=world_size, + ) + comm = PyNcclCommunicator(pg, device=0) + + patch = SparseWeightPatch( + name="test.weight", + indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"), + values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"), + ) + NCCLWeightTransferEngine.trainer_send_sparse_weights( + iter([patch]), + NCCLTrainerSendWeightsArgs(group=comm), + ) + torch.accelerator.synchronize() + return True + + +@ray.remote(num_gpus=1) +def inference_receive_sparse_tensor( + master_address: str, + master_port: int, + world_size: int, +) -> dict: + """Inference task that receives sparse patches via NCCLWeightTransferEngine.""" + from unittest.mock import MagicMock + + import torch + + from vllm.config.parallel import ParallelConfig + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + NCCLWeightTransferInitInfo, + NCCLWeightTransferUpdateInfo, + ) + + config = WeightTransferConfig(backend="nccl") + parallel_config = MagicMock(spec=ParallelConfig) + parallel_config.rank = 0 + parallel_config.world_size = 1 + parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = 0 + + engine = NCCLWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) + engine.init_transfer_engine( + NCCLWeightTransferInitInfo( + master_address=master_address, + master_port=master_port, + rank_offset=1, + world_size=world_size, + ) + ) + + target = torch.zeros(30, dtype=torch.float32, device="cuda") + + def apply_sparse_patches(patches: list[SparseWeightPatch]): + for patch in patches: + target.index_copy_(0, patch.indices.to(torch.long), patch.values) + + update_info = NCCLWeightTransferUpdateInfo( + names=["test.weight"], + dtype_names=["float32"], + shapes=[[30]], + num_updates_list=[3], + update_kind="sparse_flat", + ) + engine.receive_sparse_weights(update_info, apply_sparse_patches) + torch.accelerator.synchronize() + + expected = torch.zeros(30, dtype=torch.float32, device="cuda") + expected[[1, 7, 25]] = torch.tensor( + [10.0, 20.0, 30.0], dtype=torch.float32, device="cuda" + ) + success = torch.equal(target, expected) + engine.shutdown() + return { + "success": success, + "selected_values": target[[1, 7, 25]].cpu().tolist(), + } + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 2, + reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.", +) +def test_nccl_sparse_weight_transfer_between_processes(): + """Test NCCL sparse weight transfer from trainer to inference process.""" + ray.init(ignore_reinit_error=True) + + master_address = "127.0.0.1" + master_port = get_open_port() + world_size = 2 + + inference_future = inference_receive_sparse_tensor.remote( + master_address, master_port, world_size + ) + trainer_future = trainer_broadcast_sparse_tensor.remote( + master_address, master_port, world_size + ) + + trainer_result, result = ray.get([trainer_future, inference_future]) + + assert trainer_result, "Trainer should complete successfully" + assert result["success"], ( + "Sparse weight transfer failed. " + f"Received selected values: {result['selected_values']}" + ) + + # --- Unit Tests: IPCWeightTransferUpdateInfo Validation --- @@ -461,9 +678,101 @@ class TestIPCWeightTransferUpdateInfoValidation: ipc_handles=ipc_handles, ) - def test_missing_ipc_handles_raises(self): - """Test that omitting ipc_handles raises TypeError.""" - with pytest.raises(TypeError): + def test_sparse_update_kind_rejected(self): + """Test that IPC backend rejects sparse update metadata.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + dummy_tensor = torch.ones(10, 10, device="cuda:0") + ipc_handle = reduce_tensor(dummy_tensor) + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid) + ipc_handles = [{gpu_uuid: ipc_handle}] + + with pytest.raises(NotImplementedError, match="dense updates"): + IPCWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + num_updates_list=[1], + ipc_handles=ipc_handles, + update_kind="sparse_flat", + ) + + def test_sparse_methods_not_supported(self): + """Test that IPC engine inherits sparse rejection from the base class.""" + config = WeightTransferConfig(backend="ipc") + parallel_config = create_mock_parallel_config() + engine = IPCWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) + + with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"): + engine.receive_sparse_weights(MagicMock(), lambda _: None) + with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"): + engine.trainer_send_sparse_weights( + iter([]), + {"mode": "http", "url": "http://localhost:8000"}, + ) + + def test_valid_update_info_from_pickled(self, monkeypatch): + """Test creating IPCWeightTransferUpdateInfo from pickled handles.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + dummy_tensor = torch.ones(10, 10, device="cuda:0") + ipc_handle = reduce_tensor(dummy_tensor) + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid) + ipc_handles = [{gpu_uuid: ipc_handle}] + + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8") + + info = IPCWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + ipc_handles_pickled=pickled, + ) + assert info.ipc_handles == ipc_handles + assert info.ipc_handles_pickled is None + + def test_pickled_requires_insecure_serialization_flag(self, monkeypatch): + """Test that pickled handles are rejected unless env flag is enabled.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0") + + with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"): + IPCWeightTransferUpdateInfo( + names=[], + dtype_names=[], + shapes=[], + ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"), + ) + + def test_both_handles_and_pickled_raises(self): + """Test that providing both ipc_handles and ipc_handles_pickled raises.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + dummy_tensor = torch.ones(10, 10, device="cuda:0") + ipc_handle = reduce_tensor(dummy_tensor) + gpu_uuid = str(torch.cuda.get_device_properties(0).uuid) + ipc_handles = [{gpu_uuid: ipc_handle}] + + pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8") + + with pytest.raises(ValueError, match="Cannot specify both"): + IPCWeightTransferUpdateInfo( + names=["layer.weight"], + dtype_names=["float32"], + shapes=[[10, 10]], + ipc_handles=ipc_handles, + ipc_handles_pickled=pickled, + ) + + def test_neither_handles_nor_pickled_raises(self): + """Test that providing neither ipc_handles nor ipc_handles_pickled raises.""" + with pytest.raises(ValueError, match="must be provided"): IPCWeightTransferUpdateInfo( names=["layer.weight"], dtype_names=["float32"], @@ -558,6 +867,28 @@ class TestIPCEngineParsing: assert gpu_uuid in update_info.ipc_handles[0] assert gpu_uuid in update_info.ipc_handles[1] + def test_parse_update_info_ignores_none_pickled_handles(self): + """Test Ray/asdict payloads with a null pickled field use ipc_handles.""" + config = WeightTransferConfig(backend="ipc") + parallel_config = create_mock_parallel_config() + engine = IPCWeightTransferEngine( + config, parallel_config, MagicMock(spec=torch.nn.Module) + ) + ipc_handles = [{"gpu-uuid": ("ipc-args",)}] + + update_info = engine.parse_update_info( + { + "names": ["w1"], + "dtype_names": ["float32"], + "shapes": [[1]], + "ipc_handles": ipc_handles, + "ipc_handles_pickled": None, + } + ) + + assert isinstance(update_info, IPCWeightTransferUpdateInfo) + assert update_info.ipc_handles == ipc_handles + def test_parse_update_info_both_handles_and_pickled_raises(self): """Test that providing both ipc_handles and ipc_handles_pickled raises.""" if torch.accelerator.device_count() < 1: 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/rpc/__init__.py b/tests/entrypoints/generate/__init__.py similarity index 100% rename from tests/entrypoints/rpc/__init__.py rename to tests/entrypoints/generate/__init__.py diff --git a/tests/entrypoints/openai/generative_scoring/__init__.py b/tests/entrypoints/generate/generative_scoring/__init__.py similarity index 100% rename from tests/entrypoints/openai/generative_scoring/__init__.py rename to tests/entrypoints/generate/generative_scoring/__init__.py diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py similarity index 95% rename from tests/entrypoints/openai/generative_scoring/test_generative_scoring.py rename to tests/entrypoints/generate/generative_scoring/test_generative_scoring.py index 632c4bcc90a..d8008299229 100644 --- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py @@ -18,13 +18,13 @@ from unittest.mock import MagicMock import pytest from vllm.config.multimodal import MultiModalConfig -from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.openai.generative_scoring.serving import ( +from vllm.entrypoints.generate.generative_scoring.serving import ( GenerativeScoringItemResult, GenerativeScoringRequest, GenerativeScoringResponse, - OpenAIServingGenerativeScoring, + ServingGenerativeScoring, ) +from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.logprobs import Logprob @@ -86,13 +86,13 @@ def _create_mock_engine(): return mock_engine -def _create_serving(mock_engine) -> OpenAIServingGenerativeScoring: - """Create an OpenAIServingGenerativeScoring instance with mocks.""" +def _create_serving(mock_engine) -> ServingGenerativeScoring: + """Create an ServingGenerativeScoring instance with mocks.""" models = OpenAIServingModels( engine_client=mock_engine, base_model_paths=BASE_MODEL_PATHS, ) - return OpenAIServingGenerativeScoring(mock_engine, models, request_logger=None) + return ServingGenerativeScoring(mock_engine, models, request_logger=None) def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput: @@ -186,7 +186,7 @@ class TestProbabilityComputation: self, label_logprobs, apply_softmax, should_sum_to_one ): """Test probability computation for softmax and true probability modes.""" - serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring) + serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring) probs = serving._compute_probabilities( label_logprobs, apply_softmax=apply_softmax ) @@ -211,7 +211,7 @@ class TestProbabilityComputation: def test_score_formula(self): """Test the score formula: P(token[0]) / (P(token[0]) + P(token[1])).""" - serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring) + serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring) # With logprobs -0.5 and -2.0, softmax gives higher prob to first token logprobs = {9454: -0.5, 2753: -2.0} diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py similarity index 99% rename from tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py rename to tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py index 64a59b270f1..4fe8dbe791b 100644 --- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py +++ b/tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py @@ -8,7 +8,7 @@ Tests verify the full HTTP request/response flow using RemoteOpenAIServer. import pytest import requests -from ....utils import RemoteOpenAIServer +from tests.utils import RemoteOpenAIServer MODEL_NAME = "Qwen/Qwen3-0.6B" diff --git a/tests/entrypoints/openai/chat_completion/test_audio_in_video.py b/tests/entrypoints/openai/chat_completion/test_audio_in_video.py index 61ee91eab4d..14ac7ce0439 100644 --- a/tests/entrypoints/openai/chat_completion/test_audio_in_video.py +++ b/tests/entrypoints/openai/chat_completion/test_audio_in_video.py @@ -14,8 +14,12 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer MODEL_NAME = "Qwen/Qwen2.5-Omni-3B" -@pytest.fixture +@pytest.fixture(scope="module") def server(): + # Use module scope so the server is started once and shared across all + # tests in this file. Starting a new vLLM server per test on XPU can + # cause the second server startup to hang silently and exceed the + # wait-for-server timeout, resulting in RuntimeError. args = [ "--max-model-len", "16384", 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_chat_with_tool_reasoning.py b/tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py deleted file mode 100644 index 295b5588941..00000000000 --- a/tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import openai # use the official client for correctness check -import pytest -import pytest_asyncio - -from tests.utils import RemoteOpenAIServer - -# a reasoning and tool calling model -MODEL_NAME = "Qwen/QwQ-32B" - - -@pytest.fixture(scope="module") -def server(): - args = [ - "--max-model-len", - "8192", - "--enforce-eager", - "--reasoning-parser", - "deepseek_r1", - "--enable-auto-tool-choice", - "--tool-call-parser", - "hermes", - ] - - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: - yield remote_server - - -@pytest_asyncio.fixture -async def client(server): - async with server.get_async_client() as async_client: - yield async_client - - -TOOLS = [ - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "city": { - "type": "string", - "description": "The city to find the weather for, e.g. " - "'San Francisco'", - }, - "state": { - "type": "string", - "description": "the two-letter abbreviation for the state that " - "the city is in, e.g. 'CA' which would mean 'California'", - }, - "unit": { - "type": "string", - "description": "The unit to fetch the temperature in", - "enum": ["celsius", "fahrenheit"], - }, - }, - "required": ["city", "state", "unit"], - }, - }, - } -] - -MESSAGES = [ - {"role": "user", "content": "Hi! How are you doing today?"}, - {"role": "assistant", "content": "I'm doing well! How can I help you?"}, - { - "role": "user", - "content": "Can you tell me what the temperate will be in Dallas, " - "in fahrenheit?", - }, -] - -FUNC_NAME = "get_current_weather" -FUNC_ARGS = """{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}""" - - -def extract_reasoning_and_calls(chunks: list): - reasoning = "" - tool_call_idx = -1 - arguments = [] - function_names = [] - for chunk in chunks: - if chunk.choices[0].delta.tool_calls: - tool_call = chunk.choices[0].delta.tool_calls[0] - if tool_call.index != tool_call_idx: - tool_call_idx = chunk.choices[0].delta.tool_calls[0].index - arguments.append("") - function_names.append("") - - if tool_call.function: - if tool_call.function.name: - function_names[tool_call_idx] = tool_call.function.name - - if tool_call.function.arguments: - arguments[tool_call_idx] += tool_call.function.arguments - else: - if hasattr(chunk.choices[0].delta, "reasoning"): - reasoning += chunk.choices[0].delta.reasoning - return reasoning, arguments, function_names - - -# test streaming -@pytest.mark.asyncio -async def test_chat_streaming_of_tool_and_reasoning(client: openai.AsyncOpenAI): - stream = await client.chat.completions.create( - model=MODEL_NAME, - messages=MESSAGES, - tools=TOOLS, - temperature=0.0, - stream=True, - ) - - chunks = [] - async for chunk in stream: - chunks.append(chunk) - - reasoning, arguments, function_names = extract_reasoning_and_calls(chunks) - assert len(reasoning) > 0 - assert len(function_names) > 0 and function_names[0] == FUNC_NAME - assert len(arguments) > 0 and arguments[0] == FUNC_ARGS - - -# test full generate -@pytest.mark.asyncio -async def test_chat_full_of_tool_and_reasoning(client: openai.AsyncOpenAI): - tool_calls = await client.chat.completions.create( - model=MODEL_NAME, - messages=MESSAGES, - tools=TOOLS, - temperature=0.0, - stream=False, - ) - - assert len(tool_calls.choices[0].message.reasoning) > 0 - assert tool_calls.choices[0].message.tool_calls[0].function.name == FUNC_NAME - assert tool_calls.choices[0].message.tool_calls[0].function.arguments == FUNC_ARGS diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 45fae821af3..7c0a46a4e63 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1449,91 +1449,6 @@ class TestServingChatWithHarmony: ], ) - @pytest.mark.asyncio - async def test_tools_and_reasoning( - self, serving_chat, stream, weather_tools, weather_messages_start - ): - tools = weather_tools - messages = list(weather_messages_start) - - # Test the Harmony messages for the first turn's input - req = ChatCompletionRequest(model=MODEL_NAME, messages=messages, tools=tools) - input_messages, _ = ( - serving_chat.openai_serving_render._make_request_with_harmony(req) - ) - verify_harmony_messages( - input_messages, - [ - {"role": "system"}, - {"role": "developer", "tool_definitions": ["get_weather"]}, - {"role": "user", "content": messages[0]["content"]}, - ], - ) - - # Test the Chat Completion response for the first turn's output - reasoning_str = "I'll call get_weather." - tool_args_str = '{"location": "Paris"}' - response_str = ( - f"<|channel|>analysis<|message|>{reasoning_str}<|end|>" - "<|start|>assistant to=functions.get_weather<|channel|>commentary" - f"<|constrain|>json<|message|>{tool_args_str}<|call|>" - ) - response = await self.generate_response_from_harmony_str( - serving_chat, req, response_str, stream=stream - ) - verify_chat_response( - response, - reasoning=reasoning_str, - tool_calls=[("get_weather", tool_args_str)], - ) - - tool_call = response.choices[0].message.tool_calls[0] - - # Add the output messages from the first turn as input to the second turn - for choice in response.choices: - messages.append(choice.message.model_dump(exclude_none=True)) - - # Add our tool output message - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "content": "20 degrees Celsius", - }, - ) - - # Test the Harmony messages for the second turn's input - req_2 = ChatCompletionRequest(model=MODEL_NAME, messages=messages, tools=tools) - input_messages_2, _ = ( - serving_chat.openai_serving_render._make_request_with_harmony(req_2) - ) - verify_harmony_messages( - input_messages_2, - [ - {"role": "system"}, - {"role": "developer"}, - {"role": "user"}, - { - "role": "assistant", - "channel": "analysis", - "content": reasoning_str, - }, - { - "role": "assistant", - "channel": "commentary", - "recipient": "functions.get_weather", - "content": tool_args_str, - }, - { - "role": "tool", - "author_name": "functions.get_weather", - "channel": "commentary", - "recipient": "assistant", - "content": "20 degrees Celsius", - }, - ], - ) - @pytest.mark.asyncio async def test_multi_turn_tools_and_reasoning( self, serving_chat, stream, weather_tools, weather_messages_start @@ -2020,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(), @@ -2044,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): @@ -2069,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/chat_completion/test_serving_chat_stream_harmony.py b/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py index 0a0802a7939..1c058adaf0a 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py @@ -121,7 +121,9 @@ class TestExtractHarmonyStreamingDelta: token_states = [ TokenState( - channel=channel, recipient="functions.get_weather", text=args_text + channel=channel, + recipient="functions.get_weather", + text=args_text, ) ] @@ -168,7 +170,11 @@ class TestExtractHarmonyStreamingDelta: parser = MockStreamableParser(messages=messages) token_states = [ - TokenState(channel="commentary", recipient="functions.tool2", text="args") + TokenState( + channel="commentary", + recipient="functions.tool2", + text="args", + ) ] delta_message, _ = extract_harmony_streaming_delta( @@ -199,75 +205,6 @@ class TestExtractHarmonyStreamingDelta: assert delta_message.content == delta_text assert tools_streamed is False - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_new_tool_call_without_functions_prefix( - self, mock_make_tool_call_id, channel - ): - mock_make_tool_call_id.return_value = "call_bare123" - parser = MockStreamableParser() - - token_states = [TokenState(channel=channel, recipient="get_weather", text="")] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert len(delta_message.tool_calls) == 1 - tool_call = delta_message.tool_calls[0] - assert tool_call.id == "call_bare123" - assert tool_call.type == "function" - assert tool_call.function.name == "get_weather" - assert tool_call.function.arguments == "" - assert tool_call.index == 0 - assert tools_streamed is True - - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - def test_tool_call_argument_streaming_without_functions_prefix(self, channel): - parser = MockStreamableParser() - args_text = '{"location": "Paris"}' - - token_states = [ - TokenState(channel=channel, recipient="get_weather", text=args_text) - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="get_weather", - include_reasoning=False, - ) - - assert delta_message is not None - tool_call = delta_message.tool_calls[0] - assert tool_call.id is None - assert tool_call.function.arguments == args_text - assert tool_call.index == 0 - assert tools_streamed is True - - def test_tool_call_index_from_previous_messages_without_functions_prefix(self): - messages = [ - MockMessage(channel="commentary", recipient="tool1"), - ] - parser = MockStreamableParser(messages=messages) - - token_states = [ - TokenState(channel="commentary", recipient="tool2", text="args") - ] - - delta_message, _ = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="tool2", - include_reasoning=False, - ) - - assert delta_message.tool_calls[0].index == 1 - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") def test_new_tool_call_dotted_function_name(self, mock_make_tool_call_id, channel): 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_dp_supervisor.py b/tests/entrypoints/openai/test_dp_supervisor.py index 0b678b22625..9967e6d86d0 100644 --- a/tests/entrypoints/openai/test_dp_supervisor.py +++ b/tests/entrypoints/openai/test_dp_supervisor.py @@ -21,7 +21,10 @@ import asyncio import contextlib import os import signal +import subprocess +import tempfile import time +from pathlib import Path from types import SimpleNamespace import aiohttp @@ -35,6 +38,7 @@ from vllm.entrypoints.openai.dp_supervisor import ( DPSupervisor, _build_vllm_dp_server_args, infer_multi_port_external_lb_start_rank, + validate_multi_port_external_lb_args, ) from vllm.logger import init_logger @@ -75,6 +79,8 @@ def _make_unit_args(**overrides) -> argparse.Namespace: "ssl_keyfile": None, "ssl_certfile": None, "ssl_ca_certs": None, + "ssl_cert_reqs": 0, + "ssl_ciphers": None, "node_rank": 1, "tensor_parallel_size": 1, "pipeline_parallel_size": 1, @@ -108,6 +114,8 @@ def _make_args(**overrides) -> argparse.Namespace: ssl_keyfile=None, ssl_certfile=None, ssl_ca_certs=None, + ssl_cert_reqs=0, + ssl_ciphers=None, node_rank=0, tensor_parallel_size=1, pipeline_parallel_size=1, @@ -118,6 +126,33 @@ def _make_args(**overrides) -> argparse.Namespace: return argparse.Namespace(**base) +def _generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]: + """Generate a self-signed certificate for HTTPS lifecycle tests.""" + cert_file = cert_dir / "cert.pem" + key_file = cert_dir / "key.pem" + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:2048", + "-keyout", + str(key_file), + "-out", + str(cert_file), + "-days", + "1", + "-nodes", + "-subj", + "/CN=localhost", + ], + check=True, + capture_output=True, + ) + return cert_file, key_file + + # --------------------------------------------------------------------------- # Unit tests # --------------------------------------------------------------------------- @@ -141,6 +176,15 @@ def test_build_multi_port_external_lb_child_args_sets_external_rank_server(): assert child_args.api_server_count == 1 +def test_validate_multi_port_external_lb_args_allows_ssl(): + args = _make_unit_args( + ssl_keyfile="/tmp/server.key", + ssl_certfile="/tmp/server.crt", + ssl_ca_certs="/tmp/ca.crt", + ) + validate_multi_port_external_lb_args(args) + + def test_aggregates_health(): supervisor = DPSupervisor(_make_unit_args()) supervisor._is_ready = True @@ -236,10 +280,18 @@ class MockVLLMServer: Health state is toggled by the test via set_healthy(). """ - def __init__(self, port: int, drain_seconds: float = 0.0) -> None: + def __init__( + self, + port: int, + drain_seconds: float = 0.0, + ssl_keyfile: str | None = None, + ssl_certfile: str | None = None, + ) -> None: self.port = port self._healthy = False self._drain_seconds = drain_seconds + self._ssl_keyfile = ssl_keyfile + self._ssl_certfile = ssl_certfile self._server: uvicorn.Server | None = None self._serve_task: asyncio.Task | None = None @@ -274,6 +326,8 @@ class MockVLLMServer: port=self.port, log_level="warning", lifespan="off", + ssl_keyfile=self._ssl_keyfile, + ssl_certfile=self._ssl_certfile, ) self._server = uvicorn.Server(config) @@ -312,7 +366,11 @@ class MockVLLMServer: def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]): logger.info("Launching mock vLLM on port %s", child_args.port) - mock_vllm = MockVLLMServer(port=child_args.port) + mock_vllm = MockVLLMServer( + port=child_args.port, + ssl_keyfile=child_args.ssl_keyfile, + ssl_certfile=child_args.ssl_certfile, + ) asyncio.run(mock_vllm.start()) @@ -320,7 +378,12 @@ def launch_mock_vllm_with_drain( child_args: argparse.Namespace, env_updates: dict[str, str] ): logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port) - mock_vllm = MockVLLMServer(port=child_args.port, drain_seconds=10.0) + mock_vllm = MockVLLMServer( + port=child_args.port, + drain_seconds=10.0, + ssl_keyfile=child_args.ssl_keyfile, + ssl_certfile=child_args.ssl_certfile, + ) asyncio.run(mock_vllm.start()) @@ -329,15 +392,16 @@ def launch_mock_vllm_with_drain( # --------------------------------------------------------------------------- -async def _poll_supervisor_health(expected_status: int) -> bool: +async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool: """ Poll GET /health on the supervisor until expected_status is seen. A connection error is treated as 503-equivalent when expected_status != 200. """ - url = f"http://127.0.0.1:{_SUPERVISOR_PORT}/health" + scheme = "https" if use_ssl else "http" + url = f"{scheme}://127.0.0.1:{_SUPERVISOR_PORT}/health" async with aiohttp.ClientSession() as session: try: - async with session.get(url) as resp: + async with session.get(url, ssl=False if use_ssl else None) as resp: if resp.status != expected_status: print(f"expected: {expected_status=}, got: {resp.status=}") return False @@ -349,12 +413,15 @@ async def _poll_supervisor_health(expected_status: int) -> bool: return True -async def _poll_until_api_server_running(port: int, retries: int = 10) -> None: - url = f"http://127.0.0.1:{port}/health" +async def _poll_until_api_server_running( + port: int, retries: int = 10, use_ssl: bool = False +) -> None: + scheme = "https" if use_ssl else "http" + url = f"{scheme}://127.0.0.1:{port}/health" async with aiohttp.ClientSession() as session: for _ in range(retries): try: - async with session.get(url) as resp: + async with session.get(url, ssl=False if use_ssl else None) as resp: if resp.status != 200: return await asyncio.sleep(1.0) @@ -363,22 +430,34 @@ async def _poll_until_api_server_running(port: int, retries: int = 10) -> None: await asyncio.sleep(1.0) -async def _set_healthy(port: int) -> None: - url = f"http://127.0.0.1:{port}/set_healthy" - async with aiohttp.ClientSession() as session, session.get(url) as resp: +async def _set_healthy(port: int, use_ssl: bool = False) -> None: + scheme = "https" if use_ssl else "http" + url = f"{scheme}://127.0.0.1:{port}/set_healthy" + async with ( + aiohttp.ClientSession() as session, + session.get(url, ssl=False if use_ssl else None) as resp, + ): assert resp.status == 200 -async def _set_unhealthy(port: int) -> None: - url = f"http://127.0.0.1:{port}/set_unhealthy" - async with aiohttp.ClientSession() as session, session.get(url) as resp: +async def _set_unhealthy(port: int, use_ssl: bool = False) -> None: + scheme = "https" if use_ssl else "http" + url = f"{scheme}://127.0.0.1:{port}/set_unhealthy" + async with ( + aiohttp.ClientSession() as session, + session.get(url, ssl=False if use_ssl else None) as resp, + ): assert resp.status == 200 -async def _kill_server(port: int) -> None: - url = f"http://127.0.0.1:{port}/kill" +async def _kill_server(port: int, use_ssl: bool = False) -> None: + scheme = "https" if use_ssl else "http" + url = f"{scheme}://127.0.0.1:{port}/kill" try: - async with aiohttp.ClientSession() as session, session.get(url) as resp: + async with ( + aiohttp.ClientSession() as session, + session.get(url, ssl=False if use_ssl else None) as resp, + ): assert resp.status != 200 except Exception as e: assert isinstance(e, aiohttp.ClientConnectorError) @@ -455,6 +534,34 @@ async def test_basic_lifecycle(monkeypatch): print("everything was cleaned up!") +@pytest.mark.asyncio +async def test_basic_lifecycle_with_ssl(monkeypatch): + with tempfile.TemporaryDirectory() as cert_dir: + cert_file, key_file = _generate_self_signed_cert(Path(cert_dir)) + args = _make_args( + ssl_keyfile=str(key_file), + ssl_certfile=str(cert_file), + ) + + vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)] + + async with _run_supervisor(args, monkeypatch) as (supervisor, _task): + assert await _poll_supervisor_health(503, use_ssl=True) + assert not supervisor.is_ready + + for port in vllm_server_ports: + assert await _poll_supervisor_health(503, use_ssl=True) + assert not supervisor.is_ready + await _poll_until_api_server_running(port, use_ssl=True) + + for port in vllm_server_ports: + await _set_healthy(port, use_ssl=True) + await asyncio.sleep(1.0) + + assert await _poll_supervisor_health(200, use_ssl=True) + assert supervisor.is_ready + + @pytest.mark.asyncio async def test_failed_startup(monkeypatch): """ 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/vllm/entrypoints/openai/generate/__init__.py b/tests/entrypoints/serve/dev/__init__.py similarity index 100% rename from vllm/entrypoints/openai/generate/__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/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 6c626986507..1dd89afcf80 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -48,6 +48,7 @@ class MockUpdateInfo(WeightTransferUpdateInfo): names: list[str] | None = None dtype_names: list[str] | None = None shapes: list[list[int]] | None = None + num_updates_list: list[int] | None = None class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]): @@ -87,6 +88,15 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo # (In real implementation, this would receive and load actual weights) load_weights([]) + def receive_sparse_weights( + self, + update_info: MockUpdateInfo, + apply_patches: Callable[[list], None], + ) -> None: + MockWeightTransferEngine.receive_weights_called = True + MockWeightTransferEngine.last_update_info = update_info + apply_patches([]) + def shutdown(self) -> None: MockWeightTransferEngine.shutdown_called = True @@ -198,8 +208,6 @@ def test_update_weights_calls_engine(): llm.init_weight_transfer_engine( WeightTransferInitRequest(init_info={"test_param": "init"}) ) - - # Start weight update (required before update_weights) llm.start_weight_update(is_checkpoint_format=True) # Call update_weights @@ -232,14 +240,67 @@ def test_update_weights_calls_engine(): assert dtypes == test_dtypes assert shapes == test_shapes - # Finish weight update + llm.finish_weight_update() + + +@create_new_process_for_each_test() +def test_update_weights_passes_sparse_metadata(): + """Test sparse update metadata is forwarded unchanged to the engine.""" + if torch.accelerator.device_count() < 1: + pytest.skip("Need at least 1 GPU for this test") + + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + with patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ): + llm = LLM( + model=MODEL_NAME, + enforce_eager=True, + load_format="dummy", + tensor_parallel_size=1, + weight_transfer_config=WeightTransferConfig(backend="nccl"), + ) + + llm.init_weight_transfer_engine( + WeightTransferInitRequest(init_info={"test_param": "init"}) + ) + llm.start_weight_update(is_checkpoint_format=False) + + llm.update_weights( + WeightTransferUpdateRequest( + update_info={ + "names": ["layer.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[100]], + "num_updates_list": [3], + "update_kind": "sparse_flat", + } + ) + ) + + def check_sparse_update_called(self): + engine = self.weight_transfer_engine + if not engine.receive_weights_called: + return None + info = engine.last_update_info + return ( + info.update_kind, + info.num_updates_list, + ) + + results = llm.collective_rpc(check_sparse_update_called) + for result in results: + assert result == ("sparse_flat", [3]) + llm.finish_weight_update() @create_new_process_for_each_test() def test_full_weight_transfer_flow(): - """Test the complete weight transfer flow: - init -> start -> update -> finish.""" + """Test the complete weight transfer flow: init -> start -> update -> finish.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") 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/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index ea52a2d3398..fdd00cfa27a 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -224,10 +224,6 @@ class Config: info = expert_info(self.fused_experts_type) return info.blocked_quantization_support - def supports_expert_map(self): - info = expert_info(self.fused_experts_type) - return info.supports_expert_map - def supports_apply_weight_on_input(self): info = prepare_finalize_info(self.prepare_finalize_type) return info.supports_apply_weight_on_input @@ -326,6 +322,15 @@ class Config: if self.needs_mori() and not has_mori(): # noqa: SIM103 return False, "Needs MoRI, but MoRI not available." + try: + if not self.fused_experts_type._supports_current_device(): + return ( + False, + f"{self.fused_experts_type} not supported on the current device.", + ) + except NotImplementedError: + pass + return True, None @@ -471,7 +476,7 @@ class RankTensors: topk_ids = topk_ids.to(device=device) expert_map = None - if config.world_size > 1 and config.supports_expert_map(): + if config.world_size > 1: expert_map = torch.full( (global_num_experts,), fill_value=-1, dtype=torch.int32 ) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 7c3bde2eafa..78ee8084d90 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -67,7 +67,6 @@ class ExpertInfo: activation_format: mk.FusedMoEActivationFormat supported_dtypes: list[torch.dtype | str] blocked_quantization_support: bool - supports_expert_map: bool needs_matching_quant: bool = False needs_deep_gemm: bool = False needs_aiter: bool = False @@ -129,7 +128,6 @@ def register_experts( activation_format: mk.FusedMoEActivationFormat, supported_dtypes: list[torch.dtype | str], blocked_quantization_support: bool, - supports_expert_map: bool, needs_matching_quant: bool = False, needs_deep_gemm: bool = False, needs_aiter: bool = False, @@ -142,7 +140,6 @@ def register_experts( activation_format, supported_dtypes, blocked_quantization_support, - supports_expert_map, needs_matching_quant, needs_deep_gemm, needs_aiter, @@ -176,7 +173,6 @@ register_experts( batched_format, common_float_types, blocked_quantization_support=True, - supports_expert_map=False, needs_matching_quant=True, ) @@ -185,7 +181,6 @@ register_experts( standard_format, common_float_and_int_types, blocked_quantization_support=True, - supports_expert_map=True, needs_matching_quant=True, ) @@ -194,7 +189,6 @@ register_experts( batched_format, common_float_and_int_types, blocked_quantization_support=True, - supports_expert_map=True, ) # Disable on blackwell for now @@ -260,7 +254,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability nvfp4_types + fp8_types, blocked_quantization_support=True, # Note: this is a hack to get it to run for now - supports_expert_map=True, ) else: FlashInferCutlassMoEPrepareAndFinalize = None @@ -294,7 +287,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability standard_format, nvfp4_types, blocked_quantization_support=False, - supports_expert_map=True, ) if has_aiter(): @@ -307,7 +299,6 @@ if has_aiter(): standard_format, fp8_types, blocked_quantization_support=True, - supports_expert_map=True, needs_aiter=True, ) else: @@ -319,7 +310,6 @@ if has_deep_gemm() and is_deep_gemm_supported(): batched_format, fp8_types, blocked_quantization_support=True, - supports_expert_map=False, needs_matching_quant=False, needs_deep_gemm=True, ) @@ -328,7 +318,6 @@ if has_deep_gemm() and is_deep_gemm_supported(): standard_format, fp8_types, blocked_quantization_support=True, - supports_expert_map=True, needs_matching_quant=False, needs_deep_gemm=True, ) @@ -337,7 +326,6 @@ if has_deep_gemm() and is_deep_gemm_supported(): standard_format, common_float_and_int_types, blocked_quantization_support=True, - supports_expert_map=True, needs_matching_quant=True, needs_deep_gemm=True, ) @@ -353,14 +341,12 @@ if cutlass_fp8_supported(): standard_format, fp8_types, blocked_quantization_support=False, - supports_expert_map=False, ) register_experts( CutlassBatchedExpertsFp8, batched_format, fp8_types, blocked_quantization_support=False, - supports_expert_map=False, ) else: CutlassBatchedExpertsFp8 = None @@ -376,7 +362,6 @@ if cutlass_fp4_supported(): standard_format, nvfp4_types, blocked_quantization_support=True, - supports_expert_map=False, ) else: CutlassExpertsFp4 = None 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/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index c7295f3ed6e..0c0e1d61f90 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -227,7 +227,7 @@ def is_nyi_config(config: Config) -> bool: ) == 1 return unsupported_quant_config - return not info.supports_expert_map + return False def generate_valid_test_cases( diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index 4089e9bc468..d957cefed4d 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 @@ -16,7 +17,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 ): @@ -37,7 +40,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, @@ -77,7 +80,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 @@ -99,8 +103,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), @@ -115,8 +119,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( @@ -126,7 +130,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, @@ -157,7 +161,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 @@ -216,7 +221,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 @@ -244,7 +250,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, @@ -286,7 +292,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 @@ -305,8 +312,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) @@ -320,11 +327,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" @@ -356,7 +363,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/kernels/quantization/test_rdna3_w4a16.py b/tests/kernels/quantization/test_rdna3_w4a16.py new file mode 100644 index 00000000000..b70a2b9a86e --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_w4a16.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100). + +Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with +GPTQ-format checkpoint parameters, runs ``process_weights_after_loading`` +(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares +the result against an fp32 reference dequant-and-matmul. + +The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only +built for gfx11; tests are skipped elsewhere. + +Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 + MPLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 + RDNA3W4A16LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 + pack_quantized_values_into_int32, +) +from vllm.model_executor.parameter import ( # noqa: E402 + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 +from vllm.utils.torch_utils import set_random_seed # noqa: E402 + +device = "cuda" + +WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8 +PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32 + +# Skip everything in this module unless we are on the only architecture the +# kernel is built/registered for. +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ), + reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", +) + + +# --------------------------------------------------------------------------- +# Reference implementation +# --------------------------------------------------------------------------- + + +def _reference( + x_mk: torch.Tensor, + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + group_size: int, + bias: torch.Tensor | None, +) -> torch.Tensor: + """fp32 reference for the RDNA3 W4A16 op. + + x_mk: [M, K] fp16/bf16 activations. + q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15]. + scales_gn: [K//G, N] per-group scales (act dtype). + zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None + for the symmetric path (kernel synthesizes stored zero = 7). + group_size: G. + + The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective + zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8). + """ + K, N = q_int4_kn.shape + s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N] + if zeros_gn is None: + z_full = torch.full( + (K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32 + ) + else: + z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32) + w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N] + out = x_mk.to(torch.float32) @ w_fp # [M, N] + if bias is not None: + out = out + bias.to(torch.float32) + return out.to(x_mk.dtype) + + +# --------------------------------------------------------------------------- +# Layer construction (GPTQ checkpoint format) +# --------------------------------------------------------------------------- + + +def _build_layer( + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + dtype: torch.dtype, +) -> torch.nn.Module: + """Build a dummy layer carrying GPTQ-format params, as the loader would.""" + no_loader = lambda *args, **kwargs: None # noqa: E731 + + # qweight: int4 packed along K into int32 -> [K//8, N]. + qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0) + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.register_parameter( + "qweight", + PackedvLLMParameter( + data=qweight, + weight_loader=no_loader, + input_dim=0, + output_dim=1, + packed_dim=0, + packed_factor=PACK_FACTOR, + ), + ) + layer.register_parameter( + "scales", + GroupQuantScaleParameter( + data=scales_gn.to(dtype), + weight_loader=no_loader, + input_dim=0, + output_dim=1, + ), + ) + if zeros_gn is not None: + # qzeros: int4 packed along N into int32 -> [K//G, N//8]. + qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1) + layer.register_parameter( + "qzeros", + PackedvLLMParameter( + data=qzeros, + weight_loader=no_loader, + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=PACK_FACTOR, + ), + ) + return layer + + +def _run_kernel( + x_mk: torch.Tensor, + q_int4_kn: torch.Tensor, + scales_gn: torch.Tensor, + zeros_gn: torch.Tensor | None, + group_size: int, + bias: torch.Tensor | None, + dtype: torch.dtype, +) -> torch.Tensor: + K, N = q_int4_kn.shape + has_zp = zeros_gn is not None + + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=WEIGHT_TYPE, + act_type=dtype, + group_size=group_size, + zero_points=has_zp, + has_g_idx=False, + ) + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok, f"can_implement rejected a supported config: {reason}" + + layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype) + kernel = RDNA3W4A16LinearKernel( + config, + w_q_param_name="qweight", + w_s_param_name="scales", + w_zp_param_name="qzeros" if has_zp else None, + w_gidx_param_name=None, + ) + kernel.process_weights_after_loading(layer) + return kernel.apply_weights(layer, x_mk, bias=bias) + + +# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights +# to fp32 and accumulates in fp32, so it matches the reference almost exactly +# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2 +# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered +# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order +# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We +# compare on the relative Frobenius norm rather than elementwise, since the +# bit-trick noise produces large *relative* errors on individual near-zero +# outputs that carry negligible absolute weight. +_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2} + + +def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype): + rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to( + torch.float32 + ).norm() + tol = _REL_L2_TOL[dtype] + assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}" + + +# --------------------------------------------------------------------------- +# Forward correctness +# --------------------------------------------------------------------------- + + +# (M, K, N, group_size). M spans the scalar decode path (small M) and the +# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's +# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0). +MKNG_SHAPES = [ + (1, 128, 128, 128), # single group, decode + (2, 256, 256, 128), # two groups + (8, 256, 512, 64), # M=8 scalar, smaller group + (16, 512, 256, 128), # M=16 -> WMMA path for bf16 + (32, 512, 512, 64), # larger prefill +] + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"]) +@pytest.mark.parametrize( + "M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES] +) +def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init): + set_random_seed(0) + assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0 + + groups = K // G + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + scales_gn = ( + 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01 + ).to(dtype) + zeros_gn = ( + torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32) + if has_zp + else None + ) + + out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype) + ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None) + + assert out.shape == (M, N) and out.dtype == dtype + _assert_close(out, ref, dtype) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"]) +def test_rdna3_w4a16_bias(dtype, M, dist_init): + """Bias is added on both the scalar (M=1) and WMMA (M=32) paths.""" + set_random_seed(0) + K, N, G = 512, 256, 128 + groups = K // G + + x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + scales_gn = ( + 0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01 + ).to(dtype) + bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype) + + out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype) + ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias) + + _assert_close(out, ref, dtype) diff --git a/tests/kernels/quantization/test_rdna3_w4a16_selection.py b/tests/kernels/quantization/test_rdna3_w4a16_selection.py new file mode 100644 index 00000000000..b53e1663781 --- /dev/null +++ b/tests/kernels/quantization/test_rdna3_w4a16_selection.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel. + +Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ +config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of +``TritonW4A16LinearKernel`` in the ROCm priority list), and that +``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not +support so selection falls through to the next kernel. + +Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True) + +from vllm.model_executor.kernels.linear import ( # noqa: E402 + choose_mp_linear_kernel, +) +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402 + MPLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402 + RDNA3W4A16LinearKernel, +) +from vllm.platforms.rocm import on_gfx1100 # noqa: E402 +from vllm.scalar_type import scalar_types # noqa: E402 + +WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8 + +# The kernel is only selectable when running on gfx1100 with the custom op +# compiled in; otherwise can_implement rejects and selection falls through. +gfx1100_only = pytest.mark.skipif( + not ( + on_gfx1100() + and hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ), + reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in", +) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_selection_prefers_rdna3(dtype): + """A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100.""" + config = MPLinearLayerConfig( + full_weight_shape=(1024, 256), + partition_weight_shape=(1024, 256), + weight_type=WEIGHT_TYPE, + act_type=dtype, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel" + + +@gfx1100_only +@pytest.mark.parametrize( + "weight_type,group_size,N,full_k,expected_ok", + [ + (scalar_types.uint4b8, 128, 256, 1024, True), # nominal: supported + (scalar_types.uint4b8, -1, 256, 1024, False), # channelwise unsupported + (scalar_types.uint4b8, 128, 252, 1024, False), # N not a multiple of 8 + (scalar_types.uint4b8, 96, 256, 1024, False), # group does not divide K + (scalar_types.uint8b128, 128, 256, 1024, False), # wrong quant type + ], + ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"], +) +def test_can_implement(weight_type, group_size, N, full_k, expected_ok): + """can_implement gates on quant type, group size, and N divisibility.""" + config = MPLinearLayerConfig( + full_weight_shape=(full_k, N), + partition_weight_shape=(full_k, N), + weight_type=weight_type, + act_type=torch.float16, + group_size=group_size, + zero_points=False, + has_g_idx=False, + ) + ok, reason = RDNA3W4A16LinearKernel.can_implement(config) + assert ok is expected_ok, reason diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py new file mode 100644 index 00000000000..f855eb7aa17 --- /dev/null +++ b/tests/kernels/test_fp32_router_gemm.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for fp32_router_gemm kernel: activation×weight→fp32, H=3072, E=256. + +Correctness baseline: torch.matmul in float64. +""" + +import pytest +import torch + +from vllm._custom_ops import fp32_router_gemm + +NUM_EXPERTS = 256 +HIDDEN_DIM = 3072 +# Absolute tolerance for fp32 kernel vs float64 reference +ATOL_FP32 = 2e-4 +ATOL_BF16 = 2e-2 # bf16 activation has lower precision + + +def _requires_sm90(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + major, minor = torch.cuda.get_device_capability() + if major * 10 + minor < 90: + pytest.skip(f"fp32_router_gemm requires SM90+, got SM{major}{minor}") + + +def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + """Reference: F.linear in float32 on GPU.""" + return torch.nn.functional.linear(mat_a.float(), mat_b.float()) + + +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32]) +def test_fp32_activation(num_tokens: int): + """fp32 activation → fp32 output should match reference closely.""" + _requires_sm90() + torch.manual_seed(42) + device = torch.device("cuda") + mat_a = torch.randn(num_tokens, HIDDEN_DIM, dtype=torch.float32, device=device) + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device) + + out = fp32_router_gemm(mat_a, mat_b) + ref = _ref(mat_a, mat_b) + + assert out.shape == (num_tokens, NUM_EXPERTS) + assert out.dtype == torch.float32 + torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0) + + +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32]) +def test_bf16_activation(num_tokens: int): + """bf16 activation → fp32 output should match reference within bf16 error.""" + _requires_sm90() + torch.manual_seed(42) + device = torch.device("cuda") + mat_a_bf16 = torch.randn( + num_tokens, HIDDEN_DIM, dtype=torch.bfloat16, device=device + ) + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device) + + out = fp32_router_gemm(mat_a_bf16, mat_b) + ref = _ref(mat_a_bf16, mat_b).to(device) + + assert out.shape == (num_tokens, NUM_EXPERTS) + assert out.dtype == torch.float32 + torch.testing.assert_close(out, ref, atol=ATOL_BF16, rtol=0) + + +def test_output_shape_and_dtype(): + """Basic shape and dtype checks.""" + _requires_sm90() + device = torch.device("cuda") + mat_a = torch.randn(4, HIDDEN_DIM, dtype=torch.float32, device=device) + mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device) + out = fp32_router_gemm(mat_a, mat_b) + assert out.shape == (4, NUM_EXPERTS) + assert out.dtype == torch.float32 + assert out.device.type == "cuda" diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index e7d4cde43f1..0e0e3769f49 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -340,22 +340,17 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult): hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1 rms_eps = hc_eps = 1e-6 - out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16) - out.fill_(float("nan")) - - result = torch.ops.vllm.hc_head_fused_kernel_tilelang( + out = torch.ops.vllm.hc_head_fused_kernel_tilelang( residual, fn, hc_scale, hc_base, - out, - hidden_size, rms_eps, hc_eps, - hc_mult, ) - assert result is None + assert out.shape == (num_tokens, hidden_size) + assert out.dtype == torch.bfloat16 assert not torch.isnan(out).any() out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) 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/lora/test_utils.py b/tests/lora/test_utils.py index bec12eeeb48..4a9ba66069d 100644 --- a/tests/lora/test_utils.py +++ b/tests/lora/test_utils.py @@ -175,7 +175,7 @@ def test_get_adapter_absolute_path_local_existing(mock_abspath, mock_exist): assert get_adapter_absolute_path(path) == absolute_path -@patch("huggingface_hub.snapshot_download") +@patch("huggingface_hub.HfApi.snapshot_download") @patch("os.path.exists") def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_download): # Hugging Face model identifier @@ -186,7 +186,7 @@ def test_get_adapter_absolute_path_huggingface(mock_exist, mock_snapshot_downloa assert get_adapter_absolute_path(path) == absolute_path -@patch("huggingface_hub.snapshot_download") +@patch("huggingface_hub.HfApi.snapshot_download") @patch("os.path.exists") def test_get_adapter_absolute_path_huggingface_error( mock_exist, mock_snapshot_download diff --git a/tests/model_executor/layers/test_pooler_methods.py b/tests/model_executor/layers/test_pooler_methods.py new file mode 100644 index 00000000000..cb8533cacb8 --- /dev/null +++ b/tests/model_executor/layers/test_pooler_methods.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for sequence and token pooling methods and their factories.""" + +from dataclasses import dataclass +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.layers.pooler.seqwise.methods import ( + CLSPool, + LastPool, + MeanPool, + get_seq_pooling_method, +) +from vllm.model_executor.layers.pooler.tokwise.methods import ( + AllPool, + StepPool, + get_tok_pooling_method, +) +from vllm.pooling_params import PoolingParams +from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates + +_CPU = torch.device("cpu") + + +def _make_pooling_cursor( + prompt_lens: list[int], + *, + num_scheduled_tokens: list[int] | None = None, + seq_lens: list[int] | None = None, + device: torch.device = _CPU, +) -> PoolingCursor: + """Build a PoolingCursor from a list of per-sequence prompt lengths.""" + prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.long) + if num_scheduled_tokens is None: + num_scheduled_tokens_cpu = prompt_lens_cpu.clone() + else: + num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, dtype=torch.long) + if seq_lens is None: + seq_lens_cpu = prompt_lens_cpu.clone() + else: + seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.long) + + cumsum = torch.zeros(len(prompt_lens) + 1, dtype=torch.long, device=device) + torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:]) + + return PoolingCursor( + first_token_indices_gpu=cumsum[: len(prompt_lens)].to(device), + last_token_indices_gpu=(cumsum[1:] - 1).to(device), + prompt_lens_cpu=prompt_lens_cpu, + seq_lens_cpu=seq_lens_cpu, + num_scheduled_tokens_cpu=num_scheduled_tokens_cpu, + ) + + +def _make_metadata( + prompt_lens: list[int], + *, + tasks: list[str] | None = None, + token_ids: list[list[int]] | None = None, + pooling_params: list[PoolingParams] | None = None, + num_scheduled_tokens: list[int] | None = None, + seq_lens: list[int] | None = None, + device: torch.device = _CPU, +) -> PoolingMetadata: + """Build a minimal PoolingMetadata for testing pooling methods.""" + n_seqs = len(prompt_lens) + if tasks is None: + tasks = ["embed"] * n_seqs + if pooling_params is None: + pooling_params = [PoolingParams(task=t) for t in tasks] + + prompt_lens_tensor = torch.tensor(prompt_lens, dtype=torch.long) + + prompt_token_ids_cpu = None + prompt_token_ids = None + if token_ids is not None: + max_len = max(len(t) for t in token_ids) + padded = [t + [0] * (max_len - len(t)) for t in token_ids] + prompt_token_ids_cpu = torch.tensor(padded, dtype=torch.long) + prompt_token_ids = prompt_token_ids_cpu.to(device) + + cursor = _make_pooling_cursor( + prompt_lens, + num_scheduled_tokens=num_scheduled_tokens, + seq_lens=seq_lens, + device=device, + ) + + pooling_states = [PoolingStates() for _ in range(n_seqs)] + + return PoolingMetadata( + prompt_lens=prompt_lens_tensor, + prompt_token_ids=prompt_token_ids, + prompt_token_ids_cpu=prompt_token_ids_cpu, + pooling_params=pooling_params, + pooling_states=pooling_states, + pooling_cursor=cursor, + ) + + +# --------------------------------------------------------------------------- +# CLSPool +# --------------------------------------------------------------------------- +class TestCLSPool: + def test_extracts_first_token(self): + hidden = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]] + ) + metadata = _make_metadata([2, 3]) + pooler = CLSPool() + out = pooler(hidden, metadata) + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]]) + assert torch.equal(out, expected) + + def test_rejects_partial_prefill(self): + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + metadata = _make_metadata([3], num_scheduled_tokens=[2]) + pooler = CLSPool() + with pytest.raises(AssertionError, match="partial prefill"): + pooler(hidden, metadata) + + +# --------------------------------------------------------------------------- +# LastPool +# --------------------------------------------------------------------------- +class TestLastPool: + def test_extracts_last_token(self): + hidden = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]] + ) + metadata = _make_metadata([2, 3]) + pooler = LastPool() + out = pooler(hidden, metadata) + expected = torch.tensor([[3.0, 4.0], [9.0, 10.0]]) + assert torch.equal(out, expected) + + def test_partial_prefill_extracts_last_scheduled(self): + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + metadata = _make_metadata([4], num_scheduled_tokens=[2]) + pooler = LastPool() + out = pooler(hidden, metadata) + expected = torch.tensor([[3.0, 4.0]]) + assert torch.equal(out, expected) + + +# --------------------------------------------------------------------------- +# MeanPool +# --------------------------------------------------------------------------- +class TestMeanPool: + def test_computes_mean(self): + hidden = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [10.0, 20.0]], dtype=torch.float32 + ) + metadata = _make_metadata([2, 1]) + pooler = MeanPool() + out = pooler(hidden, metadata) + expected = torch.tensor([[2.0, 3.0], [10.0, 20.0]], dtype=torch.float32) + assert torch.allclose(out, expected, atol=1e-5) + + def test_single_token_is_identity(self): + hidden = torch.tensor([[5.0, 10.0]], dtype=torch.float32) + metadata = _make_metadata([1]) + pooler = MeanPool() + out = pooler(hidden, metadata) + assert torch.allclose(out, hidden, atol=1e-5) + + def test_uniform_values_return_same(self): + hidden = torch.full((4, 3), 7.0, dtype=torch.float32) + metadata = _make_metadata([4]) + pooler = MeanPool() + out = pooler(hidden, metadata) + expected = torch.full((1, 3), 7.0, dtype=torch.float32) + assert torch.allclose(out, expected, atol=1e-5) + + def test_multiple_sequences(self): + hidden = torch.tensor( + [ + [0.0, 0.0], + [2.0, 4.0], + [4.0, 8.0], + [10.0, 10.0], + ], + dtype=torch.float32, + ) + metadata = _make_metadata([3, 1]) + pooler = MeanPool() + out = pooler(hidden, metadata) + expected = torch.tensor([[2.0, 4.0], [10.0, 10.0]], dtype=torch.float32) + assert torch.allclose(out, expected, atol=1e-5) + + def test_empty_batch(self): + hidden = torch.empty((0, 8), dtype=torch.float32) + metadata = _make_metadata([]) + pooler = MeanPool() + out = pooler(hidden, metadata) + assert out.shape == (0, 8) + + def test_rejects_partial_prefill(self): + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) + metadata = _make_metadata([3], num_scheduled_tokens=[2]) + pooler = MeanPool() + with pytest.raises(AssertionError, match="partial prefill"): + pooler(hidden, metadata) + + def test_chunked_accumulation(self): + hidden = torch.arange(20, dtype=torch.float32).reshape(5, 4) + metadata = _make_metadata([3, 2]) + pooler = MeanPool() + with patch( + "vllm.model_executor.layers.pooler.seqwise.methods" + "._MEAN_POOL_ACCUMULATION_CHUNK_BYTES", + 16, + ): + out = pooler(hidden, metadata) + expected_seq0 = hidden[:3].float().mean(dim=0, keepdim=True) + expected_seq1 = hidden[3:].float().mean(dim=0, keepdim=True) + expected = torch.cat([expected_seq0, expected_seq1], dim=0) + assert torch.allclose(out, expected, atol=1e-5) + + def test_upcasts_to_float32(self): + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float16) + metadata = _make_metadata([2]) + pooler = MeanPool() + out = pooler(hidden, metadata) + assert out.dtype == torch.float32 + expected = torch.tensor([[2.0, 3.0]], dtype=torch.float32) + assert torch.allclose(out, expected, atol=1e-2) + + +# --------------------------------------------------------------------------- +# get_seq_pooling_method factory +# --------------------------------------------------------------------------- +class TestGetSeqPoolingMethod: + def test_cls(self): + assert isinstance(get_seq_pooling_method("CLS"), CLSPool) + + def test_last(self): + assert isinstance(get_seq_pooling_method("LAST"), LastPool) + + def test_mean(self): + assert isinstance(get_seq_pooling_method("MEAN"), MeanPool) + + def test_unknown_raises(self): + with pytest.raises(NotImplementedError, match="UNKNOWN"): + get_seq_pooling_method("UNKNOWN") + + +# --------------------------------------------------------------------------- +# AllPool +# --------------------------------------------------------------------------- + + +@dataclass +class _FakeSchedulerConfig: + enable_chunked_prefill: bool = False + + +@dataclass +class _FakeVllmConfig: + scheduler_config: _FakeSchedulerConfig + + +class TestAllPool: + @staticmethod + def _make_all_pool(*, chunked: bool = False) -> AllPool: + fake_config = _FakeVllmConfig( + scheduler_config=_FakeSchedulerConfig( + enable_chunked_prefill=chunked, + ), + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=fake_config, + ): + return AllPool() + + def test_splits_by_sequence(self): + pooler = self._make_all_pool() + hidden = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]] + ) + metadata = _make_metadata([2, 3]) + out = pooler(hidden, metadata) + assert len(out) == 2 + assert torch.equal(out[0], hidden[:2]) + assert torch.equal(out[1], hidden[2:]) + + def test_single_sequence(self): + pooler = self._make_all_pool() + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + metadata = _make_metadata([3]) + out = pooler(hidden, metadata) + assert len(out) == 1 + assert torch.equal(out[0], hidden) + + def test_chunked_prefill_returns_none_for_unfinished(self): + pooler = self._make_all_pool(chunked=True) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + metadata = _make_metadata( + [4], + num_scheduled_tokens=[2], + seq_lens=[2], + ) + out = pooler(hidden, metadata) + assert len(out) == 1 + assert out[0] is None + + def test_chunked_prefill_returns_concat_when_finished(self): + pooler = self._make_all_pool(chunked=True) + + chunk1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + metadata1 = _make_metadata( + [4], + num_scheduled_tokens=[2], + seq_lens=[2], + ) + out1 = pooler(chunk1, metadata1) + assert out1[0] is None + + chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]]) + metadata2 = _make_metadata( + [4], + num_scheduled_tokens=[2], + seq_lens=[4], + ) + metadata2.pooling_states = metadata1.pooling_states + out2 = pooler(chunk2, metadata2) + assert out2[0] is not None + expected = torch.cat([chunk1, chunk2], dim=0) + assert torch.equal(out2[0], expected) + + def test_chunked_prefill_single_shot_matches_non_chunked(self): + pooler = self._make_all_pool(chunked=True) + hidden = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]] + ) + metadata = _make_metadata([2, 3]) + out = pooler(hidden, metadata) + assert len(out) == 2 + assert torch.equal(out[0], hidden[:2]) + assert torch.equal(out[1], hidden[2:]) + + def test_chunked_prefill_mixed_finished_unfinished(self): + pooler = self._make_all_pool(chunked=True) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + metadata = _make_metadata( + [2, 4], + num_scheduled_tokens=[2, 1], + seq_lens=[2, 1], + ) + out = pooler(hidden, metadata) + assert len(out) == 2 + assert torch.equal(out[0], hidden[:2]) + assert out[1] is None + + +# --------------------------------------------------------------------------- +# StepPool +# --------------------------------------------------------------------------- +class TestStepPool: + @staticmethod + def _make_step_pool(*, chunked: bool = False) -> StepPool: + fake_config = _FakeVllmConfig( + scheduler_config=_FakeSchedulerConfig( + enable_chunked_prefill=chunked, + ), + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=fake_config, + ): + return StepPool() + + def test_filters_by_step_tag_id(self): + pooler = self._make_step_pool() + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]) + token_ids = [[10, 99, 10, 20]] + params = [PoolingParams(task="token_classify", step_tag_id=10)] + metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]]) + assert torch.equal(out[0], expected) + + def test_filters_by_returned_token_ids(self): + pooler = self._make_step_pool() + hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + token_ids = [[10, 20]] + params = [PoolingParams(task="token_classify", returned_token_ids=[0, 2])] + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + expected = torch.tensor([[1.0, 3.0], [4.0, 6.0]]) + assert torch.equal(out[0], expected) + + def test_no_filtering_without_params(self): + pooler = self._make_step_pool() + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + token_ids = [[10, 20]] + params = [PoolingParams(task="token_classify")] + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + assert torch.equal(out[0], hidden) + + def test_combined_step_tag_and_returned_token_ids(self): + pooler = self._make_step_pool() + hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + token_ids = [[99, 10, 99]] + params = [ + PoolingParams( + task="token_classify", + step_tag_id=10, + returned_token_ids=[0, 2], + ) + ] + metadata = _make_metadata([3], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + expected = torch.tensor([[4.0, 6.0]]) + assert torch.equal(out[0], expected) + + def test_step_tag_id_no_match_returns_empty(self): + pooler = self._make_step_pool() + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + token_ids = [[10, 20]] + params = [PoolingParams(task="token_classify", step_tag_id=999)] + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + assert out[0].shape == (0, 2) + + def test_chunked_prefill_propagates_none_for_unfinished(self): + pooler = self._make_step_pool(chunked=True) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + token_ids = [[10, 20, 30, 40]] + params = [PoolingParams(task="token_classify", step_tag_id=10)] + metadata = _make_metadata( + [4], + token_ids=token_ids, + pooling_params=params, + num_scheduled_tokens=[2], + seq_lens=[2], + ) + out = pooler(hidden, metadata) + assert len(out) == 1 + assert out[0] is None + + def test_chunked_prefill_filters_when_finished(self): + pooler = self._make_step_pool(chunked=True) + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]) + token_ids = [[10, 99, 10, 20]] + params = [PoolingParams(task="token_classify", step_tag_id=10)] + metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params) + out = pooler(hidden, metadata) + assert len(out) == 1 + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]]) + assert torch.equal(out[0], expected) + + def test_requires_token_ids_update(self): + pooler = self._make_step_pool() + update = pooler.get_pooling_updates("token_classify") + assert update.requires_token_ids is True + + +# --------------------------------------------------------------------------- +# get_tok_pooling_method factory +# --------------------------------------------------------------------------- +class TestGetTokPoolingMethod: + def test_all(self): + fake_config = _FakeVllmConfig( + scheduler_config=_FakeSchedulerConfig( + enable_chunked_prefill=False, + ), + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=fake_config, + ): + assert isinstance(get_tok_pooling_method("ALL"), AllPool) + + def test_step(self): + fake_config = _FakeVllmConfig( + scheduler_config=_FakeSchedulerConfig( + enable_chunked_prefill=False, + ), + ) + with patch( + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config", + return_value=fake_config, + ): + assert isinstance(get_tok_pooling_method("STEP"), StepPool) + + def test_unknown_raises(self): + with pytest.raises(NotImplementedError, match="UNKNOWN"): + get_tok_pooling_method("UNKNOWN") 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/models/multimodal/pooling/test_phi3v.py b/tests/models/multimodal/pooling/test_phi3v.py index 2794b0b2937..285ded375da 100644 --- a/tests/models/multimodal/pooling/test_phi3v.py +++ b/tests/models/multimodal/pooling/test_phi3v.py @@ -8,6 +8,8 @@ from PIL import Image from vllm.assets.base import get_vllm_public_assets from vllm.assets.image import VLM_IMAGES_DIR +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner from ....utils import large_gpu_test @@ -37,6 +39,18 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( MODELS = ["TIGER-Lab/VLM2Vec-Full"] +SPECIAL_TOKEN_IMAGE_PROMPT = ( + "\n<|user|>\n <|image_1|>\n\t " + "Represent the given image for classification<|end|>" + "\n<|assistant|>\n" +) + + +def _get_cherry_blossom_image() -> Image.Image: + return Image.open( + get_vllm_public_assets(filename="cherry_blossom.jpg", s3_prefix=VLM_IMAGES_DIR) + ) + def _run_test( hf_runner: type[HfRunner], @@ -123,19 +137,6 @@ def test_models_image( input_texts_images = [ (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) ] - # add cases for special_tokens - input_texts_images.append( - ( - "\n<|user|>\n <|image_1|>\n\t " - "Represent the given image for classification<|end|>" - "\n<|assistant|>\n", - Image.open( - get_vllm_public_assets( - filename="cherry_blossom.jpg", s3_prefix=VLM_IMAGES_DIR - ) - ), - ) - ) input_texts = [text for text, _ in input_texts_images] input_images = [image for _, image in input_texts_images] @@ -147,3 +148,48 @@ def test_models_image( model, dtype=dtype, ) + + +@pytest.mark.core_model +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("dtype", ["half"]) +def test_models_image_special_tokens_processing( + model: str, + dtype: str, +) -> None: + model_config = ModelConfig( + model, + runner="pooling", + trust_remote_code=True, + dtype=dtype, + max_model_len=1024, + ) + processor = MULTIMODAL_REGISTRY.create_processor(model_config) + image = _get_cherry_blossom_image() + + processed_inputs = processor( + SPECIAL_TOKEN_IMAGE_PROMPT, + mm_items=processor.info.parse_mm_data({"image": image}), + hf_processor_mm_kwargs={}, + ) + + hf_processor = processor.info.get_hf_processor() + hf_inputs = hf_processor( + SPECIAL_TOKEN_IMAGE_PROMPT, + images=image, + return_tensors="pt", + ) + + image_token_id = hf_processor.get_special_image_token_id() + hf_prompt_token_ids = [ + image_token_id if token_id < 0 else token_id + for token_id in hf_inputs["input_ids"][0].tolist() + ] + + prompt_token_ids = processed_inputs["prompt_token_ids"] + + assert prompt_token_ids == hf_prompt_token_ids + assert prompt_token_ids.count(image_token_id) == hf_prompt_token_ids.count( + image_token_id + ) + assert prompt_token_ids.count(image_token_id) > 0 diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 5afcab9f324..12c5071978f 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -180,6 +180,7 @@ def test_model_tensor_schema(model_id: str): dummy_hf_overrides, model_arch=model_arch, exist_overrides=model_info.hf_overrides, + use_original_num_layers=getattr(model_info, "use_original_num_layers", False), ) # ROCm: Detect if model uses AWQ quantization and set appropriate dtype diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index fe9ddd2f6ba..783f1773d21 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -22,7 +22,14 @@ import pytest from packaging import version from vllm.platforms import current_platform -from vllm.platforms.rocm import on_gfx950 + +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx950 +else: + + def on_gfx950() -> bool: + return False + MODEL_ACCURACIES = { # Full quantization: attention linears and MoE linears diff --git a/tests/models/registry.py b/tests/models/registry.py index 226eb80aca1..36e201eac8c 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -363,7 +363,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "IQuestLoopCoderForCausalLM": _HfExamplesInfo( "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct", trust_remote_code=True ), - "JAISLMHeadModel": _HfExamplesInfo("inceptionai/jais-13b-chat"), "Jais2ForCausalLM": _HfExamplesInfo( "inceptionai/Jais-2-8B-Chat", min_transformers_version="4.58" ), @@ -523,6 +522,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"), "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"), "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"), + "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"), "Qwen3NextForCausalLM": _HfExamplesInfo( "Qwen/Qwen3-Next-80B-A3B-Instruct", extras={"tiny-random": "tiny-random/qwen3-next-moe"}, @@ -936,7 +936,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { trust_remote_code=True, hf_overrides={"architectures": ["GLM4VForCausalLM"]}, ), - "Glm4vForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.1V-9B-Thinking"), + "Glm4vForConditionalGeneration": _HfExamplesInfo( + "zai-org/GLM-4.1V-9B-Thinking", + extras={"4.6V": "zai-org/GLM-4.6V-Flash"}, + ), "Glm4vMoeForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.5V"), "GlmOcrForConditionalGeneration": _HfExamplesInfo( "zai-org/GLM-OCR", @@ -1369,6 +1372,16 @@ _MULTIMODAL_EXAMPLE_MODELS = { "StepVLForConditionalGeneration": _HfExamplesInfo( "stepfun-ai/Step3-VL-10B", trust_remote_code=True ), + "Step3p7ForConditionalGeneration": _HfExamplesInfo( + "stepfun-ai/Step-3.7-Flash", + trust_remote_code=True, + use_original_num_layers=True, + # The MoE config lives in the nested ``text_config``, so the overrides + # must be nested too. Use 4 layers to initialize at least one MoE layer + # and shrink ``moe_num_experts`` (a non-standard key not handled by + # ``dummy_hf_overrides``) to avoid OOM during init. + hf_overrides={"text_config": {"num_hidden_layers": 4, "moe_num_experts": 8}}, + ), "UltravoxModel": _HfExamplesInfo( "fixie-ai/ultravox-v0_5-llama-3_2-1b", trust_remote_code=True, 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_gemma4_chat_template.py b/tests/renderers/test_gemma4_chat_template.py index a4a0b41d053..ac13c0d4d5f 100644 --- a/tests/renderers/test_gemma4_chat_template.py +++ b/tests/renderers/test_gemma4_chat_template.py @@ -343,3 +343,72 @@ class TestGemma4ChatTemplate: assert '<|"|>Alice<|"|>' in result assert "active:true" in result assert "count:42" in result + + def test_tool_response_with_multimodal_content(self, gemma4_template): + """Multimodal placeholders in tool messages are emitted after the + tool_response block.""" + messages = [ + {"role": "user", "content": "Download the image and describe it."}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "download_image", + "arguments": '{"url": "https://example.com/x.png"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "text", "text": "Image downloaded successfully."}, + {"type": "image"}, + ], + }, + ] + result = _render(gemma4_template, messages, add_generation_prompt=True) + assert "<|tool_response>" in result + assert "response:download_image{" in result + assert "" in result + assert "<|image|>" in result + + def test_tool_response_with_all_modalities(self, gemma4_template): + """All multimodal types (image, audio, video) in a single tool + response are rendered.""" + messages = [ + {"role": "user", "content": "Process media"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": { + "name": "process", + "arguments": "{}", + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "c1", + "content": [ + {"type": "text", "text": "Results."}, + {"type": "image"}, + {"type": "audio"}, + {"type": "video"}, + ], + }, + ] + result = _render(gemma4_template, messages, add_generation_prompt=True) + assert "<|image|>" in result + assert "<|audio|>" in result + assert "<|video|>" in result 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/tool_use/test_chat_completion_request_validations.py b/tests/tool_use/test_chat_completion_request_validations.py index d832feda7f5..7adf4beb9d8 100644 --- a/tests/tool_use/test_chat_completion_request_validations.py +++ b/tests/tool_use/test_chat_completion_request_validations.py @@ -116,3 +116,68 @@ def test_no_reasoning_fields_unchanged(): assistant_msg = request.messages[1] assert assistant_msg.get("reasoning") is None assert "reasoning_content" not in assistant_msg + + +SAMPLE_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + }, + }, +} + + +def test_structured_outputs_with_named_tool_choice_rejected(): + """structured_outputs cannot be combined with a named tool_choice.""" + with pytest.raises( + ValueError, + match="structured outputs or tools, not both", + ): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "Hello"}], + "model": "facebook/opt-125m", + "tools": [SAMPLE_TOOL], + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"}, + }, + "structured_outputs": {"json": {"type": "object"}}, + } + ) + + +def test_structured_outputs_with_auto_tool_choice_allowed(): + """structured_outputs with tool_choice 'auto' should be allowed.""" + request = ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "Hello"}], + "model": "facebook/opt-125m", + "tools": [SAMPLE_TOOL], + "tool_choice": "auto", + "structured_outputs": {"json": {"type": "object"}}, + } + ) + assert request.tool_choice == "auto" + + +def test_multiple_structured_outputs_rejected(): + """Only one kind of structured output constraint is allowed.""" + with pytest.raises( + ValueError, + match="You can only use one kind of constraints", + ): + ChatCompletionRequest.model_validate( + { + "messages": [{"role": "user", "content": "Hello"}], + "model": "facebook/opt-125m", + "structured_outputs": { + "json": {"type": "object"}, + "regex": ".*", + }, + } + ) diff --git a/tests/utils_/test_import_utils.py b/tests/utils_/test_import_utils.py index d42685b3fc9..464f209f0f2 100644 --- a/tests/utils_/test_import_utils.py +++ b/tests/utils_/test_import_utils.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + import pytest -from vllm.utils.import_utils import PlaceholderModule +from vllm.utils.import_utils import PlaceholderModule, _has_module def _raises_module_not_found(): @@ -44,3 +46,58 @@ def test_placeholder_module_error_handling(): with _raises_module_not_found(): # Test conflict with internal __module attribute _ = placeholder_attr.module + + +class TestHasModule: + """Tests for _has_module with trial import verification.""" + + def setup_method(self): + # Clear the @cache between tests so each test gets a fresh call + _has_module.cache_clear() + + def test_returns_true_for_importable_stdlib_module(self): + assert _has_module("json") is True + + def test_returns_false_for_nonexistent_module(self): + assert _has_module("nonexistent_module_xyz_12345") is False + + def test_returns_false_when_find_spec_succeeds_but_import_fails(self): + """Simulate a native extension whose shared library is missing. + + ``find_spec`` finds the package on disk, but the actual import + raises ``ImportError`` (e.g. missing ``libcudart.so``). + """ + 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=ImportError( + "libcudart.so.12: cannot open shared object file" + ), + ), + ): + assert _has_module("fake_native_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. + """ + with patch( + "vllm.utils.import_utils.importlib.util.find_spec", + side_effect=ModuleNotFoundError("No module named 'fake_parent'"), + ): + assert _has_module("fake_parent.child") is False + + def test_result_is_cached(self): + """Verify the @cache decorator prevents repeated imports.""" + _has_module("json") # prime the cache + + with patch("vllm.utils.import_utils.importlib.util.find_spec") as mock_spec: + result = _has_module("json") # should hit cache + mock_spec.assert_not_called() + assert result is True diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 9f6fa85da58..0f615fb8c47 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -10,6 +10,33 @@ from vllm.config import ParallelConfig from vllm.utils import numa_utils +@pytest.fixture(autouse=True) +def _disable_pct_by_default(monkeypatch): + """Force PCT detection OFF unless a test opts in via ``_patch_pct_gates``. + + The CI / dev machines themselves can be Xeon 6776P with PCT enabled, so a + plain ``cache_clear`` would let the gate auto-detect ``True`` from the + live filesystem and silently re-route "baseline" tests through the PCT + path. Stub ``/proc/cpuinfo`` and ``acpi_cppc/highest_perf`` to a state + that fails the gate; ``_patch_pct_gates`` re-stubs on top when needed. + """ + from io import StringIO + + real_open = open + + def _no_pct_open(path, *args, **kwargs): + if path == numa_utils._PROC_CPUINFO_PATH: + return StringIO("processor\t: 0\nmodel name\t: Generic Test CPU\n") + if path == numa_utils._PCT_HIGHEST_PERF_PATH: + raise OSError("PCT disabled by autouse fixture") + return real_open(path, *args, **kwargs) + + monkeypatch.setattr("builtins.open", _no_pct_open) + numa_utils._pct_sku_config.cache_clear() + yield + numa_utils._pct_sku_config.cache_clear() + + def _make_config(**parallel_kwargs): parallel_defaults = dict( numa_bind=False, @@ -31,7 +58,7 @@ def _make_config(**parallel_kwargs): def test_get_numactl_args_with_node_binding(): vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1]) assert ( - numa_utils._get_numactl_args(vllm_config, local_rank=1) + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) == "--cpunodebind=1 --membind=1" ) @@ -43,7 +70,300 @@ def test_get_numactl_args_with_cpu_binding(): numa_bind_cpus=["0-3", "4-7"], ) assert ( - numa_utils._get_numactl_args(vllm_config, local_rank=1) + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) + == "--physcpubind=4-7 --membind=1" + ) + + +def _patch_pct_gates( + monkeypatch, + *, + model_match: bool, + highest_perf: int | None, + cpulist: str | None = "0-31,64-95", + cpulist_by_node: dict[int, str | None] | None = None, + sku: str = "6776P", +): + """Force `_pct_sku_config` and node cpulist read to deterministic state. + + ``cpulist`` is the default returned for any node not present in + ``cpulist_by_node``. ``cpulist_by_node`` lets a test return different + cpulists for different NUMA nodes (e.g. node 0 vs node 1). ``sku`` lets + the test pick which Granite Rapids SKU appears in the fake + ``/proc/cpuinfo`` ``model name`` (only used when ``model_match=True``). + """ + import pathlib + from io import StringIO + + import regex as re + + cpuinfo = ( + f"processor\t: 0\nmodel name\t: Intel(R) Xeon(R) Platinum {sku} CPU @ 2.40GHz\n" + if model_match + else "processor\t: 0\nmodel name\t: Intel(R) Xeon(R) Platinum 8480+\n" + ) + + real_open = open + + def fake_open(path, *args, **kwargs): + if path == numa_utils._PROC_CPUINFO_PATH: + return StringIO(cpuinfo) + if path == numa_utils._PCT_HIGHEST_PERF_PATH: + if highest_perf is None: + raise OSError("missing") + return StringIO(f"{highest_perf}\n") + return real_open(path, *args, **kwargs) + + real_read_text = pathlib.Path.read_text + cpulist_by_node = cpulist_by_node or {} + + def fake_read_text(self, *args, **kwargs): + path_str = str(self) + if path_str.endswith("/cpulist") and "/sys/devices/system/node" in path_str: + match = re.search(r"/node(\d+)/cpulist$", path_str) + if match: + node_id = int(match.group(1)) + if node_id in cpulist_by_node: + val = cpulist_by_node[node_id] + if val is None: + raise OSError(f"missing cpulist for node{node_id}") + return val + if cpulist is None: + raise OSError("missing cpulist") + return cpulist + return real_read_text(self, *args, **kwargs) + + monkeypatch.setattr("builtins.open", fake_open) + monkeypatch.setattr("pathlib.Path.read_text", fake_read_text) + numa_utils._pct_sku_config.cache_clear() + + +def test_pct_binding_filters_cpus(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) + assert numa_utils._maybe_get_pct_cpu_binding([0]) == [0, 1, 16, 17, 64, 65, 80, 81] + + +@pytest.mark.parametrize( + "sku,expected_cpus", + [ + # 64-core SKUs (stride 16): cpus from "0-31,64-95" with cpu_id % 16 + # in (0, 1) -> 0, 1, 16, 17, 64, 65, 80, 81. + ("6776P", [0, 1, 16, 17, 64, 65, 80, 81]), + ("6774P", [0, 1, 16, 17, 64, 65, 80, 81]), + # 72-core SKU (stride 18): cpus from "0-31,64-95" with cpu_id % 18 + # in (0, 1) -> 0, 1, 18, 19, 72, 73, 90, 91. + ("6962P", [0, 1, 18, 19, 72, 73, 90, 91]), + ], +) +def test_pct_binding_fires_on_every_capable_sku(monkeypatch, sku, expected_cpus): + """Each SKU in ``_PCT_CAPABLE_SKUS`` engages the gate at its own + expected ``highest_perf`` and uses its own priority-core stride.""" + sku_config = numa_utils._PCT_CAPABLE_SKUS[sku] + _patch_pct_gates( + monkeypatch, + model_match=True, + highest_perf=sku_config.highest_perf, + sku=sku, + ) + assert numa_utils._maybe_get_pct_cpu_binding([0]) == expected_cpus + + +def test_pct_binding_fails_closed_when_sku_perf_mismatch(monkeypatch): + """6962P with 6776P's highest_perf (46 vs expected 44) must fail closed.""" + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46, sku="6962P") + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_pct_binding_disabled_when_cpu_model_mismatch(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=False, highest_perf=46) + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_pct_binding_disabled_when_highest_perf_does_not_match(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=42) + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_pct_binding_disabled_when_files_missing(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=None) + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_pct_binding_returns_none_when_node_cpulist_filter_empty(monkeypatch): + _patch_pct_gates( + monkeypatch, + model_match=True, + highest_perf=46, + cpulist="2-15,18-31", + ) + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_pct_binding_returns_none_when_node_cpulist_missing(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46, cpulist=None) + assert numa_utils._maybe_get_pct_cpu_binding([0]) is None + + +def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) + vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1]) + assert ( + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) + == "--physcpubind=0,1,16,17,64,65,80,81 --membind=1" + ) + + +def test_get_numactl_args_engine_core_baseline_single_node_shard(): + """Baseline (no PCT): single-NUMA shard -> single-node bind.""" + vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1]) + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == "--cpunodebind=0 --membind=0" + ) + + +def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes(): + """Baseline (no PCT): a TP=4 shard spanning both NUMA nodes -> bind to both.""" + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 1, 1], + tensor_parallel_size=4, + ) + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == "--cpunodebind=0,1 --membind=0,1" + ) + + +def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch): + """PCT: EngineCore for a multi-NUMA shard binds to the union of priority + cores across all shard nodes, so worker `--physcpubind` is always a + subset of EngineCore's `cpus_allowed`.""" + _patch_pct_gates( + monkeypatch, + model_match=True, + highest_perf=46, + cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"}, + ) + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 1, 1], + tensor_parallel_size=4, + ) + assert numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) == ( + "--physcpubind=" + "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209" + " --membind=0,1" + ) + + +def test_get_numactl_args_engine_core_pct_dp_shard_picks_local_nodes(monkeypatch): + """With DP=2, each shard's EngineCore binds only to its own NUMA nodes.""" + _patch_pct_gates( + monkeypatch, + model_match=True, + highest_perf=46, + cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"}, + ) + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 1, 1], + tensor_parallel_size=2, + data_parallel_rank_local=1, + ) + # Shard 1 owns gpu_indices 2 and 3 -> nodes [1, 1] -> {1}. + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == "--physcpubind=64,65,80,81,192,193,208,209 --membind=1" + ) + + +def test_get_numactl_args_engine_core_pct_external_launcher_spans_local_nodes( + monkeypatch, +): + """external_launcher (or multi-node-within-DP, or Ray) hits the + fallback branch. EngineCore must still span every local NUMA node + so it can mp-spawn its local workers without ``--physcpubind`` + strict-validation failures.""" + _patch_pct_gates( + monkeypatch, + model_match=True, + highest_perf=46, + cpulist_by_node={0: "0-31,128-159", 1: "64-95,192-223"}, + ) + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1], + distributed_executor_backend="external_launcher", + tensor_parallel_size=8, + ) + assert numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) == ( + "--physcpubind=" + "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209" + " --membind=0,1" + ) + + +def test_get_numactl_args_engine_core_baseline_multi_node_within_dp_spans_locals(): + """Multi-node-within-DP fallback: bind EngineCore to all local NUMA + nodes that the visible ``numa_bind_nodes`` reference.""" + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 0, 0, 1, 1, 1, 1], + nnodes_within_dp=2, + tensor_parallel_size=8, + ) + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == "--cpunodebind=0,1 --membind=0,1" + ) + + +def test_get_numactl_args_engine_core_skips_user_cpu_list(monkeypatch): + """EngineCore ignores ``--numa-bind-cpus``. + + Those are per-worker lists; binding EngineCore to any of them would + shrink its ``cpus_allowed`` below the strict-superset workers' + ``--physcpubind`` spawns need. We fall back to ``--cpunodebind`` over + the shard's NUMA nodes instead. PCT auto-detect is also bypassed when + the user is explicit (its priority-core union may not be a superset + of the user's per-worker cores).""" + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 0, 1, 1], + numa_bind_cpus=["0-3", "4-7", "64-67", "68-71"], + tensor_parallel_size=4, + ) + assert ( + numa_utils._get_numactl_enginecore_args( + vllm_config.parallel_config, local_rank=0 + ) + == "--cpunodebind=0,1 --membind=0,1" + ) + + +def test_get_numactl_args_user_cpus_override_pct(monkeypatch): + _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) + vllm_config = _make_config( + numa_bind=True, + numa_bind_nodes=[0, 1], + numa_bind_cpus=["0-3", "4-7"], + ) + assert ( + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) == "--physcpubind=4-7 --membind=1" ) @@ -57,7 +377,7 @@ def test_get_numactl_args_uses_dp_offset(): tensor_parallel_size=2, ) assert ( - numa_utils._get_numactl_args(vllm_config, local_rank=1) + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) == "--cpunodebind=1 --membind=1" ) @@ -66,7 +386,21 @@ def test_get_numactl_args_requires_detectable_nodes(monkeypatch): vllm_config = _make_config(numa_bind=True) monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", lambda: None) with pytest.raises(RuntimeError): - numa_utils._get_numactl_args(vllm_config, local_rank=0) + numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=0) + + +def test_configure_subprocess_rejects_unknown_process_kind(): + """configure_subprocess only knows 'worker' and 'EngineCore'; anything + else must raise ValueError instead of silently routing to the worker + path.""" + vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0]) + with ( + pytest.raises(ValueError, match="process_kind"), + numa_utils.configure_subprocess( + vllm_config, local_rank=0, process_kind="bogus" + ), + ): + pass def test_log_numactl_show(monkeypatch): 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/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 2ba140707ea..61134a4f5a2 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -87,8 +87,10 @@ class _MockModel(SupportsEncoderCudaGraph): def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: return EncoderCudaGraphConfig( modalities=["image"], - input_key_by_modality={"image": "pixel_values"}, - buffer_keys=["dummy_buf"], + buffer_keys=[ + "pixel_values", + "dummy_buf", + ], out_hidden_size=32, ) @@ -107,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 @@ -178,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 @@ -269,9 +276,6 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph): def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: return EncoderCudaGraphConfig( modalities=["image"], - input_key_by_modality={ - "image": "pixel_values", - }, buffer_keys=["dummy_buf"], out_hidden_size=_HIDDEN, ) @@ -350,11 +354,10 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph): n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE) dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype) return EncoderCudaGraphCaptureInputs( - mm_kwargs={ + values={ "pixel_values": dummy_pixel_values, - "image_grid_thw": grid_config, + "dummy_buf": dummy_buf, }, - buffers={"dummy_buf": dummy_buf}, ) def prepare_encoder_cudagraph_replay_buffers( @@ -367,14 +370,18 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph): n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE) p = next(self.parameters()) dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype) - return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf}) + return EncoderCudaGraphReplayBuffers( + values={ + "pixel_values": mm_kwargs["pixel_values"], + "dummy_buf": dummy_buf, + } + ) def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + values: dict[str, torch.Tensor], ) -> torch.Tensor: - return self._forward(mm_kwargs["pixel_values"]) + return self._forward(values["pixel_values"]) def encoder_eager_forward( self, @@ -407,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 @@ -465,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 --- @@ -473,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): @@ -551,10 +568,6 @@ class SimpleMockViTVideoModel(SimpleMockViTModel): def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: return EncoderCudaGraphConfig( modalities=["image", "video"], - input_key_by_modality={ - "image": "pixel_values", - "video": "pixel_values_videos", - }, buffer_keys=["dummy_buf"], out_hidden_size=_HIDDEN, ) @@ -654,11 +667,10 @@ class SimpleMockViTVideoModel(SimpleMockViTModel): n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE) dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype) return EncoderCudaGraphCaptureInputs( - mm_kwargs={ + values={ "pixel_values": dummy_pixel_values, - "image_grid_thw": grid_config, + "dummy_buf": dummy_buf, }, - buffers={"dummy_buf": dummy_buf}, ) def prepare_encoder_cudagraph_replay_buffers( @@ -670,14 +682,18 @@ class SimpleMockViTVideoModel(SimpleMockViTModel): n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE) p = next(self.parameters()) dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype) - return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf}) + return EncoderCudaGraphReplayBuffers( + values={ + "pixel_values": self._get_pixel_values(mm_kwargs), + "dummy_buf": dummy_buf, + } + ) def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + values: dict[str, torch.Tensor], ) -> torch.Tensor: - return self._forward(self._get_pixel_values(mm_kwargs)) + return self._forward(values["pixel_values"]) def encoder_eager_forward( self, @@ -718,14 +734,6 @@ class TestGetInputModality: } assert model.get_input_modality(mm_kwargs) == "video" - def test_video_model_config_has_both_modalities(self): - model = SimpleMockViTVideoModel() - cfg = model.get_encoder_cudagraph_config() - assert "image" in cfg.modalities - assert "video" in cfg.modalities - assert cfg.input_key_by_modality["image"] == "pixel_values" - assert cfg.input_key_by_modality["video"] == "pixel_values_videos" - # --------------------------------------------------------------------------- # GPU tests — video capture, replay, fallback, and mixed image+video @@ -749,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/executor/test_vllm_net_devices.py b/tests/v1/executor/test_vllm_net_devices.py new file mode 100644 index 00000000000..366dd001f29 --- /dev/null +++ b/tests/v1/executor/test_vllm_net_devices.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.v1.executor.vllm_net_devices import normalize_pci + + +@pytest.mark.parametrize( + "addr, expected", + [ + ("0000:3f:00.0", (0, 0x3F, 0, 0)), + ("0001:00:00.0", (1, 0, 0, 0)), + ("00000001:00:00.0", (1, 0, 0, 0)), + ("0009:00:00.0", (9, 0, 0, 0)), + ("0105:00:00.0", (0x105, 0, 0, 0)), + ("0000:0a:1f.7", (0, 0x0A, 0x1F, 7)), + ], +) +def test_normalize_pci_full_domain(addr, expected): + assert normalize_pci(addr) == expected + + +@pytest.mark.parametrize( + "addr, expected", + [ + ("01:00.0", (0, 1, 0, 0)), + ("3f:00.0", (0, 0x3F, 0, 0)), + ("40:00.0", (0, 0x40, 0, 0)), + ("ff:1f.7", (0, 0xFF, 0x1F, 7)), + ], +) +def test_normalize_pci_short_form(addr, expected): + assert normalize_pci(addr) == expected + + +def test_normalize_pci_case_insensitive(): + assert normalize_pci("0A:1F.7") == normalize_pci("0a:1f.7") + + +def test_normalize_pci_strips_whitespace(): + assert normalize_pci(" 0001:00:00.0 ") == (1, 0, 0, 0) + + +def test_normalize_pci_strips_0x_prefix(): + assert normalize_pci("0x0001:00:00.0") == (1, 0, 0, 0) + + +def test_normalize_pci_missing_function_raises(): + with pytest.raises(ValueError, match="missing function suffix"): + normalize_pci("0001:00:00") + + +def test_normalize_pci_invalid_function_char_raises(): + with pytest.raises(ValueError, match="invalid PCI function"): + normalize_pci("0001:00:00.z") + + +def test_normalize_pci_too_many_segments_raises(): + with pytest.raises(ValueError, match="invalid PCI BDF"): + normalize_pci("a:b:c:d.0") + + +def test_normalize_pci_bus_out_of_range_raises(): + with pytest.raises(ValueError, match="out of range"): + normalize_pci("0000:1ff:00.0") + + +def test_normalize_pci_device_out_of_range_raises(): + with pytest.raises(ValueError, match="out of range"): + normalize_pci("0000:00:20.0") + + +def test_normalize_pci_empty_string_raises(): + with pytest.raises(ValueError): + normalize_pci("") 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 0cfa1fc168f..20c230a4c2a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -24,7 +24,9 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.kv_offload.base import ( OffloadingEvent, OffloadingManager, + OffloadPolicy, ReqContext, + RequestOffloadingContext, get_offload_block_hash, ) from vllm.v1.request import RequestStatus @@ -748,12 +750,8 @@ class TestSlidingWindowLookup: @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bool): - """With do_remote_decode=True, after loading prefix blocks from CPU, - all blocks must be re-stored — not just the newly computed ones. - - This supports P/D disaggregation where the prefill instance offloads the - complete KV cache so a remote decode node can consume it.""" +def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool): + """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits.""" gpu_block_size = 4 block_size_factor = 3 offloaded_block_size = gpu_block_size * block_size_factor @@ -780,12 +778,13 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo # Reset GPU prefix cache so the next request must load from CPU. runner.scheduler.reset_prefix_cache() - # New request with do_remote_decode=True and 2 offloaded blocks. - # The first offloaded block matches what we stored in CPU. - runner.new_request( - token_ids=[0] * offloaded_block_size * 2, - kv_transfer_params={"do_remote_decode": True}, + # Manager returns REQUEST_LEVEL for the next request. + runner.manager.on_new_request.return_value = RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL ) + + # New request with 2 offloaded blocks; first matches what's in CPU. + runner.new_request(token_ids=[0] * offloaded_block_size * 2) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) @@ -957,9 +956,10 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): ) def setup(r, max_offload_tokens): - r.new_request(token_ids=[0] * offloaded_block_size * 3) - req = r.scheduler.requests[str(r.req_id)] - req.kv_transfer_params = {"max_offload_tokens": max_offload_tokens} + r.new_request( + token_ids=[0] * offloaded_block_size * 3, + kv_transfer_params={"max_offload_tokens": max_offload_tokens}, + ) r.manager.prepare_store.side_effect = ( lambda keys, req_context: generate_store_output(keys) ) @@ -1029,6 +1029,56 @@ def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): ) +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_offload_prompt_only(request_runner, async_scheduling: bool): + """offload_prompt_only=True offloads prompt blocks but never decode blocks. + + Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill + 4 more offloaded blocks. The flag clamps the offloadable token count to the + prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever + eligible for store; the decode blocks (offsets >= 6) are skipped. + + The request is intentionally not terminated (no EOS): a store is only + flushed when a request finishes (or is preempted), so without a finish + there is nothing to flush and the assertion stays free of flush-timing + subtleties. The decode steps are still enough for the prompt store to + complete and show up in expected_stored. + """ + gpu_block_size = 4 + block_size_factor = 3 + offloaded_block_size = gpu_block_size * block_size_factor # 12 + num_prompt_blocks = 2 + num_decode_blocks = 4 + prompt_offsets = (0, 1, 2, 3, 4, 5) + + runner = request_runner( + block_size=gpu_block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + extra_config_overrides={"offload_prompt_only": True}, + ) + + runner.manager.prepare_store.side_effect = ( + lambda keys, req_context: generate_store_output(keys) + ) + + runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks) + runner.run( + decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks), + expected_stored=prompt_offsets, + ) + + # Timing-independent guard: only the prompt's blocks were ever offered for + # store. If decode blocks leaked through, more keys would appear here. + offered_keys = { + key + for call in runner.manager.prepare_store.call_args_list + for key in call.args[0] + } + assert len(offered_keys) == num_prompt_blocks + + def test_flush_all_jobs_when_no_requests_remain(request_runner): """When all tracked requests are finished, build_connector_meta flushes all pending jobs since there will be no future step to complete them.""" @@ -1255,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 ad4d1ae2492..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 @@ -43,6 +44,7 @@ from vllm.v1.kv_offload.base import ( OffloadingSpec, OffloadKey, PrepareStoreOutput, + RequestOffloadingContext, make_offload_key, ) from vllm.v1.kv_offload.worker.worker import ( @@ -119,6 +121,7 @@ class MockOffloadingSpec(OffloadingSpec): self.manager.lookup.return_value = 0 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) self.manager.lookup.return_value = False + self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingHandler() def get_manager(self) -> OffloadingManager: @@ -169,6 +172,7 @@ class RequestRunner: block_size_factor: int = 1, async_scheduling: bool = True, kv_cache_groups: list[KVCacheGroupSpec] | None = None, + extra_config_overrides: dict[str, Any] | None = None, ): assert block_size_factor == 1 or kv_cache_groups is None, ( "block_size_factor > 1 requires all groups to have the same " @@ -192,9 +196,13 @@ class RequestRunner: extra_config: dict[str, Any] = { "spec_name": "MockOffloadingSpec", "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils", # noqa: E501 + # Preserve legacy behavior for tests; new opt-in tests override. + "offload_prompt_only": False, } if block_size_factor > 1: extra_config["block_size"] = block_size * block_size_factor + if extra_config_overrides: + extra_config.update(extra_config_overrides) vllm_config.kv_transfer_config = KVTransferConfig( kv_connector="OffloadingConnector", @@ -223,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( @@ -591,6 +604,7 @@ def request_runner(): async_scheduling, block_size_factor=1, kv_cache_groups=None, + extra_config_overrides=None, ): runner = RequestRunner( block_size=block_size, @@ -598,6 +612,7 @@ def request_runner(): block_size_factor=block_size_factor, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, + extra_config_overrides=extra_config_overrides, ) runners.append(runner) return runner 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_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 78269bfe40a..2a5c96a46e5 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib.util -import subprocess import uuid from unittest.mock import MagicMock, patch @@ -48,18 +47,6 @@ def _make_test_kv_cache_config() -> KVCacheConfig: aiter_available = importlib.util.find_spec("aiter") is not None mori_available = importlib.util.find_spec("mori") is not None - -def _rdma_available() -> bool: - """Check if RDMA devices are available.""" - try: - result = subprocess.run(["ibv_devinfo"], capture_output=True, text=True) - return "No IB devices found" not in result.stderr - except FileNotFoundError: - return False - - -rdma_available = _rdma_available() - pytestmark = pytest.mark.skipif( not (current_platform.is_rocm() and mori_available), reason="MoRIIOs are only available on ROCm with aiter package installed", @@ -224,11 +211,14 @@ def create_vllm_config( cache_dtype="auto", enable_prefix_caching=True, ) + # These tests exercise connector setup, not real RDMA transfer (MoRI wrapper is + # mocked), so we can use any backend without affecting test validity. Use xGMI to + # avoid requiring RNICs in CI. kv_transfer_config = KVTransferConfig( kv_connector="MoRIIOConnector", kv_role=role, enable_permute_local_kv=enable_permute_local_kv, - kv_connector_extra_config={"read_mode": read_mode}, + kv_connector_extra_config={"read_mode": read_mode, "backend": "xgmi"}, ) return VllmConfig( scheduler_config=scheduler_config, @@ -417,7 +407,6 @@ def test_read_mode_loads_remote_block_ids(): @pytest.mark.skipif( not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend" ) -@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available") def test_register_kv_caches(mock_parallel_groups): """Test that MoRIIOConnector.register_kv_caches correctly registers kv caches.""" ROLE = "kv_consumer" @@ -517,7 +506,6 @@ def test_register_kv_caches(mock_parallel_groups): @pytest.mark.skipif( not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend" ) -@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available") def test_moriio_handshake_returns_metadata(mock_parallel_groups): """MoRIIO handshake socket returns valid agent metadata over ZMQ.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py new file mode 100644 index 00000000000..0760d7141ec --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MultiConnector wrapping NixlConnector + SimpleCPUOffloadConnector. + +Verifies scheduler-level behavior: HMA support detection, load delegation +(first-wins), store-to-all, and connector metadata aggregation. +""" + +import uuid +from collections import defaultdict +from unittest.mock import patch + +import pytest +import torch + +from tests.v1.kv_connector.unit.utils import ( + create_model_runner_output, + create_request, + create_scheduler, + create_vllm_config, +) +from vllm import SamplingParams +from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( + MultiConnector, + MultiKVConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( + NixlConnectorMetadata, +) +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import get_request_block_hasher +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + SlidingWindowSpec, +) +from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.request import Request +from vllm.v1.simple_kv_offload.metadata import ( + SimpleCPUOffloadMetadata, + SimpleCPUOffloadWorkerMetadata, +) + +NIXL_WRAPPER_PATCH = ( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper" +) + + +class FakeNixlWrapper: + """Minimal mock of NixlWrapper for testing without NIXL hardware. + + Duplicated from test_nixl_connector.py to avoid importing that module + (which has heavy dependencies like ray). + """ + + AGENT_METADATA = b"fake_agent_metadata" + REMOTE_AGENT_NAME = "remote_agent" + + def __init__(self, agent_name: str, *args, **kwargs): + self._cycles_before_xfer_done = 0 + self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(lambda: 0) + + def get_reg_descs(self, caches_data, memory_type: str) -> list: + return [str(uuid.uuid4()) for _ in caches_data] + + def register_memory(self, descs, backends) -> None: + pass + + def deregister_memory(self, descs) -> None: + pass + + def get_xfer_descs(self, blocks_data, memory_type: str) -> list: + return [str(uuid.uuid4()) for _ in blocks_data] + + def prep_xfer_dlist(self, agent_name: str, descs: list) -> int: + return uuid.uuid4().int + + def get_agent_metadata(self) -> bytes: + return self.AGENT_METADATA + + def add_remote_agent(self, agent_metadata: bytes) -> str: + return self.REMOTE_AGENT_NAME + + def get_new_notifs(self) -> dict[str, list[bytes]]: + return {} + + def check_xfer_state(self, handle: int) -> str: + if self._check_xfer_state_cycles[handle] >= self._cycles_before_xfer_done: + return "DONE" + self._check_xfer_state_cycles[handle] += 1 + return "PROC" + + def release_xfer_handle(self, handle: int) -> None: + pass + + def release_dlist_handle(self, handle: int) -> None: + pass + + def remove_remote_agent(self, agent: str) -> None: + pass + + def send_notif(self, agent_name: str, notif_msg: bytes) -> None: + pass + + def make_prepped_xfer(self, *args, **kwargs) -> int: + return uuid.uuid4().int + + def transfer(self, handle: int) -> str: + return "PROC" + + def get_xfer_telemetry(self, handle: int) -> dict: + return {} + + def set_cycles_before_xfer_done(self, cycles: int): + pass + + +BLOCK_SIZE = 16 +NUM_KV_HEADS = 1 +HEAD_SIZE = 16 +DTYPE = torch.float16 +_BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize + + +def _make_kv_cache_config( + num_blocks: int = 100, + swa_enabled: bool = False, +) -> KVCacheConfig: + """Build KVCacheConfig with non-empty kv_cache_tensors. + + SimpleCPUOffloadConnector requires kv_cache_tensors with real sizes + (used by _derive_cpu_config). + """ + fa_layers = ["layer0", "layer2"] + groups = [ + KVCacheGroupSpec( + fa_layers, + FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ), + ) + ] + tensors = [ + KVCacheTensor( + size=_BYTES_PER_BLOCK * num_blocks, + shared_by=fa_layers, + ) + ] + if swa_enabled: + sw_layers = ["layer1", "layer3"] + groups.append( + KVCacheGroupSpec( + sw_layers, + SlidingWindowSpec( + block_size=BLOCK_SIZE, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + sliding_window=128, + ), + ) + ) + tensors.append( + KVCacheTensor( + size=_BYTES_PER_BLOCK * num_blocks, + shared_by=sw_layers, + ) + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=groups, + ) + + +def _multi_connector_config(swa_enabled: bool = False): + """Return (vllm_config, kv_cache_config) for a MultiConnector test.""" + kv_cache_config = _make_kv_cache_config(swa_enabled=swa_enabled) + vllm_config = create_vllm_config( + kv_connector="MultiConnector", + kv_connector_extra_config={ + "connectors": [ + { + "kv_connector": "NixlConnector", + "kv_role": "kv_both", + }, + { + "kv_connector": "SimpleCPUOffloadConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cpu_bytes_to_use": 1 << 30, + }, + }, + ], + }, + ) + return vllm_config, kv_cache_config + + +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper) +def test_nixl_wins_load_over_cpu_offload(): + """When NixlConnector (index 0) has matched tokens from a remote prefill, it should + win the load: Nixl metadata tracks the recv while CPU offload metadata has no load + scheduled.""" + vllm_config, kv_cache_config = _multi_connector_config() + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) + mc = scheduler.connector + assert isinstance(mc, MultiConnector) + + request = create_request( + request_id=1, + num_tokens=BLOCK_SIZE * 3, + do_remote_prefill=True, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(request) + sched_out = scheduler.schedule() + + assert mc._requests_to_connector[request.request_id] == 0 + + meta = sched_out.kv_connector_metadata + assert isinstance(meta, MultiKVConnectorMetadata) + assert len(meta.metadata) == 2 + + nixl_meta = meta.metadata[0] + assert isinstance(nixl_meta, NixlConnectorMetadata) + # nixl is tracking the request + assert request.request_id in nixl_meta.reqs_to_recv + + cpu_meta = meta.metadata[1] + assert isinstance(cpu_meta, SimpleCPUOffloadMetadata) + assert not cpu_meta.load_gpu_blocks + + +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper) +def test_cpu_offload_wins_when_nixl_has_no_match(): + """When NixlConnector returns 0 matched tokens and SimpleCPUOffloadConnector has a + CPU cache hit, the CPU offload connector (index 1) wins the load.""" + vllm_config, kv_cache_config = _multi_connector_config() + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) + mc = scheduler.connector + assert isinstance(mc, MultiConnector) + + req1 = create_request( + request_id=10, + num_tokens=BLOCK_SIZE * 3, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(req1) + sched_out1 = scheduler.schedule() + + # build_connector_meta runs before _update_after_schedule, so num_computed_tokens + # is still 0 during the first schedule and the store sees no confirmed blocks. + # Simulate one model step so the second schedule triggers the store. + model_output = create_model_runner_output(reqs=[req1]) + scheduler.update_from_output(sched_out1, model_output) + + sched_out2 = scheduler.schedule() + meta2 = sched_out2.kv_connector_metadata + assert isinstance(meta2, MultiKVConnectorMetadata) + cpu_meta = meta2.metadata[1] + assert isinstance(cpu_meta, SimpleCPUOffloadMetadata) + assert cpu_meta.store_event >= 0, "Expected a store event on the second schedule" + + cpu_connector = mc._connectors[1] + worker_meta = SimpleCPUOffloadWorkerMetadata( + completed_store_events={ + cpu_meta.store_event: cpu_connector.scheduler_manager._expected_worker_count + }, + ) + output = KVConnectorOutput( + finished_recving=set(), + kv_connector_worker_meta=worker_meta, + ) + cpu_connector.update_connector_output(output) + + req2 = Request( + request_id="id-cpu-offload-hit", + prompt_token_ids=req1.prompt_token_ids, + sampling_params=SamplingParams(max_tokens=16), + pooling_params=None, + mm_features=None, + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), + ) + + hit_tokens, is_async = mc.get_num_new_matched_tokens(req2, num_computed_tokens=0) + assert hit_tokens is not None and hit_tokens > 0 + assert mc._requests_to_connector[req2.request_id] == 1 + assert is_async is True + + +@pytest.mark.parametrize("swa_enabled", [False, True], ids=["fa_only", "fa_sw"]) +@patch(NIXL_WRAPPER_PATCH, FakeNixlWrapper) +def test_request_finished_no_async_save(swa_enabled: bool): + """A normal request (no P/D) produces no async save from either connector. + MultiConnector returns (False, None) via both request_finished and + request_finished_all_groups, and cleans up _requests_to_connector.""" + from vllm.v1.request import RequestStatus + + vllm_config, kv_cache_config = _multi_connector_config(swa_enabled=swa_enabled) + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config) + mc = scheduler.connector + assert isinstance(mc, MultiConnector) + + request = create_request( + request_id=40, + num_tokens=BLOCK_SIZE * 2, + block_size=BLOCK_SIZE, + ) + scheduler.add_request(request) + scheduler.schedule() + request.status = RequestStatus.FINISHED_STOPPED + + block_ids = (list(range(2)), list(range(2))) if swa_enabled else (list(range(2)),) + async_save, txfer_params = mc.request_finished_all_groups(request, block_ids) + + assert async_save is False + assert txfer_params is None + assert len(mc._extra_async_saves) == 0 + assert request.request_id not in mc._requests_to_connector 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 81% rename from tests/v1/kv_offload/test_tiering_offloading.py rename to tests/v1/kv_offload/tiering/test_tiering_offloading.py index 7359a39384f..5a7c11787d9 100644 --- a/tests/v1/kv_offload/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -19,7 +19,9 @@ import torch from vllm.v1.kv_offload.base import ( OffloadKey, + OffloadPolicy, ReqContext, + RequestOffloadingContext, make_offload_key, ) from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager @@ -121,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)) @@ -183,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. @@ -200,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: @@ -236,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 @@ -269,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)) @@ -287,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) @@ -326,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. @@ -352,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 @@ -387,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() @@ -412,6 +419,77 @@ class TestTieringOffloadingManager: job_metadata = self.secondary_tier1.submit_store.call_args.args[0] assert job_metadata.req_context is ctx + def test_on_new_request_lifecycle(self, manager_setup): + """Policy defaults to BLOCK_LEVEL, escalates when a tier requests it, + and is cleaned up on on_request_finished.""" + # Default: all tiers return BLOCK_LEVEL + ctx = ReqContext(req_id="req_policy_lifecycle") + result = self.manager.on_new_request(ctx) + assert result.policy == OffloadPolicy.BLOCK_LEVEL + self.manager.on_request_finished(ctx) + + # Escalate: tier1 requests REQUEST_LEVEL + self.secondary_tier1.on_new_request = ( + lambda req_context: RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL + ) + ) + + ctx = ReqContext(req_id="req_policy_lifecycle_2") + result = self.manager.on_new_request(ctx) + assert result.policy == OffloadPolicy.REQUEST_LEVEL + assert ctx.req_id in self.manager._request_level_tiers + + # Cleanup + self.manager.on_request_finished(ctx) + assert ctx.req_id not in self.manager._request_level_tiers + + def test_prepare_store_cascades_existing_blocks_to_request_level_tiers( + self, manager_setup + ): + """prepare_store cascades hit blocks to request-level tiers only.""" + # Store some blocks to primary first + existing_blocks = to_keys(range(3)) + result = self.manager.prepare_store(existing_blocks, _CTX) + assert result is not None + self.manager.complete_store(existing_blocks, _CTX, success=True) + # Drain cascade completions + self._simulate_on_schedule_end() + + # Make tier1 request-level, tier2 stays block-level + self.secondary_tier1.on_new_request = ( + lambda req_context: RequestOffloadingContext( + policy=OffloadPolicy.REQUEST_LEVEL + ) + ) + + ctx = ReqContext(req_id="req_cascade") + self.manager.on_new_request(ctx) + + # Spy on submit_store + self.secondary_tier1.submit_store = MagicMock( + wraps=self.secondary_tier1.submit_store + ) + self.secondary_tier2.submit_store = MagicMock( + wraps=self.secondary_tier2.submit_store + ) + + # Call prepare_store with existing + new blocks + new_blocks = to_keys(range(3, 5)) + all_blocks = existing_blocks + new_blocks + result = self.manager.prepare_store(all_blocks, ctx) + assert result is not None + assert set(result.keys_to_store) == set(new_blocks) + + # Only tier1 (request-level) should get existing blocks cascaded now. + # New blocks are cascaded to ALL tiers later via complete_store(). + self.secondary_tier1.submit_store.assert_called_once() + job_metadata = self.secondary_tier1.submit_store.call_args.args[0] + assert set(job_metadata.keys) == set(existing_blocks) + + # tier2 (block-level) does not get existing blocks here. + self.secondary_tier2.submit_store.assert_not_called() + class TestTieringOffloadingWithoutSecondaryTiers: """Test TieringOffloadingManager with no secondary tiers (backward compat).""" 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/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 7d488aaabf2..a80fddc9235 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -5,6 +5,7 @@ import torch from torch import Generator from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch DEVICE_TYPE = current_platform.device_type @@ -151,7 +152,7 @@ def test_flashinfer_sampler(): # ============================================================================= -@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available on this platform") class TestTritonTopkTopp: """Tests for the Triton top-k/top-p kernel.""" 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/spec_decode/test_acceptance_length.py b/tests/v1/spec_decode/test_acceptance_length.py index 62ff100fdbf..90e3821e2f1 100644 --- a/tests/v1/spec_decode/test_acceptance_length.py +++ b/tests/v1/spec_decode/test_acceptance_length.py @@ -39,6 +39,8 @@ class Eagle3ModelConfig: marks: list = field(default_factory=list) # Custom relative tolerance (defaults to DEFAULT_RTOL if None) rtol: float | None = None + # ROCm-specific test configuration + rocm_expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list) # Model configurations for EAGLE3 acceptance length tests. @@ -69,6 +71,7 @@ EAGLE3_MODEL_CONFIGS = [ # FLASHINFER incompatible: gpt-oss-20b uses sink attention which # FLASHINFER does not support ("sink setting not supported") excluded_backends={AttentionBackendEnum.FLASHINFER}, + rocm_expected_acceptance_lengths_per_pos=[0.7040, 0.4820, 0.3350], ), Eagle3ModelConfig( verifier="Qwen/Qwen3-VL-30B-A3B-Instruct-FP8", @@ -99,16 +102,14 @@ EXCLUDED_BACKENDS = {AttentionBackendEnum.FLEX_ATTENTION} def get_available_attention_backends() -> list[str]: + if current_platform.is_rocm(): + return ["auto"] + # Check if get_valid_backends is actually defined in the platform class # (not just returning None from __getattr__) get_valid_backends = getattr(current_platform.__class__, "get_valid_backends", None) if get_valid_backends is None: - if current_platform.is_rocm(): - # ROCm uses Triton as its default attention backend since - # Flash Attention is not supported. - return ["TRITON_ATTN"] - else: - return ["FLASH_ATTN"] + return ["FLASH_ATTN"] device_capability = current_platform.get_device_capability() if device_capability is None: @@ -167,6 +168,8 @@ def get_mt_bench_prompts( disable_shuffle=False, skip_chat_template=False, trust_remote_code=False, + enable_multimodal_chat=False, + request_id_prefix="", ) samples = get_samples(args, tokenizer) prompt_ids = [ @@ -233,9 +236,12 @@ def test_eagle3_acceptance_length( monkeypatch: pytest.MonkeyPatch, ): # Skip if this backend is incompatible with the model - backend_enum = AttentionBackendEnum[attention_backend] - if backend_enum in model_config.excluded_backends: - pytest.skip(f"{attention_backend} is incompatible with {model_config.id}") + attention_config = None + if attention_backend != "auto": + backend_enum = AttentionBackendEnum[attention_backend] + if backend_enum in model_config.excluded_backends: + pytest.skip(f"{attention_backend} is incompatible with {model_config.id}") + attention_config = {"backend": attention_backend} with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") @@ -247,11 +253,16 @@ def test_eagle3_acceptance_length( "model": model_config.drafter, "num_speculative_tokens": num_spec_tokens, }, - attention_config={"backend": attention_backend}, + attention_config=attention_config, tensor_parallel_size=tp_size, gpu_memory_utilization=0.7, disable_log_stats=False, max_model_len=DEFAULT_MAX_MODEL_LEN, + # Qwen/Qwen3-30B-A3B-FP8 with TP=4 needs EP + # https://github.com/vllm-project/vllm/issues/25292 + enable_expert_parallel=( + tp_size == 4 and "Qwen3-VL" in model_config.verifier + ), ) as vllm_runner: tokenizer = vllm_runner.llm.get_tokenizer() prompt_ids = get_mt_bench_prompts(tokenizer, DEFAULT_NUM_PROMPTS) @@ -272,6 +283,11 @@ def test_eagle3_acceptance_length( expected = model_config.expected_acceptance_length actual_per_pos = results["acceptance_lengths_per_pos"] expected_per_pos = model_config.expected_acceptance_lengths_per_pos + if ( + current_platform.is_rocm() + and model_config.rocm_expected_acceptance_lengths_per_pos + ): + expected_per_pos = model_config.rocm_expected_acceptance_lengths_per_pos rel_error = abs(actual_acceptance_length - expected) / expected @@ -294,14 +310,14 @@ def test_eagle3_acceptance_length( zip(actual_per_pos, expected_per_pos) ): if exp > 0: - pos_rel_error = abs(actual - exp) / exp - assert pos_rel_error <= rtol, ( + min_expected = exp * (1 - rtol) + assert actual >= min_expected, ( f"Per-position acceptance length regression at pos {pos} " f"for {model_config.id}!\n" f" Expected: {exp:.3f}\n" f" Actual: {actual:.3f}\n" - f" Relative error: {pos_rel_error:.2%} " - f"(tolerance: {rtol:.2%})" + f" Minimum: {min_expected:.3f}\n" + f" Tolerance: rtol={rtol:.2%}" ) print( 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/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index ffb4d7f474b..1a1352249c3 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -7,6 +7,7 @@ from unittest.mock import Mock import numpy as np import pytest import torch +import torch.nn as nn import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module from vllm.config import ( @@ -22,6 +23,7 @@ from vllm.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) +from vllm.distributed.weight_transfer.base import SparseWeightPatch from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.platforms import current_platform @@ -39,6 +41,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheTensor, ) +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker.gpu_input_batch import InputBatch @@ -279,7 +282,8 @@ def test_sample_tokens_receives_pp_sampled_ids_only_on_non_last_rank( lambda: SimpleNamespace(world_size=world_size, is_last_rank=is_last_rank), ) - assert GPUModelRunner.sample_tokens(runner, None) is None + output = GPUModelRunner.sample_tokens(runner, None) + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) assert receive_calls == expected_calls @@ -298,7 +302,8 @@ def test_sample_tokens_skips_pp_group_lookup_without_async_scheduling( pytest.fail, ) - assert GPUModelRunner.sample_tokens(runner, None) is None + output = GPUModelRunner.sample_tokens(runner, None) + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) def test_select_common_block_size_no_valid_option(): @@ -781,6 +786,73 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler(): assert torch.equal(passed_draft_probs, expected_draft_probs) +def test_apply_sparse_weight_patches_updates_only_selected_entries(): + class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(6, dtype=torch.float32)) + + runner = object.__new__(GPUModelRunner) + runner.model = DummyModel() + + runner.apply_sparse_weight_patches( + [ + SparseWeightPatch( + name="weight", + indices=torch.tensor([1, 4], dtype=torch.int32), + values=torch.tensor([3.5, -2.0], dtype=torch.float32), + ) + ] + ) + + expected = torch.tensor([0.0, 3.5, 0.0, 0.0, -2.0, 0.0], dtype=torch.float32) + assert torch.equal(runner.get_model().weight.data, expected) + + +def test_apply_sparse_weight_patches_rejects_mismatched_lengths(): + class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + + runner = object.__new__(GPUModelRunner) + runner.model = DummyModel() + + with pytest.raises(ValueError, match="matching lengths"): + runner.apply_sparse_weight_patches( + [ + SparseWeightPatch( + name="weight", + indices=torch.tensor([1, 2], dtype=torch.int32), + values=torch.tensor([1.0], dtype=torch.float32), + ) + ] + ) + + +def test_apply_sparse_weight_patches_rejects_non_contiguous_param(): + class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter( + torch.arange(12, dtype=torch.float32).view(3, 4).t() + ) + + runner = object.__new__(GPUModelRunner) + runner.model = DummyModel() + + with pytest.raises(NotImplementedError, match="contiguous params"): + runner.apply_sparse_weight_patches( + [ + SparseWeightPatch( + name="weight", + indices=torch.tensor([1], dtype=torch.int32), + values=torch.tensor([1.0], dtype=torch.float32), + ) + ] + ) + + def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config): torch.set_default_dtype(torch.float16) layer_0 = "model.layers.0.self_attn.attn" 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 d68ff83c407..1db07baf93d 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -7,6 +7,7 @@ from typing import Any import torch +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.worker.gpu import eplb_utils as eplb from vllm.v1.worker.gpu import model_runner as mrv2 @@ -76,7 +77,10 @@ def _make_runner(**overrides: Any) -> Any: runner.max_num_reqs = 8 runner.max_num_tokens = 16 runner.decode_query_len = 1 - runner.kv_connector = SimpleNamespace(set_disabled=lambda *_: None) + runner.kv_connector = SimpleNamespace( + set_disabled=lambda *_: None, + post_forward=lambda *_, **__: None, + ) runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device) runner.pooling_runner = None runner.execute_model_state = None @@ -163,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, @@ -171,17 +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() - assert mrv2.GPUModelRunner.sample_tokens(runner, None) is None - assert events == ["postprocess", "eplb"] + 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 == ["receive", "postprocess_num_computed_tokens", "eplb"] diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py new file mode 100644 index 00000000000..dba0f658542 --- /dev/null +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.config.parallel import ParallelConfig +from vllm.config.weight_transfer import WeightTransferConfig +from vllm.distributed.weight_transfer.base import SparseWeightPatch +from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine +from vllm.v1.worker.gpu_worker import Worker + + +def _make_nccl_engine() -> NCCLWeightTransferEngine: + parallel_config = MagicMock(spec=ParallelConfig) + parallel_config.rank = 0 + parallel_config.world_size = 1 + parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = 0 + return NCCLWeightTransferEngine( + WeightTransferConfig(backend="nccl"), + parallel_config, + MagicMock(spec=torch.nn.Module), + ) + + +def test_update_weights_sparse_dispatches_to_sparse_receive(monkeypatch): + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + worker = object.__new__(Worker) + worker.device = "cpu" + worker.parallel_config = SimpleNamespace(world_size=1) + worker.weight_transfer_engine = _make_nccl_engine() + worker._weight_update_active = True + worker._is_checkpoint_format = False + + applied_patches = [] + + def apply_sparse_weight_patches(patches): + applied_patches.extend(patches) + + worker.model_runner = SimpleNamespace( + apply_sparse_weight_patches=apply_sparse_weight_patches, + ) + + received_kinds = [] + + def receive_sparse_weights(update_info, apply_patches): + received_kinds.append(update_info.update_kind) + apply_patches( + [ + SparseWeightPatch( + name="layer.weight", + indices=torch.tensor([1], dtype=torch.int32), + values=torch.tensor([2.0], dtype=torch.float32), + ) + ] + ) + + worker.weight_transfer_engine.receive_sparse_weights = receive_sparse_weights + + Worker.update_weights( + worker, + { + "names": ["layer.weight"], + "dtype_names": ["float32"], + "shapes": [[4]], + "num_updates_list": [1], + "update_kind": "sparse_flat", + }, + ) + + assert received_kinds == ["sparse_flat"] + assert len(applied_patches) == 1 + assert torch.equal(applied_patches[0].indices, torch.tensor([1], dtype=torch.int32)) + + +def test_update_weights_sparse_rejects_tp_or_pp(monkeypatch): + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + worker = object.__new__(Worker) + worker.device = "cpu" + worker.parallel_config = SimpleNamespace(world_size=2) + worker.weight_transfer_engine = _make_nccl_engine() + worker._weight_update_active = True + worker._is_checkpoint_format = False + worker.model_runner = SimpleNamespace(apply_sparse_weight_patches=lambda _: None) + + with pytest.raises(NotImplementedError, match="TP=1 and PP=1"): + Worker.update_weights( + worker, + { + "names": ["layer.weight"], + "dtype_names": ["float32"], + "shapes": [[4]], + "num_updates_list": [1], + "update_kind": "sparse_flat", + }, + ) + assert worker._weight_update_active is False + assert worker._is_checkpoint_format is True + + +def test_update_weights_sparse_rejects_checkpoint_format(monkeypatch): + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + worker = object.__new__(Worker) + worker.device = "cpu" + worker.parallel_config = SimpleNamespace(world_size=1) + worker.weight_transfer_engine = _make_nccl_engine() + worker._weight_update_active = True + worker._is_checkpoint_format = True + worker.model_runner = SimpleNamespace(model=MagicMock()) + + with pytest.raises(ValueError, match="start_weight_update"): + Worker.update_weights( + worker, + { + "names": ["layer.weight"], + "dtype_names": ["float32"], + "shapes": [[4]], + "num_updates_list": [1], + "update_kind": "sparse_flat", + }, + ) + assert worker._weight_update_active is False + assert worker._is_checkpoint_format is True + + +def test_update_weights_resets_state_when_update_info_is_invalid(monkeypatch): + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + worker = object.__new__(Worker) + worker.device = "cpu" + worker.parallel_config = SimpleNamespace(world_size=1) + worker.weight_transfer_engine = _make_nccl_engine() + worker._weight_update_active = True + worker._is_checkpoint_format = False + + with pytest.raises(ValueError, match="cannot be empty"): + Worker.update_weights( + worker, + { + "names": [], + "dtype_names": [], + "shapes": [], + "num_updates_list": [], + "update_kind": "sparse_flat", + }, + ) + assert worker._weight_update_active is False + assert worker._is_checkpoint_format is True 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/_aiter_ops.py b/vllm/_aiter_ops.py index 2c9b939fa2f..5a8b690433c 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2394,6 +2394,7 @@ class rocm_aiter_ops: alibi_slopes: torch.Tensor | None = None, return_lse: bool = False, out: torch.Tensor | None = None, + sink_ptr: torch.Tensor | None = None, ): """ Flash attention with variable length sequences. @@ -2422,6 +2423,7 @@ class rocm_aiter_ops: alibi_slopes=alibi_slopes, return_lse=return_lse, out=out, + sink_ptr=sink_ptr, ) @staticmethod diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 84f944df2bf..f12d128f083 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -650,6 +650,51 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None torch.ops._C.gptq_shuffle(q_weight, q_perm, bit) +def gptq_gemm_rdna3( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, +) -> torch.Tensor: + return torch.ops._rocm_C.gptq_gemm_rdna3( + a, b_q_weight, b_qzeros, b_scales, b_g_idx, use_v2_format + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3"): + + @register_fake("_rocm_C::gptq_gemm_rdna3") + def _gptq_gemm_rdna3_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + +if hasattr(torch.ops, "_rocm_C") and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3_wmma"): + + @register_fake("_rocm_C::gptq_gemm_rdna3_wmma") + def _gptq_gemm_rdna3_wmma_fake( + a: torch.Tensor, + b_q_weight: torch.Tensor, + b_qzeros: torch.Tensor, + b_scales: torch.Tensor, + b_g_idx: torch.Tensor, + use_v2_format: bool, + ) -> torch.Tensor: + return torch.empty( + (a.size(0), b_q_weight.size(1)), dtype=a.dtype, device=a.device + ) + + if hasattr(torch.ops._C, "allspark_w8a16_gemm"): @register_fake("_C::allspark_w8a16_gemm") @@ -2367,6 +2412,31 @@ def dsv3_router_gemm( return output +def fp32_router_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, +) -> torch.Tensor: + output = torch.empty( + hidden_states.shape[0], + router_weight.shape[0], + device=hidden_states.device, + dtype=torch.float32, + ) + torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight) + return output + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"): + + @register_fake("_C::fp32_router_gemm") + def fp32_router_gemm_fake( + output: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + ) -> None: + return + + def topk_softmax( topk_weights: torch.Tensor, topk_ids: torch.Tensor, @@ -3719,29 +3789,6 @@ def matmul_mxf4_bf16_tn( return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) -if hasattr(torch.ops._qutlass_C, "matmul_ada_mxf4_bf16_tn"): - - @register_fake("_qutlass_C::matmul_ada_mxf4_bf16_tn") - def _fake_matmul_ada_mxf4_bf16_tn( - a: torch.Tensor, - b: torch.Tensor, - a_sf: torch.Tensor, - b_sf: torch.Tensor, - alpha: torch.Tensor, - ): - return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16) - - -def matmul_ada_mxf4_bf16_tn( - a: torch.Tensor, - b: torch.Tensor, - a_sf: torch.Tensor, - b_sf: torch.Tensor, - alpha: torch.Tensor, -) -> torch.Tensor: - return torch.ops._qutlass_C.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha) - - if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"): @register_fake("_qutlass_C::fusedQuantizeMxQuest") diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 233c8fb632f..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__) @@ -114,8 +115,34 @@ def _gdn_attention_core_xpu_impl( attn_metadata = attn_metadata_raw[self.prefix] assert isinstance(attn_metadata, GDNAttentionMetadata) - # TODO: xpu does not support speculative decoding yet - assert attn_metadata.spec_sequence_masks is None # type: ignore[attr-defined] + num_actual_tokens = attn_metadata.num_actual_tokens + num_accepted_tokens = attn_metadata.num_accepted_tokens + + num_prefills = attn_metadata.num_prefills + num_decodes = attn_metadata.num_decodes + num_spec_decodes = attn_metadata.num_spec_decodes + + has_initial_state = attn_metadata.has_initial_state + + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc + non_spec_token_indx = attn_metadata.non_spec_token_indx + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + non_spec_state_indices_tensor = ( + non_spec_state_indices_tensor.contiguous() + if non_spec_state_indices_tensor is not None + else None + ) + + spec_query_start_loc = attn_metadata.spec_query_start_loc + spec_token_indx = attn_metadata.spec_token_indx + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501 + + spec_sequence_masks = attn_metadata.spec_sequence_masks + if spec_sequence_masks is not None: + if non_spec_token_indx is not None: + non_spec_token_indx = non_spec_token_indx.to(torch.int32) + if spec_token_indx is not None: + spec_token_indx = spec_token_indx.to(torch.int32) conv_weights = self.conv1d.weight.view( self.conv1d.weight.size(0), self.conv1d.weight.size(2) @@ -137,12 +164,18 @@ def _gdn_attention_core_xpu_impl( activation=self.activation, A_log=self.A_log, dt_bias=self.dt_bias, - num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined] - num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined] - has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined] - non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, # type: ignore[attr-defined] - non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, # type: ignore[attr-defined] - num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined] + num_prefills=num_prefills, # type: ignore[attr-defined] + num_decodes=num_decodes, # type: ignore[attr-defined] + num_spec_decodes=num_spec_decodes, # type: ignore[attr-defined] + has_initial_state=has_initial_state, # type: ignore[attr-defined] + non_spec_query_start_loc=non_spec_query_start_loc, # type: ignore[attr-defined] + non_spec_token_indx=non_spec_token_indx, # type: ignore[attr-defined] + non_spec_state_indices_tensor=non_spec_state_indices_tensor, # type: ignore[attr-defined] + spec_query_start_loc=spec_query_start_loc, # type: ignore[attr-defined] + spec_token_indx=spec_token_indx, # type: ignore[attr-defined] + spec_state_indices_tensor=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, # type: ignore[attr-defined] + num_actual_tokens=num_actual_tokens, # type: ignore[attr-defined] tp_size=self.tp_size, reorder_input=not self.gqa_interleaved_layout, ) @@ -306,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 @@ -366,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 @@ -508,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/assets/video.py b/vllm/assets/video.py index 9ec2e4d1677..72cd196c68f 100644 --- a/vllm/assets/video.py +++ b/vllm/assets/video.py @@ -7,10 +7,10 @@ from typing import Any, ClassVar, Literal import numpy as np import numpy.typing as npt -from huggingface_hub import hf_hub_download from PIL import Image from vllm.multimodal.media.audio import load_audio_pyav +from vllm.transformers_utils.repo_utils import hf_api from .base import get_cache_dir @@ -27,7 +27,7 @@ def download_video_asset(filename: str) -> str: video_path = video_directory / filename video_path_str = str(video_path) if not video_path.exists(): - video_path_str = hf_hub_download( + video_path_str = hf_api().hf_hub_download( repo_id="raushan-testing-hf/videos-test", filename=filename, repo_type="dataset", diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 70545a3b495..59e2aa578c3 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -31,7 +31,6 @@ from typing import Any, cast import numpy as np import pybase64 as base64 -from huggingface_hub import snapshot_download from PIL import Image from typing_extensions import deprecated @@ -45,7 +44,9 @@ 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 from vllm.utils.import_utils import PlaceholderModule from vllm.utils.mistral import is_mistral_tokenizer @@ -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 @@ -3334,7 +3415,10 @@ class MMVUDataset(HuggingFaceDataset): self._remote_path_root = ( f"https://huggingface.co/datasets/{self.hf_name}/resolve/main" ) - self._local_path_root = snapshot_download(self.hf_name, repo_type="dataset") + self._local_path_root = hf_api().snapshot_download( + self.hf_name, + repo_type="dataset", + ) def sample( self, 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/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 5406f611e87..569fac667eb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1101,6 +1101,14 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass): return False return bool(compile_range.end <= self.max_token_num) + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + self.matched_count = self.pm_pass.apply(graph) + VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count + logger.debug( + "%s Replaced %s patterns", self.__class__.__name__, self.matched_count + ) + def __del__(self) -> None: if getattr(self, "disabled", True): return diff --git a/vllm/compilation/passes/fusion/collective_fusion.py b/vllm/compilation/passes/fusion/collective_fusion.py index 29d79c9b92c..3658877c67a 100644 --- a/vllm/compilation/passes/fusion/collective_fusion.py +++ b/vllm/compilation/passes/fusion/collective_fusion.py @@ -917,13 +917,13 @@ class AsyncTPPass(VllmFusionPatternMatcherPass): AllGatherScaledMMPattern(self.model_dtype, self.device).register( self.pm_pass ) - - CutlassScaledMMReduceScatterPattern(self.model_dtype, self.device).register( - self.pm_pass - ) - AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register( - self.pm_pass - ) + if hasattr(torch.ops._C, "cutlass_scaled_mm"): + CutlassScaledMMReduceScatterPattern( + self.model_dtype, self.device + ).register(self.pm_pass) + AllGatherCutlassScaledMMPattern(self.model_dtype, self.device).register( + self.pm_pass + ) with suppress(ImportError): import vllm.utils.flashinfer # noqa: F401 if hasattr(torch.ops.vllm, "bmm_fp8"): 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/parallel.py b/vllm/config/parallel.py index 960e1d343c0..f32ecef1482 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -78,7 +78,7 @@ class EPLBConfig: """ Interval for logging the balancedness. """ - use_async: bool = False + use_async: bool = True """ Whether to use non-blocking EPLB. """ @@ -267,7 +267,12 @@ class ParallelConfig: """num of nodes for multi-node distributed inference when distributed_executor_backend is mp.""" numa_bind: bool = False - """Enable NUMA binding for GPU worker subprocesses.""" + """Enable NUMA binding for GPU worker subprocesses. + + By default, workers are pinned to their GPU's NUMA-local CPUs and + memory; on PCT-capable Xeons they also auto-bind to the SKU's + PCT priority cores. + """ numa_bind_nodes: list[int] | None = None """NUMA node to bind each GPU worker to. diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 47d35f4ff4b..e388987d6d4 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -485,7 +485,16 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]} ) - if hf_config.model_type == "step3p5": + if hf_config.model_type in ("step3p5", "step3p7") or hf_config.architectures[ + 0 + ] in ("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration"): + quantization_config = getattr(hf_config, "quantization_config", None) + hf_config = getattr(hf_config, "text_config", hf_config) + if ( + quantization_config is not None + and getattr(hf_config, "quantization_config", None) is None + ): + hf_config.update({"quantization_config": quantization_config}) hf_config.model_type = "step3p5_mtp" n_predict = getattr(hf_config, "num_nextn_predict_layers", 1) hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]}) @@ -705,22 +714,16 @@ class SpeculativeConfig: MTPModelTypes ): self.method = "mtp" - if self.num_speculative_tokens > 1: + if ( + self.num_speculative_tokens > 1 + and self.draft_model_config.hf_config.model_type + != "step3p5_mtp" + ): logger.warning( "Enabling num_speculative_tokens > 1 will run " "multiple times of forward on same MTP layer" ",which may result in lower acceptance rate" ) - elif self.draft_model_config.hf_config.model_type in ( - "longcat_flash_mtp" - ): - self.method = "longcat_flash_mtp" - if self.num_speculative_tokens > 1: - logger.warning( - "LongCat MTP models only have " - "one layer. Might need some code changes " - "to support multiple layers." - ) elif self.method == "draft_model": pass else: @@ -1056,6 +1059,14 @@ class SpeculativeConfig: == "gemma4_mtp" ) + def use_step3p5_mtp(self) -> bool: + return ( + self.method == "mtp" + and self.draft_model_config is not None + and getattr(self.draft_model_config.hf_config, "model_type", None) + == "step3p5_mtp" + ) + def use_eagle(self) -> bool: return self.method in ("eagle", "eagle3", "mtp", "dflash") diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4dacf4980ef..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): @@ -121,7 +127,13 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool: def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool: - """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed.""" + """Enable if TP > 1, PP == 1, Hopper/Blackwell, and flashinfer installed. + + Gated off for PP > 1: the fused op's GPU-side peer-signal spin-wait + assumes byte-identical kernel launches across TP peers, but concurrent + independent warmup of multiple TP subgroups lets ranks pick divergent + FlashInfer launch configs and deadlock. + """ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer @@ -134,6 +146,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool: return ( cfg.parallel_config.tensor_parallel_size > 1 + and cfg.parallel_config.pipeline_parallel_size == 1 and current_platform.is_cuda() and has_flashinfer() and ( @@ -480,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/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..a17f0b5f5ff 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 @@ -233,7 +233,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 +262,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 +273,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 c7d10a3e393..6ee827fa17e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -29,7 +29,9 @@ from vllm.v1.kv_offload.base import ( OffloadingManager, OffloadingSpec, OffloadKey, + OffloadPolicy, ReqContext, + RequestOffloadingContext, get_offload_block_hash, make_offload_key, ) @@ -92,6 +94,7 @@ class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] block_size_factor: int num_workers: int + offload_prompt_only: bool @classmethod def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig": @@ -154,6 +157,7 @@ class SchedulerOffloadConfig(NamedTuple): for idx, gpu_block_size in enumerate(spec.gpu_block_size) ), block_size_factor=spec.block_size_factor, + offload_prompt_only=spec.offload_prompt_only, ) @@ -172,8 +176,9 @@ class RequestGroupState: class RequestOffloadState: config: SchedulerOffloadConfig req: Request + req_context: ReqContext + offloading_context: RequestOffloadingContext group_states: tuple[RequestGroupState, ...] = field(init=False) - req_context: ReqContext = field(init=False) # upper bound on tokens to offload for this request; None means no cap max_offload_tokens: int | None = None # number of hits in the GPU cache @@ -186,10 +191,6 @@ class RequestOffloadState: self.group_states = tuple( RequestGroupState() for _ in self.config.kv_group_configs ) - self.req_context = ReqContext( - req_id=self.req.request_id, - kv_transfer_params=self.req.kv_transfer_params, - ) params = self.req.kv_transfer_params # NOTE: This field is experimental and subject to change in the future. @@ -248,6 +249,13 @@ class RequestOffloadState: ) +def _create_req_context(req: Request) -> ReqContext: + return ReqContext( + req_id=req.request_id, + kv_transfer_params=req.kv_transfer_params, + ) + + class OffloadingConnectorScheduler: """Implementation of Scheduler side methods""" @@ -278,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 = ( @@ -512,6 +522,18 @@ class OffloadingConnectorScheduler: return num_hit_tokens + def on_new_request(self, request: Request) -> None: + """Called when a new request is added to the scheduler.""" + req_context = _create_req_context(request) + offloading_context = self.manager.on_new_request(req_context) + req_status = RequestOffloadState( + config=self.config, + req=request, + req_context=req_context, + offloading_context=offloading_context, + ) + self._req_status[request.request_id] = req_status + def get_num_new_matched_tokens( self, request: Request, num_computed_tokens: int ) -> tuple[int | None, bool]: @@ -534,24 +556,15 @@ class OffloadingConnectorScheduler: - `True` if tokens will be loaded asynchronously (between scheduler steps). """ - is_new_request = False - if req_status := self._req_status.get(request.request_id): - # make sure block IDs are cleared - for group_state in req_status.group_states: - group_state.block_ids.clear() - else: - is_new_request = True - req_status = RequestOffloadState(config=self.config, req=request) - self._req_status[request.request_id] = req_status + req_status = self._req_status[request.request_id] + for group_state in req_status.group_states: + group_state.block_ids.clear() req_status.update_offload_keys() req_status.num_locally_computed_tokens = num_computed_tokens num_hit_tokens = self._lookup(req_status) - if is_new_request: - req_status.update_num_hit_blocks( - num_computed_tokens + (num_hit_tokens or 0) - ) + req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) @@ -568,9 +581,6 @@ class OffloadingConnectorScheduler: num_locally_computed_tokens = req_status.num_locally_computed_tokens num_cached_tokens = num_locally_computed_tokens + num_external_tokens - params = req_status.req_context.kv_transfer_params - do_remote_decode = params is not None and params.get("do_remote_decode") - keys_to_load: list[OffloadKey] = [] dst_block_ids: list[int] = [] # per group @@ -581,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 @@ -624,23 +638,12 @@ class OffloadingConnectorScheduler: group_sizes.append(num_pending_gpu_blocks) block_indices.append(num_locally_computed_gpu_blocks) - if not do_remote_decode: - # For P/D prefill requests (do_remote_decode=True), we do - # NOT skip saving the hit prefix, as we need to stream the - # entire KV cache so a remote decode node can consume it. + # Skip prefix-hit blocks for block-level policy; for + # request-level, next_stored_block_idx stays at 0 so all + # blocks (including hits) are offloaded. + 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 @@ -664,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 @@ -704,6 +738,15 @@ class OffloadingConnectorScheduler: if max_offload_tokens is not None: num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens) + # Skip decode-phase blocks: clamp to the prompt length so only + # prefill (prompt) blocks become eligible for store. next_stored_idx + # never advances past this boundary, so decode blocks are never + # queued in this or any later step. + if self.config.offload_prompt_only: + num_offloadable_tokens = min( + num_offloadable_tokens, req.num_prompt_tokens + ) + # Filter out blocks skipped due to sliding window attention / SSM # or unreachable by the load path's alignment constraints. new_offload_keys: list[OffloadKey] = [] @@ -718,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 @@ -797,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 @@ -858,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: @@ -866,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. @@ -881,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): @@ -950,6 +1007,12 @@ class OffloadingConnectorScheduler: # TODO(orozery): possibly kickoff offload for last block # which may have been deferred due to async scheduling req_status = self._req_status.get(request.request_id) + + req_context = ( + req_status.req_context if req_status else _create_req_context(request) + ) + self.manager.on_request_finished(req_context) + if req_status is None: return False, None if not req_status.transfer_jobs: @@ -990,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/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 6c75bda0c4c..20888c71f84 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -124,6 +124,10 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): return self.connector_worker.build_connector_worker_meta() return None + def on_new_request(self, request: "Request") -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int | None, bool]: 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/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 6e99adde1ca..eda209c3f6b 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -4,8 +4,8 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterator -from dataclasses import dataclass, field -from typing import Any, Generic, TypeVar +from dataclasses import KW_ONLY, dataclass, field +from typing import Any, Generic, Literal, TypeVar import torch @@ -28,7 +28,44 @@ class WeightTransferInitInfo(ABC): # noqa: B024 class WeightTransferUpdateInfo(ABC): # noqa: B024 """Base class for backend-specific weight update info.""" - pass + _: KW_ONLY + update_kind: Literal["dense", "sparse_flat"] = "dense" + """Weight update format.""" + num_updates_list: list[int] | None = None + """Number of sparse entries to receive for each parameter in ``names``.""" + + def __post_init__(self) -> None: + if self.update_kind not in ("dense", "sparse_flat"): + raise ValueError(f"Unsupported update_kind: {self.update_kind}") + if self.update_kind == "dense": + if self.num_updates_list is not None: + raise ValueError( + "Sparse metadata is only supported for `update_kind='sparse_flat'`" + ) + return + + if self.num_updates_list is None: + raise ValueError("`num_updates_list` is required for sparse updates") + if len(self.num_updates_list) == 0: + raise ValueError("`num_updates_list` cannot be empty for sparse updates") + if any(num_updates < 0 for num_updates in self.num_updates_list): + raise ValueError("Sparse `num_updates_list` entries must be non-negative") + + names = getattr(self, "names", None) + if names is not None and len(self.num_updates_list) != len(names): + raise ValueError( + f"`num_updates_list` should be of the same size as `names`: " + f"got {len(self.num_updates_list)} and {len(names)}" + ) + + +@dataclass +class SparseWeightPatch: + """A sparse in-place patch for one existing parameter.""" + + name: str + indices: torch.Tensor + values: torch.Tensor # API-level request classes (accept dicts for backend-agnostic serialization) @@ -150,6 +187,16 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): """ raise NotImplementedError + def receive_sparse_weights( + self, + update_info: TUpdateInfo, + apply_patches: Callable[[list[SparseWeightPatch]], None], + ) -> None: + """Receive sparse weight patches from the trainer.""" + raise NotImplementedError( + f"{self.__class__.__name__} does not support sparse weight updates" + ) + @abstractmethod def shutdown(self) -> None: """ @@ -184,3 +231,11 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]): >>> engine.trainer_send_weights(param_iter, trainer_args) """ raise NotImplementedError + + @staticmethod + def trainer_send_sparse_weights( + _iterator: Iterator[SparseWeightPatch], + _trainer_args: dict[str, Any] | Any, + ) -> None: + """Send sparse weight patches from trainer to inference workers.""" + raise NotImplementedError("Sparse weight updates are not supported") diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index b138c7dd937..a77aab751ff 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -74,10 +74,12 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): names: list[str] dtype_names: list[str] shapes: list[list[int]] - ipc_handles: list[dict[str, tuple]] | dict[str, tuple] + ipc_handles: list[dict[str, tuple]] | dict[str, tuple] | None = None """IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args. For non-packed mode: list of per-parameter handle dicts. For packed mode: single handle dict for the packed buffer.""" + ipc_handles_pickled: str | None = None + """Base64-encoded pickled IPC handles, used for HTTP transport.""" tensor_sizes: list[int] | None = None """Per-parameter sizes in bytes within the packed buffer. Required when packed=True, unused otherwise.""" @@ -85,6 +87,29 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo): """Whether this update uses packed tensor format.""" def __post_init__(self): + super().__post_init__() + if self.update_kind != "dense": + raise NotImplementedError("IPC weight transfer only supports dense updates") + + if self.ipc_handles_pickled is not None: + if self.ipc_handles is not None: + raise ValueError( + "Cannot specify both `ipc_handles` and `ipc_handles_pickled`" + ) + + if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION: + raise ValueError( + "Refusing to deserialize `ipc_handles_pickled` without " + "VLLM_ALLOW_INSECURE_SERIALIZATION=1" + ) + + self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled)) + self.ipc_handles_pickled = None + + if self.ipc_handles is None: + raise ValueError( + "Either `ipc_handles` or `ipc_handles_pickled` must be provided" + ) num_params = len(self.names) if len(self.dtype_names) != num_params: raise ValueError( @@ -153,8 +178,9 @@ class IPCWeightTransferEngine( Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the payload is deserialized via ``pickle.loads``. """ - if "ipc_handles_pickled" in update_dict: - if "ipc_handles" in update_dict: + pickled = update_dict.pop("ipc_handles_pickled", None) + if pickled is not None: + if update_dict.get("ipc_handles") is not None: raise ValueError( "Cannot specify both `ipc_handles` and `ipc_handles_pickled`" ) @@ -165,7 +191,6 @@ class IPCWeightTransferEngine( "VLLM_ALLOW_INSECURE_SERIALIZATION=1" ) - pickled = update_dict.pop("ipc_handles_pickled") update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled)) return super().parse_update_info(update_dict) diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 3b04a5f65ba..674f5b524da 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( + SparseWeightPatch, WeightTransferEngine, WeightTransferInitInfo, WeightTransferUpdateInfo, @@ -81,6 +82,7 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): def __post_init__(self): """Validate that all lists have the same length.""" + super().__post_init__() num_params = len(self.names) if len(self.dtype_names) != num_params: raise ValueError( @@ -92,6 +94,13 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): f"`shapes` should be of the same size as `names`: " f"got {len(self.shapes)} and {len(self.names)}" ) + if self.update_kind == "dense": + return + + if self.packed: + raise ValueError( + "`update_kind='sparse_flat'` cannot be combined with `packed=True`" + ) class NCCLWeightTransferEngine( @@ -178,6 +187,11 @@ class NCCLWeightTransferEngine( "NCCL weight transfer not initialized. " "Call init_transfer_engine() first." ) + if update_info.update_kind != "dense": + raise ValueError( + "Sparse updates must use `receive_sparse_weights`, not " + "`receive_weights`" + ) if update_info.packed: # Build iterator of (name, (shape, dtype)) from update_info @@ -209,6 +223,42 @@ class NCCLWeightTransferEngine( load_weights([(name, weight)]) del weight + def receive_sparse_weights( + self, + update_info: NCCLWeightTransferUpdateInfo, + apply_patches: Callable[[list[SparseWeightPatch]], None], + ) -> None: + """Receive sparse flat-index patches from trainer via NCCL.""" + if self.model_update_group is None: + raise RuntimeError( + "NCCL weight transfer not initialized. " + "Call init_transfer_engine() first." + ) + if update_info.update_kind != "sparse_flat": + raise ValueError("Sparse receive path requires `update_kind='sparse_flat'`") + assert update_info.num_updates_list is not None + + for name, dtype_name, num_updates in zip( + update_info.names, + update_info.dtype_names, + update_info.num_updates_list, + ): + dtype = getattr(torch, dtype_name) + device = torch.accelerator.current_device_index() + indices = torch.empty(num_updates, dtype=torch.int32, device=device) + values = torch.empty(num_updates, dtype=dtype, device=device) + self.model_update_group.broadcast( + indices, src=0, stream=torch.cuda.current_stream() + ) + self.model_update_group.broadcast( + values, src=0, stream=torch.cuda.current_stream() + ) + apply_patches( + [SparseWeightPatch(name=name, indices=indices, values=values)] + ) + del indices + del values + def shutdown(self) -> None: if self.model_update_group is not None: # Clean up the communicator by removing the reference @@ -272,6 +322,27 @@ class NCCLWeightTransferEngine( stream=args.stream or torch.cuda.current_stream(), ) + @staticmethod + def trainer_send_sparse_weights( + iterator: Iterator[SparseWeightPatch], + trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs, + ) -> None: + """Broadcast sparse flat-index patches from trainer to vLLM workers.""" + if isinstance(trainer_args, dict): + args = NCCLTrainerSendWeightsArgs(**trainer_args) + else: + args = trainer_args + + if args.packed: + raise ValueError( + "Sparse NCCL updates cannot be combined with `packed=True`" + ) + + stream = args.stream or torch.cuda.current_stream() + for patch in iterator: + args.group.broadcast(patch.indices, src=args.src, stream=stream) + args.group.broadcast(patch.values, src=args.src, stream=stream) + @staticmethod def trainer_init( init_info: NCCLWeightTransferInitInfo | dict, 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/generate/api_router.py b/vllm/entrypoints/generate/api_router.py similarity index 95% rename from vllm/entrypoints/openai/generate/api_router.py rename to vllm/entrypoints/generate/api_router.py index 84a7fddeabe..713e2566bc5 100644 --- a/vllm/entrypoints/openai/generate/api_router.py +++ b/vllm/entrypoints/generate/api_router.py @@ -41,6 +41,10 @@ def register_generate_api_routers(app: FastAPI): register_anthropic_api_router(app) + from .generative_scoring.api_router import register_generative_scoring_api_router + + register_generative_scoring_api_router(app) + async def init_generate_state( engine_client: "EngineClient", @@ -185,3 +189,11 @@ async def init_generate_state( if "generate" in supported_tasks else None ) + + from .generative_scoring.serving import ServingGenerativeScoring + + state.serving_generative_scoring = ServingGenerativeScoring( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + ) diff --git a/vllm/entrypoints/openai/generate/factories.py b/vllm/entrypoints/generate/factories.py similarity index 100% rename from vllm/entrypoints/openai/generate/factories.py rename to vllm/entrypoints/generate/factories.py diff --git a/vllm/entrypoints/openai/generative_scoring/__init__.py b/vllm/entrypoints/generate/generative_scoring/__init__.py similarity index 100% rename from vllm/entrypoints/openai/generative_scoring/__init__.py rename to vllm/entrypoints/generate/generative_scoring/__init__.py diff --git a/vllm/entrypoints/openai/generative_scoring/api_router.py b/vllm/entrypoints/generate/generative_scoring/api_router.py similarity index 65% rename from vllm/entrypoints/openai/generative_scoring/api_router.py rename to vllm/entrypoints/generate/generative_scoring/api_router.py index ed0a81d149c..e6918b7f03b 100644 --- a/vllm/entrypoints/openai/generative_scoring/api_router.py +++ b/vllm/entrypoints/generate/generative_scoring/api_router.py @@ -1,34 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from http import HTTPStatus -from typing import TYPE_CHECKING from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.responses import JSONResponse -from vllm.entrypoints.openai.engine.protocol import ErrorResponse -from vllm.entrypoints.openai.generative_scoring.serving import ( +from vllm.entrypoints.generate.generative_scoring.serving import ( GenerativeScoringResponse, - OpenAIServingGenerativeScoring, + ServingGenerativeScoring, ) +from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.utils import load_aware_call, with_cancellation from vllm.logger import init_logger -if TYPE_CHECKING: - from argparse import Namespace - - from starlette.datastructures import State - - from vllm.engine.protocol import EngineClient - from vllm.entrypoints.logger import RequestLogger - router = APIRouter() logger = init_logger(__name__) -def generative_scoring(request: Request) -> OpenAIServingGenerativeScoring | None: +def generative_scoring(request: Request) -> ServingGenerativeScoring | None: return request.app.state.serving_generative_scoring @@ -51,7 +42,7 @@ async def create_generative_scoring(raw_request: Request): raw_body = await raw_request.json() - from vllm.entrypoints.openai.generative_scoring.serving import ( + from vllm.entrypoints.generate.generative_scoring.serving import ( GenerativeScoringRequest, ) @@ -68,20 +59,3 @@ async def create_generative_scoring(raw_request: Request): def register_generative_scoring_api_router(app: FastAPI): app.include_router(router) - - -async def init_generative_scoring_state( - engine_client: "EngineClient", - state: "State", - args: "Namespace", - request_logger: "RequestLogger | None", -): - from vllm.entrypoints.openai.generative_scoring.serving import ( - OpenAIServingGenerativeScoring, - ) - - state.serving_generative_scoring = OpenAIServingGenerativeScoring( - engine_client, - state.openai_serving_models, - request_logger=request_logger, - ) diff --git a/vllm/entrypoints/openai/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py similarity index 99% rename from vllm/entrypoints/openai/generative_scoring/serving.py rename to vllm/entrypoints/generate/generative_scoring/serving.py index fd8f89cadad..0592d0b29af 100644 --- a/vllm/entrypoints/openai/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -142,7 +142,7 @@ class GenerativeScoringResponse(OpenAIBaseModel): # ============================================================================ -class OpenAIServingGenerativeScoring(OpenAIServing): +class ServingGenerativeScoring(OpenAIServing): """Serving class for generative scoring computation. This class handles computing the probability of specified token IDs diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index c8e9b8c08f0..802d7a6d796 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -873,14 +873,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): ) def start_weight_update(self, is_checkpoint_format: bool = True) -> None: - """ - Start a new weight update. - - Args: - is_checkpoint_format: Whether incoming weights are in checkpoint - format (need layerwise processing) or kernel format (direct - copy). - """ + """Start a new weight update.""" self.llm_engine.collective_rpc( "start_weight_update", kwargs={"is_checkpoint_format": is_checkpoint_format}, @@ -902,9 +895,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): ) def finish_weight_update(self) -> None: - """ - Finish the current weight update. - """ + """Finish the current weight update.""" self.llm_engine.collective_rpc("finish_weight_update") def __repr__(self) -> str: diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 461128ed905..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,8 +193,13 @@ 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.openai.generate.api_router import ( + from vllm.entrypoints.generate.api_router import ( register_generate_api_routers, ) @@ -208,24 +211,12 @@ 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, ) elastic_ep_attach_router(app) - from vllm.entrypoints.openai.generative_scoring.api_router import ( - register_generative_scoring_api_router, - ) - - register_generative_scoring_api_router(app) - if "generate" in supported_tasks or "render" in supported_tasks: from vllm.entrypoints.serve.render.api_router import ( attach_router as attach_render_router, @@ -402,18 +393,12 @@ async def init_app_state( ) if "generate" in supported_tasks: - from vllm.entrypoints.openai.generate.api_router import init_generate_state + from vllm.entrypoints.generate.api_router import init_generate_state await init_generate_state( engine_client, state, args, request_logger, supported_tasks ) - from vllm.entrypoints.openai.generative_scoring.api_router import ( - init_generative_scoring_state, - ) - - await init_generative_scoring_state(engine_client, state, args, request_logger) - if "transcription" in supported_tasks or "realtime" in supported_tasks: from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 73ecb3f35a1..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") @@ -740,20 +745,21 @@ class ChatCompletionRequest(OpenAIBaseModel): ) # you can only use one kind of constraints for structured outputs if count > 1: - raise ValueError( + raise VLLMValidationError( "You can only use one kind of constraints for structured " - "outputs ('json', 'regex' or 'choice')." + "outputs ('json', 'regex' or 'choice').", ) # you can only either use structured outputs or tools, not both - if count > 1 and data.get("tool_choice", "none") not in ( + if count > 0 and data.get("tool_choice", "none") not in ( "none", "auto", "required", ): - raise ValueError( + raise VLLMValidationError( "You can only either use constraints for structured outputs " - "or tools, not both." + "or tools, not both.", ) + validate_structured_outputs_structural_tag(structured_outputs_kwargs) return data @model_validator(mode="before") @@ -784,17 +790,21 @@ class ChatCompletionRequest(OpenAIBaseModel): if "tool_choice" in data and data["tool_choice"] is not None: # ensure that if "tool choice" is specified, tools are present if "tools" not in data or data["tools"] is None: - raise ValueError("When using `tool_choice`, `tools` must be set.") + raise VLLMValidationError( + "When using `tool_choice`, `tools` must be set.", + parameter="tool_choice", + ) # make sure that tool choice is either a named tool # OR that it's set to "auto" or "required" if data["tool_choice"] not in ["auto", "required"] and not isinstance( data["tool_choice"], dict ): - raise ValueError( + raise VLLMValidationError( f"Invalid value for `tool_choice`: {data['tool_choice']}! " 'Only named tools, "none", "auto" or "required" ' - "are supported." + "are supported.", + parameter="tool_choice", ) # ensure that if "tool_choice" is specified as an object, @@ -807,29 +817,33 @@ class ChatCompletionRequest(OpenAIBaseModel): valid_tool = False function = data["tool_choice"].get("function") if not isinstance(function, dict): - raise ValueError( + raise VLLMValidationError( f"Invalid value for `function`: `{function}` in " - f"`tool_choice`! {correct_usage_message}" + f"`tool_choice`! {correct_usage_message}", + parameter="tool_choice.function", ) if "name" not in function: - raise ValueError( + raise VLLMValidationError( f"Expected field `name` in `function` in " - f"`tool_choice`! {correct_usage_message}" + f"`tool_choice`! {correct_usage_message}", + parameter="tool_choice.function.name", ) function_name = function["name"] if not isinstance(function_name, str) or len(function_name) == 0: - raise ValueError( + raise VLLMValidationError( f"Invalid `name` in `function`: `{function_name}`" - f" in `tool_choice`! {correct_usage_message}" + f" in `tool_choice`! {correct_usage_message}", + parameter="tool_choice.function.name", ) for tool in data["tools"]: if tool["function"]["name"] == function_name: valid_tool = True break if not valid_tool: - raise ValueError( + raise VLLMValidationError( "The tool specified in `tool_choice` does not match any" - " of the specified `tools`" + " of the specified `tools`", + parameter="tool_choice", ) return data @@ -837,9 +851,9 @@ class ChatCompletionRequest(OpenAIBaseModel): @classmethod def check_generation_prompt(cls, data): if data.get("continue_final_message") and data.get("add_generation_prompt"): - raise ValueError( + raise VLLMValidationError( "Cannot set both `continue_final_message` and " - "`add_generation_prompt` to True." + "`add_generation_prompt` to True.", ) return data @@ -849,8 +863,9 @@ class ChatCompletionRequest(OpenAIBaseModel): if data.get("cache_salt") is not None and ( not isinstance(data["cache_salt"], str) or not data["cache_salt"] ): - raise ValueError( - "Parameter 'cache_salt' must be a non-empty string if provided." + raise VLLMValidationError( + "Parameter 'cache_salt' must be a non-empty string if provided.", + parameter="cache_salt", ) return data @@ -970,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 8f6d76d7852..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, @@ -57,7 +54,6 @@ from vllm.entrypoints.openai.engine.serving import ( ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.parser.harmony_utils import ( - get_stop_tokens_for_assistant_actions, get_streamable_parser_for_assistant, parse_chat_output, ) @@ -66,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 @@ -158,13 +154,6 @@ class OpenAIServingChat(OpenAIServing): else getattr(mc, "override_generation_config", {}).get("max_new_tokens") ) self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss" - if self.use_harmony: - if "stop_token_ids" not in self.default_sampling_params: - self.default_sampling_params["stop_token_ids"] = [] - self.default_sampling_params["stop_token_ids"].extend( - get_stop_tokens_for_assistant_actions() - ) - self.tool_call_id_type = get_tool_call_id_type(self.model_config) # NOTE(woosuk): While OpenAI's chat completion API supports browsing @@ -368,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, @@ -389,7 +386,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, - reasoning_parser, + parser, ) def get_chat_request_role(self, request: ChatCompletionRequest) -> str: @@ -723,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 @@ -813,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: @@ -1012,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 @@ -1121,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 @@ -1543,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 cb793a41563..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") @@ -447,8 +453,9 @@ class CompletionRequest(OpenAIBaseModel): ) if prompt_is_empty and embeds_is_empty: - raise ValueError( - "Either prompt or prompt_embeds must be provided and non-empty." + raise VLLMValidationError( + "Either prompt or prompt_embeds must be provided and non-empty.", + parameter="prompt", ) return data @@ -459,8 +466,9 @@ class CompletionRequest(OpenAIBaseModel): if data.get("cache_salt") is not None and ( not isinstance(data["cache_salt"], str) or not data["cache_salt"] ): - raise ValueError( - "Parameter 'cache_salt' must be a non-empty string if provided." + raise VLLMValidationError( + "Parameter 'cache_salt' must be a non-empty string if provided.", + parameter="cache_salt", ) return data diff --git a/vllm/entrypoints/openai/dp_supervisor.py b/vllm/entrypoints/openai/dp_supervisor.py index 2dff91fa794..13444015ecc 100644 --- a/vllm/entrypoints/openai/dp_supervisor.py +++ b/vllm/entrypoints/openai/dp_supervisor.py @@ -55,9 +55,9 @@ def validate_multi_port_external_lb_args(args: argparse.Namespace) -> None: raise ValueError( "Error: --data-parallel-multi-port-external-lb does not support --uds" ) - if any((args.ssl_keyfile, args.ssl_certfile, args.ssl_ca_certs)): + if bool(args.ssl_keyfile) != bool(args.ssl_certfile): raise ValueError( - "Error: --data-parallel-multi-port-external-lb does not support HTTPS yet" + "Error: --ssl-keyfile and --ssl-certfile must be provided together" ) if args.api_server_count not in (None, 1): raise ValueError( @@ -151,7 +151,8 @@ def _child_base_url(args: argparse.Namespace, port: int) -> str: host = "127.0.0.1" elif host == "::": host = "::1" - return f"http://{host}:{port}" + scheme = "https" if args.ssl_keyfile and args.ssl_certfile else "http" + return f"{scheme}://{host}:{port}" def _join_processes_with_timeout(processes: list[BaseProcess], timeout: float) -> None: @@ -178,7 +179,15 @@ async def _probe_endpoint( """ for iteration in range(conn_err_failure_threshold): try: - async with session.get(_child_base_url(args, port) + path) as response: + probe_ssl = None + if args.ssl_keyfile and args.ssl_certfile: + # Probes target node-local child servers over loopback, so skip + # certificate verification to avoid SAN/hostname mismatches for + # localhost/127.0.0.1 deployments. + probe_ssl = False + async with session.get( + _child_base_url(args, port) + path, ssl=probe_ssl + ) as response: # vLLM returns 503 on EngineDeadError, so we should return # immediately if vLLM responds with a non-200 status code. return response.status == HTTPStatus.OK @@ -272,6 +281,11 @@ class DPSupervisor: host=host, port=self.supervisor_port, log_level=self.args.uvicorn_log_level, + ssl_keyfile=self.args.ssl_keyfile, + ssl_certfile=self.args.ssl_certfile, + ssl_ca_certs=self.args.ssl_ca_certs, + ssl_cert_reqs=self.args.ssl_cert_reqs, + ssl_ciphers=self.args.ssl_ciphers, ) supervisor_server = uvicorn.Server(config) supervisor_server_task = asyncio.create_task( 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/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 7dc3704cea9..e76fa38d3c3 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -365,10 +365,6 @@ def render_for_completion(messages: list[Message]) -> list[int]: return token_ids -def get_stop_tokens_for_assistant_actions() -> list[int]: - return get_encoding().stop_tokens_for_assistant_actions() - - def get_streamable_parser_for_assistant() -> StreamableParser: return StreamableParser(get_encoding(), role=Role.ASSISTANT) 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/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 370ed9a825d..30a92066365 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -372,8 +372,6 @@ class ResponsesRequest(OpenAIBaseModel): if (frequency_penalty := self.frequency_penalty) is None: frequency_penalty = default_sampling_params.get("frequency_penalty", 0.0) - stop_token_ids = default_sampling_params.get("stop_token_ids") - # Structured output structured_outputs = self.structured_outputs @@ -409,7 +407,6 @@ class ResponsesRequest(OpenAIBaseModel): top_k=top_k, max_tokens=max_tokens, logprobs=self.top_logprobs if self.is_include_output_logprobs() else None, - stop_token_ids=stop_token_ids, stop=stop, frequency_penalty=frequency_penalty, presence_penalty=presence_penalty, diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 3c80e3d015d..eee02707a97 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -46,7 +46,6 @@ from vllm.entrypoints.openai.engine.serving import ( from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.parser.harmony_utils import ( get_developer_message, - get_stop_tokens_for_assistant_actions, get_system_message, get_user_message, has_custom_tools, @@ -222,13 +221,6 @@ class OpenAIServingResponses(OpenAIServing): "For gpt-oss, we ignore --enable-auto-tool-choice " "and always enable tool use." ) - # OpenAI models have two EOS-like tokens: <|return|> and <|call|>. - # We need to add them to the stop token ids. - if "stop_token_ids" not in self.default_sampling_params: - self.default_sampling_params["stop_token_ids"] = [] - self.default_sampling_params["stop_token_ids"].extend( - get_stop_tokens_for_assistant_actions() - ) self.tool_call_id_type = get_tool_call_id_type(self.model_config) @@ -468,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() @@ -716,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, ) @@ -1041,7 +1030,10 @@ class OpenAIServingResponses(OpenAIServing): # Use parser to extract and create response output items if self.parser: - parser = self.parser(tokenizer, request.tools) + chat_template_kwargs = self._effective_chat_template_kwargs(request) + parser = self.parser( + tokenizer, request.tools, chat_template_kwargs=chat_template_kwargs + ) return parser.extract_response_outputs( model_output=final_output.text, model_output_token_ids=final_output.token_ids, @@ -1378,7 +1370,15 @@ class OpenAIServingResponses(OpenAIServing): ], ) -> AsyncGenerator[StreamingResponsesResponse, None]: processor = SimpleStreamingEventProcessor() - parser = self.parser(tokenizer, request.tools) if self.parser else None + parser = ( + self.parser( + tokenizer, + request.tools, + chat_template_kwargs=self._effective_chat_template_kwargs(request), + ) + if self.parser + else None + ) def _get_logprobs( output: CompletionOutput, @@ -1408,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/sagemaker/api_router.py b/vllm/entrypoints/sagemaker/api_router.py index b3b11cd07b4..00dd7db2818 100644 --- a/vllm/entrypoints/sagemaker/api_router.py +++ b/vllm/entrypoints/sagemaker/api_router.py @@ -11,9 +11,9 @@ from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response from vllm.config import ModelConfig +from vllm.entrypoints.generate.factories import get_generate_invocation_types from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.engine.serving import OpenAIServing -from vllm.entrypoints.openai.generate.factories import get_generate_invocation_types from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.base.serving import PoolingServingBase from vllm.entrypoints.pooling.factories import get_pooling_invocation_types 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 3a3934f3cdf..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 @@ -278,6 +277,8 @@ if TYPE_CHECKING: VLLM_XPU_ENABLE_XPU_GRAPH: bool = False VLLM_XPU_USE_SAMPLER_KERNEL: bool = True VLLM_LORA_ENABLE_DUAL_STREAM: bool = False + VLLM_GPU_NIC_PCIE_MAPPING: str = "" + VLLM_NIC_SELECTION_VARS: str = "" def get_default_cache_root(): @@ -1013,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") @@ -1973,6 +1971,13 @@ environment_variables: dict[str, Callable[[], Any]] = { # If set to 1, use Python spinloop extension to poll in a more efficient # way when using the mp backend. "VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))), + # Comma-separated GPU_BDF=NIC_BDF pairs for RDMA NIC selection. + # Must be set together with VLLM_NIC_SELECTION_VARS. + "VLLM_GPU_NIC_PCIE_MAPPING": lambda: os.getenv("VLLM_GPU_NIC_PCIE_MAPPING", ""), + # Comma-separated list of env vars to set from the GPU-NIC mapping. + # Each entry is VAR_NAME or VAR_NAME: (suffix appended to + # RDMA device name). Must be set together with VLLM_GPU_NIC_PCIE_MAPPING. + "VLLM_NIC_SELECTION_VARS": lambda: os.getenv("VLLM_NIC_SELECTION_VARS", ""), } diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 7b68c1f952a..d5c9a1a6ff8 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -4,7 +4,6 @@ import os from typing import TYPE_CHECKING -import huggingface_hub import regex as re from huggingface_hub.utils import HfHubHTTPError, HFValidationError from torch import nn @@ -37,6 +36,7 @@ from vllm.lora.layers import ( from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.utils import get_moe_expert_mapping, get_packed_modules_mapping +from vllm.transformers_utils.repo_utils import hf_api if TYPE_CHECKING: from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -340,7 +340,9 @@ def get_adapter_absolute_path(lora_path: str) -> str: error_log = "Error downloading the ModelScope model" else: # Otherwise, we assume the path is a Hugging Face Hub repo. - download_fn = lambda: huggingface_hub.snapshot_download(repo_id=lora_path) + download_fn = lambda: hf_api().snapshot_download( + repo_id=lora_path, + ) download_exceptions = (HfHubHTTPError, HFValidationError) error_log = "Error downloading the HuggingFace model" diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index cfeb40fecf6..39d2e86d3c3 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -51,6 +51,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import ( from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( MarlinLinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( + RDNA3W4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( TritonW4A16LinearKernel, ) @@ -58,6 +61,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import ( XPUW4A8IntLinearKernel, XPUwNa16LinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.zentorch import ( + ZentorchWNA16LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4 import ( MxFp4LinearKernel, MxFp4LinearLayerConfig, @@ -155,8 +161,12 @@ 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 ( + ZentorchInt8ScaledMMLinearKernel, +) from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms import PlatformEnum, current_platform @@ -254,7 +264,7 @@ def _filter_kernels_by_backend( # in priority/performance order (when available) _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = { - PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel], + PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel], PlatformEnum.CUDA: [ CutlassInt8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, @@ -308,6 +318,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ CPUFp8BlockScaledMMKernel, ], PlatformEnum.XPU: [ + XPUFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], } @@ -339,6 +350,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { TritonW4A16LinearKernel, ], PlatformEnum.ROCM: [ + RDNA3W4A16LinearKernel, TritonW4A16LinearKernel, ConchLinearKernel, ExllamaLinearKernel, @@ -349,6 +361,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { ], PlatformEnum.CPU: [ Dynamic4bitLinearKernel, + ZentorchWNA16LinearKernel, CPUWNA16LinearKernel, ], } @@ -1019,6 +1032,8 @@ __all__ = [ "RowWiseTorchFP8ScaledMMLinearKernel", "ROCmFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", + "ZentorchInt8ScaledMMLinearKernel", + "ZentorchWNA16LinearKernel", "MPLinearKernel", "MPLinearLayerConfig", "AllSparkLinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py index 4d659b36042..c0b8c35bbd5 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py @@ -29,6 +29,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearKernel, MPLinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( + RDNA3W4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( TritonW4A16LinearKernel, ) @@ -36,6 +39,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import ( XPUW4A8IntLinearKernel, XPUwNa16LinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.zentorch import ( + ZentorchWNA16LinearKernel, +) __all__ = [ "MPLinearKernel", @@ -48,7 +54,9 @@ __all__ = [ "ExllamaLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", + "RDNA3W4A16LinearKernel", "TritonW4A16LinearKernel", "XPUW4A8IntLinearKernel", "XPUwNa16LinearKernel", + "ZentorchWNA16LinearKernel", ] diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py new file mode 100644 index 00000000000..268728f4bf6 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16. + +Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16 +support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu`` +and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``. + +Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls +through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300). +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + pack_quantized_values_into_int32, +) +from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types + +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig + + +class RDNA3W4A16LinearKernel(MPLinearKernel): + SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8] + + @classmethod + def get_min_capability(cls) -> int: + # ROCm gates via on_gfx1100() in can_implement. + return 60 + + @classmethod + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "RDNA3 W4A16 kernel is ROCm-only" + + from vllm.platforms.rocm import on_gfx1100 + + if not on_gfx1100(): + return False, "RDNA3 W4A16 kernel requires gfx1100" + + # The HIP op is registered by the C++ extension; if a user is running + # against a vLLM build that doesn't include it (e.g. partial rebuild), + # fall through gracefully to the next kernel in the registry. + if not ( + hasattr(torch.ops, "_rocm_C") + and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3") + ): + return ( + False, + "torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension", + ) + + if c.act_type not in (torch.float16, torch.bfloat16): + return False, "RDNA3 W4A16 kernel only supports fp16 and bf16" + + if c.weight_type not in cls.SUPPORTED_QUANT_TYPES: + return ( + False, + f"Quant type ({c.weight_type}) not supported by " + f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}", + ) + + if c.group_size <= 0: + return ( + False, + "RDNA3 W4A16 kernel does not support channelwise quantization", + ) + + if c.full_weight_shape[0] % c.group_size != 0: + return ( + False, + f"Group size ({c.group_size}) does not evenly divide K " + f"({c.full_weight_shape[0]})", + ) + + # Output features must be a multiple of the pack factor (8 nibbles per + # int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly + # against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns. + if c.partition_weight_shape[1] % 8 != 0: + return ( + False, + "Output features must be a multiple of 8 for the RDNA3 " + "W4A16 kernel (qzeros packing)", + ) + + if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]: + return ( + False, + "Act-order with TP-partitioned input features is not " + "supported by the RDNA3 W4A16 kernel", + ) + + return True, None + + # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) ----- + + def process_weights_after_loading(self, layer: torch.nn.Module): + c = self.config + device = getattr(layer, self.w_q_name).device + + # Synthesize zero points if the checkpoint doesn't carry them. + if not c.zero_points: + self.w_zp_name = "qzeros" + groups = c.partition_weight_shape[0] // c.group_size + out_features = c.partition_weight_shape[1] + + if c.weight_type.has_bias(): + # GPTQv1 quirk: the kernel adds 1 to the stored zero, so we + # encode (bias - 1) here. See exllama.py for the link to the + # documentation of this checkpoint-format wart. + zeros = torch.full( + (groups, out_features), + c.weight_type.bias - 1, + dtype=torch.int32, + device=device, + ) + else: + raise NotImplementedError( + "RDNA3 W4A16 kernel: zero-bias 4-bit quant requires " + "explicit zero points (GPTQv1 +1 quirk)." + ) + zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1) + setattr( + layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False) + ) + + # Act-order: convert g_idx to the inverse permutation array exllama + # expects (kernel reads a[perm[k]] instead of using groups indirected + # by g_idx[k]). + if c.has_g_idx: + + def transform_w_g_idx(x): + return torch.argsort(x).to(torch.int) + + self._transform_param(layer, self.w_gidx_name, transform_w_g_idx) # type: ignore + else: + self.w_gidx_name = "g_idx" + empty_g_idx = torch.nn.Parameter( + torch.empty((0,), dtype=torch.int, device=device), + requires_grad=False, + ) + setattr(layer, self.w_gidx_name, empty_g_idx) + + def transform_w_q(x): + assert isinstance(x, BasevLLMParameter) + assert self.w_gidx_name is not None + g_idx = getattr(layer, self.w_gidx_name) + + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0) + x_cont = x.data.contiguous() + # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in + # the same shuffled int32 layout and uses the (qa & 0x000F000F) + # bit-trick on top. + ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits) + return x_cont + + def transform_w_s(x): + assert isinstance(x, BasevLLMParameter) + permute_param_layout_(x, input_dim=0, output_dim=1) + x.data = x.data.contiguous() + # Keep scales in the activation dtype (fp16 OR bf16) — the kernel + # branches on dtype internally. + return x.to(dtype=c.act_type) + + self._transform_param(layer, self.w_q_name, transform_w_q) + self._transform_param(layer, self.w_s_name, transform_w_s) + + # ----- Forward -------------------------------------------------------- + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + c = self.config + + x_2d = x.reshape(-1, x.shape[-1]) + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) + + w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer) + + assert w_zp is not None, "Zero points are required by RDNA3 W4A16" + assert w_g_idx is not None, "g_idx tensor (possibly empty) required" + + output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False) + + if bias is not None: + output.add_(bias) + return output.reshape(out_shape) 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/mixed_precision/zentorch.py b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py new file mode 100644 index 00000000000..c3e8b17cc9a --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs. + +Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed +``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector +falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``. +""" + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types + +from .cpu import CPUWNA16LinearKernel +from .MPLinearKernel import MPLinearLayerConfig + +logger = init_logger(__name__) + + +def _import_unpack_from_int32(): + """Import compressed-tensors' ``unpack_from_int32`` across versions.""" + try: + from compressed_tensors.compressors.pack_quantized.helpers import ( + unpack_from_int32, + ) + except ImportError: + from compressed_tensors.compressors.quantized_compressors.pack_quantized import ( # type: ignore[import-not-found] # noqa: E501 + unpack_from_int32, + ) + return unpack_from_int32 + + +class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): + """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``.""" + + @classmethod + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: + ok, reason = super().can_implement(c) + if not ok: + return ok, reason + + if not current_platform.is_zen_cpu(): + return False, "ZentorchWNA16 requires an AMD Zen CPU." + + if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]): + return ( + False, + "torch.ops.zentorch.{zentorch_woq_repack_weight, " + "zentorch_woq_linear} are not registered.", + ) + + if c.has_g_idx: + return False, "ZentorchWNA16 does not support activation re-ordering." + return True, None + + def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool: + """Eligibility predicate for the zentorch W4A16 GPTQ fast path. + + Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()`` + with ``layer`` untouched). + """ + if ( + self.w_gidx_name is not None + and getattr(layer, self.w_gidx_name, None) is not None + ) or (getattr(self.config, "has_g_idx", False)): + return False + + weight_packed = getattr(layer, self.w_q_name, None) + weight_scale = getattr(layer, self.w_s_name, None) + if weight_packed is None or weight_scale is None: + return False + + bits = self.config.weight_type.mantissa + pack_factor = torch.iinfo(weight_packed.dtype).bits // bits + # 4-bit -> 8 values per int32; + if pack_factor != 8: + return False + + # GPTQ-only. AWQ packs along the output dim instead. + in_dim = getattr(weight_packed, "input_dim", None) + pk_dim = getattr(weight_packed, "packed_dim", None) + if in_dim is None or pk_dim is None or in_dim != pk_dim: + return False + + is_ct_format = in_dim == pk_dim == 1 + if not is_ct_format: + return False + + if weight_packed.dim() != 2 or weight_scale.dim() != 2: + return False + + # 4-bit -> 8 values per int32; in_features must be divisible by num_groups. + in_features = weight_packed.shape[1] * 8 + num_groups = weight_scale.shape[1] + return num_groups > 0 and in_features % num_groups == 0 + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Repack CT GPTQ weights into the zentorch WOQ layout. + + Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading`` + via ``super()`` when the layer doesn't satisfy + ``_zentorch_woq_eligible``. + + On success, ``layer._zentorch_processed_weights`` is set to ``True`` + """ + if getattr(layer, "_zentorch_processed_weights", False): + return + + if not self._zentorch_woq_eligible(layer): + logger.info_once( + "[zen_cpu] ZentorchWNA16 fast path not eligible for this " + "layer (AWQ pack layout, g_idx, or non-int32 storage); " + "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)." + ) + super().process_weights_after_loading(layer) + return + + if (not self.config.zero_points) and (self.w_zp_name is not None): + setattr(layer, self.w_zp_name, None) + + if (not self.config.has_g_idx) and (self.w_gidx_name is not None): + setattr(layer, self.w_gidx_name, None) + + weight_q = getattr(layer, self.w_q_name) + weight_s = getattr(layer, self.w_s_name) + weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q + weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s + + bits = self.config.weight_type.mantissa + pack_factor = torch.iinfo(weight_packed.dtype).bits // bits + out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1] + in_features = weight_packed.shape[1] * pack_factor + original_shape = torch.Size([out_features, in_features]) + unpack_from_int32 = _import_unpack_from_int32() + repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default + + weight_unpacked = unpack_from_int32( + weight_packed, + bits, + original_shape, + packed_dim=weight_q.packed_dim, + ) + + zp_param = ( + getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None + ) + needs_unsigned_offset = self.config.weight_type == scalar_types.uint4 + + if needs_unsigned_offset: + weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15) + repacked = repack_op(weight_unpacked.to(torch.int8).contiguous()) + + if zp_param is None: + zp_tc = None + else: + zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param + zp = unpack_from_int32( + zp_tensor, + bits, + (out_features, num_groups), + packed_dim=zp_param.packed_dim, + ) + if needs_unsigned_offset: + zp = (zp.to(torch.int32) + 8).clamp(0, 15) + zp_tc = zp.to(torch.int8).t().contiguous() + + layer._zentorch_woq_packed = repacked.t() + layer._zentorch_woq_scale = weight_scale.t().contiguous() + layer._zentorch_woq_zero_point = zp_tc + + for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name): + if param_name is None: + continue + param = getattr(layer, param_name, None) + if param is None: + continue + if hasattr(param, "data"): + param.data = torch.empty(0) + else: + setattr(layer, param_name, torch.empty(0)) + + layer._zentorch_kind = "compressed_tensors_w4a16_gptq" + layer._zentorch_processed_weights = True + logger.info_once( + "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ " + "(weight_type=%s, has_zp=%s)", + self.config.weight_type, + zp_tc is not None, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if getattr(layer, "_zentorch_processed_weights", False): + return torch.ops.zentorch.zentorch_woq_linear.default( + x, + layer._zentorch_woq_packed, + layer._zentorch_woq_scale, + layer._zentorch_woq_zero_point, + bias, + ) + return super().apply_weights(layer, x, bias) + + +__all__ = ["ZentorchWNA16LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index f8f12f7b0cb..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,12 @@ 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, +) __all__ = [ "FP8ScaledMMLinearKernel", @@ -58,6 +64,8 @@ __all__ = [ "RowWiseTorchFP8ScaledMMLinearKernel", "ROCmFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", + "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", "CPUFp8BlockScaledMMKernel", + "XPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index 9e65edb851e..b52d2c5b101 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -312,7 +312,7 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): ) -> torch.Tensor: out_dtype = self.config.out_dtype if self.is_hopper: - return torch.ops.vllm.padded_cutlass( + return torch.ops.vllm.dynamic_padded_cutlass( A, B, As, @@ -320,14 +320,14 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): list(self.weight_group_shape), out_dtype, ) - else: - return ops.cutlass_scaled_mm( - A, - B.T, - out_dtype=out_dtype, - scale_a=As, - scale_b=Bs.T, - ) + + return ops.cutlass_scaled_mm( + A, + B.T, + out_dtype=out_dtype, + scale_a=As, + scale_b=Bs.T, + ) def cutlass_scaled_mm( @@ -397,8 +397,56 @@ def _padded_cutlass_fake( ) +def _dynamic_padded_cutlass( + qx: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + block_size: list[int], + output_dtype: torch.dtype, +) -> torch.Tensor: + def run_padded( + qx: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + return _padded_cutlass( + qx, weight, x_scale, weight_scale, block_size, output_dtype + ) + + def run_direct( + qx: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + return cutlass_scaled_mm( + qx, weight, x_scale, weight_scale, block_size, output_dtype + ) + + if torch.compiler.is_compiling(): + return torch.cond( + qx.shape[0] % 4 != 0, + run_padded, + run_direct, + (qx, weight, x_scale, weight_scale), + ) + + if qx.shape[0] % 4 != 0: + return run_padded(qx, weight, x_scale, weight_scale) + + return run_direct(qx, weight, x_scale, weight_scale) + + direct_register_custom_op( "padded_cutlass", _padded_cutlass, fake_impl=_padded_cutlass_fake, ) + +direct_register_custom_op( + "dynamic_padded_cutlass", + _dynamic_padded_cutlass, + fake_impl=_padded_cutlass_fake, +) 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/kernels/linear/scaled_mm/zentorch.py b/vllm/model_executor/kernels/linear/scaled_mm/zentorch.py new file mode 100644 index 00000000000..c434c9d465f --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/zentorch.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs. + +Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic +oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or +``can_implement`` rejects a layer, the selector falls through to the next +kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``. +""" + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op +from vllm.model_executor.layers.quantization.utils import replace_parameter +from vllm.platforms import current_platform + +from .ScaledMMLinearKernel import ( + Int8ScaledMMLinearKernel, + Int8ScaledMMLinearLayerConfig, +) + +logger = init_logger(__name__) + + +class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cpu(): + return False, "requires CPU." + if not current_platform.is_zen_cpu(): + return False, "requires AMD Zen CPU." + if not has_zentorch_op(["zentorch_dynamic_qlinear"]): + return ( + False, + "torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.", + ) + return True, None + + @classmethod + def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + if c.is_static_input_scheme: + return False, "requires dynamic activation quantization." + if not c.input_symmetric: + return False, "requires symmetric activation quantization." + if not c.is_channelwise: + return False, "requires per-channel weight quantization." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prepare weights for ``zentorch_dynamic_qlinear``. + + Keeps weight in [N, K] layout (int8, contiguous) and converts the + per-channel weight scale to bf16 with shape ``(N,)``. + """ + w_q_name, w_s_name, _, _, _ = self.layer_param_names + weight = getattr(layer, w_q_name) + n = weight.shape[0] + replace_parameter( + layer, + w_q_name, + torch.nn.Parameter(weight.data.contiguous(), requires_grad=False), + ) + + weight_scale = getattr(layer, w_s_name) + ws = weight_scale.data + if ws.dim() == 2 and ws.shape[-1] == 1: + ws = ws.squeeze(-1) + ws = ws.to(torch.bfloat16).contiguous() + assert ws.shape == (n,), ( + f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}" + ) + + replace_parameter( + layer, + w_s_name, + torch.nn.Parameter(ws, requires_grad=False), + ) + logger.info_once( + "[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)" + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + w_q_name, w_s_name, _, _, _ = self.layer_param_names + return torch.ops.zentorch.zentorch_dynamic_qlinear( + x, + getattr(layer, w_q_name), + getattr(layer, w_s_name), + bias, + zentorch_op_name="zentorch::zentorch_dynamic_qlinear", + ) diff --git a/vllm/model_executor/kernels/linear/zentorch_utils.py b/vllm/model_executor/kernels/linear/zentorch_utils.py new file mode 100644 index 00000000000..310aed579ef --- /dev/null +++ b/vllm/model_executor/kernels/linear/zentorch_utils.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gates zentorch CPU linear dispatch on platform/op availability.""" + +from __future__ import annotations + +import torch + +from vllm.platforms import current_platform + +__all__ = ["has_zentorch_op"] + + +def has_zentorch_op(op_names: list[str]) -> bool: + """Return ``True`` when running on Zen CPU with all named ops registered.""" + if not op_names: + raise ValueError("has_zentorch_op requires at least one op name") + if not current_platform.is_zen_cpu(): + return False + ns = getattr(torch.ops, "zentorch", None) + if ns is None: + return False + return all(hasattr(ns, op_name) for op_name in op_names) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index d76123bb762..e0007141d53 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -29,7 +29,7 @@ def _tilelang_hc_prenorm_gemm( n_thr: int = 512, n_splits: int = 1, ) -> None: - from vllm._tilelang_ops import ( + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( hc_prenorm_gemm_block_m_tilelang, hc_prenorm_gemm_tilelang, ) @@ -126,7 +126,7 @@ def mhc_pre_tilelang( comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 layer_input: shape (..., hidden_size), dtype torch.bfloat16 """ - from vllm._tilelang_ops import ( + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( compute_num_split, mhc_pre_big_fuse_tilelang, mhc_pre_big_fuse_with_norm_tilelang, @@ -306,7 +306,9 @@ def mhc_post_tilelang( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - from vllm._tilelang_ops import mhc_post_tilelang as _mhc_post_kernel + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_post_tilelang as _mhc_post_kernel, + ) out = torch.empty_like(residual) _mhc_post_kernel( @@ -353,7 +355,7 @@ def mhc_fused_post_pre_tilelang( layer_input_cur: shape (..., hidden_size) """ - from vllm._tilelang_ops import ( + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( compute_num_split, mhc_fused_tilelang, mhc_post_tilelang, @@ -608,21 +610,22 @@ def _mhc_post_tilelang_fake( return torch.empty_like(residual) -def _hc_head_fused_kernel_tilelang( +def hc_head_fused_kernel_tilelang( hs_flat: torch.Tensor, fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, - out: torch.Tensor, - hidden_size: int, rms_eps: float, hc_eps: float, - hc_mult: int, -) -> None: - """Fill pre-allocated `out` (T, H) in-place with the hc_head result.""" - if hs_flat.shape[0] == 0: - return - from vllm._tilelang_ops import hc_head_fuse_tilelang +) -> torch.Tensor: + """Apply the fused hc_head kernel and return the (T, H) bf16 result.""" + num_tokens, hc_mult, hidden_size = hs_flat.shape + out = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device + ) + if num_tokens == 0: + return out + from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang hc_head_fuse_tilelang( hs_flat, @@ -635,6 +638,21 @@ def _hc_head_fused_kernel_tilelang( hc_eps, hc_mult, ) + return out + + +def _hc_head_fused_kernel_tilelang_fake( + hs_flat: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, +) -> torch.Tensor: + num_tokens, _, hidden_size = hs_flat.shape + return torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device + ) direct_register_custom_op( @@ -659,6 +677,7 @@ direct_register_custom_op( direct_register_custom_op( op_name="hc_head_fused_kernel_tilelang", - op_func=_hc_head_fused_kernel_tilelang, - mutates_args=["out"], + op_func=hc_head_fused_kernel_tilelang, + mutates_args=[], + fake_impl=_hc_head_fused_kernel_tilelang_fake, ) diff --git a/vllm/_tilelang_ops.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py similarity index 100% rename from vllm/_tilelang_ops.py rename to vllm/model_executor/kernels/mhc/tilelang_kernels.py diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 71fd297a7ed..140e071c746 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -11,12 +11,12 @@ Sq as Q sequence length Skv as KV sequence length MLA has two possible ways of computing, a data-movement friendly approach and a -compute friendly approach, we generally want to use the compute friendly -approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1) -and the data-movement friendly approach for "decode" (i.e. the ratio -Sq / Skv is "large"). +compute friendly approach. We generally want to use the compute friendly +approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near +1) and the data-movement friendly approach for "decode" (i.e. the ratio +Sq / Skv is small). -NOTE what we deem small and large is currently determined by if its labelled +NOTE what we deem small and large is currently determined by if it is labelled prefill or decode by the scheduler, but this is something we should probably tune. @@ -96,7 +96,7 @@ NOTE: in the actual code, Runtime q_c = h_t @ W_DQ q_nope = (q_c @ W_UQ).view(-1, N, P) -ql_nope = einsum("snh,lnh->snl", q, W_UK) +ql_nope = einsum("snh,lnh->snl", q_nope, W_UK) q_pe = RoPE(q_c @ W_QR).view(Sq, N, R) new_kv_c = h_t @ W_DKV new_k_pe = RoPE(h_t @ W_KR) @@ -115,7 +115,7 @@ spda_o = scaled_dot_product_attention( ) o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV) -return o.view(-1, N * V) @ self.num_heads @ W_O +return o.view(-1, N * V) @ W_O ## Chunked Prefill diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index da9d5d3e234..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: @@ -869,19 +874,23 @@ def nvfp4_w4a16_moe_quant_config( def int4_w4a16_moe_quant_config( w1_scale: torch.Tensor, w2_scale: torch.Tensor, - w1_zp: torch.Tensor | None, - w2_zp: torch.Tensor | None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int4 weights. """ group_shape = GroupShape(*block_shape) if block_shape is not None else None return FusedMoEQuantConfig( - _a1=FusedMoEQuantDesc(shape=group_shape), - _a2=FusedMoEQuantDesc(shape=group_shape), - _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp), - _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp), + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), + _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias), + _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias), ) @@ -922,19 +931,21 @@ def fp8_w8a16_moe_quant_config( def int8_w8a16_moe_quant_config( w1_scale: torch.Tensor, w2_scale: torch.Tensor, - w1_zp: torch.Tensor | None, - w2_zp: torch.Tensor | None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int8 weights. """ group_shape = GroupShape(*block_shape) if block_shape is not None else None return FusedMoEQuantConfig( - _a1=FusedMoEQuantDesc(shape=group_shape), - _a2=FusedMoEQuantDesc(shape=group_shape), + _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale), + _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), ) @@ -965,47 +976,6 @@ def int4_w4afp8_moe_quant_config( ) -def awq_marlin_moe_quant_config( - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - w1_zp: torch.Tensor | None, - w2_zp: torch.Tensor | None, - weight_bits: int, - group_size: int, - w1_bias: torch.Tensor | None = None, - w2_bias: torch.Tensor | None = None, - a1_gscale: torch.Tensor | None = None, - a2_gscale: torch.Tensor | None = None, -) -> FusedMoEQuantConfig: - """ - Construct a quant config for awq marlin quantization. - - a1_gscale / a2_gscale are optional global scales applied to activation - quantization scales when Marlin runs with 8-bit activations. - """ - from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape - - w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) - - # Activations are NOT quantized for AWQ (fp16/bf16) - a_shape = w_shape # Same as weight shape for alignment - - # Determine weight dtype - if weight_bits == 4: - weight_dtype = "int4" - elif weight_bits == 8: - weight_dtype = torch.int8 - else: - raise ValueError(f"Unsupported weight_bits: {weight_bits}") - - return FusedMoEQuantConfig( - _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a1_gscale), - _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a2_gscale), - _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias), - _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias), - ) - - def biased_moe_quant_config( w1_bias: torch.Tensor | None, w2_bias: torch.Tensor | None, 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/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 3906a7e057c..cc2adc31fcd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -248,9 +248,6 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): ) -> bool: return True - def supports_expert_map(self) -> bool: - return False # Expert parallelism not yet supported - @property def expects_unquantized_inputs(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index 7bd383b9cda..c8611217a18 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -316,9 +316,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular): def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def supports_packed_ue8m0_act_scales(self) -> bool: """ DeepGemm supports packed ue8m0 activation scales format in devices == sm100 diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 54b264ef772..84740fc0570 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -100,9 +100,6 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): ) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def apply( self, hidden_states: torch.Tensor, @@ -256,9 +253,6 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): ) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def apply( self, hidden_states: torch.Tensor, 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 28a7d283b4b..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, ] @@ -378,7 +379,8 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular): topk_ids, activation, global_num_experts, - expert_map, + # the fp8 cutlass experts use their own expert map. + None, self.w1_scale, self.w2_scale, a1q_scale, @@ -418,9 +420,6 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base): or moe_parallel_config.use_fi_nvl_one_sided_kernels ) - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # topk weights and reduction are fused in moe_unpermute cuda kernel return TopKWeightAndReduceNoOP() @@ -460,9 +459,6 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.BatchedExperts - def supports_expert_map(self) -> bool: - return False - def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype: return self.out_dtype if self.out_dtype is not None else act_dtype @@ -724,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, ] @@ -741,9 +739,6 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() @@ -1038,9 +1033,6 @@ class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() @@ -1340,9 +1332,6 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular): def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return True - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # topk weights and reduction are fused in moe_unpermute cuda kernel return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index e3e15e31618..3b354dd3ef1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -164,9 +164,6 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): or moe_parallel_config.use_fi_nvl_one_sided_kernels ) - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() @@ -388,9 +385,6 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): or moe_parallel_config.use_fi_nvl_one_sided_kernels ) - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/fallback.py b/vllm/model_executor/layers/fused_moe/experts/fallback.py index 40741d52af5..639b2bf2668 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fallback.py +++ b/vllm/model_executor/layers/fused_moe/experts/fallback.py @@ -92,16 +92,6 @@ class FallbackExperts(mk.FusedMoEExpertsModular, ABC): moe_parallel_config ) and fallback_cls._supports_parallel_config(moe_parallel_config) - def supports_expert_map(self) -> bool: - assert ( - self.experts.supports_expert_map() - == self.fallback_experts.supports_expert_map() - ) - return ( - self.experts.supports_expert_map() - and self.fallback_experts.supports_expert_map() - ) - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: e_war = self.experts.finalize_weight_and_reduce_impl() fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl() 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/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py index 5eaaf46739f..253d1dae711 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py @@ -89,9 +89,6 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular): def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # Let PrepareAndFinalize::finalize() decide the impl. return TopKWeightAndReduceDelegate() diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 2310982792f..b512d51c135 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -98,9 +98,6 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): ) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index b891583e3ef..fd9446c2a22 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -207,9 +207,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 0e31331e726..1f5724ac39c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -555,9 +555,6 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular): "This method should not be called." ) - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # Let PrepareAndFinalize::finalize() decide the impl. return TopKWeightAndReduceDelegate() @@ -799,9 +796,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular): def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return True - def supports_expert_map(self) -> bool: - return False - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # Let PrepareAndFinalize::finalize() decide the impl. return TopKWeightAndReduceDelegate() diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 8874228a142..53623f13254 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -156,9 +156,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): ) -> bool: return True - def supports_expert_map(self) -> bool: - return True - @staticmethod def _supports_current_device() -> bool: platform = current_platform diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 98265abf7c8..03bf925fbd9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -608,9 +608,6 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular): def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: return True - def supports_expert_map(self) -> bool: - return True - def moe_problem_size( self, a1: torch.Tensor, @@ -1036,9 +1033,6 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): ) -> bool: return True - def supports_expert_map(self) -> bool: - return True - @property def expects_unquantized_inputs(self) -> bool: return True 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 a358796a7d0..64c68018f36 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -44,6 +44,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticChannelSym, kFp8StaticTensorSym, kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt8Static, kMxfp4Static, kMxfp8Static, @@ -566,8 +569,9 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): quant_config.use_mxfp4_w4a16 or quant_config.use_nvfp4_w4a16 or quant_config.use_int4_w4a16 + or quant_config.use_int8_w8a16 or quant_config.use_fp8_w8a16 - ), "Supports only {mxfp,nvfp,int}4_w4a16 or fp8_w8a16" + ), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16" self.w13_g_idx = w13_g_idx self.w2_g_idx = w2_g_idx self.w13_g_idx_sort_indices = w13_g_idx_sort_indices @@ -608,6 +612,9 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): kNvfp4Static, kInt4Static, kInt8Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, ] return weight_key in SUPPORTED_W @@ -640,6 +647,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): if self.w1_zp is not None or self.w2_zp is not None: return scalar_types.uint4.id return scalar_types.uint4b8.id + elif self.quant_config.use_int8_w8a16: + return scalar_types.uint8b128.id elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16: return scalar_types.float4_e2m1f.id elif ( @@ -681,9 +690,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): """Marlin-based fused MoE expert implementation.""" - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() @@ -915,9 +921,6 @@ class BatchedMarlinExperts(MarlinExpertsBase): is_k_full=is_k_full, ) - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceDelegate() diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index b272a458b17..8415ac02784 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -441,9 +441,6 @@ class AiterExperts(mk.FusedMoEExpertsModular): or moe_parallel_config.use_fi_nvl_one_sided_kernels ) - def supports_expert_map(self): - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 45b95b34102..25dd0584de0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -125,9 +125,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): def _supports_batch_invariance(): return True - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 02b7450a5c9..592a1513d75 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): ) -> bool: return True - def supports_expert_map(self) -> bool: - return False - @property def expects_unquantized_inputs(self) -> bool: return True 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 046d3e006af..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 @@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase: or moe_parallel_config.use_ag_rs_all2all_kernels ) and not moe_parallel_config.enable_eplb - def supports_expert_map(self) -> bool: - return False - class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): """ @@ -260,7 +257,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - return True + return router_logits_dtype in [torch.bfloat16, torch.float32] @staticmethod def _supports_routing_method( @@ -282,6 +279,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.SigmoidRenorm, + RoutingMethodType.Sigmoid, RoutingMethodType.MiniMax2, RoutingMethodType.Simulated, ] @@ -293,6 +291,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.SigmoidRenorm, + RoutingMethodType.Sigmoid, RoutingMethodType.MiniMax2, RoutingMethodType.Simulated, ] @@ -320,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 @@ -336,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, @@ -353,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, @@ -364,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, @@ -396,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_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 1e2fff8eb66..43f800343c7 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_expert_map(self) -> bool: - return False - @property def expects_unquantized_inputs(self) -> bool: return False @@ -248,9 +245,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula # routing is done externally, so accept any routing method. return True - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py new file mode 100644 index 00000000000..a65873aca49 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static32, +) +from vllm.platforms import current_platform + + +class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + """ + FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface + (fused router + experts). + + Wraps flashinfer_trtllm_mxint4_moe(). + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.local_num_experts = moe_config.num_local_experts + self.ep_rank = moe_config.ep_rank + self.routing_method = moe_config.routing_method + + @staticmethod + def _supports_current_device() -> bool: + from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501 + is_flashinfer_mxint4_moe_available, + ) + + p = current_platform + return ( + p.is_cuda() + and p.is_device_capability_family(100) + and is_flashinfer_mxint4_moe_available() + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kInt4Static32, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + # FlashInfer MxInt4 uses a fused SwiGLU activation. + return activation == MoEActivation.SWIGLUOAI + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + or moe_parallel_config.use_ag_rs_all2all_kernels + ) and not moe_parallel_config.enable_eplb + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Llama4, + RoutingMethodType.Simulated, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + if router_logits_dtype == torch.float32: + # DeepSeekV3 routing handles float32 logits internally. + # Simulated routing generates synthetic decisions. + return routing_method in ( + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Simulated, + ) + return True + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + def supports_expert_map(self) -> bool: + return False + + @property + def expects_unquantized_inputs(self) -> bool: + # The kernel handles quantization internally. + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501 + flashinfer_trtllm_mxint4_moe, + ) + + assert self.w1_scale is not None + assert self.w2_scale is not None + return flashinfer_trtllm_mxint4_moe( + x=hidden_states, + router_logits=router_logits, + w13_weight_packed=w1, + w13_weight_scale=self.w1_scale, + w2_weight_packed=w2, + w2_weight_scale=self.w2_scale, + global_num_experts=global_num_experts, + top_k=self.topk, + intermediate_size_per_partition=self.intermediate_size_per_partition, + local_num_experts=self.local_num_experts, + ep_rank=self.ep_rank, + num_expert_group=num_expert_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routing_method_type=self.routing_method, + ) 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 ee56b50acd2..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 @@ -179,9 +179,6 @@ class TrtLlmNvFp4ExpertsBase: 300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts) ) - def supports_expert_map(self) -> bool: - return False - class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular): """ @@ -363,7 +360,7 @@ class TrtLlmNvFp4ExpertsMonolithic( router_logits_dtype: torch.dtype | None, routing_method: RoutingMethodType, ) -> bool: - return True + return router_logits_dtype in [torch.bfloat16, torch.float32] def apply( self, @@ -396,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/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py index 4f77954c2a5..82969dd8e25 100644 --- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py @@ -29,9 +29,20 @@ if current_platform.is_xpu(): def prepare_fp8_moe_layer_for_xpu( w13: torch.Tensor, + w13_scale: torch.Tensor, w2: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - return w13.transpose(-1, -2).contiguous(), w2.transpose(-1, -2).contiguous() + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if w13_scale is not None and w13_scale.ndim == 3: + w13_scale = w13_scale.transpose(-1, -2).contiguous() + if w2_scale is not None and w2_scale.ndim == 3: + w2_scale = w2_scale.transpose(-1, -2).contiguous() + return ( + w13.transpose(-1, -2).contiguous(), + w13_scale, + w2.transpose(-1, -2).contiguous(), + w2_scale, + ) class XPUExperts(mk.FusedMoEExpertsModular): @@ -75,6 +86,7 @@ class XPUExperts(mk.FusedMoEExpertsModular): return activation in [ MoEActivation.SILU, MoEActivation.GELU, + MoEActivation.GELU_TANH, MoEActivation.SWIGLUOAI, MoEActivation.RELU2_NO_MUL, ] @@ -95,9 +107,6 @@ class XPUExperts(mk.FusedMoEExpertsModular): ] return (weight_key, activation_key) in SUPPORTED_W_A - def supports_expert_map(self) -> bool: - return True - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index e13ef52c01d..49957c8f5e3 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1518,7 +1518,6 @@ def fused_experts( def _get_config_quant_dtype( use_fp8_w8a8: bool, use_int8_w8a8: bool, - ocp_mx_scheme: str | None, ) -> None | torch.dtype | str: """ Get the quantization type based on the quantization strategy flags. @@ -1529,18 +1528,8 @@ def _get_config_quant_dtype( """ if use_fp8_w8a8: return current_platform.fp8_dtype() - elif use_int8_w8a8: + if use_int8_w8a8: return torch.int8 - elif ocp_mx_scheme == "w_mxfp4_a_mxfp4": - return "mxfp4" - elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e3m2", "w_mxfp6_e3m2_a_mxfp6_e3m2"}: - return "mxfp6_e3m2" - elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e2m3", "w_mxfp6_e2m3_a_mxfp6_e2m3"}: - return "mxfp6_e2m3" - elif ocp_mx_scheme in {"w_mxfp4", "w_mxfp6_e3m2", "w_mxfp6_e2m3"}: - return torch.bfloat16 - elif ocp_mx_scheme in {"w_mxfp4_a_fp8", "w_mxfp6_e3m2_a_fp8", "w_mxfp6_e2m3_a_fp8"}: - return torch.float8_e4m3fn return None @@ -1615,7 +1604,6 @@ def fused_experts_impl( quant_dtype = _get_config_quant_dtype( use_fp8_w8a8=use_fp8_w8a8, use_int8_w8a8=use_int8_w8a8, - ocp_mx_scheme=None, ) get_config_func = functools.partial( diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index 3ebb63b0057..dd21ff58fc3 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -34,11 +34,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): super().__init__(moe_kernel.moe_config) self.moe_quant_config = old_quant_method.moe_quant_config self.moe_kernel = moe_kernel - self.disable_expert_map = getattr( - old_quant_method, - "disable_expert_map", - not self.moe_kernel.supports_expert_map(), - ) self.old_quant_method = old_quant_method logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__) @@ -103,7 +98,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): activation=layer.activation, global_num_experts=layer.global_num_experts, apply_router_weight_on_input=layer.apply_router_weight_on_input, - expert_map=None if self.disable_expert_map else layer.expert_map, + expert_map=layer.expert_map, shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index b5d8f5fd016..4ff43ce21b8 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -727,16 +727,19 @@ class FusedMoE(PluggableLayer): # Only narrow if the loaded_weight is not a scalar (0-dim tensor) # and we're not loading the full weight if not load_full and loaded_weight.ndim > 0: - # Handle padding: loaded_weight might be smaller than shard_size on last - # TP rank - start_offset = shard_size * tp_rank + # When the parameter has been padded (e.g. MXFP4 rounding up + # intermediate_size_per_partition), shard_size is the padded + # size. Compute the offset into the checkpoint weight using + # the *unpadded* per-rank size so that every TP rank lands at + # the correct slice. + tp_size = self.moe_config.moe_parallel_config.tp_size + loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size + start_offset = loaded_per_rank * tp_rank available = loaded_weight.shape[shard_dim] - start_offset if available <= 0: # If there is no available weight to load for this TP rank - # (can happen on last TP rank with padding), we can skip - # loading and return early return - narrow_size = min(shard_size, available) + narrow_size = min(loaded_per_rank, available) loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size) # Narrow parameter and load. # w1, gate_proj: Load into first logical weight of w13. @@ -765,21 +768,18 @@ class FusedMoE(PluggableLayer): ): # Index the loaded weight for tp sharding. # down_proj: "RowParallel" so tp sharding on input_dim - # Narrow parameter and load. - shard_size = expert_data.shape[shard_dim] # Only narrow if the loaded_weight is not a scalar (0-dim tensor) # and we're not loading the full weight if not load_full and loaded_weight.ndim > 0: - # Handle padding: loaded_weight might be smaller than shard_size on last - # TP rank - start_offset = shard_size * tp_rank + # Same padding fix as _load_w13: use unpadded per-rank size. + tp_size = self.moe_config.moe_parallel_config.tp_size + loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size + start_offset = loaded_per_rank * tp_rank available = loaded_weight.shape[shard_dim] - start_offset if available <= 0: # If there is no available weight to load for this TP rank - # (can happen on last TP rank with padding), we can skip - # loading and return early return - narrow_size = min(shard_size, available) + narrow_size = min(loaded_per_rank, available) loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size) # w2, down_proj: Load into only logical weight of w2. hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim) @@ -974,6 +974,19 @@ class FusedMoE(PluggableLayer): # this is needed for compressed-tensors only loaded_weight = loaded_weight.to(param.data.device) + # ModelOpt NVFP4 stores w13 input scales as two logical shards. + # The generic assignment below would broadcast w1/w3 into the + # whole expert row, so the second shard would overwrite the first. + if ( + "ModelOpt" in quant_method_name + and param.data.ndim == 2 + and shard_id in ("w1", "w3") + ): + scale_expert_id = global_expert_id if use_global_sf else expert_id + scale_shard_id = 0 if shard_id == "w1" else 1 + param.data[scale_expert_id][scale_shard_id] = loaded_weight.reshape(()) + return True if return_success else None + if ( "compressed" in quant_method_name.lower() and param.data[expert_id] != 1 diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 6fbc1bffaac..9c3ecee9f9b 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -751,13 +751,6 @@ class FusedMoEExperts(ABC): """ return False - @abstractmethod - def supports_expert_map(self) -> bool: - """ - A flag indicating whether or not this class supports expert maps - """ - raise NotImplementedError - def supports_packed_ue8m0_act_scales(self) -> bool: """ A flag indicating whether or not this class can process packed ue8m0 @@ -1567,12 +1560,6 @@ class FusedMoEKernel: == self.fused_experts.activation_format() ) - def supports_expert_map(self) -> bool: - """ - A flag indicating whether or not this class supports expert maps. - """ - return self.fused_experts.supports_expert_map() - def output_is_reduced(self) -> bool: """ Indicates whether or not the output of fused MoE kernel diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index df430689436..ad9fb509a1f 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -74,6 +74,9 @@ class MoEPermuteScratch: self.sort_workspace = torch.empty( sorter_size, dtype=torch.int8, device=self.device ) + # torch.device("cuda") in config, after initialized, + # will be changed to cuda:{index}, so we need to refresh here. + self.device = self.token_expert_indices.device def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None: n_token, n_hidden = hidden_states.shape diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index bbbad3e8271..0a2e3846dd9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -483,7 +483,9 @@ def convert_to_fp8_moe_kernel_format( prepare_fp8_moe_layer_for_xpu, ) - w13, w2 = prepare_fp8_moe_layer_for_xpu(w13, w2) + w13, w13_scale, w2, w2_scale = prepare_fp8_moe_layer_for_xpu( + w13, w13_scale, w2, w2_scale + ) elif fp8_backend == Fp8MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( prepare_fp8_moe_layer_for_cpu, 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 42b62f8e345..6ad60d62e97 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -2,9 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import sys from enum import Enum -from typing import TYPE_CHECKING +from typing import Any import torch +from compressed_tensors.quantization import ( + QuantizationArgs, +) import vllm._custom_ops as ops import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -12,10 +15,16 @@ from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, + int4_w4a16_moe_quant_config, + int8_w8a16_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( BatchedMarlinExperts, MarlinExperts, + MarlinExpertsBase, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( + TrtLlmMxint4ExpertsMonolithic, ) from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -23,22 +32,20 @@ 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, ) from vllm.platforms import current_platform -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig - from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig - logger = init_logger(__name__) class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" + FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" @@ -47,26 +54,17 @@ def backend_to_kernel_cls( ) -> list[type[mk.FusedMoEExperts]]: """Return the experts class for the given backend, or None for NONE.""" if backend == WNA16MoEBackend.MARLIN: - from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( - MarlinExperts, - ) - return [MarlinExperts] - elif backend == WNA16MoEBackend.BATCHED_MARLIN: - from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( - BatchedMarlinExperts, - ) - return [BatchedMarlinExperts] - + elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: + return [TrtLlmMxint4ExpertsMonolithic] elif backend == WNA16MoEBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, ) return [XPUExpertsWNA16] - else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -77,23 +75,25 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: """ if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] - return [ + + _AVAILABLE_BACKENDS = [ + WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] + return _AVAILABLE_BACKENDS def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, - weight_bits: int, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. Args: config: the shared ``FusedMoEConfig`` for this layer. - weight_bits: quantization bit-width (4 or 8). 8-bit weights are not - supported by the modular Marlin kernel, so ``NONE`` is returned. + weight_key: The QuantKey describing the weight quantization. + Must have int4 or int8 type. Returns: A tuple of (``WNA16MoEBackend``, experts class or ``None``). @@ -156,15 +156,55 @@ def select_wna16_moe_backend( ) +def make_wna16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + group_size: int, + num_bits: int, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + w1_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, + a1_gscale: torch.Tensor | None = None, + a2_gscale: torch.Tensor | None = None, +) -> FusedMoEQuantConfig: + """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE.""" + if num_bits == 4: + return int4_w4a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + w1_bias=w1_bias, + w2_bias=w2_bias, + block_shape=[0, group_size], + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, + ) + else: + assert num_bits == 8 + return int8_w8a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=w1_zp, + w2_zp=w2_zp, + w1_bias=w1_bias, + w2_bias=w2_bias, + block_shape=[0, group_size], + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, + ) + + def make_wna16_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEExperts] | None, - is_k_full: bool, - w13_g_idx: torch.Tensor | None, - w2_g_idx: torch.Tensor | None, - w13_g_idx_sort_indices: torch.Tensor | None, - w2_g_idx_sort_indices: torch.Tensor | None, + experts_cls: type[mk.FusedMoEExperts], + is_k_full: bool = False, + w13_g_idx: torch.Tensor | None = None, + w2_g_idx: torch.Tensor | None = None, + w13_g_idx_sort_indices: torch.Tensor | None = None, + w2_g_idx_sort_indices: torch.Tensor | None = None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -174,16 +214,37 @@ def make_wna16_moe_kernel( XPUExpertsWNA16, ) - assert experts_cls in (MarlinExperts, BatchedMarlinExperts, XPUExpertsWNA16) + # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts + # and BatchedMarlinExperts + assert experts_cls in ( + MarlinExperts, + BatchedMarlinExperts, + TrtLlmMxint4ExpertsMonolithic, + XPUExpertsWNA16, + ) + + is_monolithic = experts_cls.is_monolithic() prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, quant_config=moe_quant_config, routing_tables=routing_tables, allow_new_interface=True, + use_monolithic=is_monolithic, ) assert prepare_finalize is not None - assert isinstance(prepare_finalize, mk.FusedMoEPrepareAndFinalizeModular) + + logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") + + extra_args: dict[str, Any] = {} + if issubclass(experts_cls, MarlinExpertsBase): + extra_args = { + "w13_g_idx": w13_g_idx, + "w2_g_idx": w2_g_idx, + "w13_g_idx_sort_indices": w13_g_idx_sort_indices, + "w2_g_idx_sort_indices": w2_g_idx_sort_indices, + "is_k_full": is_k_full, + } if experts_cls is XPUExpertsWNA16: assert ( @@ -199,30 +260,20 @@ def make_wna16_moe_kernel( elif ( prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts ): - assert experts_cls == BatchedMarlinExperts max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None - experts = BatchedMarlinExperts( + experts = experts_cls( max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), moe_config=moe_config, quant_config=moe_quant_config, - w13_g_idx=w13_g_idx, - w2_g_idx=w2_g_idx, - w13_g_idx_sort_indices=w13_g_idx_sort_indices, - w2_g_idx_sort_indices=w2_g_idx_sort_indices, - is_k_full=is_k_full, + **extra_args, ) else: - assert experts_cls == MarlinExperts - experts = MarlinExperts( + experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, - w13_g_idx=w13_g_idx, - w2_g_idx=w2_g_idx, - w13_g_idx_sort_indices=w13_g_idx_sort_indices, - w2_g_idx_sort_indices=w2_g_idx_sort_indices, - is_k_full=is_k_full, + **extra_args, ) return mk.FusedMoEKernel( @@ -236,10 +287,74 @@ def make_wna16_moe_kernel( # --------------------------------------------------------------------------- +def _process_weights_flashinfer( + w13_qweight: torch.Tensor, + w2_qweight: torch.Tensor, + w13_scales: torch.Tensor, + w2_scales: torch.Tensor, + w13_g_idx: torch.Tensor, + w2_g_idx: torch.Tensor, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor, # w13_g_idx + torch.Tensor, # w2_g_idx + torch.Tensor | None, # w13_g_idx_sort_indices + torch.Tensor | None, # w2_g_idx_sort_indices + torch.Tensor | None, # w13_qzeros + torch.Tensor | None, # w2_qzeros + torch.Tensor | None, # w13_input_global_scale + torch.Tensor | None, # w2_input_global_scale + torch.Tensor | None, # w13_bias + torch.Tensor | None, # w2_bias +]: + """Flashinfer (TRT-LLM MXINT4) weight post-processing. + + Steps + ----- + 1. Transform weights/scales via ``prepare_static_weights_for_trtllm_mxint4_moe``. + 2. Return transformed tensors, passing through g_idx/bias unchanged. + """ + from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( + prepare_static_weights_for_trtllm_mxint4_moe, + ) + + dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( + w13_qweight, + w13_scales, + w2_qweight, + w2_scales, + ) + + return ( + dict_weights_mxint4["gemm1_weights"], + dict_weights_mxint4["gemm2_weights"], + dict_weights_mxint4["gemm1_scales"], + dict_weights_mxint4["gemm2_scales"], + w13_g_idx, + w2_g_idx, + None, + None, + None, + None, + None, + None, + w13_bias, + w2_bias, + ) + + def _process_weights_marlin( layer: torch.nn.Module, - quant_config: "AutoGPTQConfig", input_dtype: torch.dtype | None, + num_bits: int, + pack_factor: int, + group_size: int, + actorder: str | None, w13_qweight: torch.Tensor, w2_qweight: torch.Tensor, w13_scales: torch.Tensor, @@ -292,6 +407,7 @@ def _process_weights_marlin( # --- FP8 weight / scale adjustment --- if input_dtype == torch.float8_e4m3fn: + # NOTE: for non-zp quantization format only marlin_w13_qweight = ops.marlin_int4_fp8_preprocess(w13_qweight, inplace=False) marlin_w2_qweight = ops.marlin_int4_fp8_preprocess(w2_qweight, inplace=False) marlin_w13_scales = w13_scales.data * 512 @@ -303,7 +419,7 @@ def _process_weights_marlin( marlin_w2_scales = w2_scales # --- Process act_order (g_idx) --- - if quant_config.desc_act: + if actorder == "group": num_experts = w13_g_idx.shape[0] w13_g_idx_sort_indices = torch.empty_like(w13_g_idx) w2_g_idx_sort_indices = torch.empty_like(w2_g_idx) @@ -314,6 +430,8 @@ def _process_weights_marlin( w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32) w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]] w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]] + w13_g_idx = w13_sorted_g_idx + w2_g_idx = w2_sorted_g_idx else: num_experts = w13_g_idx.shape[0] device = w13_g_idx.device @@ -338,17 +456,17 @@ def _process_weights_marlin( marlin_w13_qweight = ops.gptq_marlin_moe_repack( marlin_w13_qweight, w13_g_idx_sort_indices, - marlin_w13_qweight.shape[1] * quant_config.pack_factor, + marlin_w13_qweight.shape[1] * pack_factor, marlin_w13_qweight.shape[2], - quant_config.quant_type.size_bits, + num_bits, is_a_8bit=is_a_8bit, ) marlin_w2_qweight = ops.gptq_marlin_moe_repack( marlin_w2_qweight, w2_g_idx_sort_indices, - marlin_w2_qweight.shape[1] * quant_config.pack_factor, + marlin_w2_qweight.shape[1] * pack_factor, marlin_w2_qweight.shape[2], - quant_config.quant_type.size_bits, + num_bits, is_a_8bit=is_a_8bit, ) @@ -357,19 +475,15 @@ def _process_weights_marlin( s=marlin_w13_scales, size_k=layer.intermediate_size_per_partition, size_n=marlin_w13_scales.shape[2], - group_size=quant_config.group_size, + group_size=group_size, is_a_8bit=is_a_8bit, ) + group_size_or_pack_factor = group_size if group_size != -1 else pack_factor marlin_w2_scales = marlin_moe_permute_scales( s=marlin_w2_scales, - size_k=marlin_w2_scales.shape[1] - * ( - quant_config.group_size - if quant_config.group_size != -1 - else quant_config.pack_factor - ), + size_k=marlin_w2_scales.shape[1] * group_size_or_pack_factor, size_n=marlin_w2_scales.shape[2], - group_size=quant_config.group_size, + group_size=group_size, is_a_8bit=is_a_8bit, ) @@ -383,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) @@ -409,7 +540,9 @@ def _process_weights_marlin( def _process_awq_weights_marlin( layer: torch.nn.Module, - quant_config: "AWQMarlinConfig", + weight_bits: int, + pack_factor: int, + group_size: int, input_dtype: torch.dtype | None, w13_qweight: torch.Tensor, w2_qweight: torch.Tensor, @@ -475,16 +608,16 @@ def _process_awq_weights_marlin( w13_qweight, w13_g_idx_sort_indices, size_k=w13_qweight.shape[1], - size_n=w13_qweight.shape[2] * quant_config.pack_factor, - num_bits=quant_config.weight_bits, + size_n=w13_qweight.shape[2] * pack_factor, + num_bits=weight_bits, is_a_8bit=is_a_8bit, ) marlin_w2_qweight = ops.awq_marlin_moe_repack( w2_qweight, w2_g_idx_sort_indices, size_k=w2_qweight.shape[1], - size_n=w2_qweight.shape[2] * quant_config.pack_factor, - num_bits=quant_config.weight_bits, + size_n=w2_qweight.shape[2] * pack_factor, + num_bits=weight_bits, is_a_8bit=is_a_8bit, ) @@ -492,7 +625,7 @@ def _process_awq_weights_marlin( s=w13_scales, size_k=layer.intermediate_size_per_partition, size_n=w13_scales.shape[2], - group_size=quant_config.group_size, + group_size=group_size, is_a_8bit=is_a_8bit, ) if input_dtype == torch.int8 and layer.num_groups_w13 > 1: @@ -504,7 +637,7 @@ def _process_awq_weights_marlin( s=w2_scales, size_k=layer.intermediate_size_per_partition, size_n=w2_scales.shape[2], - group_size=quant_config.group_size, + group_size=group_size, is_a_8bit=is_a_8bit, ) if input_dtype == torch.int8 and layer.num_groups_w2 > 1: @@ -515,15 +648,15 @@ def _process_awq_weights_marlin( marlin_w13_qzeros = moe_awq_to_marlin_zero_points( w13_qzeros, size_k=w13_qzeros.shape[1], - size_n=w13_qzeros.shape[2] * quant_config.pack_factor, - num_bits=quant_config.weight_bits, + size_n=w13_qzeros.shape[2] * pack_factor, + num_bits=weight_bits, is_a_8bit=is_a_8bit, ) marlin_w2_qzeros = moe_awq_to_marlin_zero_points( w2_qzeros, size_k=w2_qzeros.shape[1], - size_n=w2_qzeros.shape[2] * quant_config.pack_factor, - num_bits=quant_config.weight_bits, + size_n=w2_qzeros.shape[2] * pack_factor, + num_bits=weight_bits, is_a_8bit=is_a_8bit, ) @@ -616,7 +749,7 @@ def _process_weights_xpu( def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, - quant_config: QuantizationConfig, + quant_config: QuantizationConfig | QuantizationArgs | None, input_dtype: torch.dtype | None, w13: torch.Tensor, w2: torch.Tensor, @@ -669,9 +802,16 @@ def convert_to_wna16_moe_kernel_format( if isinstance(quant_config, AWQMarlinConfig): if w13_qzeros is None or w2_qzeros is None: raise ValueError("AWQ Marlin MoE requires zero-point tensors.") + + weight_bits = quant_config.weight_bits + pack_factor = quant_config.pack_factor + group_size = quant_config.group_size + return _process_awq_weights_marlin( layer, - quant_config, + weight_bits, + pack_factor, + group_size, input_dtype, w13, w2, @@ -682,19 +822,30 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - - if not isinstance(quant_config, AutoGPTQConfig): + elif isinstance(quant_config, AutoGPTQConfig): + num_bits = quant_config.quant_type.size_bits + pack_factor = quant_config.pack_factor + group_size = quant_config.group_size + actorder = "group" if quant_config.desc_act else None + elif isinstance(quant_config, QuantizationArgs): + num_bits = quant_config.num_bits + pack_factor = 32 // quant_config.num_bits + group_size = quant_config.group_size + actorder = quant_config.actorder + else: raise TypeError( - "Marlin WNA16 MoE backend requires AutoGPTQConfig or " - "AWQMarlinConfig, got " - f"{type(quant_config).__name__}." + "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or " + f"QuantizationArgs, got {type(quant_config).__name__}." ) if w13_g_idx is None or w2_g_idx is None: raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") return _process_weights_marlin( layer, - quant_config, input_dtype, + num_bits, + pack_factor, + group_size, + actorder, w13, w2, w13_scale, @@ -706,7 +857,19 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) + elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: + return _process_weights_flashinfer( + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + w13_bias, + w2_bias, + ) elif backend == WNA16MoEBackend.XPU: + assert quant_config is not None ( w13_xpu, w2_xpu, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index a01718c5879..a506eaffd07 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -622,16 +622,15 @@ def select_deepseek_v4_mxfp4_moe_backend( assert last_error is not None raise last_error - # DeepSeek-V4 on ROCm is more accurate with the unfused Triton MXFP4 path - # than the default AITER path. Prefer Triton-unfused for this routing mode, - # while keeping AITER as a fallback if Triton-unfused rejects the config. + # DeepSeek-V4 on ROCm: prefer AITER FlyDSL MoE (better perf + accuracy + # after shuffle/TP-offset fixes), with Triton-unfused as fallback. if ( current_platform.is_rocm() and config.routing_method == RoutingMethodType.DeepseekV4 ): priority_backends = [ - Mxfp4MoeBackend.TRITON_UNFUSED, Mxfp4MoeBackend.AITER_MXFP4_BF16, + Mxfp4MoeBackend.TRITON_UNFUSED, ] else: priority_backends = _get_priority_backends() @@ -1421,7 +1420,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: - from vllm._aiter_ops import rocm_aiter_ops + from vllm._aiter_ops import rocm_aiter_ops # noqa: F401 if w13_bias is not None: w13_bias = w13_bias.data.to(torch.float32) @@ -1430,45 +1429,31 @@ def convert_weight_to_mxfp4_moe_kernel_format( e, n, k = w13_weight.shape - w13_weight.view(torch.uint8).copy_( - w13_weight.data.view(torch.uint8) - .view(e, n // 2, 2, k) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, k) + # No de-interleave: standard _load_w13 already produces + # [gate_all, up_all] layout. Use aiter-native shuffle functions + # (matching aiter/ops/flydsl/test_flydsl_moe_a4w4.py pattern). + from aiter.ops.shuffle import shuffle_weight as _shuf_w + from aiter.utility.fp4_utils import e8m0_shuffle as _e8m0_shuf + + # w13 (gate+up, stage1): shuffle_weight with layout (16,16) + w13_weight = torch.nn.Parameter( + _shuf_w(w13_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)), + requires_grad=False, ) - w13_weight_scale.data = ( - w13_weight_scale.data.view(e, n // 2, 2, -1) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, -1) + shuffled_w13_scale = _e8m0_shuf( + w13_weight_scale.view(-1, w13_weight_scale.shape[-1]) ) - w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2) - w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2) - - w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True) - shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( - w13_weight_scale.view(-1, w13_weight_scale.shape[-1]), - num_experts, - True, + # w2 (down-proj, stage2): same shuffle as w13 for a4w4 fp4x2 + # (tuning script uses shuffle_weight((16,16)) + e8m0_shuffle for both) + w2_weight = torch.nn.Parameter( + _shuf_w(w2_weight.data.view(torch.float4_e2m1fn_x2), (16, 16)), + requires_grad=False, ) - - w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False) - shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4( - w2_weight_scale.view(-1, w2_weight_scale.shape[-1]), - num_experts, - False, + shuffled_w2_scale = _e8m0_shuf( + w2_weight_scale.view(-1, w2_weight_scale.shape[-1]) ) - if w13_bias is not None: - w13_bias = ( - w13_bias.data.view(-1, n // 2, 2) - .permute(0, 2, 1) - .contiguous() - .view(-1, n) - ) - return ( w13_weight, w2_weight, 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/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index ffbb4c4a7d3..89f3843cc50 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -17,7 +17,7 @@ def _quantize_and_setup_dispatch( a1: torch.Tensor, quant_config: FusedMoEQuantConfig, defer_input_quant: bool = False, -) -> tuple[torch.Tensor, list[torch.Tensor] | None]: +) -> tuple[torch.Tensor, list[torch.Tensor] | None, torch.Tensor | None]: # Defer input quantization to the MoE kernel. if defer_input_quant: a1q = a1 @@ -33,7 +33,7 @@ def _quantize_and_setup_dispatch( # which makes the scales tensor different shape than # the hidden states, breaking the A2A kernel. So, we # delay the swizzling until after the A2A. - a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input( + a1q, a1q_scale = moe_kernel_quantize_input( a1, input_sf, quant_dtype=quant_config.quant_dtype, @@ -49,7 +49,7 @@ def _quantize_and_setup_dispatch( skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0 scales = None if skip_gather_scales else [a1q_scale] - return a1q, scales + return a1q, scales, a1q_scale def _unwrap_scale_and_prepare_for_moe( @@ -129,7 +129,9 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular ) a1 = a1 * topk_weights.to(a1.dtype) - a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) + a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch( + a1, quant_config, defer_input_quant + ) # When LoRA is active, dispatch the per-token LoRA id along with # hidden_states so every rank receives the correct mapping for the @@ -164,7 +166,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular if extra_tensors is None: assert len(res) == 3 a1q, topk_weights, topk_ids = res - a1q_scale = None + a1q_scale = a1q_scale_orig else: assert len(res) == 4 a1q, topk_weights, topk_ids, gathered_extras = res @@ -178,7 +180,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular gathered_extras, quant_config ) else: - a1q_scale = None + a1q_scale = a1q_scale_orig return a1q, a1q_scale, None, topk_ids, topk_weights @@ -249,7 +251,9 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono ) -> mk.PrepareMonolithicResultType: """Quantize and Dispatch Router Logits.""" - a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) + a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch( + a1, quant_config, defer_input_quant + ) res = get_ep_group().dispatch_router_logits( a1q, @@ -261,7 +265,7 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono if scales is None: assert len(res) == 2 a1q, router_logits = res - a1q_scale = None + a1q_scale = a1q_scale_orig else: assert len(res) == 3 a1q, router_logits, scales = res 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/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index a868c9c8487..0a57a6f4dfe 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -3,18 +3,22 @@ import torch from torch.nn.parameter import Parameter +import vllm._custom_ops as ops from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op @PluggableLayer.register("gate_linear") class GateLinear(ReplicatedLinear): - """MoE gate linear layer with three-tier GEMM dispatch: + """MoE gate linear layer with multi-tier GEMM dispatch: - 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims) - 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype) - 3. F.linear via ReplicatedLinear (ultimate fallback) + 1. DSV3 specialized kernel (SM90+, fp32 out, M<=16, H=7168, E=256/384) + 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, + M<=32, H=3072, E=256) + 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 4. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -25,6 +29,11 @@ class GateLinear(ReplicatedLinear): DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] DSV3_SUPPORTED_HIDDEN_SIZES = [7168] + # Dimensions supported by the fp32 specialized kernel + FP32_SUPPORTED_NUM_EXPERTS = [256] + FP32_SUPPORTED_HIDDEN_SIZES = [3072] + FP32_MAX_TOKENS = 32 + def __init__( self, input_size: int, @@ -43,7 +52,7 @@ class GateLinear(ReplicatedLinear): ) # If fp32 compute is required and no specialized kernel is available, - # store weights in fp32 so Tier 3 computes in fp32 natively. + # store weights in fp32 so the fallback linear path computes in fp32. if force_fp32_compute and not can_use_specialized_kernels: params_dtype = torch.float32 @@ -65,6 +74,16 @@ class GateLinear(ReplicatedLinear): and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES ) + # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight) + self.allow_fp32_router_gemm = ( + not bias + and self.weight.dtype == torch.float32 + and current_platform.is_cuda() + and is_hopper_or_blackwell + and output_size in self.FP32_SUPPORTED_NUM_EXPERTS + and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES + ) + # cuBLAS bf16→fp32 eligibility self.allow_cublas_router_gemm = ( self.allow_specialized_router_gemm @@ -92,8 +111,6 @@ class GateLinear(ReplicatedLinear): def forward( self, x: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: - import vllm._custom_ops as ops - # Tier 1: DSV3 specialized kernel if self.allow_dsv3_router_gemm and x.shape[0] <= 16: output = ops.dsv3_router_gemm( @@ -103,15 +120,56 @@ class GateLinear(ReplicatedLinear): ) return output, None - # Tier 2: cuBLAS bf16→fp32 + # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32) + # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph + # capture does not freeze the runtime num_tokens branch. + if self.allow_fp32_router_gemm and x.dtype in ( + torch.float32, + torch.bfloat16, + ): + output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight) + return output, None + + # Tier 3: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return output, None - # Tier 3: F.linear (ReplicatedLinear) + # Tier 4: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) output, output_bias = super().forward(x) if self.out_dtype is not None and output.dtype != self.out_dtype: output = output.to(self.out_dtype) return output, output_bias + + +_FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS + + +def fp32_router_gemm_dispatch_impl( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + """ + Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, + otherwise fall back to F.linear. + This must be wrapped in a custom op because our torch.compile integration + does not support runtime dispatching on num_tokens. + """ + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: + return ops.fp32_router_gemm(x, weight) + else: + return torch.nn.functional.linear(x.float(), weight) + + +def fp32_router_gemm_dispatch_fake( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32) + + +direct_register_custom_op( + op_name="fp32_router_gemm_dispatch", + op_func=fp32_router_gemm_dispatch_impl, + fake_impl=fp32_router_gemm_dispatch_fake, +) 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/linear.py b/vllm/model_executor/layers/linear.py index dbd5577ee03..e50a0e6b002 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -50,9 +50,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [ "AWQLinearMethod", "AutoGPTQLinearMethod", "Fp8LinearMethod", - "MarlinLinearMethod", - "GPTQMarlin24LinearMethod", - "TPUInt8LinearMethod", "FBGEMMFp8LinearMethod", "ModelOptFp8LinearMethod", "ModelOptFp8PcPtLinearMethod", 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/mhc.py b/vllm/model_executor/layers/mhc.py index b720fa1f6fe..5249481293a 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -243,21 +243,13 @@ class HCHeadOp(CustomOp): hc_mult, hidden_size = hidden_states.shape[-2:] outer_shape = hidden_states.shape[:-2] hs_flat = hidden_states.view(-1, hc_mult, hidden_size) - num_tokens = hs_flat.shape[0] - - out = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device - ) - torch.ops.vllm.hc_head_fused_kernel_tilelang( + out = torch.ops.vllm.hc_head_fused_kernel_tilelang( hs_flat, hc_fn, hc_scale, hc_base, - out, - hidden_size, rms_norm_eps, hc_eps, - hc_mult, ) return out.view(*outer_shape, hidden_size) @@ -273,25 +265,24 @@ class HCHeadOp(CustomOp): hc_mult, hidden_size = hidden_states.shape[-2:] outer_shape = hidden_states.shape[:-2] hs_flat = hidden_states.view(-1, hc_mult, hidden_size) - num_tokens = hs_flat.shape[0] - - out = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device - ) if HAS_TILELANG: - torch.ops.vllm.hc_head_fused_kernel_tilelang( + out = torch.ops.vllm.hc_head_fused_kernel_tilelang( hs_flat, hc_fn, hc_scale, hc_base, - out, - hidden_size, rms_norm_eps, hc_eps, - hc_mult, ) else: + num_tokens = hs_flat.shape[0] + out = torch.empty( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device=hidden_states.device, + ) torch.ops.vllm.hc_head_triton( hs_flat, hc_fn, diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 2506c62d390..0e83f80aebd 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -46,7 +46,6 @@ QuantizationMethods = Literal[ QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) DEPRECATED_QUANTIZATION_METHODS = [ - "tpu_int8", "fbgemm_fp8", "fp_quant", ] diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 85ee4061ada..1821fd5c7f7 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -483,7 +483,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): weight_key = QuantKey(quant_type, scale) self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( - moe, weight_key, quant_config.weight_bits + moe, + weight_key, ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 14ebd920fd9..81c0fcb331e 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( convert_to_wna16_moe_kernel_format, make_wna16_moe_kernel, + make_wna16_moe_quant_config, select_wna16_moe_backend, ) from vllm.model_executor.layers.linear import ( @@ -521,7 +522,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase): self.input_dtype = None self.use_marlin = True self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( - moe, kInt4Static, quant_config.weight_bits + moe, + kInt4Static, ) def create_weights( @@ -706,15 +708,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase): ) def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: - from vllm.model_executor.layers.fused_moe.config import ( - awq_marlin_moe_quant_config, - ) - - return awq_marlin_moe_quant_config( + return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, - weight_bits=self.quant_config.weight_bits, group_size=self.quant_config.group_size, + num_bits=self.quant_config.weight_bits, w1_zp=getattr(layer, "w13_qzeros", None) if self.quant_config.zero_point else None, 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_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index da5d85e4abc..14ef8bf614c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -405,8 +405,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): topk_ids, activation=layer.activation, global_num_experts=layer.global_num_experts, - # TODO(rob): investigate the disable_expert_map introduced by: - # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, shared_experts=shared_experts, 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 c56962422bc..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 @@ -1,16 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import enum -from enum import Enum +from typing import Any import torch from compressed_tensors.quantization import ( QuantizationArgs, ) -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( RoutedExperts, @@ -19,40 +16,36 @@ from vllm.model_executor.layers.fused_moe import ( from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( - BatchedMarlinExperts, - MarlinExperts, - fused_marlin_moe, +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + convert_to_wna16_moe_kernel_format, + make_wna16_moe_kernel, + make_wna16_moe_quant_config, + select_wna16_moe_backend, ) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, ) from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa WNA16_SUPPORTED_TYPES_MAP, -) -from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( - flashinfer_trtllm_mxint4_moe, - is_flashinfer_mxint4_moe_available, - prepare_static_weights_for_trtllm_mxint4_moe, + WNA16_ZP_SUPPORTED_TYPES_MAP, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, - marlin_act_int8_process_scales, marlin_make_workspace_new, - marlin_moe_permute_scales, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static32GroupScale, + kInt4StaticGroupScale, + kInt8StaticGroupScale, ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs logger = init_logger(__name__) -class GPTQMarlinState(Enum): - REPACK = enum.auto() - READY = enum.auto() - - class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): def __init__( self, @@ -64,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 @@ -74,20 +65,33 @@ 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) - self.use_flashinfer_mxint4_moe = ( - is_flashinfer_mxint4_moe_available() - and self.group_size == 32 - and weight_quant.num_bits == 4 - ) - self.kernel_backend = ( - "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin" - ) - logger.info_once( - f"Using {self.kernel_backend} backend for WNA16 MoE " - f"(group_size={self.group_size}, num_bits={self.num_bits})", + + if self.num_bits == 4: + if self.group_size == 32: + scale = kInt4Static32GroupScale + else: + scale = kInt4StaticGroupScale + elif self.num_bits == 8: + assert self.group_size == -1 + scale = kInt8StaticGroupScale + else: + raise ValueError( + "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now." + ) + + 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( + config=self.moe, + weight_key=weight_key, ) def get_weight_shape( @@ -103,17 +107,18 @@ 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 shape_map = { "w13_weight": { "Flashinfer": ( @@ -139,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, @@ -155,8 +169,16 @@ 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, + ), + }, } - return shape_map[weight_name][self.kernel_backend] + backend_key = "Flashinfer" if is_flashinfer else "Marlin" + return shape_map[weight_name][backend_key] def create_weights( self, @@ -172,7 +194,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims - is_transposed = self.kernel_backend != "Flashinfer" + is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM extra_weight_attrs.update( {"is_transposed": is_transposed, "quant_method": self.strategy} ) @@ -261,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 ) @@ -319,200 +374,112 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): layer.a13_scale = None layer.a2_scale = None - layer.marlin_state = GPTQMarlinState.REPACK def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - num_experts = layer.w13_weight_g_idx.shape[0] - device = layer.w13_weight_g_idx.device - if self.kernel_backend == "Flashinfer": - dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( - layer.w13_weight_packed, - layer.w13_weight_scale, - layer.w2_weight_packed, - layer.w2_weight_scale, - ) - replace_parameter( - layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"] - ) - replace_parameter( - layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"] - ) - replace_parameter( - layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"] - ) - replace_parameter( - layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"] - ) - return None - - is_a_8bit = ( - self.marlin_input_dtype is not None - and self.marlin_input_dtype.itemsize == 1 + # Process weights using the shared oracle infrastructure + is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM + ( + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, + w13_g_idx_processed, + w2_g_idx_processed, + w13_g_idx_sort_indices, + w2_g_idx_sort_indices, + w13_qzeros, + w2_qzeros, + w13_input_global_scale, + w2_input_global_scale, + _, # w13_bias + _, # w2_bias + ) = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.weight_quant, + input_dtype=self.marlin_input_dtype, + w13=layer.w13_weight_packed, + w2=layer.w2_weight_packed, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_qzeros=getattr(layer, "w13_weight_zero_point", None), + w2_qzeros=getattr(layer, "w2_weight_zero_point", None), ) - if self.marlin_input_dtype == torch.float8_e4m3fn: - # NOTE: for non-zp quantization format only - ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True) - ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True) - layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512 - layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512 + # Replace common parameters + replace_parameter(layer, "w13_weight_packed", w13_qweight) + replace_parameter(layer, "w2_weight_packed", w2_qweight) + replace_parameter(layer, "w13_weight_scale", w13_scales) + replace_parameter(layer, "w2_weight_scale", w2_scales) - # when running models with grouped act order, - # resort to g_idx values provided in checkpoint - if self.actorder == "group": - w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx) - w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx) - w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx) - w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx) + if not self.symmetric: + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) - for e in range(num_experts): - w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to( - torch.int32 - ) - w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to( - torch.int32 - ) - w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][ - w13_g_idx_sort_indices[e] - ] - w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]] - - replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx) - replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx) + # Marlin-specific parameters (not needed for Flashinfer) + if not is_flashinfer: + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - else: - layer.w13_weight_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w2_weight_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, + # Register input global scales if present + if w13_input_global_scale is not None: + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + if w2_input_global_scale is not None: + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + + layer.workspace = marlin_make_workspace_new( + layer.w13_weight_g_idx.device, 4 ) - marlin_w13_qweight = ops.gptq_marlin_moe_repack( - layer.w13_weight_packed, - layer.w13_g_idx_sort_indices, - layer.w13_weight_packed.shape[1] * self.packed_factor, - layer.w13_weight_packed.shape[2], - self.num_bits, - is_a_8bit=is_a_8bit, + # Alias packed weights to w13_weight/w2_weight for the modular kernel interface + layer.w13_weight = layer.w13_weight_packed + layer.w2_weight = layer.w2_weight_packed + + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + + # Add Marlin-specific arguments + marlin_args: dict[str, Any] = {} + if not is_flashinfer: + marlin_args = { + "w13_g_idx": layer.w13_weight_g_idx, + "w2_g_idx": layer.w2_weight_g_idx, + "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices, + "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices, + "is_k_full": self.is_k_full, + } + + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + **marlin_args, ) - replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight) - - marlin_w2_qweight = ops.gptq_marlin_moe_repack( - layer.w2_weight_packed, - layer.w2_g_idx_sort_indices, - layer.w2_weight_packed.shape[1] * self.packed_factor, - layer.w2_weight_packed.shape[2], - self.num_bits, - is_a_8bit=is_a_8bit, - ) - replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight) - - # Repack scales - marlin_w13_scales = marlin_moe_permute_scales( - s=layer.w13_weight_scale, - size_k=layer.w13_weight_packed.shape[2], - size_n=layer.w13_weight_scale.shape[2], - group_size=self.group_size, - is_a_8bit=is_a_8bit, - ) - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1: - marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales( - marlin_w13_scales - ) - layer.register_parameter( - "w13_input_global_scale", - torch.nn.Parameter(w13_input_global_scale, requires_grad=False), - ) - replace_parameter(layer, "w13_weight_scale", marlin_w13_scales) - - marlin_w2_scales = marlin_moe_permute_scales( - s=layer.w2_weight_scale, - size_k=layer.w2_weight_scale.shape[1] - * (self.group_size if self.group_size != -1 else self.packed_factor), - size_n=layer.w2_weight_scale.shape[2], - group_size=self.group_size, - is_a_8bit=is_a_8bit, - ) - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1: - marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales( - marlin_w2_scales - ) - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) - - layer.workspace = marlin_make_workspace_new(device, 4) def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - if self.num_bits != 4: - return None - return int4_w4a16_moe_quant_config( + return make_wna16_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, self.group_size], + 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 select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - assert self.num_bits == 4, "only supporting w4" - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]]) - assert self.moe_quant_config is not None - if ( - prepare_finalize.activation_format - == mk.FusedMoEActivationFormat.BatchedExperts - ): - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank() - assert max_num_tokens_per_rank is not None - return BatchedMarlinExperts( - max_num_tokens=max_num_tokens_per_rank, - num_dispatchers=prepare_finalize.num_dispatchers(), - moe_config=self.moe, - quant_config=self.moe_quant_config, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, - ) - else: - return MarlinExperts( - moe_config=self.moe, - quant_config=self.moe_quant_config, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, - ) - - @property - def is_monolithic(self) -> bool: - return self.kernel_backend == "Flashinfer" - def apply_monolithic( self, layer: RoutedExperts, @@ -520,23 +487,21 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, ) -> torch.Tensor: - assert self.kernel_backend == "Flashinfer" - return flashinfer_trtllm_mxint4_moe( - x=x, - router_logits=router_logits, - w13_weight_packed=layer.w13_weight_packed, - w13_weight_scale=layer.w13_weight_scale, - w2_weight_packed=layer.w2_weight_packed, - w2_weight_scale=layer.w2_weight_scale, + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, global_num_experts=layer.global_num_experts, - top_k=layer.top_k, - intermediate_size_per_partition=layer.intermediate_size_per_partition, - local_num_experts=layer.local_num_experts, - ep_rank=layer.ep_rank, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, num_expert_group=layer.num_expert_group, topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, - routing_method_type=layer.routing_method_type, + routed_scaling_factor=layer.routed_scaling_factor, ) def apply( @@ -548,29 +513,18 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - assert self.kernel_backend == "Marlin" - return fused_marlin_moe( + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, - layer.w13_weight_packed, - layer.w2_weight_packed, - None, - None, - layer.w13_weight_scale, - layer.w2_weight_scale, + layer.w13_weight, + layer.w2_weight, topk_weights, topk_ids, - input_global_scale1=getattr(layer, "w13_input_global_scale", None), - input_global_scale2=getattr(layer, "w2_input_global_scale", None), - quant_type_id=self.quant_type.id, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, activation=layer.activation, + global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - g_idx1=layer.w13_weight_g_idx, - g_idx2=layer.w2_weight_g_idx, - sort_indices1=layer.w13_g_idx_sort_indices, - sort_indices2=layer.w2_g_idx_sort_indices, - workspace=layer.workspace, - input_dtype=self.marlin_input_dtype, - is_k_full=self.is_k_full, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, ) 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/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index de224684788..eabaf62be78 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -192,7 +192,11 @@ class ModelOptQuantConfigBase(QuantizationConfig): # exclude_modules config. But need to keep them for loading quantized # checkpoints generated by older versions. Then check substring matching # for patterns not caught by exact match - if "vision_tower" in prefix or "vision_model" in prefix: + if ( + "vision_tower" in prefix + or "vision_model" in prefix + or "vit_large_projector" in prefix + ): return UnquantizedLinearMethod() # now, the layer is quantized, handle it here @@ -2340,6 +2344,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): if key.startswith(prefix_dot): return info["quant_algo"].upper() + # FusedMoE expert prefix is e.g. "...moe.experts", while ModelOpt's + # quantized_layers entries use "...moe.gate_proj" / "...moe.up_proj". + if prefix.endswith(".experts"): + parent_dot = prefix.rsplit(".experts", 1)[0] + "." + for key, info in self.quantized_layers.items(): + if key.startswith(parent_dot): + return info["quant_algo"].upper() + return None @staticmethod 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 29fa8bcbacb..ba1016a4fb9 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -178,6 +178,16 @@ kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True) kInt8StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) 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 81526415ff2..bc2504b09c5 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -10,7 +10,6 @@ from typing import Any import numpy as np import torch -from huggingface_hub import HfApi from packaging import version from torch import nn from transformers.utils import SAFE_WEIGHTS_INDEX_NAME @@ -48,6 +47,7 @@ from vllm.model_executor.utils import ( set_weight_attrs, ) from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import set_default_torch_dtype logger = init_logger(__name__) @@ -97,8 +97,7 @@ class BitsAndBytesModelLoader(BaseModelLoader): if weight_files: return model_name_or_path, weight_files, pattern else: - hf_api = HfApi() - repo_files = hf_api.list_repo_files(repo_id=model_name_or_path) + repo_files = hf_api().list_repo_files(repo_id=model_name_or_path) for pattern in allowed_patterns: matching_files = fnmatch.filter(repo_files, pattern) if matching_files: @@ -745,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/model_loader/gguf_loader.py b/vllm/model_executor/model_loader/gguf_loader.py index 6148caa9874..2db5efd0e5b 100644 --- a/vllm/model_executor/model_loader/gguf_loader.py +++ b/vllm/model_executor/model_loader/gguf_loader.py @@ -8,7 +8,6 @@ import gguf import regex as re import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoModelForCausalLM, AutoModelForImageTextToText from vllm.config import ModelConfig, VllmConfig @@ -27,6 +26,7 @@ from vllm.model_executor.model_loader.weight_utils import ( gguf_quant_weights_iterator_multi, ) from vllm.transformers_utils.gguf_utils import detect_gguf_multimodal +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.torch_utils import set_default_torch_dtype if TYPE_CHECKING: @@ -57,7 +57,7 @@ class GGUFModelLoader(BaseModelLoader): # repo id/filename.gguf if "/" in model_name_or_path and model_name_or_path.endswith(".gguf"): repo_id, filename = model_name_or_path.rsplit("/", 1) - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=repo_id, filename=filename, revision=model_config.revision, diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index 37d37d55f54..736b2134604 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -16,7 +16,6 @@ from typing import TYPE_CHECKING, Any, ClassVar import regex as re import torch -from huggingface_hub import snapshot_download from torch import nn from torch.utils._python_dispatch import TorchDispatchMode from transformers import PretrainedConfig @@ -26,6 +25,7 @@ from vllm.config import ModelConfig, ParallelConfig, VllmConfig, set_current_vll from vllm.logger import init_logger from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.platforms import current_platform +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.import_utils import PlaceholderModule @@ -629,7 +629,7 @@ def serialize_extra_artifacts( ) with tempfile.TemporaryDirectory() as tmpdir: - snapshot_download( + hf_api().snapshot_download( served_model_name, local_dir=tmpdir, ignore_patterns=[ diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 972271e8c30..cbb191ebb62 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -23,7 +23,6 @@ import huggingface_hub.constants import numpy as np import regex as re import torch -from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download from safetensors.torch import load, load_file, safe_open, save_file from tqdm.auto import tqdm from transformers.utils import SAFE_WEIGHTS_INDEX_NAME @@ -46,6 +45,7 @@ from vllm.model_executor.model_loader.ep_weight_filter import ( ) from vllm.platforms import current_platform from vllm.tracing import instrument +from vllm.transformers_utils.repo_utils import hf_api, hf_fs from vllm.utils.import_utils import PlaceholderModule try: @@ -373,7 +373,7 @@ def get_quant_config( if not is_local: # Download the config files. with get_lock(model_config.model, load_config.download_dir): - hf_folder = snapshot_download( + hf_folder = hf_api().snapshot_download( model_config.model, revision=model_config.revision, allow_patterns="*.json", @@ -431,7 +431,7 @@ def get_sparse_attention_config( if not is_local: # Download the config files. with get_lock(model_name_or_path, load_config.download_dir): - hf_folder = snapshot_download( + hf_folder = hf_api().snapshot_download( model_name_or_path, revision=model_config.revision, allow_patterns="*.json", @@ -534,7 +534,7 @@ def download_weights_from_hf( # Attempt to reduce allow_patterns to a single pattern # so we only have to call snapshot_download once. try: - fs = HfFileSystem() + fs = hf_fs() file_list = fs.ls( os.path.join(model_name_or_path, subfolder or ""), detail=False, @@ -546,7 +546,7 @@ def download_weights_from_hf( # unnecessary files (e.g., from subdirectories like "original/"). index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}" if "*.safetensors" in allow_patterns and index_file in file_list: - index_path = hf_hub_download( + index_path = hf_api().hf_hub_download( repo_id=model_name_or_path, filename=SAFE_WEIGHTS_INDEX_NAME, cache_dir=cache_dir, @@ -582,7 +582,7 @@ def download_weights_from_hf( with get_lock(model_name_or_path, cache_dir): start_time = time.perf_counter() for allow_pattern in allow_patterns: - hf_folder = snapshot_download( + hf_folder = hf_api().snapshot_download( model_name_or_path, allow_patterns=allow_pattern, ignore_patterns=ignore_patterns, @@ -631,7 +631,7 @@ def download_safetensors_index_file_from_hf( with get_lock(model_name_or_path, cache_dir): try: # Download the safetensors index file. - hf_hub_download( + hf_api().hf_hub_download( repo_id=model_name_or_path, filename=index_file, cache_dir=cache_dir, diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 56e119207da..a45d0ca81de 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -130,7 +130,12 @@ class BailingAttention(nn.Module): prefix=f"{prefix}.dense", ) - rotary_dim = getattr(config, "rotary_dim", self.head_dim) + rotary_dim = getattr(config, "rotary_dim", None) + if rotary_dim is None: + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + rotary_dim = int(self.head_dim * partial_rotary_factor) + if rotary_dim is None: + rotary_dim = self.head_dim config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim self.rotary_emb = get_rope( diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 37f94c687a2..b8987a99872 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -35,7 +35,7 @@ from .deepseek_v2 import ( _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) -from .utils import maybe_prefix +from .utils import get_pp_missing_layer_names, maybe_prefix logger = init_logger(__name__) @@ -267,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ), ) + pp_missing_layer_names = get_pp_missing_layer_names(self) params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() _pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer @@ -282,7 +283,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): name = self._rewrite_spec_layer_name(spec_layer, name) if _try_load_fp8_indexer_wk( - name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, ): continue diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index a268ca1974e..e80d00437c7 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -105,6 +105,7 @@ from .interfaces import ( ) from .utils import ( PPMissingLayer, + get_pp_missing_layer_names, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, @@ -304,10 +305,8 @@ class DeepseekV2MoE(nn.Module): self.is_rocm_aiter_moe_enabled and self.gate.e_score_correction_bias is not None ): - # AITER biased_grouped_topk requires the correction bias dtype to - # match the router logits. Keep DeepSeek's correction bias in fp32 - # by requesting fp32 router logits for this routing path. - self.gate.set_out_dtype(torch.float32) + # Accumulates in fp32; avoids bf16->fp32 cast. + self.gate.set_out_dtype(self.gate.weight.dtype) if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled: self.shared_experts = None @@ -744,7 +743,9 @@ class Indexer(nn.Module): return self.indexer_op(hidden_states, q_fp8, k, weights) -def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): +def _try_load_fp8_indexer_wk( + name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names +): """ We fuse the WK and weights_proj projections, but in some checkpoints WK is stored in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16. @@ -760,6 +761,12 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): return False # WK is not in FP8 format, ignore. # Buffer this tensor (weight or scale) until both have arrived. layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer" + fused_name = f"{layer_prefix}.wk_weights_proj.weight" + if any( + name.startswith(missing_layer_name) + for missing_layer_name in pp_missing_layer_names + ): + return True entry = buf.setdefault(layer_prefix, {}) entry["weight" if is_weight else "scale"] = tensor if "weight" not in entry or "scale" not in entry: @@ -777,7 +784,6 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): ) # Load the dequantized weight into shard 0 of the fused buffer. - fused_name = f"{layer_prefix}.wk_weights_proj.weight" param = params_dict[fused_name] param.weight_loader(param, weight_bf16, 0) loaded_params.add(fused_name) @@ -1381,6 +1387,7 @@ class DeepseekV2Model(nn.Module): num_redundant_experts=self.num_redundant_experts, ) + pp_missing_layer_names = get_pp_missing_layer_names(self) params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: @@ -1396,7 +1403,12 @@ class DeepseekV2Model(nn.Module): ) if _try_load_fp8_indexer_wk( - name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + name, + loaded_weight, + _pending_wk_fp8, + params_dict, + loaded_params, + pp_missing_layer_names, ): continue 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/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index c294ffc6f9a..122855400d9 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -501,6 +501,7 @@ class Gemma4MTP(nn.Module): config = vllm_config.speculative_config.draft_model_config.hf_config text_config = _get_text_config(config) self.config = config + self._stable_full_lm_head_weight: torch.Tensor | None = None self.model = Gemma4MultiTokenPredictor( vllm_config=vllm_config, @@ -567,6 +568,8 @@ class Gemma4MTP(nn.Module): ) def _get_full_lm_head_weight(self) -> torch.Tensor: + if self._stable_full_lm_head_weight is not None: + return self._stable_full_lm_head_weight lm_head_weight = self.lm_head.weight tp_size = get_tensor_model_parallel_world_size() if tp_size > 1: @@ -574,7 +577,11 @@ class Gemma4MTP(nn.Module): lm_head_weight, dim=0, ) - return lm_head_weight[: self.masked_embedding.vocab_size] + lm_head_weight = lm_head_weight[: self.masked_embedding.vocab_size] + if tp_size > 1: + lm_head_weight = lm_head_weight.contiguous() + self._stable_full_lm_head_weight = lm_head_weight + return lm_head_weight def compute_logits( self, @@ -599,5 +606,6 @@ class Gemma4MTP(nn.Module): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + self._stable_full_lm_head_weight = None loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 20caa8672de..c6cb6ab103a 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -36,7 +36,9 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +import transformers from einops import rearrange +from packaging.version import Version from transformers import BatchFeature, Glm4vProcessor from transformers.models.glm4v.configuration_glm4v import ( Glm4vTextConfig, @@ -122,6 +124,15 @@ logger = init_logger(__name__) # For profile run _MAX_FRAMES_PER_VIDEO = 600 +TRANSFORMERS_WITH_GA = Version(transformers.__version__) >= Version("5.10.0.dev0") + + +def _to_video_metadata(metadata: Mapping[str, Any]) -> VideoMetadata: + return VideoMetadata( + **{k: metadata[k] for k in metadata if k != "do_sample_frames"} + ) + + # === Vision Inputs === # @@ -832,6 +843,46 @@ class Glm4vProcessingInfo(BaseProcessingInfo): def get_video_processor(self, **kwargs: object) -> Glm4vVideoProcessor: return self.get_hf_processor(**kwargs).video_processor + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int] | None: + processor = self.get_hf_processor() + if isinstance(processor, Glm4vProcessor): + return None + + result: dict[str, int] = {} + + if mm_counts.get("image", 0) > 0: + result["image"] = self.get_max_image_tokens() + + if mm_counts.get("video", 0) > 0: + video_processor = self.get_video_processor() + max_pixels = video_processor.size["longest_edge"] + + vision_config = self.get_hf_config().vision_config + temporal_patch_size = vision_config.temporal_patch_size + patch_size = vision_config.patch_size + merge_size = vision_config.spatial_merge_size + + max_vision_tokens = max_pixels // ( + temporal_patch_size * patch_size**2 * merge_size**2 + ) + + # GLMGA supports up to 640 frames (max_frames). + max_grid_t = 640 // temporal_patch_size + + tokenizer = self.get_tokenizer() + max_ts_tokens = max( + len(tokenizer.encode(f"{t:.1f} seconds", add_special_tokens=False)) + for t in range(min(max_grid_t, 300)) + ) + + result["video"] = max_vision_tokens + max_grid_t * (2 + max_ts_tokens) + 2 + + return result + def get_data_parser(self): return MultiModalDataParser( video_needs_metadata=True, @@ -1104,12 +1155,83 @@ class Glm4vProcessingInfo(BaseProcessingInfo): selected_timestamps.append(timestamps_list[idx]) return selected_timestamps + def _is_glmga_model(self, processor: object) -> bool: + """Detect GLMGA variant via its Glmga sub-processors.""" + for attr in ("image_processor", "video_processor"): + sub = getattr(processor, attr, None) + if sub and "Glmga" in type(sub).__name__: + return True + return False + + def _get_video_second_idx_glmga( + self, metadata: dict[str, Any], total_frames: int + ) -> list[int]: + """Fixed fps=2 frame selection matching GlmgaVideoProcessor.sample_frames.""" + video_processor = self.get_video_processor() + + video_fps = metadata["fps"] + meta_frames = metadata.get("total_num_frames", total_frames) + max_frame_idx = meta_frames - 1 + duration = metadata.get("duration", round(max_frame_idx / video_fps) + 1) + + do_sample_frames = metadata.get("do_sample_frames", True) + if not do_sample_frames: + frame_indices = metadata["frames_indices"] + else: + target_fps = 2 + max_frames = getattr(video_processor, "max_frames", 640) + extract_t = int(duration * target_fps) + extract_t = min(extract_t, max_frames) + + duration_per_frame = 1 / video_fps + timestamps = [i * duration_per_frame for i in range(meta_frames)] + + if meta_frames < extract_t: + frame_indices = [ + math.floor(i * meta_frames / extract_t) for i in range(extract_t) + ] + else: + frame_indices = [] + current_second = 0.0 + inv_fps = 1 / target_fps + for frame_index in range(meta_frames): + if timestamps[frame_index] >= current_second: + current_second += inv_fps + frame_indices.append(frame_index) + if current_second >= duration - inv_fps: + break + + if len(frame_indices) < extract_t: + if len(frame_indices) == 0: + start, end = 0, max(meta_frames - 1, 0) + else: + start, end = frame_indices[0], frame_indices[-1] + frame_indices = np.linspace(start, end, extract_t, dtype=int).tolist() + elif len(frame_indices) > extract_t: + frame_indices = np.linspace( + 0, meta_frames - 1, extract_t, dtype=int + ).tolist() + + seen, uniq = set(), [] + for idx in frame_indices: + if idx not in seen: + seen.add(idx) + uniq.append(idx) + + if len(uniq) & 1: + uniq.append(uniq[-1]) + + frame_indices = uniq + full_second_idxs = [int(idx / video_fps) for idx in frame_indices] + timestamps_list = full_second_idxs[::2] + return list(timestamps_list) + def _construct_video_placeholder( self, video_array: np.ndarray, metadata: dict[str, Any], grid_thw: torch.Tensor, - ) -> str: + ) -> list[int]: hf_processor = self.get_hf_processor() tokenizer = self.get_tokenizer() image_processor = hf_processor.image_processor @@ -1122,11 +1244,13 @@ class Glm4vProcessingInfo(BaseProcessingInfo): merge_length = image_processor.merge_size**2 assert isinstance(grid_thw, torch.Tensor) - timestamps = ( - self._get_video_second_idx_glm4v(metadata, len(video_array)) - if isinstance(hf_processor, Glm4vProcessor) - else self._get_video_second_idx_glm46v(metadata, len(video_array)) - ) + + if isinstance(hf_processor, Glm4vProcessor): + timestamps = self._get_video_second_idx_glm4v(metadata, len(video_array)) + elif self._is_glmga_model(hf_processor): + timestamps = self._get_video_second_idx_glmga(metadata, len(video_array)) + else: + timestamps = self._get_video_second_idx_glm46v(metadata, len(video_array)) timestamp_format = ( "{}" if isinstance(hf_processor, Glm4vProcessor) else "{:.1f} seconds" @@ -1139,9 +1263,16 @@ class Glm4vProcessingInfo(BaseProcessingInfo): num_tokens_per_frame = int(H * W) // merge_length placeholder = [] placeholder.append(bov_token_id) + # Glm46VProcessor uses image_token_id for video frame embeddings; + # Glm4vProcessor uses video_token_id. + frame_embed_token_id = ( + hf_processor.video_token_id + if isinstance(hf_processor, Glm4vProcessor) or not TRANSFORMERS_WITH_GA + else hf_processor.image_token_id + ) for frame_idx in frames_idx_token: placeholder.append(boi_token_id) - placeholder.extend([hf_processor.video_token_id] * num_tokens_per_frame) + placeholder.extend([frame_embed_token_id] * num_tokens_per_frame) placeholder.append(eoi_token_id) placeholder.extend(frame_idx) placeholder.append(eov_token_id) @@ -1241,6 +1372,47 @@ class Glm4vDummyInputsBuilder(BaseDummyInputsBuilder[Glm4vProcessingInfo]): class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): + @staticmethod + def _get_direct_path_inputs( + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + prepared_data = dict(mm_data) + prepared_kwargs = dict(mm_kwargs) + + videos = prepared_data.get("videos") + if not (isinstance(videos, list) and len(videos) > 0): + return prepared_data, prepared_kwargs + + hf_videos = [] + hf_video_metadata = [] + for item in videos: + if isinstance(item, tuple) and len(item) == 2: + video_array, metadata = item + hf_videos.append(video_array) + if isinstance(metadata, VideoMetadata): + hf_video_metadata.append(metadata) + elif isinstance(metadata, Mapping): + hf_video_metadata.append(_to_video_metadata(metadata)) + if "do_sample_frames" in metadata: + prepared_kwargs["do_sample_frames"] = metadata[ + "do_sample_frames" + ] + elif metadata is not None: + raise TypeError( + "Video metadata must be a mapping or VideoMetadata, " + f"got {type(metadata)}" + ) + else: + hf_videos.append(item) + + prepared_data["videos"] = hf_videos + if hf_video_metadata: + prepared_data["video_metadata"] = hf_video_metadata + prepared_kwargs["return_metadata"] = True + + return prepared_data, prepared_kwargs + def _call_hf_processor( self, prompt: str, @@ -1249,11 +1421,32 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) + if not mm_data: + tokenizer = self.info.get_tokenizer() + prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) + return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") + processor = self.info.get_hf_processor(**mm_kwargs) - # GLM-4.1V use `image_token_id` as video placeholder, we need to - # replace it with `video_token_id` for video processing. So we - # separate video processing from image processing. + # Glm46VProcessor and GLMGA handle image/video placeholders together + # via the direct path. Only Glm4vProcessor (GLM-4.1V) needs the + # split-video path because it uses image_token_id as the video + # placeholder. The direct path requires transformers >= 5.5.0 + # (Glm46VProcessor / GlmgaVideoProcessor support). + use_direct_path = ( + not isinstance(processor, Glm4vProcessor) and TRANSFORMERS_WITH_GA + ) + if use_direct_path: + prepared_data, prepared_kwargs = self._get_direct_path_inputs( + mm_data, mm_kwargs + ) + return super()._call_hf_processor( + prompt=prompt, + mm_data=prepared_data, + mm_kwargs=prepared_kwargs, + tok_kwargs=tok_kwargs, + ) + if ( "videos" in mm_data and isinstance(mm_data["videos"], list) @@ -1272,19 +1465,7 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): video_mm_data = dict() video_mm_data["videos"] = [[video_array]] - - unuse_metadata = ["do_sample_frames"] - video_mm_data["video_metadata"] = [ - [ - VideoMetadata( - **{ - k: metadata[k] - for k in metadata - if k not in unuse_metadata - } - ) - ] - ] + video_mm_data["video_metadata"] = [[_to_video_metadata(metadata)]] video_outputs = super()._call_hf_processor( prompt="<|begin_of_video|><|video|><|end_of_video|>", @@ -1366,6 +1547,22 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): embed_token_id=hf_processor.video_token_id, ) + def get_video_replacement_glm46v(item_idx: int): + out_item = out_mm_kwargs["video"][item_idx] + grid_thw = out_item["video_grid_thw"].data + assert isinstance(grid_thw, torch.Tensor) + + video, metadata = mm_items["video"][item_idx] + placeholder = self.info._construct_video_placeholder( + video, metadata, grid_thw + ) + return PromptUpdateDetails.select_token_id( + placeholder, + embed_token_id=hf_processor.image_token_id, + ) + + is_glm46v = not isinstance(hf_processor, Glm4vProcessor) + return [ PromptReplacement( modality="image", @@ -1375,7 +1572,11 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]): PromptReplacement( modality="video", target="<|begin_of_video|><|video|><|end_of_video|>", - replacement=get_video_replacement_glm4v, + replacement=( + get_video_replacement_glm46v + if is_glm46v and TRANSFORMERS_WITH_GA + else get_video_replacement_glm4v + ), ), ] @@ -1436,7 +1637,7 @@ class Glm4vForConditionalGeneration( prefix=maybe_prefix(prefix, "visual"), ) - if config.model_type in ("glm4v", "glm_ocr"): + if config.model_type in ("glm4v", "glm_ocr", "glmga"): architectures = ["Glm4ForCausalLM"] elif config.model_type == "glm4v_moe": architectures = ["Glm4MoeForCausalLM"] @@ -1600,19 +1801,27 @@ class Glm4vForConditionalGeneration( hf_config = self.config spatial_merge_size = hf_config.vision_config.spatial_merge_size for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): - offset = mm_feature.mm_position.offset + embed_ranges = mm_feature.mm_position.extract_embeds_range() if mm_feature.modality == "image": t, h, w = mm_feature.data["image_grid_thw"].data.tolist() assert t == 1, f"Image must have 1 frame, got {t}" + assert len(embed_ranges) == 1 + offset, end = embed_ranges[0] + assert end - offset + 1 == h * w // spatial_merge_size**2 yield offset, t, h // spatial_merge_size, w // spatial_merge_size elif mm_feature.modality == "video": t, h, w = mm_feature.data["video_grid_thw"].data.tolist() - yield ( - offset, - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + num_tokens_per_frame = llm_grid_h * llm_grid_w + + if len(embed_ranges) == t: + for offset, end in embed_ranges: + assert end - offset + 1 == num_tokens_per_frame + yield offset, 1, llm_grid_h, llm_grid_w + else: + offset = mm_feature.mm_position.offset + yield offset, t, llm_grid_h, llm_grid_w else: raise ValueError(f"Unsupported modality: {mm_feature.modality}") diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 3d46bda7ffb..749222b0847 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1623,8 +1623,7 @@ class SupportsEncoderCudaGraph(Protocol): def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + inputs: dict[str, torch.Tensor], ) -> torch.Tensor: """Run the encoder forward pass with precomputed buffers. diff --git a/vllm/model_executor/models/jais.py b/vllm/model_executor/models/jais.py deleted file mode 100644 index cc0c1aa01ba..00000000000 --- a/vllm/model_executor/models/jais.py +++ /dev/null @@ -1,401 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from -# https://huggingface.co/inceptionai/jais-30b-chat-v3/blob/main/modeling_jais.py -# Copyright 2023 The vLLM team. -# Copyright 2023 the Jais authors and HuggingFace Inc. team. All rights -# reserved. -# Copyright 2023 Cerebras Systems. -# -# 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. -"""Inference-only Jais model compatible with HuggingFace weights.""" - -import math -from collections.abc import Iterable -from itertools import islice - -import torch -from torch import nn - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import ( - get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -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.sequence import IntermediateTensors -from vllm.transformers_utils.configs.jais import JAISConfig - -from .interfaces import SupportsPP -from .utils import ( - AutoWeightsLoader, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class SwiGLUActivation(nn.Module): - def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: - return x1 * nn.functional.silu(x2) - - -def _get_alibi_slopes(n): - def get_slopes_power_of_2(n): - start = 2 ** (-(2 ** -(math.log2(n) - 3))) - ratio = start - return [start * ratio**i for i in range(n)] - - if math.log2(n).is_integer(): - return get_slopes_power_of_2(n) - else: - closest_power_of_2 = 2 ** math.floor(math.log2(n)) - return ( - get_slopes_power_of_2(closest_power_of_2) - + _get_alibi_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] - ) - - -class JAISAttention(nn.Module): - def __init__( - self, - config: JAISConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.hidden_size = config.hidden_size - total_num_heads = config.num_attention_heads - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() - assert total_num_heads % tensor_model_parallel_world_size == 0 - self.num_heads = total_num_heads // tensor_model_parallel_world_size - self.head_dim = self.hidden_size // total_num_heads - if hasattr(config, "scale_qk_dot_by_d"): - config.mup_scale_qk_dot_by_d = config.scale_qk_dot_by_d - self.attn_scale_power = 1.0 if config.mup_scale_qk_dot_by_d else 0.5 - self.scale = self.head_dim**-self.attn_scale_power - - self.c_attn = QKVParallelLinear( - self.hidden_size, - self.head_dim, - total_num_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.c_attn", - ) - self.c_proj = RowParallelLinear( - self.hidden_size, - self.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.c_proj", - ) - - self.use_alibi = config.position_embedding_type == "alibi" - alibi_slopes = None - if self.use_alibi: - tp_rank = get_tensor_model_parallel_rank() - head_start = tp_rank * self.num_heads - head_end = (tp_rank + 1) * self.num_heads - alibi_slopes = _get_alibi_slopes(total_num_heads) - alibi_slopes = alibi_slopes[head_start:head_end] - self.attn = Attention( - self.num_heads, - self.head_dim, - scale=self.scale, - alibi_slopes=alibi_slopes, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.c_attn(hidden_states) - q, k, v = qkv.chunk(chunks=3, dim=-1) - attn_output = self.attn(q, k, v) - attn_output, _ = self.c_proj(attn_output) - return attn_output - - -class JAISMLP(nn.Module): - def __init__( - self, - intermediate_size: int, - config: JAISConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - hidden_size = config.hidden_size - self.swiglu = config.activation_function == "swiglu" - self.c_fc = ColumnParallelLinear( - hidden_size, - intermediate_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.c_fc", - ) - self.c_fc2 = ( - ColumnParallelLinear( - hidden_size, - intermediate_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.c_fc2", - ) - if self.swiglu - else None - ) - self.c_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.c_proj", - ) - - self.act = SwiGLUActivation() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.swiglu: - hidden_states2, _ = self.c_fc2(hidden_states) - hidden_states, _ = self.c_fc(hidden_states) - hidden_states = ( - self.act(hidden_states, hidden_states2) - if self.swiglu - else self.act(hidden_states) - ) - hidden_states, _ = self.c_proj(hidden_states) - return hidden_states - - -class JAISBlock(nn.Module): - def __init__( - self, - config: JAISConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - hidden_size = config.hidden_size - inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size - - self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) - self.attn = JAISAttention( - config, cache_config, quant_config, prefix=f"{prefix}.attn" - ) - self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) - self.mlp = JAISMLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp") - - def forward( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.ln_1(hidden_states) - attn_output = self.attn( - hidden_states=hidden_states, - ) - # residual connection - hidden_states = attn_output + residual - - residual = hidden_states - hidden_states = self.ln_2(hidden_states) - feed_forward_hidden_states = self.mlp(hidden_states) - # residual connection - hidden_states = residual + feed_forward_hidden_states - return hidden_states - - -@support_torch_compile -class JAISModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.config = config - assert not config.scale_attn_by_inverse_layer_idx - assert not config.reorder_and_upcast_attn - self.embed_dim = config.hidden_size - self.wte = VocabParallelEmbedding(config.vocab_size, self.embed_dim) - self.wpe = ( - nn.Embedding(config.max_position_embeddings, self.embed_dim) - if config.position_embedding_type != "alibi" - else None - ) - if hasattr(config, "embeddings_scale"): - self.embeddings_scale = config.embeddings_scale - else: - self.embeddings_scale = config.mup_embeddings_scale - - self.start_layer, self.end_layer, self.h = make_layers( - config.num_hidden_layers, - lambda prefix: JAISBlock( - config=config, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ), - prefix=f"{prefix}.h", - ) - - self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states"], config.n_embd - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.wte(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - position_ids: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> IntermediateTensors | torch.Tensor: - if get_pp_group().is_first_rank: - if inputs_embeds is None: - inputs_embeds = self.embed_input_ids(input_ids) - if self.wpe is not None: - position_embeds = self.wpe(position_ids) - hidden_states = inputs_embeds + position_embeds - else: - hidden_states = inputs_embeds - hidden_states *= torch.tensor( - float(self.embeddings_scale), dtype=hidden_states.dtype - ) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - - for layer in islice(self.h, self.start_layer, self.end_layer): - hidden_states = layer(hidden_states) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - - hidden_states = self.ln_f(hidden_states) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if ".attn.bias" in name or ".attn.masked_bias" in name: - # Skip attention mask. - # NOTE: "c_attn.bias" should not be skipped. - continue - if "relative_pe" in name: - continue - - if is_pp_missing_parameter(name, self): - continue - - param = params_dict[name] - # The HF's GPT-2 implementation uses Conv1D instead of Linear. - # Because of this, we need to transpose the weights. - # Note(zhuohan): the logic below might break quantized models. - for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]: - if conv1d_weight_name not in name: - continue - if not name.endswith(".weight"): - continue - loaded_weight = loaded_weight.t() - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -class JAISLMHeadModel(nn.Module, SupportsPP): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.quant_config = quant_config - self.transformer = JAISModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer") - ) - if self.config.tie_word_embeddings: - self.lm_head = self.transformer.wte - else: - self.lm_head = ParallelLMHead( - self.config.vocab_size, - self.config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - if hasattr(config, "width_scale"): - self.output_logits_scale = config.width_scale - else: - self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale - self.logits_processor = LogitsProcessor( - vocab_size=config.vocab_size, scale=self.output_logits_scale - ) - self.make_empty_intermediate_tensors = ( - self.transformer.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.transformer.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> IntermediateTensors | torch.Tensor: - hidden_states = self.transformer( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), - ) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py new file mode 100644 index 00000000000..bdbf0df7fd1 --- /dev/null +++ b/vllm/model_executor/models/mellum.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead + +from .qwen3_moe import ( + Qwen3MoeAttention, + Qwen3MoeDecoderLayer, + Qwen3MoeForCausalLM, + Qwen3MoeMLP, + Qwen3MoeModel, + Qwen3MoeSparseMoeBlock, +) +from .utils import PPMissingLayer, extract_layer_index, maybe_prefix + + +class MellumAttention(Qwen3MoeAttention): + """ + Differences from `Qwen3MoeAttention`: + - Supports `per_layer_sliding_window` for `Attention`. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_parameters: dict[str, Any], + max_position_embeddings: int = 8192, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + qkv_bias: bool = False, + cache_config: Any | None = None, + quant_config: Any | None = None, + prefix: str = "", + dual_chunk_attention_config: dict[str, Any] | None = None, + per_layer_sliding_window: int | None = None, + ) -> None: + nn.Module.__init__(self) + + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.dual_chunk_attention_config = dual_chunk_attention_config + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=per_layer_sliding_window, + prefix=f"{prefix}.attn", + **( + { + "layer_idx": extract_layer_index(prefix), + "dual_chunk_attention_config": dual_chunk_attention_config, + } + if dual_chunk_attention_config + else {} + ), + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + +class MellumDecoderLayer(Qwen3MoeDecoderLayer): + """ + Differences from `Qwen3MoeDecoderLayer`: + - Supports interleaved SWA and per-layer RoPE scaling. + """ + + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + nn.Module.__init__(self) + + config = vllm_config.model_config.hf_text_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.hidden_size = config.hidden_size + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + + layer_idx = extract_layer_index(prefix) + layer_type = config.layer_types[layer_idx] + if layer_type == "sliding_attention": + sliding_window = getattr(config, "sliding_window", None) + else: + sliding_window = None + rope_parameters = config.rope_parameters[layer_type] + + self.self_attn = MellumAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_parameters=rope_parameters, + max_position_embeddings=max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, "attention_bias", False), + head_dim=getattr(config, "head_dim", None), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + dual_chunk_attention_config=dual_chunk_attention_config, + per_layer_sliding_window=sliding_window, + ) + + if config.mlp_layer_types[layer_idx] == "sparse": + self.mlp = Qwen3MoeSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + else: + self.mlp = Qwen3MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + +@support_torch_compile +class MellumModel(Qwen3MoeModel): + """ + Differences from `Qwen3MoeModel`: + - Uses `MellumDecoderLayer`. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=MellumDecoderLayer, + ) + + +class MellumForCausalLM(Qwen3MoeForCausalLM): + """ + Differences from `Qwen3MoeForCausalLM`: + - Uses `MellumModel`. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + nn.Module.__init__(self) + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + if "dense" in getattr(config, "mlp_layer_types", []): + self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"] + self.model = MellumModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + self.expert_weights = [] + + self.moe_layers = [] + example_layer = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + + assert isinstance(layer, Qwen3MoeDecoderLayer) + if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + example_layer = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_layer is None: + raise RuntimeError("No MoE layer found in the model.layers.") + + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_layer.n_logical_experts + self.num_physical_experts = example_layer.n_physical_experts + self.num_local_physical_experts = example_layer.n_local_physical_experts + self.num_routed_experts = example_layer.n_routed_experts + self.num_redundant_experts = example_layer.n_redundant_experts 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/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 53cb29d1122..e5da9154150 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -43,10 +43,10 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoE, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, - ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -113,12 +113,12 @@ class MiniMaxM2MoE(nn.Module): router_logits_dtype=torch.float32, ) - self.gate = ReplicatedLinear( + self.gate = GateLinear( config.hidden_size, config.num_local_experts, bias=False, params_dtype=torch.float32, - quant_config=None, + out_dtype=torch.float32, prefix=f"{prefix}.gate", ) @@ -132,7 +132,7 @@ class MiniMaxM2MoE(nn.Module): hidden_states = hidden_states.view(-1, hidden_dim) # router_logits: (num_tokens, n_experts) - router_logits, _ = self.gate(hidden_states.to(torch.float32)) + router_logits, _ = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states=hidden_states, router_logits=router_logits ) diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 8f9c4fafe80..d4b6984afea 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -1057,7 +1057,7 @@ class Qwen2_5_VisionTransformer(nn.Module): def forward( self, x: torch.Tensor, - grid_thw: list[list[int]], + grid_thw: list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: @@ -1712,11 +1712,8 @@ class Qwen2_5_VLForConditionalGeneration( ) return EncoderCudaGraphConfig( modalities=modalities, - input_key_by_modality={ - "image": "pixel_values", - "video": "pixel_values_videos", - }, buffer_keys=[ + "pixel_values", "rotary_pos_emb_cos", "rotary_pos_emb_sin", "window_index", @@ -1930,7 +1927,7 @@ class Qwen2_5_VLForConditionalGeneration( // self.visual.patch_size ) max_seqlen_window_override = vit_merger_window_size**2 * (spatial_merge_size**2) - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_config, max_batch_size=max_batch_size, max_frames_per_batch=max_frames_per_batch, @@ -1942,14 +1939,12 @@ class Qwen2_5_VLForConditionalGeneration( # Just use image-modality dummy input_buffer for capturing, since it's also # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]). - mm_kwargs = { + values = metadata | { "pixel_values": dummy_pixel_values, - "image_grid_thw": grid_config, } return EncoderCudaGraphCaptureInputs( - mm_kwargs=mm_kwargs, - buffers=buffers, + values=values, ) def prepare_encoder_cudagraph_replay_buffers( @@ -1968,28 +1963,30 @@ class Qwen2_5_VLForConditionalGeneration( # bound and can over-pad window attention into many empty FlashAttention # CTAs. if modality == "image": - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_batch_size=max_batch_size, ) elif modality == "video": - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_frames_per_batch=max_frames_per_batch, ) else: raise AssertionError("This line should be unreachable.") - return EncoderCudaGraphReplayBuffers(buffers=buffers) + values = metadata | { + "pixel_values": self._get_pixel_values_by_modality(mm_kwargs), + } + return EncoderCudaGraphReplayBuffers(values=values) def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + values: dict[str, torch.Tensor], ) -> torch.Tensor: - pixel_values = self._get_pixel_values_by_modality(mm_kwargs) - grid_thw = self._get_grid_thw_by_modality(mm_kwargs) - return self.visual(pixel_values, grid_thw, encoder_metadata=buffers) + pixel_values = values.pop("pixel_values") + metadata = values + return self.visual(pixel_values, None, encoder_metadata=metadata) def encoder_eager_forward( self, diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 869e044c2d4..1fb5587cb3f 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -707,7 +707,7 @@ class Qwen2VisionTransformer(nn.Module): def forward( self, x: torch.Tensor, - grid_thw: torch.Tensor | list[list[int]], + grid_thw: torch.Tensor | list[list[int]] | None, *, encoder_metadata: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: @@ -715,9 +715,11 @@ class Qwen2VisionTransformer(nn.Module): x = x.to(device=self.device, dtype=self.dtype) x = self.patch_embed(x) - grid_thw_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist() - if encoder_metadata is None: + assert grid_thw is not None + grid_thw_list = ( + grid_thw if isinstance(grid_thw, list) else grid_thw.tolist() + ) encoder_metadata = self.prepare_encoder_metadata(grid_thw_list) rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"] @@ -1460,11 +1462,8 @@ class Qwen2VLForConditionalGeneration( max_frames = self.get_max_frames_per_video() return EncoderCudaGraphConfig( modalities=["image", "video"], - input_key_by_modality={ - "image": "pixel_values", - "video": "pixel_values_videos", - }, buffer_keys=[ + "pixel_values", "rotary_pos_emb_cos", "rotary_pos_emb_sin", "cu_seqlens", @@ -1622,7 +1621,7 @@ class Qwen2VLForConditionalGeneration( ) # max_seqlen.item() gets baked into the CUDA graph at capture time. - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_config, max_batch_size=max_batch_size, max_frames_per_batch=max_frames_per_batch, @@ -1632,14 +1631,12 @@ class Qwen2VLForConditionalGeneration( # Capture with image-format kwargs; pixel_values shape is compatible with # both image and video replay paths. - mm_kwargs = { + values = metadata | { "pixel_values": dummy_pixel_values, - "image_grid_thw": grid_config, } return EncoderCudaGraphCaptureInputs( - mm_kwargs=mm_kwargs, - buffers=buffers, + values=values, ) def prepare_encoder_cudagraph_replay_buffers( @@ -1652,24 +1649,28 @@ class Qwen2VLForConditionalGeneration( grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs) if modality == "image": - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_batch_size=max_batch_size, ) else: - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_frames_per_batch=max_frames_per_batch, ) - return EncoderCudaGraphReplayBuffers(buffers=buffers) + values = metadata | { + "pixel_values": self._get_pixel_values_by_modality(mm_kwargs), + } + return EncoderCudaGraphReplayBuffers(values=values) def encoder_cudagraph_forward( - self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor] + self, + values: dict[str, torch.Tensor], ) -> torch.Tensor: - pixel_values = self._get_pixel_values_by_modality(mm_kwargs) - grid_thw = self._get_grid_thw_by_modality(mm_kwargs) - return self.visual(pixel_values, grid_thw, encoder_metadata=buffers) + pixel_values = values.pop("pixel_values") + metadata = values + return self.visual(pixel_values, None, encoder_metadata=metadata) def encoder_eager_forward( self, diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 231ed646e09..25f139f26cb 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -509,7 +509,6 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) - self.config.target_layer_count = target_layer_num self.model = DFlashQwen3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index a474649cc93..28c62e59bd1 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1784,11 +1784,8 @@ class Qwen3VLForConditionalGeneration( return EncoderCudaGraphConfig( modalities=modalities, - input_key_by_modality={ - "image": "pixel_values", - "video": "pixel_values_videos", - }, buffer_keys=[ + "pixel_values", "pos_embeds", "rotary_pos_emb_cos", "rotary_pos_emb_sin", @@ -1978,7 +1975,7 @@ class Qwen3VLForConditionalGeneration( # so the capture value must cover any replay scenario. # Worst case: 1 item consuming the full budget -> # seq_len = token_budget * spatial_merge_size^2. - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_config, max_batch_size=max_batch_size, max_frames_per_batch=max_frames_per_batch, @@ -1988,14 +1985,12 @@ class Qwen3VLForConditionalGeneration( # Just use image-modality dummy input_buffer for capturing, since it's also # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]). - mm_kwargs = { + values = metadata | { "pixel_values": dummy_pixel_values, - "image_grid_thw": grid_config, } return EncoderCudaGraphCaptureInputs( - mm_kwargs=mm_kwargs, - buffers=buffers, + values=values, ) def prepare_encoder_cudagraph_replay_buffers( @@ -2008,28 +2003,30 @@ class Qwen3VLForConditionalGeneration( grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs) if modality == "image": - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_batch_size=max_batch_size, ) elif modality == "video": - buffers = self.visual.prepare_encoder_metadata( + metadata = self.visual.prepare_encoder_metadata( grid_thw_list, max_frames_per_batch=max_frames_per_batch, ) else: raise AssertionError("This line should be unreachable.") - return EncoderCudaGraphReplayBuffers(buffers=buffers) + values = metadata | { + "pixel_values": self._get_pixel_values_by_modality(mm_kwargs), + } + return EncoderCudaGraphReplayBuffers(values=values) def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + values: dict[str, torch.Tensor], ) -> torch.Tensor: - pixel_values = self._get_pixel_values_by_modality(mm_kwargs) - grid_thw = self._get_grid_thw_by_modality(mm_kwargs) - return self.visual(pixel_values, grid_thw, encoder_metadata=buffers) + pixel_values = values.pop("pixel_values") + metadata = values + return self.visual(pixel_values, None, encoder_metadata=metadata) def encoder_eager_forward( self, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5e6f106150a..d96ceeb4b50 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -146,7 +146,6 @@ _TEXT_GENERATION_MODELS = { "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), - "JAISLMHeadModel": ("jais", "JAISLMHeadModel"), "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), "JambaForCausalLM": ("jamba", "JambaForCausalLM"), "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), @@ -160,6 +159,7 @@ _TEXT_GENERATION_MODELS = { "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"), "MambaForCausalLM": ("mamba", "MambaForCausalLM"), "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"), + "MellumForCausalLM": ("mellum", "MellumForCausalLM"), "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"), "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"), "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), @@ -568,6 +568,7 @@ _MULTIMODAL_MODELS = { "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"), "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"), + "Step3p7ForConditionalGeneration": ("step3p7", "Step3p7ForConditionalGeneration"), "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), "Tarsier2ForConditionalGeneration": ( "qwen2_vl", @@ -705,6 +706,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "Phi3SmallForCausalLM": "0.9.2", "Phi4FlashForCausalLM": "0.10.2", "Phi4MultimodalForCausalLM": "0.12.0", + "JAISLMHeadModel": "0.22.0", # encoder-decoder models except whisper # have been removed for V0 deprecation. "DonutForConditionalGeneration": "0.10.2", diff --git a/vllm/model_executor/models/step3_text.py b/vllm/model_executor/models/step3_text.py index a0e7e16a9bb..dd4af6f0fec 100644 --- a/vllm/model_executor/models/step3_text.py +++ b/vllm/model_executor/models/step3_text.py @@ -345,7 +345,7 @@ class Step3TextModel(nn.Module): self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states"], config.hidden_size + ["hidden_states", "residual"], config.hidden_size ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index efe552f5549..5a28ce3d004 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -702,8 +702,10 @@ class Step3VLForConditionalGeneration( return EncoderCudaGraphConfig( modalities=["image"], - input_key_by_modality={"image": "pixel_values"}, - buffer_keys=["patch_pixel_values"], + buffer_keys=[ + "pixel_values", + "patch_pixel_values", + ], out_hidden_size=self.config.hidden_size, ) @@ -846,33 +848,27 @@ class Step3VLForConditionalGeneration( device=device, dtype=dtype, ) - # num_patches is NOT in buffers -- the per-item merge is done + # num_patches is NOT in values -- the per-item merge is done # CPU-side by finalize_encoder_cudagraph_output using the actual # batch's num_patches from mm_kwargs. - mm_kwargs = { + values = { "pixel_values": dummy_pixel_values, "patch_pixel_values": dummy_patch_pixel_values, } - buffers = { - "patch_pixel_values": dummy_patch_pixel_values, - } - return EncoderCudaGraphCaptureInputs( - mm_kwargs=mm_kwargs, - buffers=buffers, + values=values, ) def encoder_cudagraph_forward( self, - mm_kwargs: dict[str, Any], - buffers: dict[str, torch.Tensor], + values: dict[str, torch.Tensor], ) -> torch.Tensor: # Graph captures only the compute (vision model + conv projector). # Per-item merge happens CPU-side in finalize_encoder_cudagraph_output # using actual num_patches from the batch data. - pixel_values = mm_kwargs["pixel_values"] - patch_pixel_values = buffers["patch_pixel_values"] + pixel_values = values["pixel_values"] + patch_pixel_values = values["patch_pixel_values"] image_features = self._process_image_features( self._get_vision_model_output(pixel_values) @@ -974,10 +970,11 @@ class Step3VLForConditionalGeneration( EncoderCudaGraphReplayBuffers, ) - # Only patch_pixel_values lives in the buffers dict; num_patches is + # Only patch_pixel_values lives in the values dict; num_patches is # processed CPU-side by finalize_encoder_cudagraph_output. return EncoderCudaGraphReplayBuffers( - buffers={ + values={ + "pixel_values": mm_kwargs["pixel_values"], "patch_pixel_values": mm_kwargs["patch_pixel_values"], }, ) diff --git a/vllm/model_executor/models/step3p5.py b/vllm/model_executor/models/step3p5.py index cd73f6e26d9..c15cf18413b 100644 --- a/vllm/model_executor/models/step3p5.py +++ b/vllm/model_executor/models/step3p5.py @@ -637,6 +637,54 @@ class Step3p5Model(nn.Module): (f".moe.experts.{base_layer}w13_weight", ".moe.gate_proj.weight", "w1"), (f".moe.experts.{base_layer}w13_weight", ".moe.up_proj.weight", "w3"), (f".moe.experts.{base_layer}w2_weight", ".moe.down_proj.weight", "w2"), + ( + f".moe.experts.{base_layer}w13_weight_scale_2", + ".moe.gate_proj.weight_scale_2", + "w1", + ), + ( + f".moe.experts.{base_layer}w13_weight_scale_2", + ".moe.up_proj.weight_scale_2", + "w3", + ), + ( + f".moe.experts.{base_layer}w2_weight_scale_2", + ".moe.down_proj.weight_scale_2", + "w2", + ), + ( + f".moe.experts.{base_layer}w13_weight_scale", + ".moe.gate_proj.weight_scale", + "w1", + ), + ( + f".moe.experts.{base_layer}w13_weight_scale", + ".moe.up_proj.weight_scale", + "w3", + ), + ( + f".moe.experts.{base_layer}w2_weight_scale", + ".moe.down_proj.weight_scale", + "w2", + ), + # Required due to the Step3 HF model's packed expert format: + # input scales are stored as moe.{gate,up,down}_proj.input_scale + # rather than the standard per-expert format handled generically. + ( + f".moe.experts.{base_layer}w13_input_scale", + ".moe.gate_proj.input_scale", + "w1", + ), + ( + f".moe.experts.{base_layer}w13_input_scale", + ".moe.up_proj.input_scale", + "w3", + ), + ( + f".moe.experts.{base_layer}w2_input_scale", + ".moe.down_proj.input_scale", + "w2", + ), ] # New per-expert format: .moe.experts.E.gate_proj.weight_packed [out, in] @@ -756,7 +804,11 @@ class Step3p5Model(nn.Module): # Per-tensor global scales (e.g. weight_global_scale) # have shape [1] in compressed-tensors NVFP4 checkpoints. # Expand to per-expert before the iteration loop. - if ( + if loaded_weight.ndim == 0: + loaded_weight = loaded_weight.unsqueeze(0).expand( + moe_expert_num + ) + elif ( loaded_weight.shape[0] == 1 and loaded_weight.shape[0] != moe_expert_num ): diff --git a/vllm/model_executor/models/step3p7.py b/vllm/model_executor/models/step3p7.py new file mode 100644 index 00000000000..2a98ec53fc3 --- /dev/null +++ b/vllm/model_executor/models/step3p7.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Jurassic model.""" + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.linear import ColumnParallelLinear + +from .step3_vl import Step3VLForConditionalGeneration +from .step_vl import PerceptionEncoder +from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix +from .vision import run_dp_sharded_vision_model + +logger = init_logger(__name__) + + +class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.vision_model.": "vision_model.", + "model.vit_large_projector.": "vit_large_projector.", + "model.vit_large_projector": "vit_large_projector", + "model.language_model.": "language_model.model.", + "model.language_model": "language_model.model", + "model.": "language_model.model.", + "lm_head.": "language_model.lm_head.", + "lm_head": "language_model.lm_head", + }, + orig_to_new_substr={ + ".attn.in_proj_weight": ".attn.qkv_proj.weight", + ".attn.in_proj_bias": ".attn.qkv_proj.bias", + ".mlp.c_fc": ".mlp.fc1", + ".mlp.c_proj": ".mlp.fc2", + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super(Step3VLForConditionalGeneration, self).__init__() + + config = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + quant_config = vllm_config.quant_config + + self.config = config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + + with self._mark_tower_model(vllm_config, "image"): + self.vision_model = PerceptionEncoder( + config.vision_config, + get_act_fn(config.vision_config.hidden_act), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "vision_model"), + ) + self.vit_large_projector = ColumnParallelLinear( + config.vision_config.width * 4, + config.text_config.hidden_size, + bias=config.projector_bias, + gather_output=True, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "vit_large_projector"), + disable_tp=self.use_data_parallel, + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def _get_vision_model_output( + self, input_tensor: torch.Tensor | None + ) -> torch.Tensor | None: + if input_tensor is None: + return None + if self.use_data_parallel: + return run_dp_sharded_vision_model(input_tensor, self.vision_model) + return self.vision_model(input_tensor) + + def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor: + image_features, _ = self.vit_large_projector(image_features) + return image_features 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 d052b41fc54..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 @@ -25,7 +24,6 @@ from vllm.models.deepseek_v4.common.ops import ( fused_q_kv_rmsnorm, ) from vllm.utils.deep_gemm import fp8_einsum -from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum if TYPE_CHECKING: @@ -39,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 @@ -91,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, @@ -143,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, @@ -163,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) @@ -174,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, @@ -189,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 @@ -199,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) @@ -217,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. @@ -244,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, @@ -254,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, @@ -292,13 +253,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): device=hidden_states.device, ) - # Attention (inside custom op for torch.compile boundary) - torch.ops.vllm.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. @@ -334,14 +292,12 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): device=o.device, dtype=torch.bfloat16, ) - torch.ops.vllm.deepseek_v4_fp8_einsum( - o_fp8, - o_scale, - wo_a_fp8, - wo_a_scale, - z, + fp8_einsum( "bhr,hdr->bhd", - list(self._einsum_recipe), + (o_fp8, o_scale), + (wo_a_fp8, wo_a_scale), + z, + recipe=self._einsum_recipe, ) return self.wo_b(z.flatten(1)) @@ -406,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, @@ -542,67 +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) - - -def deepseek_v4_attention_fake( - hidden_states: torch.Tensor, - positions: torch.Tensor, - out: torch.Tensor, - layer_name: str, -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_attention", - op_func=deepseek_v4_attention, - mutates_args=["out"], - fake_impl=deepseek_v4_attention_fake, -) - - -def deepseek_v4_fp8_einsum( - a: torch.Tensor, - a_scale: torch.Tensor, - b: torch.Tensor, - b_scale: torch.Tensor, - out: torch.Tensor, - equation: str, - recipe: list[int], -) -> None: - fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) - - -def deepseek_v4_fp8_einsum_fake( - a: torch.Tensor, - a_scale: torch.Tensor, - b: torch.Tensor, - b_scale: torch.Tensor, - out: torch.Tensor, - equation: str, - recipe: list[int], -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_fp8_einsum", - op_func=deepseek_v4_fp8_einsum, - mutates_args=["out"], - fake_impl=deepseek_v4_fp8_einsum_fake, -) - - 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 a5f4bae842e..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,9 +15,18 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from vllm.forward_context import get_forward_context +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, + mhc_post_tilelang, + mhc_pre_tilelang, +) 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, ) @@ -29,20 +38,13 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mhc import ( - HCHeadOp, - MHCFusedPostPreOp, - MHCPostOp, - 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, ) 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, @@ -55,12 +57,11 @@ 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 -from vllm.utils.torch_utils import direct_register_custom_op class DeepseekV4MLP(nn.Module): @@ -147,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 @@ -159,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( @@ -216,10 +224,15 @@ class DeepseekV4MegaMoEExperts(nn.Module): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - 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 + 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, @@ -230,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: @@ -274,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(), @@ -308,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() @@ -334,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, @@ -349,30 +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) - torch.ops.vllm.deepseek_v4_mega_moe_experts( - hidden_states, - topk_weights, - topk_ids, - y, - self.prefix, - 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, @@ -397,51 +466,12 @@ 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] -def _deepseek_v4_mega_moe_experts_op( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - out: torch.Tensor, - layer_name: str, - activation_clamp: float | None, - fast_math: bool, -) -> None: - self = get_forward_context().no_compile_layers[layer_name] - self._run_mega_moe( - hidden_states, - topk_weights, - topk_ids, - out, - activation_clamp, - fast_math, - ) - - -def _deepseek_v4_mega_moe_experts_op_fake( - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - out: torch.Tensor, - layer_name: str, - activation_clamp: float | None, - fast_math: bool, -) -> None: - return None - - -direct_register_custom_op( - op_name="deepseek_v4_mega_moe_experts", - op_func=_deepseek_v4_mega_moe_experts_op, - mutates_args=["out"], - fake_impl=_deepseek_v4_mega_moe_experts_op_fake, -) - - class DeepseekV4MoE(nn.Module): def __init__( self, @@ -544,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, @@ -574,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, @@ -747,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 @@ -791,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, @@ -805,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, @@ -844,10 +882,6 @@ class DeepseekV4DecoderLayer(nn.Module): ): super().__init__() - # Lazy import to avoid top-level tilelang dependency. - # Registers both torch.ops.vllm.mhc_pre and mhc_post - import vllm.model_executor.layers.mhc # noqa: F401 - config = vllm_config.model_config.hf_config self.hidden_size = config.hidden_size @@ -910,42 +944,6 @@ class DeepseekV4DecoderLayer(nn.Module): ), requires_grad=False, ) - self.mhc_pre = MHCPreOp() - self.mhc_post = MHCPostOp() - self.mhc_fused_post_pre = MHCFusedPostPreOp() - - def hc_pre( - self, - x: torch.Tensor, - hc_fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, - norm_weight: torch.Tensor | None = None, - norm_eps: float = 1e-6, - ): - post_mix, res_mix, layer_input = self.mhc_pre( - residual=x, - fn=hc_fn, - hc_scale=hc_scale, - hc_base=hc_base, - rms_eps=self.rms_norm_eps, - hc_pre_eps=self.hc_eps, - hc_sinkhorn_eps=self.hc_eps, - hc_post_mult_value=self.hc_post_alpha, - sinkhorn_repeat=self.hc_sinkhorn_iters, - norm_weight=norm_weight, - norm_eps=norm_eps, - ) - return layer_input, post_mix, res_mix - - def hc_post( - self, - x: torch.Tensor, - residual: torch.Tensor, - post: torch.Tensor, - comb: torch.Tensor, - ): - return self.mhc_post(x, residual, post, comb) def forward( self, @@ -959,18 +957,23 @@ class DeepseekV4DecoderLayer(nn.Module): attn_norm_weight = self.attn_norm.weight.data attn_norm_eps = self.attn_norm.variance_epsilon if residual is None: - # Run standalone hc_pre on first layer + # Run standalone mhc_pre on first layer residual = x - x, post_mix, res_mix = self.hc_pre( + post_mix, res_mix, x = mhc_pre_tilelang( x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, norm_weight=attn_norm_weight, norm_eps=attn_norm_eps, ) else: - residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( x, residual, post_mix, @@ -989,12 +992,12 @@ class DeepseekV4DecoderLayer(nn.Module): norm_eps=attn_norm_eps, ) - # attn_norm is fused into hc_pre / mhc_fused_post_pre above. + # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above. x = self.attn(positions, x, None) ffn_norm_weight = self.ffn_norm.weight.data ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( x, residual, post_mix, @@ -1040,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)] @@ -1097,7 +1100,6 @@ class DeepseekV4Model(nn.Module): torch.empty(1, dtype=torch.float32), requires_grad=False, ) - self.hc_head_op = HCHeadOp() # Pre-hc_head residual stream buffer for the MTP draft. Stable # address (outside the cudagraph pool) so the copy_ in forward() # refreshes it correctly across captured shapes. @@ -1167,7 +1169,9 @@ class DeepseekV4Model(nn.Module): residual, ) if layer is not None: - hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) + hidden_states = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) @@ -1176,7 +1180,7 @@ class DeepseekV4Model(nn.Module): num_tokens = hidden_states.shape[0] self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) - hidden_states = self.hc_head_op( + hidden_states = hc_head_fused_kernel_tilelang( hidden_states, self.hc_head_fn, self.hc_head_scale, @@ -1343,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. @@ -1375,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/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 3db831f0261..133a96e3acd 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -24,11 +24,14 @@ from vllm.distributed import ( get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger +from vllm.model_executor.kernels.mhc.tilelang import ( + hc_head_fused_kernel_tilelang, + mhc_post_tilelang, +) from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mhc import HCHeadOp from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) @@ -122,8 +125,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): aux_stream_list=aux_stream_list, ) - self.hc_head_op = HCHeadOp() - def forward( self, input_ids: torch.Tensor, @@ -155,9 +156,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): hidden_states, residual, post_mix, res_mix = self.mtp_block( positions=positions, x=hidden_states, input_ids=None ) - hidden_states = self.mtp_block.hc_post( - hidden_states, residual, post_mix, res_mix - ) + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) # Return the flat pre-hc_head residual so it can be re-fed as the # next spec step's `previous_hidden_states` when # num_speculative_tokens > 1. hc_head is deferred to compute_logits. @@ -237,7 +236,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module): hidden_states = hidden_states.view( -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size ) - hidden_states = mtp_layer.hc_head_op( + hidden_states = hc_head_fused_kernel_tilelang( hidden_states, mtp_layer.hc_head_fn, mtp_layer.hc_head_scale, 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 0eba82126af..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 @@ -96,9 +96,16 @@ class SparseAttnCompressNormRopeStoreC4Kernel: self.quant_block = quant_block self.token_stride = token_stride self.scale_dim = scale_dim - self.num_warps = head_size // quant_block + self.elems_per_lane = 8 + self.copy_elems = 4 + self.copy_chunks = self.elems_per_lane // self.copy_elems + self.lanes_per_group = quant_block // self.elems_per_lane + self.groups_per_warp = 32 // self.lanes_per_group + self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1 + self.scale_reduce_offset = self.lanes_per_group // 2 + self.num_warps = (head_size // quant_block) // self.groups_per_warp self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 + self.tb_size = self.num_warps * 32 self.compress_ratio = compress_ratio self.overlap = overlap self.window = (1 + int(overlap)) * compress_ratio @@ -156,8 +163,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel: tid, _, _ = cute.arch.thread_idx() warp_id = cute.arch.make_warp_uniform(tid // 32) lane_id = tid % 32 - elem0 = tid * 2 - elem1 = elem0 + 1 + group_lane = lane_id % self.lanes_per_group + group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group + elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane slot_id = slot_mapping[token_idx] has_position = token_idx < positions.shape[0] @@ -201,12 +209,24 @@ class SparseAttnCompressNormRopeStoreC4Kernel: s_block_numbers[row] = block_number_i32 cute.arch.sync_threads() - max0 = -Float32.inf - max1 = -Float32.inf - sum0 = Float32(0.0) - sum1 = Float32(0.0) - product0 = Float32(0.0) - product1 = Float32(0.0) + local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(self.elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (self.copy_chunks, self.copy_elems), + stride=(self.copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) for row in cutlass.range_constexpr(self.window): pos = start + Int64(row) @@ -215,46 +235,51 @@ class SparseAttnCompressNormRopeStoreC4Kernel: block_offset = pos - block_index * block_size block_number = s_block_numbers[row].to(Int64) head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_base = ( - block_number * state_cache.stride[0] - + block_offset * state_cache.stride[1] - + head_offset - ) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(self.copy_chunks): + copy_elem = const_expr(chunk * self.copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(self.copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(self.copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(self.copy_elems,), + coord=( + col_tile + Int64(self.state_width // self.copy_elems), + ), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - score0 = state_cache.iterator[ - row_base + Int64(self.state_width) + elem0.to(Int64) - ] - kv0 = state_cache.iterator[row_base + elem0.to(Int64)] - new_max0 = cute.arch.fmax(max0, score0) - old_scale0 = cute.math.exp2( - (max0 - new_max0) * Float32(self.rcp_ln2), fastmath=True - ) - new_scale0 = cute.math.exp2( - (score0 - new_max0) * Float32(self.rcp_ln2), fastmath=True - ) - sum0 = sum0 * old_scale0 + new_scale0 - product0 = product0 * old_scale0 + kv0 * new_scale0 - max0 = new_max0 + for e in cutlass.range_constexpr(self.elems_per_lane): + chunk = const_expr(e // self.copy_elems) + copy_elem = const_expr(e % self.copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(self.rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = local_product[e] * old_scale + kv * new_scale + local_max[e] = new_max - score1 = state_cache.iterator[ - row_base + Int64(self.state_width) + elem1.to(Int64) - ] - kv1 = state_cache.iterator[row_base + elem1.to(Int64)] - new_max1 = cute.arch.fmax(max1, score1) - old_scale1 = cute.math.exp2( - (max1 - new_max1) * Float32(self.rcp_ln2), fastmath=True - ) - new_scale1 = cute.math.exp2( - (score1 - new_max1) * Float32(self.rcp_ln2), fastmath=True - ) - sum1 = sum1 * old_scale1 + new_scale1 - product1 = product1 * old_scale1 + kv1 * new_scale1 - max1 = new_max1 + x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(self.elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] - x0 = product0 / sum0 - x1 = product1 / sum1 - - local_sumsq = x0 * x0 + x1 * x1 warp_sum = local_sumsq for step in cutlass.range_constexpr(5): offset = const_expr(16 >> step) @@ -273,8 +298,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel: cute.arch.sync_threads() rrms = rrms_shared[0] - x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) - x1 = x1 * rrms * rms_norm_weight[elem1].to(Float32) + for e in cutlass.range_constexpr(self.elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) k_cache_u16 = cute.recast_tensor(k_cache, Uint16) k_cache_u32 = cute.recast_tensor(k_cache, Uint32) @@ -287,31 +313,53 @@ class SparseAttnCompressNormRopeStoreC4Kernel: + kv_offset * Int64(self.scale_dim) ) - if warp_id == self.nope_blocks: - pair_idx = lane_id + if group_idx == self.nope_blocks: compressed_pos = (position // Int64(self.compress_ratio)) * Int64( self.compress_ratio ) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x0 * cos_v - x1 * sin_v - imag = x0 * sin_v + x1 * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = packed + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(self.rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = ( + value_base + + Int64(self.nope_dim) + + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64) + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed else: - q_packed = _fp32x2_to_bf16x2(x0, x1) - q0, q1 = _bf16x2_to_fp32(q_packed) - abs0 = cute.math.absf(q0) - abs1 = cute.math.absf(q1) - local_absmax = cute.arch.fmax(abs0, abs1) + q = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + local_absmax = Float32(0.0) + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + q0, q1 = _bf16x2_to_fp32(q_packed) + q[elem] = q0 + q[elem + 1] = q1 + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), + ) absmax = local_absmax - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) + group_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - self.lanes_per_group) << 8 + | (cute.arch.WARP_SIZE - 1) + ) + for step in cutlass.range_constexpr(self.scale_reduce_steps): + offset = const_expr(self.scale_reduce_offset >> step) absmax = cute.arch.fmax( - absmax, cute.arch.shuffle_sync_bfly(absmax, offset) + absmax, + cute.arch.shuffle_sync_bfly( + absmax, + offset=offset, + mask_and_clamp=group_mask_and_clamp, + ), ) scale_raw = cute.arch.fmax( Float32(self.min_scale), @@ -320,22 +368,22 @@ class SparseAttnCompressNormRopeStoreC4Kernel: 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( - cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cute.arch.fmin( - cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to( - Int64 - ) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if lane_id == 0: - k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8) - if warp_id == 0: + for pair in cutlass.range_constexpr(self.elems_per_lane // 2): + elem = const_expr(pair * 2) + y0 = cutlass.min( + cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), + Float32(self.fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if group_lane == 0: + k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8) + if group_idx == 0: k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( 0 ) @@ -462,11 +510,11 @@ class SparseAttnCompressNormRopeStoreC4Kernel: class SparseAttnCompressKernel: head_tile = 64 - rows_per_warp = 8 + rows_per_warp = 16 row_pairs_per_warp = rows_per_warp // 2 elems_per_lane = 4 lanes_per_row = head_tile // elems_per_lane - num_warps = 16 + num_warps = 8 stats_warp_stride = num_warps + 1 tb_size = num_warps * 32 rcp_ln2 = 1.4426950408889634 @@ -715,8 +763,8 @@ class SparseAttnCompressKernel: local_warp_max = s_max[out_lane, out_elem, final_lane] global_max = local_warp_max - for step in cutlass.range_constexpr(4): - offset = const_expr(8 >> step) + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) global_max = cute.arch.fmax( global_max, cute.arch.shuffle_sync_bfly( @@ -732,8 +780,8 @@ class SparseAttnCompressKernel: ) global_sum = s_sum[out_lane, out_elem, final_lane] * scale global_product = s_product[out_lane, out_elem, final_lane] * scale - for step in cutlass.range_constexpr(4): - offset = const_expr(8 >> step) + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) global_sum += cute.arch.shuffle_sync_bfly( global_sum, offset=offset, @@ -978,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 697156a5b4d..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,7 +696,114 @@ class DynamicVideoBackend(VideoBackend): ) -@VIDEO_LOADER_REGISTRY.register("molmo2") +@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 + # does not report it (common for WebM/streaming inputs). + if source.duration: + return source + if source.original_fps > 0: + max_frame_idx = source.total_frames_num - 1 + duration = round(max_frame_idx / source.original_fps) + 1 + else: + duration = 0 + return VideoSourceMetadata( + source.total_frames_num, source.original_fps, duration + ) + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + total_frames_num = source.total_frames_num + duration = source.duration + original_fps = source.original_fps + target_fps = target.fps + max_frame_idx = source.total_frames_num - 1 + max_frames = kwargs.get("max_frames", 640) + + duration = duration or round(max_frame_idx / original_fps) + 1 + + extract_t = int(duration * target_fps) + extract_t = min(extract_t, max_frames) + + duration_per_frame = 1 / original_fps + timestamps = [i * duration_per_frame for i in range(total_frames_num)] + + if total_frames_num < extract_t: + frame_indices = [ + math.floor(i * total_frames_num / extract_t) for i in range(extract_t) + ] + else: + frame_indices = [] + current_second = 0.0 + inv_fps = 1 / target_fps + for frame_index in range(total_frames_num): + if timestamps[frame_index] >= current_second: + current_second += inv_fps + frame_indices.append(frame_index) + if current_second >= duration - inv_fps: + break + + if len(frame_indices) < extract_t: + if len(frame_indices) == 0: + start, end = 0, max(total_frames_num - 1, 0) + else: + start, end = frame_indices[0], frame_indices[-1] + frame_indices = np.linspace(start, end, extract_t, dtype=int).tolist() + elif len(frame_indices) > extract_t: + frame_indices = np.linspace( + 0, total_frames_num - 1, extract_t, dtype=int + ).tolist() + + seen, uniq = set(), [] + for idx in frame_indices: + if idx not in seen: + seen.add(idx) + uniq.append(idx) + + return uniq + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int = 2, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: Literal["opencv", "pyav"] = "opencv", + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + frames, metadata = super().load_bytes( + data, + num_frames=num_frames, + fps=fps, + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + # Ensure even frame count — matches HF's sample_frames even-padding + # and _preprocess temporal_patch_size divisibility check. + if frames.shape[0] & 1: + frames = np.concatenate([frames, frames[-1:]], axis=0) + return frames, metadata + + +@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/cpu.py b/vllm/platforms/cpu.py index 999bcfcc6db..cf4319ac722 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -408,6 +408,10 @@ class CpuPlatform(Platform): def support_hybrid_kv_cache(cls) -> bool: return True + @classmethod + def num_compute_units(cls, device_id: int = 0) -> int: + return torch.get_num_threads() + @classmethod def import_kernels(cls) -> None: if Platform.get_cpu_architecture() in (CpuArchEnum.X86,): diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 58cef2ec976..57814d29bef 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -809,6 +809,22 @@ class NvmlCudaPlatform(CudaPlatformBase): logger.warning("Failed to get NUMA nodes for GPUs: %s", e) return None + @classmethod + @with_nvml_context + def get_all_gpu_pci_bus_ids(cls) -> dict[int, str]: + """Query NVML for GPU index -> PCI bus ID mapping.""" + out: dict[int, str] = {} + for idx in range(pynvml.nvmlDeviceGetCount()): + handle = pynvml.nvmlDeviceGetHandleByIndex(idx) + pci_info = pynvml.nvmlDeviceGetPciInfo(handle) + bus_id = pci_info.busId + if isinstance(bus_id, bytes): + bus_id = bus_id.decode("utf-8") + out[idx] = bus_id.rstrip("\x00") + if not out: + raise RuntimeError("NVML returned no GPU PCI bus ID rows") + return out + @classmethod @with_nvml_context def log_warnings(cls): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index cf774b7bda9..546361229e1 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: """ @@ -381,6 +389,19 @@ class Platform: """Get the total memory of a device in bytes.""" raise NotImplementedError + @classmethod + def get_all_gpu_pci_bus_ids(cls) -> dict[int, str]: + """Return a mapping of device index to PCI bus ID string. + + Used by ``VLLM_GPU_NIC_PCIE_MAPPING`` for RDMA NIC selection. + Subclasses should override with platform-specific discovery + (e.g. pynvml for CUDA). + """ + raise NotImplementedError( + "VLLM_GPU_NIC_PCIE_MAPPING is not supported on the " + f"current platform ({cls.device_name})" + ) + @classmethod def inference_mode(cls): """A device-specific wrapper of `torch.inference_mode`. diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index c75d68954f8..89471e844d8 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -190,6 +190,9 @@ def _get_gcn_arch() -> str: _GCN_ARCH = _get_gcn_arch() _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) +_ON_GFX11 = "gfx11" in _GCN_ARCH +_ON_GFX1100 = "gfx1100" in _GCN_ARCH +_ON_GFX1151 = "gfx1151" in _GCN_ARCH _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) @@ -273,6 +276,18 @@ def on_gfx1x() -> bool: return _ON_GFX1X +def on_gfx11() -> bool: + return _ON_GFX11 + + +def on_gfx1100() -> bool: + return _ON_GFX1100 + + +def on_gfx1151() -> bool: + return _ON_GFX1151 + + def on_gfx12x() -> bool: return _ON_GFX12X 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/plugins/lora_resolvers/hf_hub_resolver.py b/vllm/plugins/lora_resolvers/hf_hub_resolver.py index c152d4c2946..57ccdc2e48d 100644 --- a/vllm/plugins/lora_resolvers/hf_hub_resolver.py +++ b/vllm/plugins/lora_resolvers/hf_hub_resolver.py @@ -3,13 +3,12 @@ import asyncio import os -from huggingface_hub import HfApi, snapshot_download - import vllm.envs as envs from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.lora.resolver import LoRAResolverRegistry from vllm.plugins.lora_resolvers.filesystem_resolver import FilesystemResolver +from vllm.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -49,7 +48,7 @@ class HfHubResolver(FilesystemResolver): return None repo_path = await asyncio.to_thread( - snapshot_download, + hf_api().snapshot_download, repo_id=maybe_repo, allow_patterns=f"{maybe_subpath}/*" if maybe_subpath != "." else "*", ) @@ -110,7 +109,10 @@ class HfHubResolver(FilesystemResolver): Args: repo_name: Name of the HF hub repo to inspect. """ - repo_files = await asyncio.to_thread(HfApi().list_repo_files, repo_id=repo_name) + repo_files = await asyncio.to_thread( + hf_api().list_repo_files, + repo_id=repo_name, + ) adapter_dirs = { os.path.dirname(name) for name in repo_files diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 41d8c0075fb..9fab3aff04e 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -104,6 +104,9 @@ class BaseRenderer(ABC, Generic[_T]): self._process_multimodal_async = make_async( self._process_multimodal, executor=self._mm_executor ) + self._safe_load_prompt_embeds_async = make_async( + safe_load_prompt_embeds, executor=self._executor + ) if mm_registry.supports_multimodal_inputs(config.model_config): mm_processor_cache = mm_registry.processor_cache_from_config(config) @@ -376,11 +379,28 @@ class BaseRenderer(ABC, Generic[_T]): return [self.render_prompt(prompt) for prompt in prompts] + async def _render_prompt_async( + self, + prompt: DictPrompt | bytes, + ) -> DictPrompt: + if isinstance(prompt, bytes): + embeds = await self._safe_load_prompt_embeds_async( + self.model_config, prompt + ) + return EmbedsPrompt(prompt_embeds=embeds) + + return prompt + async def render_prompts_async( self, prompts: Sequence[DictPrompt | bytes], ) -> list[DictPrompt]: - return self.render_prompts(prompts) + if len(prompts) == 0: + raise ValueError("You must pass at least one prompt") + + return await asyncio.gather( + *(self._render_prompt_async(prompt) for prompt in prompts) + ) @abstractmethod def render_messages( 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/sequence.py b/vllm/sequence.py index 17630623646..c90531935bb 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -3,15 +3,9 @@ """Sequence and its related classes.""" from dataclasses import dataclass -from typing import TYPE_CHECKING, Any import torch -if TYPE_CHECKING: - from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput -else: - KVConnectorOutput = Any - # cannot use msgspec.Struct here because Dynamo does not support it @dataclass @@ -19,24 +13,19 @@ class IntermediateTensors: """For all pipeline stages except the last, we need to return the hidden states and residuals to be sent to the next stage. This data structure contains the hidden states and residuals for a request. - - Each stage also needs to handle its own kv_connector_output. """ tensors: dict[str, torch.Tensor] - kv_connector_output: KVConnectorOutput | None def __init__( self, tensors: dict[str, torch.Tensor], - kv_connector_output: KVConnectorOutput | None = None, ) -> None: # manually define this function, so that # Dynamo knows `IntermediateTensors()` comes from this file. # Otherwise, dataclass will generate this function by evaluating # a string, and we will lose the information about the source file. self.tensors = tensors - self.kv_connector_output = kv_connector_output def __getitem__(self, key: str | slice): if isinstance(key, str): diff --git a/vllm/tokenizers/grok2.py b/vllm/tokenizers/grok2.py index 61fa1107e2a..612af537408 100644 --- a/vllm/tokenizers/grok2.py +++ b/vllm/tokenizers/grok2.py @@ -8,7 +8,6 @@ from collections.abc import Collection, Sequence, Set from pathlib import Path from typing import Any, Literal, overload -from huggingface_hub import hf_hub_download from huggingface_hub.utils import ( EntryNotFoundError, HfHubHTTPError, @@ -20,6 +19,7 @@ from transformers.utils import chat_template_utils as hf_chat_utils from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.logger import init_logger +from vllm.transformers_utils.repo_utils import hf_api from .protocol import TokenizerLike @@ -70,7 +70,7 @@ def _maybe_load_tokenizer_config( return {} try: - config_file = hf_hub_download( + config_file = hf_api().hf_hub_download( repo_id=repo_id, filename="tokenizer_config.json", revision=revision, @@ -208,7 +208,7 @@ class Grok2Tokenizer(TokenizerLike): repo_id = None else: vocab_file = Path( - hf_hub_download( + hf_api().hf_hub_download( repo_id=str(path_or_repo_id), filename="tokenizer.tok.json", revision=revision, diff --git a/vllm/tokenizers/kimi_audio.py b/vllm/tokenizers/kimi_audio.py index d2b0a2a557e..10f57a7f7de 100644 --- a/vllm/tokenizers/kimi_audio.py +++ b/vllm/tokenizers/kimi_audio.py @@ -10,13 +10,13 @@ from typing import Any, overload import pybase64 import tiktoken -from huggingface_hub import hf_hub_download from transformers import AddedToken, BatchEncoding from transformers.utils import chat_template_utils as hf_chat_utils from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.logger import init_logger from vllm.tokenizers.protocol import TokenizerLike +from vllm.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -78,7 +78,7 @@ class KimiAudioTokenizer(TokenizerLike): # Try to download tiktoken.model or tokenizer.model try: - vocab_path = hf_hub_download( + vocab_path = hf_api().hf_hub_download( repo_id=repo_id, filename="tiktoken.model", revision=revision, @@ -87,7 +87,7 @@ class KimiAudioTokenizer(TokenizerLike): vocab_file = Path(vocab_path) except Exception: try: - vocab_path = hf_hub_download( + vocab_path = hf_api().hf_hub_download( repo_id=repo_id, filename="tokenizer.model", revision=revision, @@ -101,7 +101,7 @@ class KimiAudioTokenizer(TokenizerLike): # Also download tokenizer_config.json if available with contextlib.suppress(Exception): - hf_hub_download( + hf_api().hf_hub_download( repo_id=repo_id, filename="tokenizer_config.json", revision=revision, diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index d72772ea00c..7578d3b43ab 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -38,7 +38,7 @@ logger = init_logger(__name__) # temporary workaround and better long term solutions are: # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better) # - Fix tokenizer_class on the hub for the affected models (best) -_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl"} +_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"} _VLLM_TOKENIZERS = { "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), @@ -249,6 +249,10 @@ def get_tokenizer( tokenizer_cls_ = tokenizer_cls tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) + if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: + from vllm.tokenizers.hf import get_cached_tokenizer + + tokenizer = get_cached_tokenizer(tokenizer) if not tokenizer.is_fast: logger.warning( "Using a slow tokenizer. This might cause a significant " 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/ernie45_tool_parser.py b/vllm/tool_parsers/ernie45_tool_parser.py index 9722dddf734..f22eaca1f80 100644 --- a/vllm/tool_parsers/ernie45_tool_parser.py +++ b/vllm/tool_parsers/ernie45_tool_parser.py @@ -34,7 +34,6 @@ class Ernie45ToolParser(ToolParser): abc\n\n\n\n\ndef\n\n """ super().__init__(tokenizer, tools) - self.current_tool_name_sent = False self.prev_tool_call_arr: list[dict] = [] self.current_tool_id = -1 self.streamed_args_for_tool: list[str] = [] diff --git a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py index 29b2a5eae27..9723ef45d24 100644 --- a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py +++ b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py @@ -38,7 +38,6 @@ class HunyuanA13BToolParser(ToolParser): # Initialize state for streaming mode self.prev_tool_calls: list[dict] = [] self.current_tool_id = -1 - self.current_tool_name_sent = False self.streamed_args: list[str] = [] # Track arguments sent for each tool # For backward compatibility with tests @@ -262,7 +261,6 @@ class HunyuanA13BToolParser(ToolParser): ) else: self.streaming_state["sent_tools"][0]["sent_name"] = True - self.current_tool_name_sent = True return delta return None @@ -306,7 +304,6 @@ class HunyuanA13BToolParser(ToolParser): ] ) self.streaming_state["sent_tools"][current_idx]["sent_name"] = True - self.current_tool_name_sent = True while len(self.streamed_args) <= current_idx: self.streamed_args.append("") return delta diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py index 809a85ce417..496deb4f2d5 100644 --- a/vllm/tool_parsers/hy_v3_tool_parser.py +++ b/vllm/tool_parsers/hy_v3_tool_parser.py @@ -246,7 +246,6 @@ class HYV3ToolParser(ToolParser): def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) - self.current_tool_name_sent: bool = False self.prev_tool_call_arr: list[dict] = [] self.current_tool_id: int = -1 self.streamed_args_for_tool: list[ diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py index 2dc262bba2e..f2fa3ce9983 100644 --- a/vllm/tool_parsers/phi4mini_tool_parser.py +++ b/vllm/tool_parsers/phi4mini_tool_parser.py @@ -47,7 +47,6 @@ class Phi4MiniJsonToolParser(ToolParser): # streaming mode self.prev_tool_call_arr: list[dict[str, Any]] = [] self.current_tool_id: int = -1 - self.current_tool_name_sent: bool = False self.streamed_args_for_tool: list[ str ] = [] # map what has been streamed for each tool so far to a list diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 0969d816902..8339c183c0f 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -12,7 +12,7 @@ from typing import Any, Literal, TypeAlias import huggingface_hub import torch -from huggingface_hub import constants, get_safetensors_metadata +from huggingface_hub import constants from packaging.version import Version from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE from transformers import GenerationConfig, PretrainedConfig @@ -43,6 +43,7 @@ from .gguf_utils import ( from .repo_utils import ( file_or_path_exists, get_hf_file_to_dict, + hf_api, list_repo_files, try_get_local_file, with_retry, @@ -113,9 +114,9 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( kimi_k25="KimiK25Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) - jais="JAISConfig", mlp_speculator="MLPSpeculatorConfig", medusa="MedusaConfig", + mellum="MellumConfig", midashenglm="MiDashengLMConfig", moondream3="Moondream3Config", eagle="EAGLEConfig", @@ -1144,7 +1145,7 @@ def try_get_safetensors_metadata( revision: str | None = None, ): get_safetensors_metadata_partial = partial( - get_safetensors_metadata, model, revision=revision + hf_api().get_safetensors_metadata, model, revision=revision ) try: @@ -1238,7 +1239,7 @@ def get_safetensors_params_metadata( # local HF cache: weights may already be cached from a prior run, and # weight loading itself uses the same cache. try: - local_dir = huggingface_hub.snapshot_download( + local_dir = hf_api().snapshot_download( repo_id=model, revision=revision, allow_patterns=["*.safetensors"], diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 4a90655dcc0..71f7723e4c8 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -46,10 +46,10 @@ _CLASS_TO_MODULE: dict[str, str] = { # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the # `FalconConfig` class from the official HuggingFace transformers library. "RWConfig": "vllm.transformers_utils.configs.falcon", - "JAISConfig": "vllm.transformers_utils.configs.jais", "LagunaConfig": "vllm.transformers_utils.configs.laguna", "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe", "MedusaConfig": "vllm.transformers_utils.configs.medusa", + "MellumConfig": "vllm.transformers_utils.configs.mellum", "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm", "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator", "Moondream3Config": "vllm.transformers_utils.configs.moondream3", @@ -115,10 +115,10 @@ __all__ = [ "HyperCLOVAXConfig", "IsaacConfig", "RWConfig", - "JAISConfig", "LagunaConfig", "Lfm2MoeConfig", "MedusaConfig", + "MellumConfig", "MiDashengLMConfig", "MLPSpeculatorConfig", "Moondream3Config", diff --git a/vllm/transformers_utils/configs/jais.py b/vllm/transformers_utils/configs/jais.py deleted file mode 100644 index 6b581bf1877..00000000000 --- a/vllm/transformers_utils/configs/jais.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2023 The OpenAI Team Authors and HuggingFace Inc. team. -# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. -# Copyright 2023 Cerebras Systems. -# -# 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. -"""JAIS configuration""" - -from transformers.configuration_utils import PretrainedConfig -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class JAISConfig(PretrainedConfig): - """ - This is the configuration class to store the configuration of a - [`JAISModel`]. It is used to instantiate a JAIS model according to the - specified arguments, defining the model architecture. - - Configuration objects inherit from [`PretrainedConfig`] and can be used - to control the model outputs. Read the documentation from - [`PretrainedConfig`] for more information. - - - Args: - vocab_size (`int`, *optional*, defaults to 50257): - Vocabulary size of the JAIS model. Defines the number of different - tokens that can be represented by the - `inputs_ids` passed when calling [`JAISModel`]. - n_positions (`int`, *optional*, defaults to 1024): - The maximum sequence length that this model might ever be used - with. Typically set this to something large just in case - (e.g., 512 or 1024 or 2048). - n_embd (`int`, *optional*, defaults to 768): - Dimensionality of the embeddings and hidden states. - n_layer (`int`, *optional*, defaults to 12): - Number of hidden layers in the Transformer encoder. - n_head (`int`, *optional*, defaults to 12): - Number of attention heads for each attention layer in the - Transformer encoder. - n_inner (`int`, *optional*, defaults to None): - Dimensionality of the inner feed-forward layers. `None` will set - it to 4 times n_embd - activation_function (`str`, *optional*, defaults to `"gelu"`): - Activation function, to be selected in the list - `["relu", "silu", "gelu", "tanh", "gelu_new", "swiglu"]`. - resid_pdrop (`float`, *optional*, defaults to 0.1): - The dropout probability for all fully connected layers in - the embeddings, encoder, and pooler. - embd_pdrop (`float`, *optional*, defaults to 0.1): - The dropout ratio for the embeddings. - attn_pdrop (`float`, *optional*, defaults to 0.1): - The dropout ratio for the attention. - layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): - The epsilon to use in the layer normalization layers. - initializer_range (`float`, *optional*, defaults to 0.02): - The standard deviation of the truncated_normal_initializer for - initializing all weight matrices. - scale_attn_weights (`bool`, *optional*, defaults to `True`): - Scale attention weights by dividing by sqrt(hidden_size).. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values - attentions (not used by all models). - scale_attn_by_inverse_layer_idx (`bool`, *optional*, default `True`): - Whether to additionally scale attention weights - by `1 / layer_idx + 1`. - reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`): - Whether to scale keys (K) prior to computing attention - (dot-product) - and upcast attention dot-product/softmax to float() when training - with mixed precision. - position_embedding_type (`str`, *optional*, defaults to `"learned"`): - Positional embedding can be either `"alibi"` or `"learned"`. - mup_width_scale (`float`, *optional*, defaults to 1.0): - muP parameter to scale learning rate and initializers. Calculated - as (`d_model,0 / d_model`), where - `d_model` is the model's width and `d_model,0` is the proxy - model's width. - mup_embeddings_scale (`float`, *optional*, defaults to 1.0): - muP parameter to scale token and position embeddings. - mup_output_alpha (`float`, *optional*, defaults to 1.0): - muP parameter to scale output logits - (`output_logits_scale = mup_output_alpha * mup_width_scale`). - mup_scale_qk_dot_by_d (`bool`, *optional*, defaults to `False`): - Scale attention weights by dividing by hidden_size instead of - sqrt(hidden_size). Need to set scale_attn_weights to `True` as - well. - alibi_scaling (`dict`, *optional*): - Dictionary containing the scaling configuration for ALiBi - embeddings. Currently only supports linear - scaling strategy. Can specify either the scaling `factor` (must be - a float greater than 1) for fixed scaling - or `train_seq_len` for dynamic scaling on input samples with - sequence length > `train_seq_len`. The expected - formats are `{"type": strategy name, "factor": scaling factor}` or - `{"type": strategy name, - "train_seq_len": training sequence length}`. - architectures (`list`, *optional*, defaults to ['JAISLMHeadModel']): - architecture names for Jais. - - Example: - - ```python - >>> from transformers import JAISConfig, JAISModel - - >>> # Initializing a JAIS configuration - >>> configuration = JAISConfig() - - >>> # Initializing a model (with random weights) from the configuration - >>> model = JAISModel(configuration) - - >>> # Accessing the model configuration - >>> configuration = model.config - ```""" - - model_type = "jais" - keys_to_ignore_at_inference = ["past_key_values"] - attribute_map = { - "hidden_size": "n_embd", - "max_position_embeddings": "n_positions", - "num_attention_heads": "n_head", - "num_hidden_layers": "n_layer", - } - - def __init__( - self, - vocab_size=50257, - n_positions=1024, - n_embd=768, - n_layer=12, - n_head=12, - n_inner=None, - activation_function="gelu_new", - resid_pdrop=0.1, - embd_pdrop=0.1, - attn_pdrop=0.1, - layer_norm_epsilon=1e-5, - initializer_range=0.02, - scale_attn_weights=True, - use_cache=True, - bos_token_id=50256, - eos_token_id=50256, - scale_attn_by_inverse_layer_idx=False, - reorder_and_upcast_attn=False, - position_embedding_type="learned", - mup_width_scale=1.0, - mup_embeddings_scale=1.0, - mup_output_alpha=1.0, - mup_scale_qk_dot_by_d=False, - alibi_scaling=None, - architectures=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.n_positions = n_positions - self.n_embd = n_embd - self.n_layer = n_layer - self.n_head = n_head - self.n_inner = n_inner - self.activation_function = activation_function - self.resid_pdrop = resid_pdrop - self.embd_pdrop = embd_pdrop - self.attn_pdrop = attn_pdrop - self.layer_norm_epsilon = layer_norm_epsilon - self.initializer_range = initializer_range - self.scale_attn_weights = scale_attn_weights - self.use_cache = use_cache - self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx - self.reorder_and_upcast_attn = reorder_and_upcast_attn - - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - - self.position_embedding_type = position_embedding_type - self.mup_width_scale = mup_width_scale - self.mup_embeddings_scale = mup_embeddings_scale - self.mup_output_alpha = mup_output_alpha - self.mup_scale_qk_dot_by_d = mup_scale_qk_dot_by_d - - self.alibi_scaling = alibi_scaling - self._alibi_scaling_validation() - if architectures is None: - architectures = ["JAISLMHeadModel"] - - super().__init__( - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - architectures=architectures, - **kwargs, - ) - - def _alibi_scaling_validation(self): - """ - Validate the `alibi_scaling` configuration. - """ - if self.alibi_scaling is None: - return - - if not isinstance(self.alibi_scaling, dict) or len(self.alibi_scaling) != 2: - raise ValueError( - "`alibi_scaling` must be a dictionary with two fields, " - "`type` and `factor` or `type` and `train_seq_len`, " - f"got {self.alibi_scaling}" - ) - alibi_scaling_type = self.alibi_scaling.get("type", None) - alibi_scaling_factor = self.alibi_scaling.get("factor", None) - alibi_dynamic_scaling = self.alibi_scaling.get("train_seq_len", None) - if alibi_scaling_type is None or alibi_scaling_type != "linear": - raise ValueError( - f"`alibi_scaling`'s type field must be 'linear', " - f"got {alibi_scaling_type}" - ) - if ( - alibi_scaling_factor is not None - and not isinstance(alibi_scaling_factor, float) - or (alibi_scaling_factor is not None and alibi_scaling_factor <= 1.0) - ): - raise ValueError( - f"`alibi_scaling`'s factor field must be a float > 1.0, " - f"got {alibi_scaling_factor}" - ) - if ( - alibi_dynamic_scaling is not None - and not isinstance(alibi_dynamic_scaling, int) - or (alibi_dynamic_scaling is not None and alibi_dynamic_scaling <= 1) - ): - raise ValueError( - f"`alibi_scaling`'s `train_seq_len` field must be an " - f"integer > 1, got {alibi_dynamic_scaling}" - ) diff --git a/vllm/transformers_utils/configs/mellum.py b/vllm/transformers_utils/configs/mellum.py new file mode 100644 index 00000000000..2bed53394b2 --- /dev/null +++ b/vllm/transformers_utils/configs/mellum.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import Qwen3MoeConfig + + +class MellumConfig(Qwen3MoeConfig): + model_type = "mellum" diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 35fa1313d1e..85452197535 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -502,6 +502,11 @@ class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0) +class Step3p5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): + def get_num_hidden_layers(self) -> int: + return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) + + class PanguUltraMoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_num_hidden_layers(self) -> int: return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) @@ -543,31 +548,32 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase): # hf_config.model_type -> convertor class MODEL_ARCH_CONFIG_CONVERTORS = { "cohere_asr": CohereAsrModelArchConfigConvertor, - "mamba": MambaModelArchConfigConvertor, - "falcon_mamba": MambaModelArchConfigConvertor, - "timm_wrapper": TerratorchModelArchConfigConvertor, - "medusa": MedusaModelArchConfigConvertor, - "zamba2": Zamba2ModelArchConfigConvertor, - "mpt": MPTModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, - "falcon": FalconModelArchConfigConvertor, - "gemma4": Gemma4ModelArchConfigConvertor, - "gemma4_text": Gemma4ModelArchConfigConvertor, - "gemma4_mtp": Gemma4MTPModelArchConfigConvertor, - "RefinedWeb": FalconModelArchConfigConvertor, - "RefinedWebModel": FalconModelArchConfigConvertor, - "nemotron-nas": NemotronNasModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, - "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, - "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, + "ernie_mtp": ErnieMTPModelArchConfigConvertor, + "falcon": FalconModelArchConfigConvertor, + "falcon_mamba": MambaModelArchConfigConvertor, + "gemma4": Gemma4ModelArchConfigConvertor, + "gemma4_mtp": Gemma4MTPModelArchConfigConvertor, + "gemma4_text": Gemma4ModelArchConfigConvertor, + "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor, + "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor, + "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor, + "mamba": MambaModelArchConfigConvertor, + "medusa": MedusaModelArchConfigConvertor, "mimo_mtp": MimoMTPModelArchConfigConvertor, "mimo_v2": MimoV2ModelArchConfigConvertor, "mimo_v2_flash": MimoV2ModelArchConfigConvertor, "mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor, "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor, - "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor, - "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor, - "ernie_mtp": ErnieMTPModelArchConfigConvertor, + "mpt": MPTModelArchConfigConvertor, + "nemotron-nas": NemotronNasModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, - "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor, + "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, + "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, + "RefinedWeb": FalconModelArchConfigConvertor, + "RefinedWebModel": FalconModelArchConfigConvertor, + "step3p5_mtp": Step3p5MTPModelArchConfigConvertor, + "timm_wrapper": TerratorchModelArchConfigConvertor, + "zamba2": Zamba2ModelArchConfigConvertor, } 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/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 704d505617f..8385057e911 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -12,8 +12,7 @@ from pathlib import Path from typing import TypeVar import huggingface_hub -from huggingface_hub import hf_hub_download, try_to_load_from_cache -from huggingface_hub import list_repo_files as hf_list_repo_files +from huggingface_hub import HfApi, try_to_load_from_cache from huggingface_hub.utils import ( EntryNotFoundError, HfHubHTTPError, @@ -24,9 +23,31 @@ from huggingface_hub.utils import ( from vllm import envs from vllm.logger import init_logger +from vllm.version import __version__ as VLLM_VERSION logger = init_logger(__name__) +_hf_api: HfApi | None = None + + +def hf_api() -> HfApi: + """Return a shared HfApi instance tagged with vLLM's library info.""" + global _hf_api + if _hf_api is None: + _hf_api = HfApi( + library_name="vllm", + library_version=VLLM_VERSION, + ) + return _hf_api + + +def hf_fs() -> "huggingface_hub.HfFileSystem": + """Return a fresh HfFileSystem tagged with vLLM's library info.""" + return huggingface_hub.HfFileSystem( + library_name="vllm", + library_version=VLLM_VERSION, + ) + _R = TypeVar("_R") @@ -80,7 +101,7 @@ def list_repo_files( revision=revision, token=os.getenv("MODELSCOPE_API_TOKEN", None), ) - return hf_list_repo_files( + return hf_api().list_repo_files( repo_id, revision=revision, repo_type=repo_type, token=token ) except huggingface_hub.errors.OfflineModeIsEnabled: @@ -215,9 +236,10 @@ def get_model_path(model: str | Path, revision: str | None = None): return snapshot_download(model_id=model, **common_kwargs) - from huggingface_hub import snapshot_download - - return snapshot_download(repo_id=model, **common_kwargs) + return hf_api().snapshot_download( + repo_id=model, + **common_kwargs, + ) def _try_download_from_hf_hub( @@ -231,7 +253,13 @@ def _try_download_from_hf_hub( if Path(model).is_dir(): return None try: - return Path(hf_hub_download(model, file_name, revision=revision)) + return Path( + hf_api().hf_hub_download( + model, + file_name, + revision=revision, + ) + ) except huggingface_hub.errors.OfflineModeIsEnabled: return None except ( 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 e97228bfa60..c37b3b6c70c 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -392,12 +392,24 @@ class LazyLoader(ModuleType): # Optional dependency detection utilities @cache def _has_module(module_name: str) -> bool: - """Return True if *module_name* can be found in the current environment. + """Return True if *module_name* can be imported in the current environment. - The result is cached so that subsequent queries for the same module incur - no additional overhead. + Uses ``importlib.util.find_spec`` as a fast pre-check, then performs a + trial import to verify that native dependencies (shared libraries, etc.) + are also satisfied. Any failure during the trial import is treated as the + module being unavailable. The result is cached so that subsequent queries + for the same module incur no additional overhead. """ - return importlib.util.find_spec(module_name) is not None + try: + if importlib.util.find_spec(module_name) is None: + return False + importlib.import_module(module_name) + except Exception: + logger.warning( + "Module %s was found but failed to import", module_name, exc_info=True + ) + return False + return True def has_deep_ep() -> bool: diff --git a/vllm/utils/numa_utils.py b/vllm/utils/numa_utils.py index 4e1addad980..6e4b4b471c1 100644 --- a/vllm/utils/numa_utils.py +++ b/vllm/utils/numa_utils.py @@ -14,7 +14,7 @@ import subprocess from contextlib import contextmanager from functools import cache from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import psutil @@ -112,6 +112,117 @@ def get_auto_numa_nodes() -> list[int] | None: return numa_nodes +# PCT (Priority Core Turbo) auto-detection workaround for Granite Rapids +# Xeon SKUs. +# +# Background: +# * The Linux kernel does not expose PCT priority-core membership via any +# unprivileged sysfs path. The official interface +# (/dev/isst_interface, used by `intel-speed-select`) is root-only, +# which is a non-starter in most production deployments (shared +# clusters, prebuilt containers, managed cloud). +# * Even recent stable kernels (e.g. 6.14, March 2025) do not yet +# preferentially schedule work on PCT priority cores, so vLLM cannot +# just "let the scheduler handle it". +# +# Empirical heuristic (DGX B300 / Xeon 6776P, the SKU we measured): +# * /proc/cpuinfo `model name` contains the SKU number. +# * cpu0 is a PCT priority core on these SKUs, so it reports the +# priority-cohort CPPC `highest_perf` (the value matches the SKU's +# "Max PCT core frequency" in 100 MHz units, e.g. 4.6 GHz -> 46). +# * Priority cores within each NUMA node satisfy `cpu_id % S in (0, 1)` +# intersected with the node's cpulist, where `S` is the SKU's logical +# CPUs per priority "group" (= total threads / 8 priority cores; 16 on +# 64-core SKUs, 18 on 72-core SKUs). +# +# SKU table: +# ``_PCT_CAPABLE_SKUS`` maps each known PCT-capable Granite Rapids part +# to a ``_PctSku(highest_perf, priority_stride)`` config: +# * highest_perf is the expected ``acpi_cppc/highest_perf`` on cpu0, +# derived from Intel ARK's "Max PCT core frequency" * 10 (CPPC max +# ratio reports in 100 MHz units). +# * priority_stride is the SKU's "Total Cores" / 4 (= total HT threads +# / 8 priority cores), used in the ``cpu_id % stride`` filter above. +# Values: +# * 6776P - 4.6 GHz, 64C/128T -> (46, 16) measured on DGX B300 +# * 6774P - 4.6 GHz, 64C/128T -> (46, 16) per Intel ARK, not measured +# * 6962P - 4.4 GHz, 72C/144T -> (44, 18) per Intel ARK, not measured +# The non-measured SKUs are listed best-effort: the gate fails closed +# (no PCT engagement) if a host's actual highest_perf doesn't match the +# table value, so adding entries is safe. If you have access to a 6962P +# or 6774P box and find a different value or cpu-id pattern, update the +# table below. +# +# This whole block is a stop-gap until the kernel exposes PCT membership +# in an unprivileged way; see the tracking issue linked from the PR. + + +class _PctSku(NamedTuple): + """Per-SKU config used by the PCT auto-detection gate.""" + + highest_perf: int + priority_stride: int + + +_PCT_CAPABLE_SKUS: dict[str, _PctSku] = { + "6776P": _PctSku(highest_perf=46, priority_stride=16), + "6774P": _PctSku(highest_perf=46, priority_stride=16), + "6962P": _PctSku(highest_perf=44, priority_stride=18), +} +_PCT_HIGHEST_PERF_PATH = "/sys/devices/system/cpu/cpu0/acpi_cppc/highest_perf" +_PROC_CPUINFO_PATH = "/proc/cpuinfo" + + +def _pct_sku_from_cpuinfo() -> _PctSku | None: + """Return the ``_PctSku`` config for this host's SKU, or None. + + Reads ``/proc/cpuinfo``'s ``model name`` and looks the SKU up in + ``_PCT_CAPABLE_SKUS``. Returns ``None`` when the host is not a known + PCT-capable Granite Rapids Xeon (or when ``/proc/cpuinfo`` is + unreadable). + """ + try: + with open(_PROC_CPUINFO_PATH) as f: + for line in f: + if not line.lstrip().lower().startswith("model name"): + continue + for sku, config in _PCT_CAPABLE_SKUS.items(): + if sku in line: + return config + except OSError: + return None + return None + + +@cache +def _pct_sku_config() -> _PctSku | None: + """Detect a PCT-capable Granite Rapids Xeon with PCT enabled. + + See the comment block above ``_PCT_CAPABLE_SKUS`` for the full context + (why we hard-code SKUs, why we read CPPC ``highest_perf``, etc.). + + Returns the matching ``_PctSku`` config when both gates hold: + * ``/proc/cpuinfo`` ``model name`` contains an SKU listed in + ``_PCT_CAPABLE_SKUS``. + * ``/sys/devices/system/cpu/cpu0/acpi_cppc/highest_perf`` matches + that SKU's expected ``highest_perf``. + Otherwise returns ``None`` and the caller falls back to the default + NUMA-node bind. + """ + sku = _pct_sku_from_cpuinfo() + if sku is None: + return None + + try: + with open(_PCT_HIGHEST_PERF_PATH) as f: + actual = int(f.read().strip()) + except (OSError, ValueError): + return None + if actual != sku.highest_perf: + return None + return sku + + def _get_gpu_index( parallel_config, local_rank: int, dp_local_rank: int | None = None ) -> int: @@ -156,55 +267,183 @@ def _get_numa_node(parallel_config, gpu_index: int) -> int: return numa_nodes[gpu_index] -def _get_cpu_binding(parallel_config, gpu_index: int) -> str | None: +def _maybe_get_pct_cpu_binding(numa_nodes: list[int]) -> list[int] | None: + """Return the union of PCT priority cores across ``numa_nodes`` (or None). + + PCT (Priority Core Turbo) lets a subset of cores boost above the rest; + we want workers and the EngineCore on those cores. The Linux kernel does + not expose PCT membership without root, so we use the empirical heuristic + documented above ``_PCT_CAPABLE_SKUS``: priority cores within each NUMA + node satisfy ``cpu_id % stride in (0, 1)`` intersected with the node's + ``cpulist``, where ``stride`` is the SKU's logical CPUs per priority + group (16 on 64-core SKUs, 18 on 72-core SKUs). Only triggers on the + SKUs in ``_PCT_CAPABLE_SKUS`` with the expected CPPC ``highest_perf`` + signal; on any other host it returns None and the caller falls back to + the default NUMA-node bind. + + Returns the sorted CPU ids as a ``list[int]``; the caller is expected + to format them for the chosen tool (e.g. comma-joined for + ``numactl --physcpubind``). + """ + sku = _pct_sku_config() + if sku is None: + return None + + from vllm.utils.cpu_resource_utils import parse_id_list + + stride = sku.priority_stride + union_cpus: set[int] = set() + for numa_node in numa_nodes: + cpulist_path = Path(f"/sys/devices/system/node/node{numa_node}/cpulist") + try: + cpulist_raw = cpulist_path.read_text().strip() + except OSError: + continue + if not cpulist_raw: + continue + try: + node_cpus = parse_id_list(cpulist_raw) + except ValueError: + continue + + priority = [cpu for cpu in node_cpus if cpu % stride in (0, 1)] + if not priority: + continue + union_cpus.update(priority) + logger.info( + "Detected PCT-capable Granite Rapids Xeon (stride=%d); " + "NUMA node %d priority cores: %s", + stride, + numa_node, + ",".join(str(c) for c in priority), + ) + + if not union_cpus: + return None + return sorted(union_cpus) + + +def _get_cpu_binding( + parallel_config, gpu_index: int, numa_nodes: list[int] +) -> str | None: + """Return the CPU list a process should be pinned to (or None).""" cpu_bindings = parallel_config.numa_bind_cpus if cpu_bindings is None: - return None + pct_cpus = _maybe_get_pct_cpu_binding(numa_nodes) + if pct_cpus is None: + return None + return ",".join(str(c) for c in pct_cpus) if gpu_index >= len(cpu_bindings): raise ValueError( f"GPU index {gpu_index} exceeds numa_bind_cpus size " f"{len(cpu_bindings)}. Ensure the binding lists cover every visible GPU." ) - return cpu_bindings[gpu_index] -def _get_numactl_args( - vllm_config: "VllmConfig", - local_rank: int, - dp_local_rank: int | None = None, - process_kind: str = "worker", -) -> str | None: - parallel_config = vllm_config.parallel_config - if not parallel_config.numa_bind: - return None - +def _get_numactl_worker_args( + parallel_config, local_rank: int, dp_local_rank: int | None = None +) -> str: + """Compute the numactl args for a single TP/PP worker subprocess.""" gpu_index = _get_gpu_index(parallel_config, local_rank, dp_local_rank) numa_node = _get_numa_node(parallel_config, gpu_index) - cpu_binding = _get_cpu_binding(parallel_config, gpu_index) + cpu_binding = _get_cpu_binding(parallel_config, gpu_index, [numa_node]) if cpu_binding is not None: - bind_arg = f"--physcpubind={cpu_binding}" logger.info( - "Binding %s subprocess (local_rank=%s, gpu_index=%s) to CPUs %s and NUMA node %s", # noqa: E501 - process_kind, + "Binding worker subprocess (local_rank=%s, gpu_index=%s) to CPUs %s and NUMA node %s", # noqa: E501 local_rank, gpu_index, cpu_binding, numa_node, ) - else: - bind_arg = f"--cpunodebind={numa_node}" - logger.info( - "Binding %s subprocess (local_rank=%s, gpu_index=%s) to NUMA node %s", - process_kind, - local_rank, - gpu_index, - numa_node, - ) + return f"--physcpubind={cpu_binding} --membind={numa_node}" - return f"{bind_arg} --membind={numa_node}" + logger.info( + "Binding worker subprocess (local_rank=%s, gpu_index=%s) to NUMA node %s", + local_rank, + gpu_index, + numa_node, + ) + return f"--cpunodebind={numa_node} --membind={numa_node}" + + +def _get_enginecore_numa_nodes( + parallel_config, dp_local_rank: int | None = None +) -> list[int]: + """Return the sorted, unique NUMA nodes of the EngineCore's DP shard.""" + numa_nodes = parallel_config.numa_bind_nodes + if numa_nodes is None: + # Trigger auto-detection (it caches into parallel_config). + _get_numa_node(parallel_config, 0) + numa_nodes = parallel_config.numa_bind_nodes + + if ( + parallel_config.distributed_executor_backend not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + ): + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + + tp_pp_world_size = ( + parallel_config.pipeline_parallel_size + * parallel_config.tensor_parallel_size + ) + shard_start = dp_local_rank * tp_pp_world_size + shard_end = min(shard_start + tp_pp_world_size, len(numa_nodes)) + shard_indices: range | tuple[int, ...] = range(shard_start, shard_end) + else: + shard_indices = range(len(numa_nodes)) + + if not shard_indices: + return [numa_nodes[0]] + return sorted({numa_nodes[i] for i in shard_indices}) + + +def _get_numactl_enginecore_args( + parallel_config, local_rank: int, dp_local_rank: int | None = None +) -> str: + """Compute the numactl args for an EngineCore subprocess. + + ``--numa-bind-cpus`` is deliberately ignored here: the user provides a + per-worker CPU list, and binding EngineCore to any of those entries + would shrink its ``cpus_allowed`` below the strict-superset that the + workers' ``--physcpubind`` spawns require. We fall back to + ``--cpunodebind=`` instead, which is always a safe + superset. PCT auto-detection still applies when the user did not pass + ``--numa-bind-cpus`` (its priority-core union across the shard nodes + is also a safe superset by construction). + """ + shard_nodes = _get_enginecore_numa_nodes(parallel_config, dp_local_rank) + membind_arg = ",".join(str(n) for n in shard_nodes) + + pct_cpus = ( + None + if parallel_config.numa_bind_cpus is not None + else _maybe_get_pct_cpu_binding(shard_nodes) + ) + + if pct_cpus is not None: + cpu_binding = ",".join(str(c) for c in pct_cpus) + logger.info( + "Binding EngineCore subprocess (local_rank=%s) to CPUs %s " + "and NUMA nodes %s", + local_rank, + cpu_binding, + membind_arg, + ) + return f"--physcpubind={cpu_binding} --membind={membind_arg}" + + logger.info( + "Binding EngineCore subprocess (local_rank=%s) to NUMA nodes %s", + local_rank, + membind_arg, + ) + return f"--cpunodebind={membind_arg} --membind={membind_arg}" def _log_numactl_show(label: str) -> bool: @@ -242,13 +481,24 @@ def configure_subprocess( process_kind: str = "worker", ): """Temporarily replace the multiprocessing executable with a numactl wrapper.""" - numactl_args = _get_numactl_args( - vllm_config, local_rank, dp_local_rank, process_kind - ) - if numactl_args is None: + parallel_config = vllm_config.parallel_config + if not parallel_config.numa_bind: yield return + if process_kind == "EngineCore": + numactl_args = _get_numactl_enginecore_args( + parallel_config, local_rank, dp_local_rank + ) + elif process_kind == "worker": + numactl_args = _get_numactl_worker_args( + parallel_config, local_rank, dp_local_rank + ) + else: + raise ValueError( + f"Unknown process_kind {process_kind!r}; expected 'worker' or 'EngineCore'." + ) + executable, debug_str = _get_numactl_executable() python_executable = os.fsdecode(multiprocessing.spawn.get_executable()) with ( 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/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index dd25e721d33..a58ecf2c651 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -375,7 +375,7 @@ class ROCMAiterMLASparseMetadataBuilder( (1, 1), dtype=torch.int32, device=self.device ) - self.req_id_per_token_buffer = torch.empty( + self.req_id_per_token_buffer = torch.zeros( (vllm_config.scheduler_config.max_num_batched_tokens,), dtype=torch.int32, device=device, @@ -458,6 +458,10 @@ class ROCMAiterMLASparseMetadataBuilder( device=device, ) + self._prev_req_extent: int = 0 + self._prev_indices_extent: int = 0 + self._prev_metadata_key: tuple | None = None + def build( self, common_prefix_len: int, @@ -470,11 +474,23 @@ class ROCMAiterMLASparseMetadataBuilder( req_id_per_token = np.repeat( np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths ) - # Zero-fill for cudagraphs - self.req_id_per_token_buffer.fill_(0) - self.paged_kv_indices.fill_(0) - self.paged_kv_indptr.fill_(0) - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( + # Only re-zero the shrink-tail. paged_kv_indptr is fully rewritten + # by the cumsum below. paged_kv_indices entries past new_indices_extent + # are never read (the attention kernel only touches the ranges + # defined by paged_kv_indptr). + new_req_extent = int(req_id_per_token.shape[0]) + new_indices_extent = num_tokens * self.topk_tokens + if self._prev_req_extent > new_req_extent: + self.req_id_per_token_buffer[new_req_extent : self._prev_req_extent].fill_( + 0 + ) + if self._prev_indices_extent > new_indices_extent: + self.paged_kv_indices[new_indices_extent : self._prev_indices_extent].fill_( + 0 + ) + self._prev_req_extent = new_req_extent + self._prev_indices_extent = new_indices_extent + self.req_id_per_token_buffer[:new_req_extent].copy_( torch.from_numpy(req_id_per_token), non_blocking=True ) query_lens = ( @@ -505,27 +521,43 @@ class ROCMAiterMLASparseMetadataBuilder( # treated as its own batch entry), so persistent metadata can always # be precomputed here. The kernel switches to the persistent # work-stealing path automatically when work_meta_data is non-None. - from aiter import get_mla_metadata_v1 - - get_mla_metadata_v1( - qo_indptr, - paged_kv_indptr, - paged_kv_last_page_len, - self._num_attention_heads, - 1, - True, - self._mla_work_meta_data, - self._mla_work_info_set, - self._mla_work_indptr, - self._mla_reduce_indptr, - self._mla_reduce_final_map, - self._mla_reduce_partial_map, - page_size=1, - kv_granularity=16, - max_seqlen_qo=1, - uni_seqlen_qo=1, - fast_mode=True, + # The output is a deterministic function of (num_tokens, max_query_len, + # num_heads, min(seq_lens, topk_tokens)); fingerprint those CPU-side + # and skip the launch when nothing changed. + num_reqs = common_attn_metadata.num_reqs + clamped_seq_lens = np.minimum( + common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(), + self.topk_tokens, ) + metadata_key = ( + num_tokens, + int(common_attn_metadata.max_query_len), + self._num_attention_heads, + clamped_seq_lens.tobytes(), + ) + if metadata_key != self._prev_metadata_key: + from aiter import get_mla_metadata_v1 + + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self._num_attention_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=True, + ) + self._prev_metadata_key = metadata_key metadata = ROCMAiterMLASparseMetadata( num_reqs=common_attn_metadata.num_reqs, diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index d0cc011fe4e..a9fa45debcf 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -709,6 +709,11 @@ class AiterFlashAttentionMetadataBuilder( class AiterFlashAttentionBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + + @classmethod + def supports_sink(cls) -> bool: + return True + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "float16", @@ -788,6 +793,7 @@ class AiterFlashAttentionImpl(AttentionImpl): logits_soft_cap: float | None = None, attn_type: AttentionType = AttentionType.DECODER, kv_sharing_target_layer_name: int | None = None, + sinks: torch.Tensor | None = None, ) -> None: self.num_heads = num_heads self.head_size = head_size @@ -806,6 +812,7 @@ class AiterFlashAttentionImpl(AttentionImpl): logits_soft_cap = 0.0 self.logits_soft_cap = logits_soft_cap self.kv_sharing_target_layer_name = kv_sharing_target_layer_name + self.sinks = sinks assert self.num_heads % self.num_kv_heads == 0 self.num_queries_per_kv = self.num_heads // self.num_kv_heads @@ -878,6 +885,7 @@ class AiterFlashAttentionImpl(AttentionImpl): alibi_slopes=self.alibi_slopes, return_lse=False, out=output, + sink_ptr=self.sinks, ) def extend_forward( @@ -927,6 +935,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, return_lse=True, + sink_ptr=self.sinks, ) assert attn_metadata.extend_metadata is not None chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata @@ -974,6 +983,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, return_lse=True, + sink_ptr=self.sinks, ) if chunked_output is None: chunked_output = suf_out @@ -1092,6 +1102,7 @@ class AiterFlashAttentionImpl(AttentionImpl): window_size=self.sliding_window, alibi_slopes=self.alibi_slopes, out=output_actual_tokens[num_decode_tokens + num_extend_tokens :], + sink_ptr=self.sinks, ) # calculate for extends @@ -1136,11 +1147,17 @@ class AiterFlashAttentionImpl(AttentionImpl): assert attn_metadata.decode_metadata is not None decode_max_query_len = attn_metadata.decode_metadata.max_query_len - # Multi-token speculative decode path. - if decode_max_query_len > 1: + # Use unified_attention for speculative decoding (multi-token), + # sliding window, or sinks + # (pa_fwd_asm and paged_attention_v1 don't support sinks) + if ( + self.sliding_window[0] != -1 + or decode_max_query_len > 1 + or self.sinks is not None + ): assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), ( - "Shuffle KV cache layout is not supported with " - "speculative decoding (multi-token decode)." + "Shuffle KV cache layout is not supported with sliding " + "window, sinks, or speculative decoding (multi-token decode)." ) if not attn_metadata.causal: from aiter.ops.triton.attention.mha_v3 import ( @@ -1207,6 +1224,7 @@ class AiterFlashAttentionImpl(AttentionImpl): q_descale=None, k_descale=layer._k_scale.expand(descale_shape), v_descale=layer._v_scale.expand(descale_shape), + sinks=self.sinks, ) return diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index f43a42f75df..984fc20ecaf 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -134,6 +134,30 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): self.unified_attention = unified_attention self.supports_quant_query_input = True + def _split_kv_cache( + self, kv_cache: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.attn_type != AttentionType.ENCODER_DECODER: + return kv_cache.unbind(1) + + # NOTE: Encoder-decoder layers can share the same raw KV allocation with + # ROCM_ATTN decoder layers, whose physical layout is K/V first. Keep + # this cross-attention path on that physical layout so block IDs do not + # alias different bytes across the shared allocation. + num_blocks, _, block_size, num_kv_heads, head_size = kv_cache.shape + block_stride = block_size * num_kv_heads * head_size + kv_cache = kv_cache.as_strided( + (2, num_blocks, block_size, num_kv_heads, head_size), + ( + num_blocks * block_stride, + block_stride, + num_kv_heads * head_size, + head_size, + 1, + ), + ) + return kv_cache.unbind(0) + def forward( self, layer: torch.nn.Module, @@ -194,7 +218,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): layer, ) - key_cache, value_cache = kv_cache.unbind(1) + key_cache, value_cache = self._split_kv_cache(kv_cache) softmax_scale = self.scale if is_quantized_kv_cache(self.kv_cache_dtype): @@ -243,7 +267,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): # For encoder attention, # we use direct Q, K, V tensors without caching return - key_cache, value_cache = kv_cache.unbind(1) + key_cache, value_cache = self._split_kv_cache(kv_cache) # Reshape the input keys and values and store them in the cache. ops.reshape_and_cache_flash( @@ -276,7 +300,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): # For encoder attention, # we use direct Q, K, V tensors without caching return - key_cache, value_cache = kv_cache.unbind(1) + key_cache, value_cache = self._split_kv_cache(kv_cache) flash_layout = True is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 7ddcc493449..332350d8380 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -998,15 +998,16 @@ def build_ragged_indices_from_dense( max_width = indices.shape[1] if indices.ndim == 2 else 0 lengths = lengths.clamp(min=0, max=max_width).contiguous() - indptr = torch.empty(indices.shape[0] + 1, dtype=torch.int32, device=indices.device) - indptr[0] = 0 + indptr = torch.zeros(indices.shape[0] + 1, dtype=torch.int32, device=indices.device) torch.cumsum(lengths, dim=0, out=indptr[1:]) if indices.numel() == 0: flat = torch.empty(0, dtype=torch.int32, device=indices.device) else: flat = torch.empty( - int(indptr[-1].item()), dtype=torch.int32, device=indices.device + indices.shape[0] * max_width, + dtype=torch.int32, + device=indices.device, ) if flat.numel() > 0: block_size = 128 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/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..8edad4c3ec2 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, @@ -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 e29919022ec..7b16d9c6f05 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -40,6 +40,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 +51,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 +60,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 +86,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) @@ -237,7 +247,11 @@ class SingleTypeKVCacheManager(ABC): cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) ) req_blocks.extend(allocated_blocks) - if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): + if type(self.kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + ): self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -265,7 +279,11 @@ class SingleTypeKVCacheManager(ABC): else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) - if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): + if type(self.kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + ): self.new_block_ids.extend(b.block_id for b in new_blocks) return new_blocks @@ -305,8 +323,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, @@ -320,11 +342,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. @@ -377,7 +402,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, @@ -397,7 +422,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. @@ -487,7 +515,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, @@ -516,7 +544,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() @@ -544,6 +572,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, @@ -552,7 +593,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, @@ -563,17 +604,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 @@ -596,7 +630,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. @@ -624,7 +658,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 @@ -638,17 +672,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: @@ -702,7 +752,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, @@ -733,7 +783,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 @@ -746,7 +796,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." @@ -862,7 +912,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, @@ -1156,7 +1206,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, diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index d8a413f4c3f..848f530ce33 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -74,8 +74,10 @@ class EngineCoreReadyResponse: max_model_len: int num_gpu_blocks: int + block_size: int dp_stats_address: str | None - dtype: str | None = None + dtype: str + vllm_version: str class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 5c8507b73ee..50fd98e1fdf 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -188,7 +188,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 @@ -534,14 +534,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,8 +1460,10 @@ 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, ) ready_payload = msgspec.msgpack.encode(ready_response) for input_socket in input_sockets: diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 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/engine/utils.py b/vllm/v1/engine/utils.py index 554e8d6f005..8a7269a7707 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -575,7 +575,9 @@ class CoreEngineActorManager: node_ip_keys = [ key for key in node_resources - if key != "node:__internal_head__" and key.startswith("node:") + if key != "node:__internal_head__" + and key.startswith("node:") + and "_group_" not in key ] assert len(node_ip_keys) == 1, ( f"Zero or multiple node IP keys found in node resources: {node_ip_keys}" @@ -654,6 +656,9 @@ class CoreEngineActorManager: if len(placement_groups) == dp_size: break + if len(placement_groups) == dp_size: + break + if len(placement_groups) < dp_size: raise ValueError( f"Not enough resources to allocate {dp_size} " 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 db21d7cee77..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 @@ -60,6 +60,7 @@ from vllm.utils.system_utils import ( ) from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.executor.abstract import Executor, FailureCallback +from vllm.v1.executor.vllm_net_devices import set_worker_net_device from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.worker_base import WorkerWrapperBase @@ -471,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). @@ -811,6 +806,9 @@ class WorkerProc: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set + set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"]) + worker = None ready_writer = kwargs.pop("ready_pipe") death_pipe = kwargs.pop("death_pipe", None) 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 92e668406f9..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 @@ -16,6 +15,7 @@ from vllm.platforms import current_platform from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.executor.abstract import Executor +from vllm.v1.executor.vllm_net_devices import set_worker_net_device from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput from vllm.v1.serial_utils import run_method from vllm.v1.worker.worker_base import WorkerWrapperBase @@ -56,6 +56,9 @@ class UniProcExecutor(Executor): shared_worker_lock=Lock(), ) + # Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set + set_worker_net_device(local_rank, self.vllm_config) + self.driver_worker.init_worker(all_kwargs=[kwargs]) self.driver_worker.init_device() @@ -73,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/executor/vllm_net_devices.py b/vllm/v1/executor/vllm_net_devices.py new file mode 100644 index 00000000000..be442f35ee0 --- /dev/null +++ b/vllm/v1/executor/vllm_net_devices.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU-to-NIC net-device mapping for RDMA transports (UCX, NVSHMEM, ...). + +Shared by both UniProcExecutor (TP=1) and MultiprocExecutor (TP>1). +All transport-specific env vars and sysfs lookups live here so executor +files only need to call ``set_worker_net_device(local_rank, vllm_config)``. + +Requires two env vars set together: +- ``VLLM_GPU_NIC_PCIE_MAPPING`` -- comma-separated GPU_BDF=NIC_BDF pairs. +- ``VLLM_NIC_SELECTION_VARS`` -- comma-separated list of env vars to set, + each optionally suffixed (e.g. ``UCX_NET_DEVICES:1,NCCL_IB_HCA:1``). +""" + +import os +from pathlib import Path + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def normalize_pci(addr: str) -> tuple[int, int, int, int]: + """Parse PCI BDF/domain-bus-device-function into comparable ints (all hex). + + Supported shapes: + - ``domain:bus:dev.fn`` -- domain width varies (e.g. ``00000001:00:00.0``, + ``0001:00:00.0``, ``0000:3f:00.0``). + - ``bus:dev.fn`` -- domain **0** (e.g. ``01:00.0``, ``40:00.0``). + + Function suffix is hex (typically ``0``--``7``). Raises ``ValueError`` if malformed. + """ + s = addr.strip().lower().replace(" ", "") + if s.startswith("0x"): + s = s[2:] + if "." not in s: + raise ValueError(f"invalid PCI BDF (missing function suffix): {addr!r}") + body, fn_s = s.rsplit(".", 1) + if not fn_s or any(c not in "0123456789abcdef" for c in fn_s): + raise ValueError(f"invalid PCI function in BDF: {addr!r}") + fn = int(fn_s, 16) + if fn > 0xFF: + raise ValueError(f"PCI function out of range: {addr!r}") + + parts = body.split(":") + if len(parts) == 2: + domain = 0 + bus = int(parts[0], 16) + device = int(parts[1], 16) + elif len(parts) == 3: + domain = int(parts[0], 16) + bus = int(parts[1], 16) + device = int(parts[2], 16) + else: + raise ValueError( + f"invalid PCI BDF (want domain:bus:dev.fn or bus:dev.fn): {addr!r}" + ) + + if bus > 0xFF or device > 0x1F: + raise ValueError(f"PCI bus or device out of range: {addr!r}") + return (domain, bus, device, fn) + + +def parse_gpu_nic_mapping( + raw: str, +) -> dict[tuple[int, int, int, int], tuple[int, int, int, int]]: + out: dict[tuple[int, int, int, int], tuple[int, int, int, int]] = {} + for segment in raw.split(","): + segment = segment.strip() + if not segment: + continue + if "=" not in segment: + raise ValueError( + "VLLM_GPU_NIC_PCIE_MAPPING: expected comma-separated" + f" gpu_bdf=nic_bdf pairs; ambiguous segment: {segment!r}" + ) + gpu_s, nic_s = segment.split("=", 1) + gpu_key = normalize_pci(gpu_s.strip()) + nic_val = normalize_pci(nic_s.strip()) + out[gpu_key] = nic_val + return out + + +def rdma_name_for_nic_pci(nic_pci: tuple[int, int, int, int]) -> str: + """Map NIC PCI BDF to sysfs RDMA name (mlx5_*, ibp*, ...). + + Under ``/sys/class/infiniband//``, ``device`` is a **symlink** to the PCI + device directory (e.g. ``.../0101:00:00.0``). We take ``Path(...).resolve().name`` + as the BDF string. + + ``VLLM_GPU_NIC_PCIE_MAPPING`` NIC keys must **normalize** (via ``normalize_pci``) + to the same tuple as this basename. + """ + ib = Path("/sys/class/infiniband") + if not ib.is_dir(): + raise RuntimeError("/sys/class/infiniband not found or not a directory") + names = sorted(p.name for p in ib.iterdir() if p.is_dir()) + for name in names: + dev_link = ib / name / "device" + if not dev_link.exists(): + continue + try: + resolved = dev_link.resolve() + except OSError: + continue + pci_name = resolved.name + try: + if normalize_pci(pci_name) == nic_pci: + return name + except ValueError: + continue + raise RuntimeError( + f"No /sys/class/infiniband device for NIC PCI {nic_pci}; have entries: {names}" + ) + + +def parse_nic_selection_vars(raw: str) -> list[tuple[str, str]]: + """Parse ``VLLM_NIC_SELECTION_VARS`` into ``(env_var_name, suffix)`` pairs. + + Each entry is ``VAR_NAME`` or ``VAR_NAME:``. The colon and + everything after it is appended verbatim to the RDMA device name. + """ + result: list[tuple[str, str]] = [] + for entry in raw.split(","): + entry = entry.strip() + if not entry: + continue + if ":" in entry: + var_name, suffix = entry.split(":", 1) + result.append((var_name, ":" + suffix)) + else: + result.append((entry, "")) + return result + + +def set_worker_gpu_nic_mapping(local_rank: int) -> None: + """Set NIC selection env vars from VLLM_GPU_NIC_PCIE_MAPPING for a worker. + + Which env vars are set is controlled by ``VLLM_NIC_SELECTION_VARS``. + """ + raw = envs.VLLM_GPU_NIC_PCIE_MAPPING.strip() + if not raw: + return + selection_raw = envs.VLLM_NIC_SELECTION_VARS.strip() + selection_vars = parse_nic_selection_vars(selection_raw) + mapping = parse_gpu_nic_mapping(raw) + pci_by_index = current_platform.get_all_gpu_pci_bus_ids() + # Translate CUDA-relative local_rank to the physical device index, + # which accounts for CUDA_VISIBLE_DEVICES narrowing (e.g. DP sharding). + physical_id = current_platform.device_id_to_physical_device_id(local_rank) + if physical_id not in pci_by_index: + raise RuntimeError( + f"No GPU PCI for physical device index {physical_id} " + f"(local_rank={local_rank}) in map " + f"(have indices {sorted(pci_by_index.keys())})" + ) + gpu_bdf = pci_by_index[physical_id] + gpu_key = normalize_pci(gpu_bdf) + if gpu_key not in mapping: + keys_fmt = ", ".join( + f"{d:04x}:{b:02x}:{dev:02x}.{fn}" + for d, b, dev, fn in sorted(mapping.keys()) + ) + raise RuntimeError( + f"No VLLM_GPU_NIC_PCIE_MAPPING entry for GPU PCI {gpu_bdf} " + f"(worker local_rank={local_rank}); mapped GPUs: {keys_fmt}" + ) + nic_pci = mapping[gpu_key] + rdma_dev = rdma_name_for_nic_pci(nic_pci) + + set_vars: list[str] = [] + for var_name, suffix in selection_vars: + value = f"{rdma_dev}{suffix}" + existing = os.environ.get(var_name, "").strip() + if existing: + value = f"{value},{existing}" + os.environ[var_name] = value + set_vars.append(f"{var_name}={value}") + + nic_fmt = f"{nic_pci[0]:04x}:{nic_pci[1]:02x}:{nic_pci[2]:02x}.{nic_pci[3]}" + logger.info( + "GPU rank %s (PCIe addr %s) mapped to NIC %s (PCIe addr %s) via env vars: %s", + local_rank, + gpu_bdf, + rdma_dev, + nic_fmt, + ", ".join(set_vars), + ) + + +def _dp_adjusted_local_rank(tp_local_rank: int, vllm_config: VllmConfig) -> int: + """Compute the node-wide GPU index accounting for data parallelism. + + On CUDA-alike platforms without env-var device isolation (the common + MP-backend path), the worker sees *all* GPUs on the node and selects + its device via ``torch.accelerator.set_device_index()`` using:: + + dp_local_rank * tp_pp_world_size + tp_local_rank + + This mirrors the adjustment in ``Worker.init_device()`` so we resolve + the correct GPU PCI address *before* the CUDA device is initialised. + """ + pc = vllm_config.parallel_config + if ( + pc.distributed_executor_backend not in ("ray", "external_launcher") + and pc.data_parallel_backend != "ray" + and pc.nnodes_within_dp == 1 + ): + dp_local_rank = pc.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = pc.data_parallel_index + tp_pp_world_size = pc.pipeline_parallel_size * pc.tensor_parallel_size + return dp_local_rank * tp_pp_world_size + tp_local_rank + return tp_local_rank + + +def set_worker_net_device(local_rank: int, vllm_config: VllmConfig) -> None: + """Top-level entry point for both UniProcExecutor and MultiprocExecutor. + + Sets NIC selection env vars from ``VLLM_GPU_NIC_PCIE_MAPPING`` and + ``VLLM_NIC_SELECTION_VARS`` if present; no-op otherwise. + """ + has_pcie_mapping = bool(envs.VLLM_GPU_NIC_PCIE_MAPPING.strip()) + has_selection_vars = bool(envs.VLLM_NIC_SELECTION_VARS.strip()) + if has_pcie_mapping and not has_selection_vars: + raise RuntimeError( + "VLLM_GPU_NIC_PCIE_MAPPING is set but VLLM_NIC_SELECTION_VARS " + "is not; both must be set together." + ) + if has_selection_vars and not has_pcie_mapping: + raise RuntimeError( + "VLLM_NIC_SELECTION_VARS is set but VLLM_GPU_NIC_PCIE_MAPPING " + "is not; both must be set together." + ) + # No-op when neither env var is present. + if not has_pcie_mapping and not has_selection_vars: + return + # Both env vars are present, so set the NIC selection env vars. + adjusted_rank = _dp_adjusted_local_rank(local_rank, vllm_config) + set_worker_gpu_nic_mapping(adjusted_rank) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index de65be1c05e..5f798f41eac 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -7,10 +7,12 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Iterator, Sequence from dataclasses import dataclass +from enum import Enum 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 @@ -49,6 +51,20 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None +class OffloadPolicy(Enum): + # Offload only newly-computed blocks as they arrive; prefix-hit + # blocks (already offloaded by a prior request) are skipped. + BLOCK_LEVEL = "block_level" + # Offload all blocks for the request, including prefix hits. + # Used by tiers that need the complete KV context for a request. + REQUEST_LEVEL = "request_level" + + +@dataclass +class RequestOffloadingContext: + policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL + + class LoadStoreSpec(ABC): """ Abstract metadata that encapsulates information allowing a worker @@ -210,6 +226,28 @@ class OffloadingManager(ABC): """ return + @abstractmethod + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """ + Called when a new request is first seen by the scheduler. + + Returns a RequestOffloadingContext indicating how this request's + blocks should be offloaded. + + Args: + req_context: per-request context. + """ + pass + + def on_request_finished(self, req_context: ReqContext) -> None: + """ + Called when a request has finished. + + Args: + req_context: per-request context. + """ + return + def take_events(self) -> Iterable[OffloadingEvent]: """ Take the offloading events from the manager. @@ -219,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 @@ -274,6 +320,7 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec): self.block_indices: Sequence[int] = block_indices @staticmethod + @override def medium() -> str: return "GPU" @@ -343,6 +390,14 @@ class OffloadingSpec(ABC): assert kv_transfer_config is not None self.extra_config = kv_transfer_config.kv_connector_extra_config + # When True, only prompt (prefill) blocks are offloaded; decode-phase + # blocks (KV generated after the prompt) are skipped. Useful when prior + # turns' generated tokens are dropped before the next turn (e.g. + # reasoning models that strip thinking). + self.offload_prompt_only: bool = bool( + self.extra_config.get("offload_prompt_only", True) + ) + parallel_config = vllm_config.parallel_config context_parallel_factor = ( parallel_config.decode_context_parallel_size 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 39e64933b5c..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, @@ -11,6 +13,7 @@ from vllm.v1.kv_offload.base import ( OffloadKey, PrepareStoreOutput, ReqContext, + RequestOffloadingContext, ) from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy @@ -94,6 +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: @@ -110,6 +118,7 @@ class CPUOffloadingManager(OffloadingManager): return None # write in-flight; caller should retry return True + @override def prepare_load( self, keys: Collection[OffloadKey], @@ -124,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: @@ -136,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], @@ -193,6 +205,7 @@ class CPUOffloadingManager(OffloadingManager): evicted_keys=to_evict, ) + @override def complete_store( self, keys: Collection[OffloadKey], @@ -223,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 @@ -234,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 7d4d6f031da..d4f0cefe5eb 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING import numpy as np -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext if TYPE_CHECKING: from vllm.v1.kv_offload.base import OffloadingSpec @@ -163,6 +163,36 @@ class SecondaryTierManager(ABC): """ return + @abstractmethod + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + """ + Called when a new request is first seen by the scheduler. + + Returns a RequestOffloadingContext expressing this tier's preference + for how blocks should be offloaded for this request. + + Args: + req_context: Per-request context. + """ + pass + + def on_request_finished(self, req_context: ReqContext) -> None: + """ + Called when a request has finished. + + Args: + req_context: per-request context. + """ + 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 933a8aa8c83..caf1d2c71b4 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -13,7 +13,9 @@ import logging from collections.abc import Iterable from typing import TYPE_CHECKING -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from typing_extensions import override + +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -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,10 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs = [] return result + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + def get_num_blocks(self) -> int: """Get the number of blocks currently stored in this tier.""" return len(self.blocks) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index e5921d6ffd8..a33de02f43d 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -21,12 +21,15 @@ 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 from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, + RequestOffloadingContext, SecondaryTierManager, ) from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block @@ -49,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__( @@ -100,11 +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( @@ -118,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( @@ -131,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. @@ -140,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 36fc27e48de..cb8de749ec7 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,10 +20,12 @@ Key Design Principles: protecting blocks from eviction until complete_read() is called """ -from collections.abc import Collection, Iterable +from collections import defaultdict +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 ( @@ -31,8 +33,10 @@ from vllm.v1.kv_offload.base import ( OffloadingEvent, OffloadingManager, OffloadKey, + OffloadPolicy, PrepareStoreOutput, ReqContext, + RequestOffloadingContext, ) from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager @@ -97,6 +101,7 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): """ return self._kv_memoryview + @override def shutdown(self) -> None: super().shutdown() self._kv_memoryview.release() @@ -147,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[ @@ -155,9 +160,16 @@ 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 + # policy. Populated in on_new_request(), + # cleaned up in on_request_finished(). + self._request_level_tiers: defaultdict[str, set[SecondaryTierManager]] = ( + defaultdict(set) + ) + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -170,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 @@ -212,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. @@ -291,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: @@ -307,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 @@ -328,6 +341,7 @@ class TieringOffloadingManager(OffloadingManager): self._pending_load_submissions.clear() + @override def prepare_load( self, keys: Collection[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: @@ -352,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. @@ -364,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. @@ -377,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: @@ -387,6 +404,9 @@ class TieringOffloadingManager(OffloadingManager): that any completed async transfers have their ref_cnt decremented before the primary tier makes eviction decisions. + For request-level tiers, blocks already present in the primary tier + are immediately cascaded via submit_store(). + Args: keys: Blocks to prepare for storing. req_context: Per-request context. @@ -400,14 +420,64 @@ class TieringOffloadingManager(OffloadingManager): # successfully transferred to secondary tiers. self._maybe_process_finished_jobs() - # Step 2: Store to primary tier + # Step 2: Store to primary tier (new blocks only). + # Cascading of these newly-stored blocks to ALL secondary tiers + # happens later in complete_store(), after the GPU→Primary transfer + # completes. primary_result = self.primary_tier.prepare_store(keys, req_context) - # Note: Secondary tier cascading will happen in complete_store() - # after the GPU→Primary transfer completes and blocks are ready. + if primary_result is None: + return None + + # Step 3: For request-level tiers, cascade blocks already in primary + request_level_tiers = self._request_level_tiers.get(req_context.req_id) + if request_level_tiers is not None: + keys_to_store_set = set(primary_result.keys_to_store) + keys_already_in_primary = tuple( + k for k in keys if k not in keys_to_store_set + ) + if keys_already_in_primary: + self._cascade_existing_blocks_to_request_level_tiers( + keys_already_in_primary, req_context, request_level_tiers + ) return primary_result + def _cascade_existing_blocks_to_request_level_tiers( + self, + keys: Sequence[OffloadKey], + req_context: ReqContext, + request_level_tiers: set[SecondaryTierManager], + ) -> None: + """ + For tiers that requested request-level policy, submit_store() for + blocks that are already present in the primary tier. + """ + # Filter out keys that are not ready in primary (e.g. in-flight) + ready_keys = tuple( + k for k in keys if self.primary_tier.lookup(k, req_context) is True + ) + if not ready_keys: + return + + for tier in request_level_tiers: + primary_blocks_spec = self.primary_tier.prepare_read( + ready_keys, req_context + ) + + job_id = self._next_job_id() + assert isinstance(primary_blocks_spec, CPULoadStoreSpec) + job_metadata = JobMetadata( + job_id=job_id, + keys=ready_keys, + block_ids=primary_blocks_spec.block_ids, + is_promotion=False, + req_context=req_context, + ) + self._transfer_jobs[job_id] = job_metadata + tier.submit_store(job_metadata) + + @override def complete_store( self, keys: Collection[OffloadKey], @@ -466,38 +536,61 @@ class TieringOffloadingManager(OffloadingManager): # Note: The async transfers are now in flight. Their completion is # tracked via get_finished_jobs() / _maybe_process_finished_jobs(). - def take_events(self) -> Iterable[OffloadingEvent]: + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: """ - End-of-step hook: flush deferred work, yield events, reset per-step state. + Query each secondary tier for its offload policy preference. - 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. + Returns REQUEST_LEVEL if ANY secondary tier wants request-level. + Only stores REQUEST_LEVEL tier decisions for use in prepare_store. + """ + for tier in self.secondary_tiers: + tier_ctx = tier.on_new_request(req_context) + if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL: + self._request_level_tiers[req_context.req_id].add(tier) + + policy = ( + OffloadPolicy.REQUEST_LEVEL + if req_context.req_id in self._request_level_tiers + else OffloadPolicy.BLOCK_LEVEL + ) + 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) + + @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 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/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 6855efd9f54..0052a35366a 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -1161,7 +1161,8 @@ class PrometheusStatLogger(AggregateStatLoggerBase): iteration_stats.num_generation_tokens ) self.histogram_iteration_tokens[engine_idx].observe( - iteration_stats.num_prompt_tokens + iteration_stats.num_generation_tokens + iteration_stats.prompt_token_stats.computed + + iteration_stats.num_generation_tokens ) for max_gen_tokens in iteration_stats.max_num_generation_tokens_iter: diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 9703dfa9e70..9f13ad939fc 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from copy import copy from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple, TypeAlias @@ -279,6 +280,19 @@ class ModelRunnerOutput: # ``None`` when ``enable_return_routed_experts`` is off. routed_experts: RoutedExpertsLists | None = None + @staticmethod + def with_kv_conn_output_only( + kv_connector_output: KVConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return ModelRunnerOutput containing the provided KVConnectorOutput, + otherwise empty. Returns None if kv_connector_output is passed as None. + """ + if kv_connector_output is None or kv_connector_output.is_empty(): + return EMPTY_MODEL_RUNNER_OUTPUT + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.kv_connector_output = kv_connector_output + return output + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): 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 99a5bd238c4..f98b12a379d 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -168,7 +168,7 @@ class TopKTopPSampler(nn.Module): The logits tensor may be updated in-place. """ - logits = apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True) + logits = apply_top_k_top_p(logits, k, p) logits_to_return = None if self.logprobs_mode == "processed_logits": logits_to_return = logits @@ -310,8 +310,9 @@ def apply_top_k_top_p( if p is None and k is None: return logits - # Keep CPU logits on the PyTorch path to avoid invoking Triton kernels. if current_platform.is_cpu(): + if HAS_TRITON: + return apply_top_k_top_p_triton(logits, k, p) return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True) if HAS_TRITON and logits.shape[0] >= 8: diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index fee87883c96..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) @@ -1038,6 +928,12 @@ def apply_top_k_top_p_triton( else: normal_cdf_to_sigma_table, percentile_to_std_table = tables + # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. + if logits.device.type == "cpu": + block_size, block_size_trunc = 256, 128 + else: + block_size, block_size_trunc = 8192, 4096 + _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, logits.stride(0), @@ -1049,8 +945,8 @@ def apply_top_k_top_p_triton( BATCH_SIZE=batch_size, MASK_VALUE=mask_value, VOCAB_SIZE=vocab_size, - BLOCK_SIZE=8192, - BLOCK_SIZE_TRUNC=4096, + BLOCK_SIZE=block_size, + BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, ) 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/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index db74044f4fd..72d0f99d07d 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -68,6 +68,8 @@ class DFlashProposer(SpecDecodeBaseProposer): # For DFlash we use the input embeddings to embed the mask token self.parallel_drafting_hidden_state_tensor = None + self.dflash_causal = self.dflash_config.get("causal", False) + @override def _create_draft_vllm_config(self) -> VllmConfig: base = super()._create_draft_vllm_config() @@ -75,7 +77,7 @@ class DFlashProposer(SpecDecodeBaseProposer): base, attention_config=replace( base.attention_config, - use_non_causal=True, + use_non_causal=not self.dflash_causal, ), ) @@ -184,7 +186,7 @@ class DFlashProposer(SpecDecodeBaseProposer): max_seq_len=cad.max_seq_len + num_query_per_req, block_table_tensor=cad.block_table_tensor, slot_mapping=query_slot_mapping, - causal=False, # Non-causal attention is required for DFlash + causal=self.dflash_causal, ) return num_query_total, token_indices_to_sample, new_cad @@ -281,20 +283,20 @@ class DFlashProposer(SpecDecodeBaseProposer): per_group, per_layer = super().build_per_group_and_layer_attn_metadata( cad, draft_index ) - for layer_name, attn_metadata in per_layer.items(): - assert getattr(attn_metadata, "causal", None) is False, ( - f"Attention metadata for layer {layer_name} does not have" - " non-causal support, which is required for DFlash." - " Consider using a different attention backend, such as FlashAttention." - ) + if not self.dflash_causal: + # Require all layers to support non-causal attention when required by DFlash + for layer_name, attn_metadata in per_layer.items(): + assert getattr(attn_metadata, "causal", None) is False, ( + f"Attention metadata for layer {layer_name} does not have" + " non-causal support, which is required for DFlash." + " Consider using a different attention backend, e.g FlashAttention." + ) return per_group, per_layer @override def _get_eagle3_use_aux_hidden_state_from_config(self): - use_aux_hidden_state = True - dflash_config = getattr( - self.draft_model_config.hf_config, "dflash_config", None - ) - if dflash_config is not None: - use_aux_hidden_state = dflash_config.get("use_aux_hidden_state", True) - return use_aux_hidden_state + return self.dflash_config.get("use_aux_hidden_state", True) + + @property + def dflash_config(self): + return getattr(self.draft_model_config.hf_config, "dflash_config", None) or {} diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 9979a051727..9a0b537175b 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1232,6 +1232,7 @@ class SpecDecodeBaseProposer: "Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Gemma4ForConditionalGeneration", + "Step3p7ForConditionalGeneration", ]: self.model.config.image_token_index = target_model.config.image_token_id elif self.get_model_name(target_model) == "PixtralForConditionalGeneration": diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py new file mode 100644 index 00000000000..ccca17a3188 --- /dev/null +++ b/vllm/v1/spec_decode/step3p5.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from copy import copy + +import torch + +from vllm.config import VllmConfig, get_layers_from_vllm_config, replace +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.models.utils import get_draft_quant_config +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.utils import PADDING_SLOT_ID +from vllm.v1.worker.utils import AttentionGroup + + +class Step3p5MTPProposer(EagleProposer): + """Step3.5 MTP proposer with per-layer draft-step selection.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + runner=None, + ): + super().__init__(vllm_config, device, runner) + self._per_group_block_tables: dict[int, torch.Tensor] = {} + self._per_group_slot_mappings: dict[int, torch.Tensor] = {} + # Slot-mapping buffers for non-primary KV cache groups (the primary + # group reuses self._slot_mapping_buffer from the base class). + self._per_group_slot_mapping_buffers: dict[int, torch.Tensor] = {} + + def set_per_group_attn_metadata( + self, + gid: int, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + self._per_group_block_tables[gid] = block_table + self._per_group_slot_mappings[gid] = slot_mapping + + def _slot_mapping_buffer_for(self, gid: int) -> torch.Tensor: + if gid == self.kv_cache_gid: + return self._slot_mapping_buffer + buf = self._per_group_slot_mapping_buffers.get(gid) + if buf is None: + buf = torch.zeros(self.max_positions, dtype=torch.int64, device=self.device) + self._per_group_slot_mapping_buffers[gid] = buf + return buf + + def _get_slot_mapping( + self, + num_tokens: int, + slot_mapping: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Per-layer slot_mapping with one buffer per KV cache group.""" + per_layer: dict[str, torch.Tensor] = {} + for attn_group in self.draft_attn_groups: + gid = attn_group.kv_cache_group_id + buf = self._slot_mapping_buffer_for(gid) + source = self._per_group_slot_mappings.get(gid, slot_mapping) + if source is not None and buf.data_ptr() != source.data_ptr(): + n = source.shape[0] + buf[:n].copy_(source) + if num_tokens > n: + buf[n:num_tokens].fill_(PADDING_SLOT_ID) + view = buf[:num_tokens] + for layer_name in attn_group.layer_names: + per_layer[layer_name] = view + return per_layer + + def _update_positions_dependent_metadata( + self, + positions: torch.Tensor, + common_attn_metadata: CommonAttentionMetadata, + batch_size: int, + input_batch_size: int, + block_size: int, + ) -> torch.Tensor: + old_positions_1d = positions[0] if self.uses_mrope else positions + positions = super()._update_positions_dependent_metadata( + positions, + common_attn_metadata, + batch_size, + input_batch_size, + block_size, + ) + # Parent already produced slot_mapping for the primary gid. + self._per_group_slot_mappings[self.kv_cache_gid] = ( + common_attn_metadata.slot_mapping + ) + # Recompute slot_mapping for the remaining gids using their own block tables. + new_positions_1d = positions[0] if self.uses_mrope else positions + exceeds = old_positions_1d + 1 >= self.max_model_len + for attn_group in self.draft_attn_groups: + gid = attn_group.kv_cache_group_id + if gid == self.kv_cache_gid: + continue + block_table = self._per_group_block_tables.get(gid) + if block_table is None: + continue + n_blocks = block_table.shape[1] + bn = (new_positions_1d // block_size).clamp(max=n_blocks - 1).to(torch.long) + block_ids = block_table[:batch_size].gather(1, bn.unsqueeze(1)).squeeze(1) + sm = block_ids * block_size + (new_positions_1d % block_size) + sm.masked_fill_(exceeds, PADDING_SLOT_ID) + buf = self._slot_mapping_buffer_for(gid) + buf[:batch_size].copy_(sm) + if input_batch_size > batch_size: + buf[batch_size:input_batch_size].fill_(PADDING_SLOT_ID) + self._per_group_slot_mappings[gid] = buf[:batch_size] + return positions + + def build_per_group_and_layer_attn_metadata( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int = 0, + ) -> tuple[list[object], dict[str, object]]: + per_group_attn_metadata: list[object] = [] + per_layer_attn_metadata: dict[str, object] = {} + # The proposer always works in unpadded shape. Per-group block tables + # registered via set_per_group_attn_metadata are stored at the model + # runner's padded shape; slice them to match cm's num_reqs. + num_reqs = common_attn_metadata.num_reqs + num_actual_tokens = common_attn_metadata.num_actual_tokens + for attn_group in self.draft_attn_groups: + gid = attn_group.kv_cache_group_id + if gid in self._per_group_block_tables: + cm = copy(common_attn_metadata) + cm.block_table_tensor = self._per_group_block_tables[gid][:num_reqs] + if gid in self._per_group_slot_mappings: + sm = self._per_group_slot_mappings[gid] + if sm.shape[0] >= num_actual_tokens: + sm = sm[:num_actual_tokens] + cm.slot_mapping = sm + else: + cm = common_attn_metadata + attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + common_attn_metadata=cm, + draft_index=draft_index, + ) + per_group_attn_metadata.append(attn_metadata) + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + return per_group_attn_metadata, per_layer_attn_metadata + + def _maybe_share_lm_head(self, target_language_model: torch.nn.Module) -> None: + """Step3.5 MTP uses the lm_head stored in each MTP layer.""" + + # The base MTP path shares target lm_head into shared_head.head. + # Step3.5 checkpoints carry per-MTP-layer shared_head weights. + return + + def _create_draft_vllm_config(self) -> VllmConfig: + base = super()._create_draft_vllm_config() + return replace( + base, + model_config=self.draft_model_config, + quant_config=get_draft_quant_config(base), + ) + + def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: + """Step3.5 MTP draft layers may span multiple KV cache groups.""" + return + + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int] | None = None, + ) -> None: + all_attn_layers = get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ) + + layer_to_gid: dict[str, int] = {} + layer_to_spec: dict[str, KVCacheSpec] = {} + for gid, group in enumerate(kv_cache_config.kv_cache_groups): + group_spec = group.kv_cache_spec + for layer_name in group.layer_names: + layer_to_gid[layer_name] = gid + if isinstance(group_spec, UniformTypeKVCacheSpecs): + if layer_name in group_spec.kv_cache_specs: + layer_to_spec[layer_name] = group_spec.kv_cache_specs[ + layer_name + ] + else: + target_layer_name = getattr( + all_attn_layers.get(layer_name), + "kv_sharing_target_layer_name", + None, + ) + if ( + target_layer_name + and target_layer_name in group_spec.kv_cache_specs + ): + layer_to_spec[layer_name] = group_spec.kv_cache_specs[ + target_layer_name + ] + else: + layer_to_spec[layer_name] = group_spec + else: + layer_to_spec[layer_name] = group_spec + + attention_groups: dict[tuple[tuple[str, str], int], AttentionGroup] = {} + for layer_name in sorted(self._draft_attn_layer_names): + if layer_name not in layer_to_spec: + continue + attn_layer = all_attn_layers[layer_name] + attn_backend = attn_layer.get_attn_backend() + spec = layer_to_spec[layer_name] + gid = layer_to_gid[layer_name] + group_key = (attn_backend.full_cls_name(), gid) + + if group_key not in attention_groups: + kernel_block_size = ( + kernel_block_sizes[gid] + if kernel_block_sizes is not None and gid < len(kernel_block_sizes) + else None + ) + attn_group = AttentionGroup( + backend=attn_backend, + layer_names=[layer_name], + kv_cache_spec=spec, + kv_cache_group_id=gid, + ) + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_size=kernel_block_size, + ) + attention_groups[group_key] = attn_group + else: + attention_groups[group_key].layer_names.append(layer_name) + + self.draft_attn_groups = list(attention_groups.values()) + if self.draft_attn_groups: + self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id + self.block_size = ( + self.draft_attn_groups[0] + .get_metadata_builder() + .kv_cache_spec.block_size + ) + else: + self.kv_cache_gid = 0 + self.block_size = kv_cache_config.kv_cache_groups[ + 0 + ].kv_cache_spec.block_size + + def _sample_draft_tokens_for_step( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + spec_step_idx: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy: + if self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states), None + logits = self.model.compute_logits( + hidden_states, spec_step_idx=spec_step_idx + ) + return logits.argmax(dim=-1), None + + logits = self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx) + return self._sample_from_logits(logits, sampling_metadata) + + def propose( + self, + target_token_ids: torch.Tensor, + target_positions: torch.Tensor, + target_hidden_states: torch.Tensor, + next_token_ids: torch.Tensor, + token_indices_to_sample: torch.Tensor | None, + common_attn_metadata: CommonAttentionMetadata, + sampling_metadata: SamplingMetadata, + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + num_rejected_tokens_gpu: torch.Tensor | None = None, + slot_mappings: dict[str, torch.Tensor] + | list[dict[str, torch.Tensor]] + | None = None, + ) -> torch.Tensor: + self._last_draft_probs = None + batch_size = common_attn_metadata.batch_size() + + num_tokens, token_indices_to_sample, common_attn_metadata = ( + self.set_inputs_first_pass( + target_token_ids=target_token_ids, + next_token_ids=next_token_ids, + target_positions=target_positions, + target_hidden_states=target_hidden_states, + token_indices_to_sample=token_indices_to_sample, + cad=common_attn_metadata, + num_rejected_tokens_gpu=num_rejected_tokens_gpu, + ) + ) + + per_group_attn_metadata, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata(common_attn_metadata) + ) + + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( + self._determine_batch_execution_and_padding(num_tokens) + ) + + model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass( + num_tokens, num_input_tokens, mm_embed_inputs + ) + model_kwargs["spec_step_idx"] = 0 + + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + slot_mapping=self._get_slot_mapping( + slot_mapping_size, common_attn_metadata.slot_mapping + ), + ): + ret_hidden_states = self.model(**model_kwargs) + if not self.model_returns_tuple(): + last_hidden_states = ret_hidden_states + hidden_states = last_hidden_states + else: + last_hidden_states, hidden_states = ret_hidden_states + + sample_hidden_states = last_hidden_states[token_indices_to_sample] + + if self.num_speculative_tokens == 1 or self.parallel_drafting: + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( + sample_hidden_states, sampling_metadata, spec_step_idx=0 + ) + if draft_probs is not None: + self._last_draft_probs = draft_probs.view( + -1, self.num_speculative_tokens, draft_probs.shape[-1] + ).contiguous() + return draft_token_ids.view(-1, self.num_speculative_tokens) + + if self.uses_mrope: + positions = self.mrope_positions[:, token_indices_to_sample] + else: + positions = self.positions[token_indices_to_sample] + hidden_states = hidden_states[token_indices_to_sample] + + if self.constant_draft_positions: + self.positions[:batch_size] = positions + + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( + sample_hidden_states, sampling_metadata, spec_step_idx=0 + ) + draft_probs_list = None if draft_probs is None else [draft_probs] + + if self.allowed_attn_types is not None: + for group_md in per_group_attn_metadata: + if not isinstance(group_md, self.allowed_attn_types): + raise ValueError( + f"Unsupported attention metadata type for speculative " + "decoding with num_speculative_tokens > 1: " + f"{type(group_md)}. Supported types are: " + f"{self.allowed_attn_types}" + ) + + draft_token_ids_list = [draft_token_ids] + + cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = ( + self._determine_batch_execution_and_padding(batch_size) + ) + + common_attn_metadata.num_actual_tokens = batch_size + common_attn_metadata.max_query_len = 1 + common_attn_metadata.query_start_loc = self.arange[: batch_size + 1] + common_attn_metadata.query_start_loc_cpu = torch.from_numpy( + self.token_arange_np[: batch_size + 1] + ).clone() + + if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None: + common_attn_metadata.seq_lens -= num_rejected_tokens_gpu + common_attn_metadata._seq_lens_cpu = None + common_attn_metadata._num_computed_tokens_cpu = None + + block_size = self.block_size + assert block_size > 0, "block_size has not been initialized." + for token_index in range(self.num_speculative_tokens - 1): + spec_step_idx = token_index + 1 + input_ids = draft_token_ids_list[-1].int() + + if not self.constant_draft_positions: + positions = self._update_positions_dependent_metadata( + positions, + common_attn_metadata, + batch_size, + input_batch_size, + block_size, + ) + + if not self.constant_draft_positions or token_index == 0: + _, per_layer_attn_metadata = ( + self.build_per_group_and_layer_attn_metadata( + common_attn_metadata, draft_index=spec_step_idx + ) + ) + + self.input_ids[:batch_size] = input_ids + self.hidden_states[:batch_size] = hidden_states + if self.supports_mm_inputs: + self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids) + + input_ids = None + inputs_embeds = self.inputs_embeds[:input_batch_size] + else: + input_ids = self.input_ids[:input_batch_size] + inputs_embeds = None + + model_kwargs = { + "input_ids": input_ids, + "positions": self._get_positions(input_batch_size), + "inputs_embeds": inputs_embeds, + "spec_step_idx": spec_step_idx, + } + if self.pass_hidden_states_to_model: + model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] + + with set_forward_context( + per_layer_attn_metadata, + self.vllm_config, + num_tokens=input_batch_size, + num_tokens_across_dp=batch_size_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + slot_mapping=self._get_slot_mapping(input_batch_size), + ): + ret_hidden_states = self.model(**model_kwargs) + if not self.model_returns_tuple(): + last_hidden_states = ret_hidden_states + hidden_states = ret_hidden_states + else: + last_hidden_states, hidden_states = ret_hidden_states + + hidden_states = hidden_states[:batch_size] + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step( + last_hidden_states[:batch_size], + sampling_metadata, + spec_step_idx=spec_step_idx, + ) + if draft_probs is not None: + assert draft_probs_list is not None + draft_probs_list.append(draft_probs) + draft_token_ids_list.append(draft_token_ids) + + draft_token_ids = torch.stack(draft_token_ids_list, dim=1) + if draft_probs_list is not None: + self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous() + return draft_token_ids diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index e39db0f23bf..583fd78ced0 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -32,23 +32,20 @@ class BudgetGraphMetadata: """Metadata for a single budget graph. CUDA graph replay pattern: - 1. Copy new batch data into input_buffer (e.g. pixel_values) - 2. Copy precomputed values into metadata_buffers - 3. Replay graph - 4. Read encoder outputs from output_buffer + * Copy precomputed values into input_buffers + * Replay graph + * Read encoder outputs from output_buffer """ token_budget: int max_batch_size: int # Max number of images/videos per batch max_frames_per_batch: int # Max total frames per batch (for video) graph: torch.cuda.CUDAGraph - # The input tensor updated before replay (e.g. pixel_values) - input_buffer: torch.Tensor # Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata). # Before replay the manager updates these in-place. By default buffers are # zeroed before slice-copying the actual values; model-specific padding # behavior is provided by EncoderCudaGraphConfig.padding_logics. - metadata_buffers: dict[str, torch.Tensor] + input_buffers: dict[str, torch.Tensor] # Output written by graph, read after replay output_buffer: torch.Tensor @@ -156,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 @@ -186,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( @@ -196,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( @@ -214,29 +222,23 @@ class EncoderCudaGraphManager: self.dtype, ) - mm_kwargs = capture_inputs.mm_kwargs - buffers = capture_inputs.buffers + values = capture_inputs.values with torch.inference_mode(): - output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers) + output = self.model.encoder_cudagraph_forward({**values}) output_buffer = torch.empty_like(output) graph = torch.cuda.CUDAGraph() - with torch.inference_mode(), torch.cuda.graph(graph): - output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers) + with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool): + output = self.model.encoder_cudagraph_forward({**values}) output_buffer.copy_(output) - # Since the image and video modalities share the same per-patch shape, - # so we can use the image dummy inputs to capture CUDA graph for both - # image and video. - input_key = self.config.input_key_by_modality["image"] self.budget_graphs[token_budget] = BudgetGraphMetadata( token_budget=token_budget, max_batch_size=self.max_batch_size, max_frames_per_batch=self.max_frames_per_batch, graph=graph, - input_buffer=mm_kwargs[input_key], - metadata_buffers=buffers, + input_buffers=values, output_buffer=output_buffer, ) @@ -273,15 +275,12 @@ class EncoderCudaGraphManager: self, mm_kwargs: dict[str, Any], token_budget: int, - replay_buffers: dict[str, torch.Tensor | None], ) -> torch.Tensor | None: """Execute budget graph. Args: mm_kwargs: Multimodal inputs for the batch. token_budget: Token budget to use. - replay_buffers: Buffer values to copy into captured buffers. - None values leave the corresponding buffer unchanged. Returns: Encoder outputs, or None if graph not captured. @@ -293,22 +292,18 @@ class EncoderCudaGraphManager: graph_meta = self.budget_graphs[token_budget] - # Copy the input tensor. Buffers are sized for the full budget; - # actual inputs may be smaller. Zero then slice-copy so padded - # positions are invisible to attention (cu_seqlens masks them out). - input_key = self.config.input_key_by_modality[ - self.model.get_input_modality(mm_kwargs) - ] - src = mm_kwargs[input_key] - n = src.shape[0] - graph_meta.input_buffer[:n].copy_(src) + replay = self.model.prepare_encoder_cudagraph_replay_buffers( + mm_kwargs, + self.max_batch_size, + self.max_frames_per_batch, + ) # Copy metadata buffers using keys from config.buffer_keys. for key in self.config.buffer_keys: - src = replay_buffers.get(key) + src = replay.values.get(key) if src is None: continue - buf = graph_meta.metadata_buffers[key] + buf = graph_meta.input_buffers[key] if src.ndim == 0: buf.copy_(src) else: @@ -430,16 +425,9 @@ class EncoderCudaGraphManager: token_budget, (token_budget - batch_out_tokens) / token_budget * 100, ) - replay = self.model.prepare_encoder_cudagraph_replay_buffers( - batch_mm_kwargs, - self.max_batch_size, - self.max_frames_per_batch, - ) # graph_hits counted inside _run_budget_graph after replay. - output = self._run_budget_graph( - batch_mm_kwargs, token_budget, replay.buffers - ) + output = self._run_budget_graph(batch_mm_kwargs, token_budget) assert output is not None self.model.postprocess_encoder_output( output, diff --git a/vllm/v1/worker/encoder_cudagraph_defs.py b/vllm/v1/worker/encoder_cudagraph_defs.py index 20f1d7c33d4..7fb08f63aaf 100644 --- a/vllm/v1/worker/encoder_cudagraph_defs.py +++ b/vllm/v1/worker/encoder_cudagraph_defs.py @@ -4,7 +4,6 @@ from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any import torch @@ -40,11 +39,6 @@ class EncoderCudaGraphConfig: modalities: list[str] """Supported modalities (e.g. ["image"]).""" - input_key_by_modality: dict[str, str] - """Per-modality input tensor key mapping, e.g. - {"image": "pixel_values", "video": "pixel_values_videos"}. - """ - buffer_keys: list[str] """Keys for the tensor buffers recorded into the CUDA graph. Before replay the manager zeros then slice-copies new data @@ -74,11 +68,7 @@ class EncoderCudaGraphCaptureInputs: Returned by ``prepare_encoder_cudagraph_capture_inputs()``. """ - mm_kwargs: dict[str, Any] - """Dummy forward inputs (model-specific keys). - For Qwen3-VL this contains pixel_values and grid_thw.""" - - buffers: dict[str, torch.Tensor] + values: dict[str, torch.Tensor] """Precomputed tensor buffers that will be recorded into the CUDA graph. The manager stores references to these exact tensor objects and copies new data into them before each @@ -94,7 +84,7 @@ class EncoderCudaGraphReplayBuffers: Keys match ``EncoderCudaGraphConfig.buffer_keys``. """ - buffers: dict[str, torch.Tensor | None] + values: dict[str, torch.Tensor | None] """Data to copy into the captured buffers before replay. ``None`` values leave the corresponding captured buffer unchanged.""" diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index c2ba12437f8..6fc55ee3203 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -199,7 +199,12 @@ def _reshape_kv_cache( if isinstance(kv_cache_spec, AttentionSpec): has_attn = True - num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size + # Use storage_block_size: it equals block_size for uncompressed + # specs but is smaller for compressed ones (DeepSeek V4), which + # store block_size tokens in block_size // compress_ratio slots. + num_blocks_per_kv_block = ( + kv_cache_spec.storage_block_size // kernel_block_size + ) kernel_num_blocks = num_blocks * num_blocks_per_kv_block kv_cache_shape = group.backend.get_kv_cache_shape( kernel_num_blocks, 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 c7a7ffe442d..0648de29859 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -9,6 +9,10 @@ import torch import torch.nn as nn from tqdm import tqdm +from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphWrapper, + is_breakable_cudagraph_enabled, +) from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode @@ -116,6 +120,13 @@ class CudaGraphManager: ) self._init_candidates() + # Breakable CUDA graph (PW CUDA graph without torch.compile) + self.use_breakable_cg = ( + is_breakable_cudagraph_enabled() + and self.cudagraph_mode.has_piecewise_cudagraphs() + ) + self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None + def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" capture_sizes = self.compilation_config.cudagraph_capture_sizes @@ -283,6 +294,19 @@ class CudaGraphManager: get_offloader().sync_prev_onload() self.graphs[desc].replay() + def init_breakable_cg_runner(self, model: nn.Module) -> None: + if self.breakable_cg_runner is None: + self.breakable_cg_runner = BreakableCUDAGraphWrapper( + model, self.vllm_config + ) + + def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any: + if not self.use_breakable_cg: + # Default: Use torch-compiled piecewise cudagraph. + return model(**model_inputs) + assert self.breakable_cg_runner is not None + return self.breakable_cg_runner(**model_inputs) + class ModelCudaGraphManager(CudaGraphManager): """CudaGraphManager with model-specific capture and hidden state management.""" @@ -316,6 +340,8 @@ class ModelCudaGraphManager(CudaGraphManager): ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]: """Capture CUDA graphs for model forward pass.""" self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs + if self.use_breakable_cg: + self.init_breakable_cg_runner(model) def create_forward_fn( desc: BatchExecutionDescriptor, @@ -370,11 +396,16 @@ class ModelCudaGraphManager(CudaGraphManager): slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, ): - model_output = model(**model_inputs) + if cg_mode == CUDAGraphMode.PIECEWISE: + # PIECEWISE graph (compiled PW or breakable, chosen inside + # run_pw_graph). + model_output = self.run_pw_graph(model, model_inputs) + else: + model_output = model(**model_inputs) if cg_mode == CUDAGraphMode.PIECEWISE: - # PW CUDA graph internally handles the model outputs. - # No need to keep track of the hidden states. + # PW CUDA graph (compiled or breakable) internally handles the + # model outputs. No need to keep track of the hidden states. return None if self.is_last_pp_rank: 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/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index 84763310984..cdacb36e583 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy from typing import TYPE_CHECKING import torch @@ -103,11 +102,7 @@ class ActiveKVConnector(KVConnector): self.pre_forward(scheduler_output) finished_req_ids = scheduler_output.finished_req_ids kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False) - if kv_connector_output is None or kv_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.kv_connector_output = kv_connector_output - return output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) def set_disabled(self, disabled: bool) -> None: # Ensure that layer-wise connector hooks aren't called when disabled. diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e989c35e933..6f039eec9d8 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 @@ -65,7 +66,10 @@ from vllm.v1.worker.gpu.block_table import ( restore_block_table_serving_state, save_block_table_serving_state, ) -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, @@ -81,7 +85,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, ) @@ -94,7 +98,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 @@ -108,6 +112,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__) @@ -134,6 +139,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 @@ -148,6 +157,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 @@ -209,6 +225,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 @@ -343,6 +366,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 @@ -398,7 +427,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( @@ -407,7 +436,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, @@ -447,11 +476,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( @@ -676,8 +716,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: @@ -698,6 +741,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 @@ -710,11 +761,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] @@ -757,13 +810,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: @@ -830,7 +889,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): @@ -882,13 +944,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, @@ -907,7 +975,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, @@ -987,12 +1059,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: @@ -1001,19 +1074,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( @@ -1026,6 +1099,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) @@ -1144,9 +1218,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 = { @@ -1166,12 +1238,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. @@ -1200,7 +1273,17 @@ class GPUModelRunner(LoRAModelRunnerMixin): skip_compiled=skip_compiled, ): self.kv_connector.pre_forward(scheduler_output) - model_output = self.model(**model_inputs) + if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE: + # Run the PIECEWISE graph (compiled PW cudagraph or breakable + # cudagraph, chosen inside run_pw_graph). cg_mode is only + # PIECEWISE after the cudagraph manager exists. + assert self.cudagraph_manager is not None + model_output = self.cudagraph_manager.run_pw_graph( + self.model, model_inputs + ) + else: + # Eager (NONE): call the raw model directly. + model_output = self.model(**model_inputs) if self.is_last_pp_rank: if self.use_aux_hidden_state_outputs: @@ -1229,9 +1312,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): if not self.is_last_pp_rank: # Non-last PP rank: return IntermediateTensors for sending. - assert output_intermediate_tensors is not None - kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - output_intermediate_tensors.kv_connector_output = kv_connector_output return output_intermediate_tensors return None @@ -1256,20 +1336,33 @@ 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) - return None + 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) + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) # Last rank: sample tokens sampler_output, num_sampled, num_rejected = self.sample( 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( @@ -1279,8 +1372,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. @@ -1306,17 +1397,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. @@ -1324,8 +1412,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: @@ -1378,18 +1470,18 @@ class GPUModelRunner(LoRAModelRunnerMixin): finished_req_ids = self.execute_model_state.finished_req_ids self.execute_model_state = None + # Post-step KV connector related operations. + kv_connector_output = self.kv_connector.post_forward(finished_req_ids) + if not self.is_last_pp_rank: - self.postprocess_pool(input_batch) - return None + self.postprocess_num_computed_tokens(input_batch) + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) assert self.pooling_runner is not None pooler_output, is_valid = self.pooling_runner.pool( hidden_states, input_batch, self.req_states ) - # Post-step KV connector related operations. - kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - # Build the model runner output. model_runner_output = ModelRunnerOutput( req_ids=input_batch.req_ids, @@ -1404,14 +1496,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/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py index 43bece01d0e..300a57ec705 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py @@ -4,7 +4,6 @@ from collections.abc import Callable import torch -from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.block_table import BlockTables @@ -19,27 +18,7 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.utils import AttentionGroup -class EagleCudaGraphManagerBase(CudaGraphManager): - """Base CudaGraphManager for Eagle with a dedicated graph pool.""" - - def __init__( - self, - vllm_config: VllmConfig, - device: torch.device, - cudagraph_mode: CUDAGraphMode, - decode_query_len: int, - ): - super().__init__(vllm_config, device, cudagraph_mode, decode_query_len) - - # Use a dedicated pool for Eagle to avoid memory overlap with the main - # model's cudagraph. The base class uses a shared global pool, but Eagle's - # internal allocations (e.g., gumbel_sample temporaries) can conflict with - # the main model's allocations when sharing the same pool. - if cudagraph_mode: - self.pool = torch.cuda.graph_pool_handle() - - -class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase): +class PrefillEagleCudaGraphManager(CudaGraphManager): """Eagle CudaGraphManager for prefill, using pre-built attention states from the target model's capture.""" @@ -74,7 +53,7 @@ class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase): super().capture(create_forward_fn, progress_bar_desc) -class DecodeEagleCudaGraphManager(EagleCudaGraphManagerBase): +class DecodeEagleCudaGraphManager(CudaGraphManager): """Eagle CudaGraphManager for decode draft generation, building its own attention metadata from scratch.""" diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 6ae3fe793bd..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). @@ -145,9 +146,6 @@ class EagleSpeculator: cudagraph_mode, decode_query_len=1, ) - # Share a single pool between prefill and decode since they never - # execute concurrently. - self.decode_cudagraph_manager.pool = self.prefill_cudagraph_manager.pool def load_model(self, target_model: nn.Module) -> None: target_attn_layer_names = get_layers_from_vllm_config( @@ -215,12 +213,22 @@ class EagleSpeculator: ) inputs_embeds = self.inputs_embeds[:num_tokens] - ret_hidden_states = self.model( + model_inputs = dict( input_ids=self.input_buffers.input_ids[:num_tokens], positions=self.input_buffers.positions[:num_tokens], hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) + if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE: + # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable), + # chosen inside run_pw_graph. + assert self.prefill_cudagraph_manager is not None + ret_hidden_states = self.prefill_cudagraph_manager.run_pw_graph( + self.model, model_inputs + ) + else: + # Eager (NONE): call the raw model directly. + ret_hidden_states = self.model(**model_inputs) if self.method == "mtp": last_hidden_states = ret_hidden_states hidden_states = ret_hidden_states @@ -409,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, @@ -431,6 +439,8 @@ class EagleSpeculator: # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. assert self.prefill_cudagraph_manager is not None + if self.prefill_cudagraph_manager.use_breakable_cg: + self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model) self.prefill_cudagraph_manager.capture( self.prefill, attn_states, @@ -485,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 7e856c1e031..9eeaf39e857 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -49,6 +49,7 @@ from vllm.distributed.parallel_state import ( is_global_first_rank, prepare_communication_buffer_for_model, ) +from vllm.distributed.weight_transfer.base import SparseWeightPatch from vllm.forward_context import ( BatchDescriptor, set_forward_context, @@ -103,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 @@ -186,6 +187,7 @@ from vllm.v1.spec_decode.ngram_proposer_gpu import ( update_ngram_gpu_tensors_incremental, update_scheduler_for_invalid_drafts, ) +from vllm.v1.spec_decode.step3p5 import Step3p5MTPProposer from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change from vllm.v1.structured_output.utils import apply_grammar_bitmask @@ -547,6 +549,7 @@ class GPUModelRunner( | MedusaProposer | ExtractHiddenStatesProposer | Gemma4Proposer + | Step3p5MTPProposer ) if self.speculative_config.method == "custom_class": self.drafter = create_custom_proposer( # type: ignore[assignment] @@ -581,6 +584,8 @@ class GPUModelRunner( ) elif self.speculative_config.use_gemma4_mtp(): self.drafter = Gemma4Proposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_step3p5_mtp(): + self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self) elif self.speculative_config.use_dflash(): self.drafter = DFlashProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True @@ -1099,11 +1104,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, @@ -2448,7 +2453,11 @@ class GPUModelRunner( else: spec_decode_common_attn_metadata = cm # Capture per-group block tables for multi-group proposers. - if self.speculative_config and isinstance(self.drafter, Gemma4Proposer): + if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer): + self.drafter.set_per_group_attn_metadata( + kv_cache_gid, cm.block_table_tensor, cm.slot_mapping + ) + elif self.speculative_config and isinstance(self.drafter, Gemma4Proposer): self.drafter.set_per_group_block_table( kv_cache_gid, cm.block_table_tensor ) @@ -3111,23 +3120,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, @@ -3201,6 +3205,44 @@ class GPUModelRunner( return self.model.unwrap() return self.model + def apply_sparse_weight_patches(self, patches: Iterable[SparseWeightPatch]) -> None: + """Apply sparse flat-index patches directly to existing model params.""" + model = self.get_model() + for patch in patches: + param = model.get_parameter(patch.name) + if not param.data.is_contiguous(): + raise NotImplementedError( + "Sparse weight updates currently require contiguous params: " + f"{patch.name}" + ) + + if patch.indices.dtype != torch.int32: + raise ValueError( + "Sparse weight updates currently require int32 indices: " + f"{patch.name}" + ) + if patch.indices.ndim != 1 or patch.values.ndim != 1: + raise ValueError( + f"Sparse weight patches must be 1D flattened updates: {patch.name}" + ) + if patch.indices.numel() != patch.values.numel(): + raise ValueError( + "`indices` and `values` must have matching lengths for " + f"{patch.name}" + ) + if patch.values.dtype != param.dtype: + raise ValueError( + f"Sparse values dtype {patch.values.dtype} does not match " + f"parameter dtype {param.dtype} for {patch.name}" + ) + + flat_param = param.data.view(-1) + flat_param.index_copy_( + 0, + patch.indices.to(device=flat_param.device, dtype=torch.long), + patch.values.to(device=flat_param.device), + ) + def get_supported_generation_tasks(self) -> list[GenerationTask]: model = self.get_model() supported_tasks = list[GenerationTask]() @@ -4270,7 +4312,6 @@ class GPUModelRunner( if not get_pp_group().is_last_rank: # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) - hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output return hidden_states @@ -4336,6 +4377,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" @@ -4346,17 +4402,9 @@ class GPUModelRunner( # receive sampled token ids from the last PP rank. if self.use_async_scheduling and not get_pp_group().is_last_rank: self._pp_receive_prev_sampled_token_ids_to_input_batch() - if not kv_connector_output: - return None # type: ignore[return-value] - # In case of PP with kv transfer, we need to pass through the # kv_connector_output - if kv_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT - - output = copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.kv_connector_output = kv_connector_output - return output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) # Unpack ephemeral state. ( @@ -4423,9 +4471,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() @@ -6246,11 +6293,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() @@ -6286,6 +6343,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): @@ -6294,24 +6387,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 @@ -6320,73 +6429,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), @@ -6404,31 +6538,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 @@ -6455,7 +6565,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 e63f50bc8dc..259cd05554c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -992,23 +992,19 @@ class Worker(WorkerBase): def start_weight_update(self, is_checkpoint_format: bool = True) -> None: """ - Start a new weight update. - - Prepares the model for receiving weights. For checkpoint format, - this initializes state for layerwise processing. For kernel format, this is - a no-op but must still be called for consistency. + Start a new weight update session. Args: is_checkpoint_format: Whether incoming weights are in checkpoint format (need layerwise processing) or kernel format (direct - copy). Stored as state for finish_weight_update. + copy / sparse patch application). """ self._check_weight_transfer_engine() if self._weight_update_active: raise RuntimeError( - "start_weight_update called while a weight update is " - "already active. Call finish_weight_update first." + "start_weight_update called while a weight update is already " + "active. Call finish_weight_update first." ) if is_checkpoint_format: @@ -1020,16 +1016,15 @@ class Worker(WorkerBase): with torch.device(self.device): initialize_layerwise_reload(model) - # Store state so update_weights/finish_weight_update can check self._is_checkpoint_format = is_checkpoint_format self._weight_update_active = True def update_weights(self, update_info: dict) -> None: """ - Receive weights from the trainer (one or more chunks). + Receive one weight update chunk from the trainer. start_weight_update must be called before update_weights and - finish_weight_update must be called after. + finish_weight_update must be called after all chunks have been sent. Args: update_info: Dictionary containing backend-specific update info @@ -1042,52 +1037,72 @@ class Worker(WorkerBase): "start_weight_update must be called before update_weights." ) - # Parse dict into backend-specific typed dataclass - typed_update_info = self.weight_transfer_engine.parse_update_info(update_info) + update_succeeded = False + try: + # Parse dict into backend-specific typed dataclass + typed_update_info = self.weight_transfer_engine.parse_update_info( + update_info + ) - model = self.model_runner.model + with torch.device(self.device): + if self._is_checkpoint_format: + if typed_update_info.update_kind != "dense": + raise ValueError( + "Sparse weight updates require " + "`start_weight_update(is_checkpoint_format=False)`." + ) - with torch.device(self.device): - if self._is_checkpoint_format: - self.weight_transfer_engine.receive_weights( - typed_update_info, - load_weights=model.load_weights, - ) - else: - # Weights are already in kernel format, copy directly - def load_weights_direct( - weights: list[tuple[str, torch.Tensor]], - ) -> None: - for name, weight in weights: - param = model.get_parameter(name) - param.copy_(weight) + model = self.model_runner.model - self.weight_transfer_engine.receive_weights( - typed_update_info, - load_weights=load_weights_direct, - ) + # Use layerwise reload pattern for checkpoint format weights + self.weight_transfer_engine.receive_weights( + typed_update_info, + load_weights=model.load_weights, + ) + elif typed_update_info.update_kind == "sparse_flat": + if self.parallel_config.world_size != 1: + raise NotImplementedError( + "Sparse weight updates currently require TP=1 and PP=1" + ) + self.weight_transfer_engine.receive_sparse_weights( + typed_update_info, + apply_patches=self.model_runner.apply_sparse_weight_patches, + ) + else: + model = self.model_runner.model - # NCCL broadcast/packed path are asynchronous. - # Sync here so the next step uses the new weights. - torch.accelerator.synchronize() + # Weights are already in kernel format, copy directly. + def load_weights_direct( + weights: list[tuple[str, torch.Tensor]], + ) -> None: + for name, weight in weights: + param = model.get_parameter(name) + param.copy_(weight) + + self.weight_transfer_engine.receive_weights( + typed_update_info, + load_weights=load_weights_direct, + ) + + # NCCL broadcast/packed path are asynchronous. + # Sync here so the next step uses the new weights. + torch.accelerator.synchronize() + update_succeeded = True + finally: + if not update_succeeded: + self._weight_update_active = False + self._is_checkpoint_format = True def finish_weight_update(self) -> None: - """ - Finish the current weight update. - - For checkpoint format, this runs layerwise postprocessing. - Uses the is_checkpoint_format state stored by start_weight_update. - """ + """Finish the current weight update session.""" self._check_weight_transfer_engine() if not self._weight_update_active: raise RuntimeError( - "start_weight_update must be called before finish_weight_update." + "finish_weight_update called without a matching start_weight_update." ) - is_checkpoint_format = self._is_checkpoint_format - - if is_checkpoint_format: + if self._is_checkpoint_format: from vllm.model_executor.model_loader.reload import ( finalize_layerwise_reload, ) @@ -1096,7 +1111,6 @@ class Worker(WorkerBase): with torch.device(self.device): finalize_layerwise_reload(model, self.model_config) - # Reset state self._weight_update_active = False self._is_checkpoint_format = True @@ -1133,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/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py index 4fc1aff94fe..797e59c0290 100644 --- a/vllm/v1/worker/kv_connector_model_runner_mixin.py +++ b/vllm/v1/worker/kv_connector_model_runner_mixin.py @@ -4,7 +4,6 @@ Define KV connector functionality mixin for model runners. """ -import copy from collections.abc import Generator from contextlib import AbstractContextManager, contextmanager, nullcontext from typing import TYPE_CHECKING @@ -20,7 +19,6 @@ from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionBackend from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig from vllm.v1.outputs import ( - EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput, ModelRunnerOutput, ) @@ -47,12 +45,7 @@ class KVConnectorModelRunnerMixin: ): pass - if kv_connector_output.is_empty(): - return EMPTY_MODEL_RUNNER_OUTPUT - - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT) - output.kv_connector_output = kv_connector_output - return output + return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) @staticmethod def maybe_get_kv_connector_output( 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