forked from Karylab-cklius/vllm
Compare commits
62
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d7cbffd5f8 | ||
|
|
fc5fda105f | ||
|
|
b07ec92faa | ||
|
|
27ffbfde8d | ||
|
|
229e01e9e1 | ||
|
|
191146dba5 | ||
|
|
f3a920a076 | ||
|
|
149daf0d72 | ||
|
|
917fdb5bf7 | ||
|
|
4b594b4aa1 | ||
|
|
7d10a4cfce | ||
|
|
910cc8543a | ||
|
|
431934522b | ||
|
|
61a09532f2 | ||
|
|
3de4b2bf3c | ||
|
|
b44311b6ef | ||
|
|
b0d7875180 | ||
|
|
53c2f20dd9 | ||
|
|
37e370fe93 | ||
|
|
2dc5a72e7e | ||
|
|
c79ff5f918 | ||
|
|
1a659a0c37 | ||
|
|
0f6cf7f628 | ||
|
|
c79ad3ae21 | ||
|
|
61c9ef986a | ||
|
|
d6dbdb9b0d | ||
|
|
06da482fb4 | ||
|
|
2f75e7f712 | ||
|
|
7c21548ce3 | ||
|
|
9df2f91232 | ||
|
|
387189c429 | ||
|
|
75576c63be | ||
|
|
16aca639b7 | ||
|
|
6049424b7e | ||
|
|
060b5f61dc | ||
|
|
ec59c1579f | ||
|
|
1750e443f2 | ||
|
|
ba18929079 | ||
|
|
0500ca6a58 | ||
|
|
a1c15bcb0f | ||
|
|
4809de7317 | ||
|
|
05781e21dd | ||
|
|
85f638a2b8 | ||
|
|
08e5067561 | ||
|
|
a7d00ec051 | ||
|
|
b8fb56d970 | ||
|
|
96a739289e | ||
|
|
60d443f738 | ||
|
|
1dca300653 | ||
|
|
fca252d59e | ||
|
|
33178f9006 | ||
|
|
b2b8f679d0 | ||
|
|
de6ec294ef | ||
|
|
61e10f0116 | ||
|
|
6e96891ba0 | ||
|
|
47f1b47a73 | ||
|
|
5aab491bc9 | ||
|
|
5812e1a66b | ||
|
|
8950394e0a | ||
|
|
7bb49be4d1 | ||
|
|
c67650f04b | ||
|
|
f890e1dbe2 |
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
|
||||
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
|
||||
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
|
||||
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_METADATA_VERSION="1"
|
||||
IMAGE_EXISTED_BEFORE_BUILD=0
|
||||
|
||||
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
|
||||
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
|
||||
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
|
||||
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
|
||||
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
|
||||
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
|
||||
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
|
||||
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
|
||||
|
||||
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
|
||||
return 0
|
||||
fi
|
||||
|
||||
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
if [[ -n "${!var:-}" ]]; then
|
||||
echo "Using provided ${var}: ${!var}"
|
||||
continue
|
||||
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
|
||||
compute_dependency_cache_keys() {
|
||||
local bake_dir=""
|
||||
local dockerfile_rocm=""
|
||||
local rixl_branch=""
|
||||
local nixl_branch=""
|
||||
local ucx_branch=""
|
||||
local rocshmem_branch=""
|
||||
local deepep_branch=""
|
||||
local rixl_material=""
|
||||
local nixl_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")
|
||||
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_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=$(
|
||||
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
|
||||
NIXL_CACHE_KEY=$(
|
||||
compose_dependency_cache_key \
|
||||
"${rixl_branch}-ucx-${ucx_branch}" \
|
||||
"${rixl_material}"
|
||||
"${nixl_branch}-ucx-${ucx_branch}" \
|
||||
"${nixl_material}"
|
||||
)
|
||||
export RIXL_CACHE_KEY
|
||||
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
|
||||
export NIXL_CACHE_KEY
|
||||
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
|
||||
fi
|
||||
|
||||
if [[ -n "${rocshmem_branch}" ]]; then
|
||||
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
|
||||
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:-}"
|
||||
nixl-rocm-ci)
|
||||
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
|
||||
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
fi
|
||||
;;
|
||||
rocshmem-rocm-ci)
|
||||
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
|
||||
|
||||
resolve_ci_base_dependency_targets() {
|
||||
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
|
||||
local rixl_ref=""
|
||||
local nixl_ref=""
|
||||
local rocshmem_ref=""
|
||||
local deepep_ref=""
|
||||
|
||||
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
|
||||
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
|
||||
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
|
||||
add_dependency_cache_target "${target}"
|
||||
fi
|
||||
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
|
||||
;;
|
||||
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}"
|
||||
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${nixl_ref}"; then
|
||||
echo "NIXL dependency cache exists: ${nixl_ref}"
|
||||
else
|
||||
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
|
||||
add_dependency_cache_target "rixl-rocm-ci"
|
||||
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
|
||||
add_dependency_cache_target "nixl-rocm-ci"
|
||||
fi
|
||||
fi
|
||||
|
||||
|
||||
@@ -56,6 +56,15 @@ fi
|
||||
export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY
|
||||
export PYTHONFAULTHANDLER
|
||||
|
||||
# The AMD pipeline template sets this for eligible pilot jobs. Capture it
|
||||
# before clear_ci_orchestration_env removes CI-only controls from test processes.
|
||||
hf_offline_retry_enabled="${VLLM_CI_HF_OFFLINE_RETRY:-0}"
|
||||
if [[ "${hf_offline_retry_enabled}" != "0" && "${hf_offline_retry_enabled}" != "1" ]]; then
|
||||
echo "VLLM_CI_HF_OFFLINE_RETRY must be 0 or 1" >&2
|
||||
exit 2
|
||||
fi
|
||||
hf_retry_command_file=""
|
||||
|
||||
# Export Python path for commands that run directly on the host. Containerized
|
||||
# tests set this to /vllm-workspace below so spawned Python processes do not
|
||||
# depend on their current working directory.
|
||||
@@ -73,6 +82,7 @@ report_docker_usage() {
|
||||
clear_ci_orchestration_env() {
|
||||
unset -v \
|
||||
VLLM_TEST_GROUP_NAME \
|
||||
VLLM_CI_HF_OFFLINE_RETRY \
|
||||
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \
|
||||
VLLM_CI_ARTIFACT_STEP \
|
||||
VLLM_TEST_CACHE \
|
||||
@@ -92,6 +102,29 @@ clear_ci_orchestration_env() {
|
||||
VLLM_ALLOW_DEPRECATED_BEAM_SEARCH
|
||||
}
|
||||
|
||||
prepare_hf_retry_command_file() {
|
||||
local command_text=$1
|
||||
|
||||
if ! hf_retry_command_file=$(mktemp "${TMPDIR:-/tmp}/vllm-hf-offline-command.XXXXXX"); then
|
||||
echo "Failed to create the Hugging Face retry command file" >&2
|
||||
return 1
|
||||
fi
|
||||
if ! printf '%s\n' "${command_text}" >"${hf_retry_command_file}"; then
|
||||
echo "Failed to write the Hugging Face retry command file" >&2
|
||||
rm -f -- "${hf_retry_command_file}"
|
||||
hf_retry_command_file=""
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# shellcheck disable=SC2329 # Called by cleanup functions registered as traps.
|
||||
cleanup_hf_retry_command_file() {
|
||||
if [[ -n "${hf_retry_command_file}" ]]; then
|
||||
rm -f -- "${hf_retry_command_file}"
|
||||
hf_retry_command_file=""
|
||||
fi
|
||||
}
|
||||
|
||||
cleanup_network() {
|
||||
local max_nodes=${NUM_NODES:-2}
|
||||
for node in $(seq 0 $((max_nodes - 1))); do
|
||||
@@ -649,6 +682,7 @@ if is_native_runtime; then
|
||||
artifact_work_dir=""
|
||||
|
||||
cleanup_native_workspace() {
|
||||
cleanup_hf_retry_command_file
|
||||
if [[ -n "${artifact_work_dir}" ]]; then
|
||||
rm -rf "${artifact_work_dir}"
|
||||
fi
|
||||
@@ -691,7 +725,13 @@ if is_native_runtime; then
|
||||
run_native_preflight || exit 1
|
||||
# Keep AMD CI orchestration variables out of vLLM's runtime environment.
|
||||
clear_ci_orchestration_env
|
||||
/bin/bash -o pipefail -c "${commands}"
|
||||
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
|
||||
prepare_hf_retry_command_file "${commands}" || exit 1
|
||||
bash "${VLLM_CI_WORKSPACE:-/vllm-workspace}/.buildkite/scripts/hf-offline-retry.sh" \
|
||||
"${hf_retry_command_file}"
|
||||
else
|
||||
/bin/bash -o pipefail -c "${commands}"
|
||||
fi
|
||||
handle_pytest_exit "$?"
|
||||
fi
|
||||
|
||||
@@ -709,6 +749,7 @@ artifact_work_dir=""
|
||||
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
|
||||
remove_docker_container() {
|
||||
cleanup_hf_retry_command_file
|
||||
if docker container inspect "${container_name}" >/dev/null 2>&1; then
|
||||
docker rm -f "${container_name}" || true
|
||||
fi
|
||||
@@ -850,6 +891,10 @@ fi
|
||||
# --- Route: multi-node vs single-node ---
|
||||
clear_ci_orchestration_env
|
||||
if is_multi_node "$commands"; then
|
||||
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
|
||||
echo "Hugging Face offline retry is not enabled for AMD multi-node jobs yet" >&2
|
||||
exit 2
|
||||
fi
|
||||
echo "--- Multi-node job detected"
|
||||
export DCKR_VER=$(docker --version | sed 's/Docker version \(.*\), build .*/\1/')
|
||||
|
||||
@@ -894,6 +939,16 @@ if is_multi_node "$commands"; then
|
||||
fi
|
||||
else
|
||||
echo "--- Single-node job"
|
||||
hf_retry_container_args=()
|
||||
container_test_command="${commands}"
|
||||
if [[ "${hf_offline_retry_enabled}" == "1" ]]; then
|
||||
prepare_hf_retry_command_file "${commands}" || exit 1
|
||||
hf_retry_container_path="/tmp/vllm-hf-offline-command.sh"
|
||||
hf_retry_container_args=(
|
||||
-v "${hf_retry_command_file}:${hf_retry_container_path}:ro"
|
||||
)
|
||||
container_test_command="bash /vllm-workspace/.buildkite/scripts/hf-offline-retry.sh ${hf_retry_container_path}"
|
||||
fi
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
docker_run_terminal_args=(-i)
|
||||
if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
|
||||
@@ -950,9 +1005,10 @@ else
|
||||
-e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
|
||||
-e "PYTORCH_ROCM_ARCH=" \
|
||||
"${standalone_merge_base_env[@]}" \
|
||||
"${hf_retry_container_args[@]}" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
|
||||
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${container_test_command}"
|
||||
|
||||
exit_code=$?
|
||||
handle_pytest_exit "$exit_code"
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
|
||||
helper="${script_dir}/hf-offline-retry.sh"
|
||||
test_root=$(mktemp -d "${TMPDIR:-/tmp}/vllm-hf-offline-retry-test.XXXXXX")
|
||||
|
||||
cleanup() {
|
||||
rm -rf -- "${test_root}"
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
assert_equal() {
|
||||
local expected=$1
|
||||
local actual=$2
|
||||
local description=$3
|
||||
|
||||
if [[ "${actual}" != "${expected}" ]]; then
|
||||
printf 'FAIL: %s\n expected: %q\n actual: %q\n' \
|
||||
"${description}" "${expected}" "${actual}" >&2
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
status_command_file="${test_root}/status-command.sh"
|
||||
cat >"${status_command_file}" <<'VLLM_TEST_COMMAND'
|
||||
if [[ "${HF_HUB_OFFLINE-}" == "1" &&
|
||||
"${TRANSFORMERS_OFFLINE-}" == "1" &&
|
||||
"${HF_DATASETS_OFFLINE-}" == "1" ]]; then
|
||||
printf '%s\n' offline >>"${ATTEMPTS_FILE}"
|
||||
exit "${OFFLINE_STATUS}"
|
||||
fi
|
||||
|
||||
if [[ -z "${HF_HUB_OFFLINE+x}" &&
|
||||
-z "${TRANSFORMERS_OFFLINE+x}" &&
|
||||
-z "${HF_DATASETS_OFFLINE+x}" ]]; then
|
||||
printf '%s\n' online >>"${ATTEMPTS_FILE}"
|
||||
exit "${ONLINE_STATUS}"
|
||||
fi
|
||||
|
||||
exit 91
|
||||
VLLM_TEST_COMMAND
|
||||
|
||||
run_status_case() {
|
||||
local name=$1
|
||||
local offline_status=$2
|
||||
local online_status=$3
|
||||
local expected_status=$4
|
||||
local expected_attempts=$5
|
||||
local attempts_file="${test_root}/${name}.attempts"
|
||||
local actual_status
|
||||
local actual_attempts
|
||||
|
||||
set +e
|
||||
ATTEMPTS_FILE="${attempts_file}" \
|
||||
OFFLINE_STATUS="${offline_status}" \
|
||||
ONLINE_STATUS="${online_status}" \
|
||||
bash "${helper}" "${status_command_file}"
|
||||
actual_status=$?
|
||||
set -e
|
||||
|
||||
actual_attempts=$(<"${attempts_file}")
|
||||
assert_equal "${expected_status}" "${actual_status}" "${name} exit status"
|
||||
assert_equal "${expected_attempts}" "${actual_attempts}" "${name} attempts"
|
||||
}
|
||||
|
||||
run_status_case offline_success 0 99 0 "offline"
|
||||
run_status_case retry_status_1 1 0 0 $'offline\nonline'
|
||||
run_status_case retry_status_2 2 0 0 $'offline\nonline'
|
||||
run_status_case retry_status_123 123 0 0 $'offline\nonline'
|
||||
run_status_case no_retry_status_42 42 0 42 "offline"
|
||||
run_status_case online_failure 1 7 7 $'offline\nonline'
|
||||
|
||||
transport_file="${test_root}/transport.out"
|
||||
transport_command_file="${test_root}/transport-command.sh"
|
||||
cat >"${transport_command_file}" <<'VLLM_TRANSPORT_COMMAND'
|
||||
cat >"${OUTPUT_FILE}" <<'INNER_PAYLOAD'
|
||||
literal:$HOME
|
||||
quotes:'single' "double"
|
||||
two command-file lines
|
||||
INNER_PAYLOAD
|
||||
printf 'expanded:%s\n' "${TRANSPORT_VALUE}" >>"${OUTPUT_FILE}"
|
||||
VLLM_TRANSPORT_COMMAND
|
||||
|
||||
OUTPUT_FILE="${transport_file}" \
|
||||
TRANSPORT_VALUE='value with spaces, $dollars, and "quotes"' \
|
||||
bash "${helper}" "${transport_command_file}"
|
||||
|
||||
transport_output=$(<"${transport_file}")
|
||||
assert_equal \
|
||||
$'literal:$HOME\nquotes:\'single\' "double"\ntwo command-file lines\nexpanded:value with spaces, $dollars, and "quotes"' \
|
||||
"${transport_output}" \
|
||||
"quotes, dollars, and newlines"
|
||||
|
||||
stdin_file="${test_root}/stdin.out"
|
||||
stdin_payload="stdin payload with \$dollars and \"quotes\""
|
||||
stdin_command_file="${test_root}/stdin-command.sh"
|
||||
cat >"${stdin_command_file}" <<'VLLM_STDIN_COMMAND'
|
||||
IFS= read -r input
|
||||
printf '%s\n' "${input}" >"${OUTPUT_FILE}"
|
||||
VLLM_STDIN_COMMAND
|
||||
|
||||
printf '%s\n' "${stdin_payload}" |
|
||||
OUTPUT_FILE="${stdin_file}" bash "${helper}" "${stdin_command_file}"
|
||||
|
||||
stdin_output=$(<"${stdin_file}")
|
||||
assert_equal \
|
||||
"${stdin_payload}" \
|
||||
"${stdin_output}" \
|
||||
"stdin preservation"
|
||||
|
||||
set +e
|
||||
bash "${helper}" >/dev/null 2>&1
|
||||
usage_status=$?
|
||||
bash "${helper}" "${test_root}/missing-command" >/dev/null 2>&1
|
||||
missing_status=$?
|
||||
set -e
|
||||
assert_equal "64" "${usage_status}" "missing argument exit status"
|
||||
assert_equal "66" "${missing_status}" "unreadable command file exit status"
|
||||
|
||||
echo "PASS: hf-offline-retry"
|
||||
@@ -0,0 +1,48 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -uo pipefail
|
||||
|
||||
if [[ $# -ne 1 ]]; then
|
||||
echo "Usage: $0 COMMAND_FILE" >&2
|
||||
exit 64
|
||||
fi
|
||||
|
||||
command_file=$1
|
||||
if [[ ! -f "${command_file}" || ! -r "${command_file}" ]]; then
|
||||
echo "Command file is not readable: ${command_file}" >&2
|
||||
exit 66
|
||||
fi
|
||||
|
||||
run_attempt() (
|
||||
if [[ "$1" == "offline" ]]; then
|
||||
export HF_HUB_OFFLINE=1
|
||||
export TRANSFORMERS_OFFLINE=1
|
||||
export HF_DATASETS_OFFLINE=1
|
||||
else
|
||||
unset HF_HUB_OFFLINE
|
||||
unset TRANSFORMERS_OFFLINE
|
||||
unset HF_DATASETS_OFFLINE
|
||||
fi
|
||||
|
||||
bash -e -o pipefail -- "${command_file}"
|
||||
)
|
||||
|
||||
echo "--- :package: Hugging Face offline attempt"
|
||||
run_attempt offline
|
||||
offline_status=$?
|
||||
|
||||
if [[ ${offline_status} -eq 0 ]]; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
case "${offline_status}" in
|
||||
1 | 2 | 123)
|
||||
echo "--- :globe_with_meridians: Hugging Face offline attempt failed; retrying online"
|
||||
run_attempt online
|
||||
online_status=$?
|
||||
exit "${online_status}"
|
||||
;;
|
||||
*)
|
||||
exit "${offline_status}"
|
||||
;;
|
||||
esac
|
||||
@@ -29,7 +29,6 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
|
||||
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
|
||||
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
|
||||
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
|
||||
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
|
||||
|
||||
log_section() {
|
||||
echo "--- $*"
|
||||
@@ -107,18 +106,6 @@ install_cargo_nextest() {
|
||||
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
|
||||
}
|
||||
|
||||
install_cargo_llvm_cov() {
|
||||
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
|
||||
local toolchain
|
||||
toolchain="$(rust_toolchain)"
|
||||
rustup component add --toolchain "$toolchain" llvm-tools-preview
|
||||
cargo binstall \
|
||||
--no-confirm \
|
||||
--force \
|
||||
--secure \
|
||||
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
|
||||
}
|
||||
|
||||
install_uv() {
|
||||
log_section "Installing uv ${UV_VERSION}"
|
||||
curl -L --proto '=https' --tlsv1.2 -sSf \
|
||||
@@ -189,41 +176,14 @@ run_tests() {
|
||||
setup_pyo3_python
|
||||
install_cargo_binstall
|
||||
install_cargo_nextest
|
||||
install_cargo_llvm_cov
|
||||
|
||||
log_section "Running cargo nextest with Rust coverage"
|
||||
mkdir -p artifacts
|
||||
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
|
||||
cargo llvm-cov clean \
|
||||
--manifest-path rust/Cargo.toml \
|
||||
--profraw-only
|
||||
|
||||
set +e
|
||||
cargo llvm-cov nextest \
|
||||
log_section "Running cargo nextest"
|
||||
cargo nextest run \
|
||||
--manifest-path rust/Cargo.toml \
|
||||
--workspace \
|
||||
--all-features \
|
||||
--locked \
|
||||
--no-fail-fast \
|
||||
--no-clean \
|
||||
--lcov \
|
||||
--output-path artifacts/rust-unit.lcov \
|
||||
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
|
||||
local coverage_rc=$?
|
||||
|
||||
local upload_rc=0
|
||||
if [[ $coverage_rc -eq 0 ]]; then
|
||||
# shellcheck source=.buildkite/scripts/rust-coverage.sh
|
||||
source .buildkite/scripts/rust-coverage.sh
|
||||
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
|
||||
upload_rc=$?
|
||||
fi
|
||||
set -e
|
||||
|
||||
if [[ $coverage_rc -ne 0 ]]; then
|
||||
return "$coverage_rc"
|
||||
fi
|
||||
return "$upload_rc"
|
||||
--no-fail-fast
|
||||
}
|
||||
|
||||
install_protoc
|
||||
|
||||
@@ -1,182 +0,0 @@
|
||||
#!/bin/sh
|
||||
|
||||
RUST_CODECOV_VERSION="v11.3.1"
|
||||
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
|
||||
|
||||
rust_coverage_repo_root() {
|
||||
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
|
||||
printf '%s\n' /vllm-workspace
|
||||
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
|
||||
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
|
||||
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
|
||||
else
|
||||
git rev-parse --show-toplevel
|
||||
fi
|
||||
}
|
||||
|
||||
rust_coverage_start() {
|
||||
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
|
||||
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
|
||||
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
|
||||
mkdir -p "$RUST_COVERAGE_DIR"
|
||||
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
|
||||
export LLVM_PROFILE_FILE
|
||||
trap rust_coverage_finalize 0
|
||||
}
|
||||
|
||||
rust_coverage_objects() {
|
||||
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
|
||||
python3 - "$rust_cov_objects_manifest" <<'PY'
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
for relative in Path(sys.argv[1]).read_text().splitlines():
|
||||
for entry in sys.path:
|
||||
path = Path(entry or ".").resolve() / relative
|
||||
if path.is_file():
|
||||
print(path)
|
||||
break
|
||||
else:
|
||||
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
|
||||
PY
|
||||
}
|
||||
|
||||
rust_coverage_collect() {
|
||||
rust_cov_collect_flag=${1:?coverage flag is required}
|
||||
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
|
||||
|
||||
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
|
||||
rust_cov_collect_primary=
|
||||
set --
|
||||
while IFS= read -r rust_cov_collect_object; do
|
||||
if [ -z "$rust_cov_collect_primary" ]; then
|
||||
rust_cov_collect_primary=$rust_cov_collect_object
|
||||
else
|
||||
set -- "$@" "--object=$rust_cov_collect_object"
|
||||
fi
|
||||
done <<EOF
|
||||
$rust_cov_collect_objects
|
||||
EOF
|
||||
|
||||
llvm-profdata merge \
|
||||
-sparse \
|
||||
"$RUST_COVERAGE_DIR"/*.profraw \
|
||||
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
|
||||
llvm-cov export \
|
||||
"$rust_cov_collect_primary" \
|
||||
"$@" \
|
||||
--format=lcov \
|
||||
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
|
||||
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
|
||||
> "$rust_cov_collect_lcov" || return 1
|
||||
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
|
||||
export RUST_COVERAGE_LCOV
|
||||
}
|
||||
|
||||
rust_coverage_upload() {
|
||||
rust_cov_upload_lcov=${1:?LCOV path is required}
|
||||
rust_cov_upload_flag=${2:?coverage flag is required}
|
||||
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
|
||||
|
||||
if [ "$(uname -m)" != "x86_64" ]; then
|
||||
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|
||||
|| return 1
|
||||
curl -fsSL \
|
||||
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
|
||||
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
|
||||
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
|
||||
| sha256sum -c - || return 1
|
||||
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
|
||||
|
||||
rust_cov_upload_slug="vllm-project/vllm"
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
|
||||
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
|
||||
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
|
||||
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
|
||||
case "$rust_cov_upload_slug" in
|
||||
*/*) ;;
|
||||
*) rust_cov_upload_slug="vllm-project/vllm" ;;
|
||||
esac
|
||||
fi
|
||||
|
||||
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
|
||||
if [ -z "${CODECOV_TOKEN:-}" ]; then
|
||||
# Codecov accepts tokenless public uploads on unprotected branch names.
|
||||
# A colon-separated prefix keeps feature-branch and fork uploads from
|
||||
# requiring a repository secret.
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
|
||||
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
|
||||
else
|
||||
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
|
||||
fi
|
||||
fi
|
||||
|
||||
set --
|
||||
set -- "$@" upload-process
|
||||
set -- "$@" --file "$rust_cov_upload_lcov"
|
||||
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
|
||||
# source-line fix scanning, which is unrelated to path mapping.
|
||||
set -- "$@" --disable-search --disable-file-fixes
|
||||
set -- "$@" --fail-on-error --git-service github
|
||||
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
|
||||
set -- "$@" --branch "$rust_cov_upload_branch"
|
||||
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
|
||||
set -- "$@" --slug "$rust_cov_upload_slug"
|
||||
set -- "$@" --flag "$rust_cov_upload_flag"
|
||||
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
|
||||
set -- "$@" --dir "$rust_cov_upload_repo_root"
|
||||
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
|
||||
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
|
||||
fi
|
||||
|
||||
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
|
||||
# E2E steps run from tests/, so execute from the repository root to resolve
|
||||
# codecov.yml and repository paths consistently.
|
||||
(
|
||||
cd "$rust_cov_upload_repo_root" || exit 1
|
||||
"$rust_cov_upload_codecov_dir/codecov" "$@"
|
||||
) >"$rust_cov_upload_log" 2>&1
|
||||
rust_cov_upload_rc=$?
|
||||
cat "$rust_cov_upload_log"
|
||||
# v11.3.1 can log API failures while returning zero even with
|
||||
# --fail-on-error. Preserve the strict CI contract explicitly.
|
||||
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
|
||||
echo "Codecov CLI reported an upload error" >&2
|
||||
rust_cov_upload_rc=1
|
||||
fi
|
||||
rm -rf "$rust_cov_upload_codecov_dir"
|
||||
return "$rust_cov_upload_rc"
|
||||
}
|
||||
|
||||
rust_coverage_finalize() {
|
||||
rust_cov_finalize_test_rc=$?
|
||||
trap - 0
|
||||
set +e
|
||||
|
||||
rust_coverage_collect "$RUST_COVERAGE_FLAG"
|
||||
rust_cov_finalize_collect_rc=$?
|
||||
|
||||
rust_cov_finalize_upload_rc=0
|
||||
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
|
||||
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
|
||||
rust_cov_finalize_upload_rc=$?
|
||||
fi
|
||||
|
||||
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
|
||||
|
||||
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
|
||||
exit "$rust_cov_finalize_test_rc"
|
||||
fi
|
||||
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
|
||||
exit "$rust_cov_finalize_collect_rc"
|
||||
fi
|
||||
exit "$rust_cov_finalize_upload_rc"
|
||||
}
|
||||
@@ -25,6 +25,9 @@
|
||||
# and $$BUILDKITE_PARALLEL_JOB_COUNT environment variables.
|
||||
# working_dir(str): specify the place where the command should execute, default to /vllm-workspace/tests
|
||||
# source_file_dependencies(list): the list of prefixes to opt-in the test for, if empty, the test will always run.
|
||||
# hf_offline_retry(bool): for single-node run-amd-test jobs, run offline first
|
||||
# and retry eligible failures online. May be set as the file default or a
|
||||
# per-step override.
|
||||
|
||||
# When adding a test
|
||||
# - If the test belongs to an existing group, add it there
|
||||
@@ -105,6 +108,8 @@
|
||||
#####################################################################################################################################
|
||||
|
||||
|
||||
# AMD-only pilot: test jobs inherit this; non-test jobs opt out below.
|
||||
hf_offline_retry: true
|
||||
steps:
|
||||
|
||||
#########################################################################################################################################
|
||||
@@ -411,6 +416,7 @@ steps:
|
||||
#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------#
|
||||
|
||||
- label: Rust Frontend Cargo Style + Clippy # TBD
|
||||
hf_offline_retry: false
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
@@ -425,6 +431,7 @@ steps:
|
||||
- bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
|
||||
|
||||
- label: Rust Frontend Cargo Tests # TBD
|
||||
hf_offline_retry: false
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
@@ -441,6 +448,7 @@ steps:
|
||||
#----------------------------------------------------------- mi250 · docker ----------------------------------------------------------#
|
||||
|
||||
- label: Docker Build Metadata (ROCm) # TBD
|
||||
hf_offline_retry: false
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
@@ -2264,6 +2272,7 @@ steps:
|
||||
#--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------#
|
||||
|
||||
- label: Ray Dependency Compatibility Check # TBD
|
||||
hf_offline_retry: false
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
dind: false
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -26,7 +26,10 @@ steps:
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- vllm/config/compilation.py
|
||||
- vllm/compilation
|
||||
- vllm/v1/worker/encoder_cudagraph.py
|
||||
- vllm/v1/worker/encoder_cudagraph_defs.py
|
||||
commands:
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 35
|
||||
timeout_in_minutes: 53
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
|
||||
@@ -39,7 +39,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -59,7 +59,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -78,7 +78,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -156,7 +156,7 @@ steps:
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
@@ -90,6 +90,7 @@ steps:
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/simple_kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/streaming_input
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/ec_connector/unit
|
||||
- tests/v1/metrics
|
||||
@@ -103,6 +104,7 @@ steps:
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s v1/streaming_input
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -143,6 +145,7 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
|
||||
@@ -46,6 +46,19 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Inkling Unit Tests (B200)
|
||||
key: inkling-unit-tests-b200
|
||||
timeout_in_minutes: 40
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- vllm/models/inkling/
|
||||
- vllm/cute_utils/
|
||||
- cmake/external_projects/tml_fa4.cmake
|
||||
- tests/models/inkling/
|
||||
commands:
|
||||
# FA4 kernel tests require SM100; the suite skips them elsewhere.
|
||||
- pytest -v -s models/inkling
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
depends_on:
|
||||
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 98
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
device: h200_35gb
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 110
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
|
||||
@@ -8,11 +8,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/benchmarks/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
@@ -28,7 +23,6 @@ steps:
|
||||
- tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
@@ -49,11 +43,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
@@ -65,7 +54,6 @@ steps:
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
@@ -84,16 +72,10 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
@@ -104,17 +86,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/tool_parsers/
|
||||
- tests/utils.py
|
||||
- tests/tool_use/
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
@@ -125,11 +101,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
@@ -140,7 +111,6 @@ steps:
|
||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||
- tests/v1/distributed/test_internal_lb_dp.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
|
||||
@@ -26,7 +26,5 @@ steps:
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
commands:
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
|
||||
|
||||
@@ -170,3 +170,19 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
|
||||
|
||||
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
|
||||
key: spec-decode-deepseek-mtp-parallel-load-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/llm_base_proposer.py
|
||||
- vllm/v1/spec_decode/eagle.py
|
||||
- vllm/v1/worker/gpu/spec_decode/eagle/
|
||||
- vllm/model_executor/models/deepseek_mtp.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
|
||||
@@ -181,6 +181,18 @@ pull_request_rules:
|
||||
add:
|
||||
- performance
|
||||
|
||||
- name: label-quantization
|
||||
description: Automatically apply quantization label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm/model_executor/layers/quantization/
|
||||
- title~=(?i)quant
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- quantization
|
||||
|
||||
- name: label-qwen
|
||||
description: Automatically apply qwen label
|
||||
conditions:
|
||||
|
||||
@@ -130,6 +130,47 @@ jobs:
|
||||
},
|
||||
],
|
||||
},
|
||||
quantization: {
|
||||
keywords: [
|
||||
{
|
||||
term: "quantization",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "quantized",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
"intel-gpu": {
|
||||
// Keyword search - matches whole words only (with word boundaries)
|
||||
keywords: [
|
||||
{
|
||||
term: "B50",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B60",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B70",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "intel gpu",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "Arc GPU",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "BMG",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
// Add more label configurations here as needed
|
||||
// example: {
|
||||
// keywords: [...],
|
||||
@@ -491,4 +532,4 @@ jobs:
|
||||
issue_number: context.issue.number,
|
||||
body: message,
|
||||
});
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
|
||||
@@ -257,4 +257,3 @@ vllm/grpc/vllm_engine_pb2.pyi
|
||||
|
||||
# Ignore generated cpu headers
|
||||
csrc/cpu/cpu_attn_dispatch_generated.h
|
||||
rust-coverage-tools/
|
||||
|
||||
@@ -8,8 +8,6 @@
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
|
||||
CARGO_LLVM_COV_VERSION="0.8.7"
|
||||
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
|
||||
|
||||
# Read the required toolchain from rust-toolchain.toml.
|
||||
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
||||
@@ -32,39 +30,4 @@ else
|
||||
PROFILE_ARG="--release"
|
||||
fi
|
||||
|
||||
rm -rf "$COVERAGE_TOOLS_DIR"
|
||||
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
|
||||
|
||||
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
|
||||
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
|
||||
# Keep instrumented intermediates isolated when local builds switch modes.
|
||||
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
|
||||
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
|
||||
cargo +"$TOOLCHAIN" install \
|
||||
--locked \
|
||||
--version "$CARGO_LLVM_COV_VERSION" \
|
||||
cargo-llvm-cov
|
||||
|
||||
eval "$(
|
||||
cargo +"$TOOLCHAIN" llvm-cov show-env \
|
||||
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
|
||||
--sh
|
||||
)"
|
||||
|
||||
# Build scripts and proc macros can run during compilation. Their profiles
|
||||
# are unrelated to runtime coverage and would otherwise pollute the tree.
|
||||
export LLVM_PROFILE_FILE=/dev/null
|
||||
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
|
||||
fi
|
||||
|
||||
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
|
||||
|
||||
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
|
||||
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
|
||||
--print target-libdir)")/bin"
|
||||
|
||||
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
|
||||
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
|
||||
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
|
||||
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
|
||||
fi
|
||||
|
||||
@@ -17,7 +17,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
fmha_sm100
|
||||
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
|
||||
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
|
||||
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
|
||||
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
-13
@@ -10,16 +10,3 @@ fixes:
|
||||
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
|
||||
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
|
||||
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
|
||||
# Map Rust sources built in the E2E image and on Buildkite agents.
|
||||
- "/workspace/rust/::rust/"
|
||||
- "/var/lib/buildkite-agent/.*/rust/::rust/"
|
||||
|
||||
flags:
|
||||
rust-unit:
|
||||
paths:
|
||||
- rust/
|
||||
carryforward: false
|
||||
rust-e2e:
|
||||
paths:
|
||||
- rust/
|
||||
carryforward: false
|
||||
|
||||
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
|
||||
batch_offset += offset;
|
||||
int32_t block_table_id = batch_offset / block_size;
|
||||
int32_t slot_id = batch_offset % block_size;
|
||||
// seq_starts may push the block index past the end of the batch's block
|
||||
// table row.
|
||||
if (block_table_id >= block_table_stride) continue;
|
||||
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
|
||||
int32_t block_id = block_table[block_table_offset];
|
||||
int64_t cache_offset =
|
||||
|
||||
@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias);
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_sigmoid(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor);
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
|
||||
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
|
||||
double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid);
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding);
|
||||
|
||||
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> topk_ids,
|
||||
|
||||
@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int end_expert,
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor)
|
||||
const double routed_scaling_factor,
|
||||
const bool* is_padding)
|
||||
{
|
||||
|
||||
using cub_kvp = cub::KeyValuePair<int, float>;
|
||||
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int expert = result_kvp.key;
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
|
||||
|
||||
const int idx = k * block_row + k_idx;
|
||||
// Return the unbiased scores for output weights
|
||||
output[idx] = inputs_after_softmax[thread_read_offset + expert];
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
|
||||
assert(indices[idx] >= 0);
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : num_experts);
|
||||
assert(is_pad_row || indices[idx] >= 0);
|
||||
source_rows[idx] = k_idx * num_rows + block_row;
|
||||
if (renormalize) {
|
||||
selected_sum += inputs_after_softmax[thread_read_offset + expert];
|
||||
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
|
||||
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor)
|
||||
const float* bias, const double routed_scaling_factor, const bool* is_padding)
|
||||
{
|
||||
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
|
||||
// single) thread per row of the input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = max_val;
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -605,7 +610,7 @@ struct TopkConstants
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
|
||||
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
|
||||
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
|
||||
{
|
||||
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
|
||||
}
|
||||
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
|
||||
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor,
|
||||
cudaStream_t stream)
|
||||
cudaStream_t stream,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const float* bias_ptr = nullptr;
|
||||
if (bias.has_value()) {
|
||||
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
|
||||
bias_ptr = bias_tensor.const_data_ptr<float>();
|
||||
}
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
|
||||
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
|
||||
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
|
||||
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
|
||||
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -806,7 +825,8 @@ void topk_softmax(
|
||||
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias)
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -825,15 +845,15 @@ void topk_softmax(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
@@ -846,7 +866,8 @@ void topk_sigmoid(
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor)
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -865,15 +886,15 @@ void topk_sigmoid(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
|
||||
@@ -80,22 +80,27 @@ __launch_bounds__(128) __global__
|
||||
OutIndType* indices, int num_rows,
|
||||
int num_experts, float routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid) {
|
||||
const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
if (warp >= num_rows) return;
|
||||
const int64_t token_id = load_index_as_int64(input_ids, warp);
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
int expert = 0;
|
||||
float weight = 0.f;
|
||||
if (lane < 6) {
|
||||
if (lane < 6 && !is_pad_row) {
|
||||
// only load and calculate for 6 experts
|
||||
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
|
||||
const float x = input[warp * num_experts + expert];
|
||||
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
|
||||
if (isnan(weight)) {
|
||||
weight = 0.f;
|
||||
}
|
||||
}
|
||||
float weight_sum = weight;
|
||||
#pragma unroll
|
||||
@@ -111,7 +116,8 @@ __launch_bounds__(128) __global__
|
||||
const int offset = warp * 6 + lane;
|
||||
output[offset] =
|
||||
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
|
||||
indices[offset] = static_cast<OutIndType>(expert);
|
||||
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
|
||||
: static_cast<OutIndType>(-1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
int num_rows, int num_experts,
|
||||
double routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
const HashIndType* tid2eid, cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
if (num_rows == 0) return;
|
||||
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
|
||||
cudaLaunchConfig_t config = {};
|
||||
@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
config.numAttrs = 1;
|
||||
const float scale = static_cast<float>(routed_scaling_factor);
|
||||
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
|
||||
num_experts, scale, input_ids, tid2eid);
|
||||
num_experts, scale, input_ids, tid2eid, is_padding);
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid) {
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
return;
|
||||
}
|
||||
const bool row_is_active = finished ? !finished[thread_row] : true;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// We finally start setting up the read pointers for each thread. First, each
|
||||
// thread jumps to the start of the row it will read.
|
||||
@@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
// to float
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = 0.f;
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, float>) {
|
||||
using VecType = AlignedArray<float, ELTS_PER_LDG>;
|
||||
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
@@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
if constexpr (USE_HASH) {
|
||||
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
|
||||
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
row_chunk[ii] = sqrtf(val);
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
}
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
@@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = static_cast<IndType>(expert);
|
||||
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
|
||||
: static_cast<IndType>(-1);
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
@@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
#endif
|
||||
return;
|
||||
} else {
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
|
||||
// Original TopK path: find top-k experts by score
|
||||
@@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert =
|
||||
expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool should_process_row =
|
||||
row_is_active && node_uses_expert && !is_pad_row;
|
||||
|
||||
// The lead thread from each sub-group will write out the final results
|
||||
// to global memory. (This will be a single) thread per row of the
|
||||
// input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
if (correction_bias != nullptr) {
|
||||
if (correction_bias != nullptr && should_process_row) {
|
||||
max_val -= correction_bias[expert];
|
||||
}
|
||||
output[idx] = max_val;
|
||||
indices[idx] =
|
||||
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
@@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
|
||||
indices, source_row, k, start_expert, end_expert,
|
||||
renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#else
|
||||
kernel<<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert,
|
||||
end_expert, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#endif
|
||||
})
|
||||
}
|
||||
@@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
|
||||
stream);
|
||||
stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream, const bool* is_padding) {
|
||||
#ifndef USE_ROCM
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (use_hash && topk == 6 && renormalize &&
|
||||
(num_experts == 256 || num_experts == 384)) {
|
||||
launchDsv4HashTopk<IndType, HashIndType>(
|
||||
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
|
||||
routed_scaling_factor, input_ids, tid2eid, stream);
|
||||
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().const_data_ptr<float>();
|
||||
@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
auto launch = [&](auto* topk_indices_ptr) {
|
||||
using OutIndType =
|
||||
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
|
||||
"is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
|
||||
"is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (tid2eid.has_value()) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(),
|
||||
"input_ids is required for hash MoE");
|
||||
@@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream);
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
|
||||
torch::headeronly::ScalarType::Int);
|
||||
@@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
|
||||
tid2eid.value().const_data_ptr<int>(), stream);
|
||||
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
|
||||
}
|
||||
} else {
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
|
||||
@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
|
||||
static_cast<const OutIndType*>(nullptr), stream);
|
||||
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -793,7 +842,8 @@ void topk_softplus_sqrt(
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
@@ -806,21 +856,22 @@ void topk_softplus_sqrt(
|
||||
dispatch_topk_softplus_sqrt_launch<float>(
|
||||
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
|
||||
token_expert_indices, num_tokens, num_experts, topk, renormalize,
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
|
||||
is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_softplus_sqrt_launch<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
|
||||
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
m.def(
|
||||
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
|
||||
"bias) -> ()");
|
||||
"bias, Tensor? is_padding) -> ()");
|
||||
|
||||
// Apply topk sigmoid to the gating outputs.
|
||||
m.def(
|
||||
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, "
|
||||
"Tensor? bias, float routed_scaling_factor) -> ()");
|
||||
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
|
||||
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
|
||||
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts. topk_ids/expert_map are optional and, when
|
||||
|
||||
+1
-12
@@ -294,9 +294,6 @@ FROM base AS rust-build
|
||||
ARG BUILD_OS
|
||||
ARG USE_SCCACHE
|
||||
ARG SCCACHE_ENDPOINT
|
||||
# Temporary default for the initial CI validation. Set this back to 0 when
|
||||
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
|
||||
ARG VLLM_RUST_COVERAGE=1
|
||||
|
||||
# Install native tools needed only for Rust/protoc builds.
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
@@ -796,7 +793,7 @@ 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.14
|
||||
ARG FLASHINFER_VERSION=0.6.15.post1
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
@@ -905,14 +902,6 @@ COPY ./vllm/collect_env.py .
|
||||
# note that this uses vllm installed by `pip`
|
||||
FROM vllm-base AS test
|
||||
|
||||
COPY --from=rust-build \
|
||||
/workspace/rust-coverage-tools/ \
|
||||
/opt/vllm-rust-coverage/
|
||||
|
||||
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
|
||||
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
|
||||
ENV LLVM_PROFILE_FILE=/dev/null
|
||||
|
||||
ADD . /vllm-workspace/
|
||||
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
+53
-31
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="39be1de8"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="bfb51733"
|
||||
# NIXL/UCX build stages
|
||||
FROM base AS build_nixl
|
||||
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
|
||||
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
|
||||
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
|
||||
ENV NIXL_HOME=/usr/local/nixl
|
||||
|
||||
# RIXL build system dependences and RDMA support
|
||||
# NIXL build system dependencies and RDMA support
|
||||
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
libgrpc-dev \
|
||||
libgrpc++-dev \
|
||||
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system meson auditwheel patchelf tomlkit
|
||||
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
|
||||
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
cd /usr/local/src && \
|
||||
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
|
||||
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
git clone ${RIXL_REPO} /opt/rixl && \
|
||||
cd /opt/rixl && \
|
||||
git checkout ${RIXL_BRANCH} && \
|
||||
git clone ${NIXL_REPO} /opt/nixl && \
|
||||
cd /opt/nixl && \
|
||||
git checkout ${NIXL_BRANCH} && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
meson setup build --prefix=${RIXL_HOME} \
|
||||
meson setup build --prefix=${NIXL_HOME} \
|
||||
-Ducx_path=${UCX_HOME} \
|
||||
-Drocm_path=${ROCM_PATH} && \
|
||||
-Dwheel_variant=rocm \
|
||||
-Dbuild_tests=false \
|
||||
-Dbuild_examples=false && \
|
||||
cd build && \
|
||||
ninja -j$(nproc) && \
|
||||
ninja install
|
||||
ninja install && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
|
||||
> /etc/ld.so.conf.d/nixl.conf && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
|
||||
>> /etc/ld.so.conf.d/nixl.conf && \
|
||||
ldconfig
|
||||
|
||||
# Generate RIXL wheel
|
||||
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
|
||||
# so configure the ROCm wheel variant directly through Meson.
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
_ucx_install_dir=${UCX_HOME} \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
RUN cd /opt/nixl && \
|
||||
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
|
||||
--python ${PYTHON_VERSION} \
|
||||
-Csetup-args=-Ducx_path=${UCX_HOME} \
|
||||
-Csetup-args=-Dwheel_variant=rocm \
|
||||
-Csetup-args=-Dbuild_tests=false \
|
||||
-Csetup-args=-Dbuild_examples=false && \
|
||||
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
|
||||
auditwheel repair \
|
||||
--exclude 'libamdhip64*' \
|
||||
--exclude 'libcore*' \
|
||||
--exclude 'libpull*' \
|
||||
/tmp/nixl_wheels/nixl_rocm*.whl \
|
||||
--plat manylinux_2_34_$(uname -m) \
|
||||
--wheel-dir /tmp/nixl_wheels/repaired && \
|
||||
./contrib/wheel_add_ucx_plugins.py \
|
||||
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
|
||||
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
|
||||
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
|
||||
/tmp/nixl_wheels/repaired/*.whl && \
|
||||
cp /tmp/nixl_wheels/repaired/*.whl /app/install
|
||||
|
||||
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
|
||||
# invalidate the slow ROCShmem build.
|
||||
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
|
||||
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
|
||||
fi
|
||||
|
||||
# Install RIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
# Install NIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
|
||||
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
|
||||
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
|
||||
|
||||
# Copy ROCShmem runtime libraries.
|
||||
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
|
||||
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
# 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
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
# ROCm profiler limits workaround.
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
&& pip uninstall -y vllm \
|
||||
&& uv pip install --system *.whl
|
||||
|
||||
# Install RIXL wheel
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
# Install NIXL ROCm wheel
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
uv pip install --system /nixl_install/*.whl
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
ARG BASE_IMAGE
|
||||
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
|
||||
# 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
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.16.post3"
|
||||
ARG AITER_BRANCH="v0.1.16.post5"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
+13
-13
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
|
||||
}
|
||||
|
||||
# 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
|
||||
# rebuilding NIXL/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" {
|
||||
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
|
||||
# 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" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "RIXL_CACHE_KEY" {
|
||||
variable "NIXL_CACHE_KEY" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
|
||||
])
|
||||
}
|
||||
|
||||
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
|
||||
# Cache functions for upstream dependency stages (NIXL/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.
|
||||
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
|
||||
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}" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_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" {
|
||||
function "get_cache_to_rocm_nixl" {
|
||||
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" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
|
||||
# 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" {
|
||||
target "nixl-rocm-ci" {
|
||||
inherits = ["_common-rocm", "_ci-rocm"]
|
||||
target = "build_rixl"
|
||||
target = "build_nixl"
|
||||
cache-from = get_cache_from_rocm_deps()
|
||||
cache-to = get_cache_to_rocm_rixl()
|
||||
cache-to = get_cache_to_rocm_nixl()
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
|
||||
# Builds only the ci_base stage (NIXL, 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.
|
||||
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
|
||||
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
|
||||
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
|
||||
]),
|
||||
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
|
||||
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
|
||||
# are cache hits even when ci_base itself needs rebuilding.
|
||||
get_cache_from_rocm_deps(),
|
||||
)
|
||||
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
|
||||
# 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"]
|
||||
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
|
||||
# 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" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ target "test-rocm" {
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
|
||||
# CI base image target - builds only the ci_base stage (NIXL, 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" {
|
||||
|
||||
@@ -46,9 +46,6 @@
|
||||
"TORCH_CUDA_ARCH_LIST": {
|
||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
|
||||
},
|
||||
"VLLM_RUST_COVERAGE": {
|
||||
"default": "1"
|
||||
},
|
||||
"MAX_JOBS": {
|
||||
"default": "2"
|
||||
},
|
||||
@@ -71,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.14"
|
||||
"default": "0.6.15.post1"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
|
||||
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
|
||||
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
|
||||
|
||||
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
|
||||
|
||||
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
|
||||
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
|
||||
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
|
||||
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
|
||||
|
||||
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
// Reo.Dev documentation tracking
|
||||
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
|
||||
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
|
||||
@@ -160,3 +160,4 @@ extra_javascript:
|
||||
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
|
||||
- mkdocs/javascript/edit_and_feedback.js
|
||||
- mkdocs/javascript/slack_and_forum.js
|
||||
- mkdocs/javascript/reo.js
|
||||
|
||||
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
|
||||
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
|
||||
# install_requires so the published wheel does not carry an unresolvable pin
|
||||
--extra-index-url https://flashinfer.ai/whl/
|
||||
flashinfer-python==0.6.14
|
||||
flashinfer-cubin==0.6.14
|
||||
flashinfer-python==0.6.15.post1
|
||||
flashinfer-cubin==0.6.15.post1
|
||||
apache-tvm-ffi==0.1.10
|
||||
tilelang==0.1.9
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.6.0
|
||||
quack-kernels>=0.4.0 # Required for tml-fa4
|
||||
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
@@ -12,4 +12,4 @@ ray[data]
|
||||
setuptools==78.1.0
|
||||
setuptools-rust>=1.9.0
|
||||
nixl==0.3.0
|
||||
tpu-inference==0.24.0
|
||||
tpu-inference==0.25.0
|
||||
|
||||
@@ -14,6 +14,10 @@ service Generate {
|
||||
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
|
||||
}
|
||||
|
||||
service Control {
|
||||
rpc Abort (AbortRequest) returns (AbortResponse) {}
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Generate Request
|
||||
// ======================================================================================
|
||||
@@ -201,3 +205,12 @@ message TokenIds {
|
||||
repeated uint32 ids = 1;
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Control
|
||||
// ======================================================================================
|
||||
|
||||
message AbortRequest {
|
||||
repeated string request_ids = 1;
|
||||
}
|
||||
|
||||
message AbortResponse {}
|
||||
|
||||
+106
-61
@@ -54,6 +54,7 @@ mod stream;
|
||||
|
||||
use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
@@ -88,6 +89,92 @@ pub fn validate_parser_overrides(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Chat request preparation shared by inference and render-only frontends.
|
||||
pub struct ChatRequestProcessor {
|
||||
backend: DynChatBackend,
|
||||
/// Effective model dtype reported by the engine.
|
||||
/// Absent for text-only frontends without an engine handshake.
|
||||
model_dtype: Option<ModelDtype>,
|
||||
}
|
||||
|
||||
impl ChatRequestProcessor {
|
||||
/// Create a processor with multimodal support using the effective model
|
||||
/// dtype reported by the engine.
|
||||
fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
model_dtype: Some(model_dtype),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a render-only processor that rejects multimodal requests.
|
||||
pub fn render_only(backend: DynChatBackend) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
model_dtype: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn finalize_rendered_prompt(
|
||||
&self,
|
||||
request: &ChatRequest,
|
||||
rendered: RenderedPrompt,
|
||||
) -> Result<(Prompt, Option<MmFeatures>)> {
|
||||
match self.model_dtype {
|
||||
Some(model_dtype) => {
|
||||
multimodal::finalize_rendered_prompt(
|
||||
request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
model_dtype,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
|
||||
None => Err(Error::UnsupportedMultimodalRenderer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Prepare one chat request without submitting it to an engine.
|
||||
pub async fn prepare(
|
||||
&self,
|
||||
mut request: ChatRequest,
|
||||
options: NewChatOutputProcessorOptions<'_>,
|
||||
) -> Result<(TextRequest, DynChatOutputProcessor)> {
|
||||
request.validate()?;
|
||||
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let reasoning_parser_kwargs =
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.is_some()
|
||||
.then(|| ReasoningParserKwargs {
|
||||
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
|
||||
});
|
||||
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
|
||||
let text_request = TextRequest {
|
||||
request_id: request.request_id,
|
||||
prompt,
|
||||
mm_features,
|
||||
sampling_params: request.sampling_params,
|
||||
decode_options: request.decode_options,
|
||||
intermediate: request.intermediate,
|
||||
priority: request.priority,
|
||||
cache_salt: request.cache_salt,
|
||||
add_special_tokens: request.add_special_tokens,
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
reasoning_parser_kwargs,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
};
|
||||
Ok((text_request, output_processor))
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured chat facade above [`TextLlm`].
|
||||
///
|
||||
/// This layer stays above raw text semantics: it takes care of chat-template
|
||||
@@ -95,9 +182,7 @@ pub fn validate_parser_overrides(
|
||||
/// request semantics such as tool calls.
|
||||
pub struct ChatLlm {
|
||||
text: TextLlm,
|
||||
backend: DynChatBackend,
|
||||
/// Effective model dtype reported by the engine.
|
||||
model_dtype: ModelDtype,
|
||||
processor: ChatRequestProcessor,
|
||||
/// Tool-call parser selection.
|
||||
tool_call_parser: ParserSelection,
|
||||
/// Reasoning parser selection.
|
||||
@@ -112,8 +197,7 @@ impl ChatLlm {
|
||||
|
||||
Self {
|
||||
text,
|
||||
backend,
|
||||
model_dtype,
|
||||
processor: ChatRequestProcessor::new(backend, model_dtype),
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
}
|
||||
@@ -140,7 +224,7 @@ impl ChatLlm {
|
||||
|
||||
/// Override the effective model dtype used for multimodal tensor encoding.
|
||||
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
|
||||
self.model_dtype = model_dtype;
|
||||
self.processor.model_dtype = Some(model_dtype);
|
||||
self
|
||||
}
|
||||
|
||||
@@ -172,57 +256,23 @@ impl ChatLlm {
|
||||
}
|
||||
|
||||
/// Render, tokenize, and submit one chat request.
|
||||
pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
|
||||
request.validate()?;
|
||||
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
|
||||
let output_processor = self.backend.new_chat_output_processor(
|
||||
&mut request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &self.tool_call_parser,
|
||||
reasoning_parser: &self.reasoning_parser,
|
||||
},
|
||||
)?;
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let reasoning_parser_kwargs =
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.is_some()
|
||||
.then(|| ReasoningParserKwargs {
|
||||
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
|
||||
});
|
||||
|
||||
let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
|
||||
&request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
self.model_dtype,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let text_request = TextRequest {
|
||||
request_id: request.request_id.clone(),
|
||||
prompt,
|
||||
mm_features,
|
||||
sampling_params: request.sampling_params,
|
||||
decode_options: request.decode_options,
|
||||
intermediate: request.intermediate,
|
||||
priority: request.priority,
|
||||
cache_salt: request.cache_salt,
|
||||
add_special_tokens: request.add_special_tokens,
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
reasoning_parser_kwargs,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
};
|
||||
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
|
||||
let (text_request, output_processor) = self
|
||||
.processor
|
||||
.prepare(
|
||||
request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &self.tool_call_parser,
|
||||
reasoning_parser: &self.reasoning_parser,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let request_id = text_request.request_id.clone();
|
||||
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
|
||||
|
||||
let structured_stream = output_processor.process(decoded_stream)?;
|
||||
|
||||
Ok(ChatEventStream::new(request.request_id, structured_stream))
|
||||
Ok(ChatEventStream::new(request_id, structured_stream))
|
||||
}
|
||||
|
||||
/// Render through the chat template and tokenize, without submitting to the engine.
|
||||
@@ -233,14 +283,9 @@ impl ChatLlm {
|
||||
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
|
||||
request.validate()?;
|
||||
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
|
||||
&request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
self.model_dtype,
|
||||
)
|
||||
.await?;
|
||||
let rendered = self.processor.backend.chat_renderer().render(&request)?;
|
||||
let (prompt, _mm_features) =
|
||||
self.processor.finalize_rendered_prompt(&request, rendered).await?;
|
||||
|
||||
let tokenizer = self.text.tokenizer();
|
||||
let token_ids = match prompt {
|
||||
|
||||
@@ -8,7 +8,7 @@ use tonic_health::ServingStatus;
|
||||
use tonic_health::server::HealthReporter;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::GenerateGrpcService;
|
||||
use super::{ControlGrpcService, GenerateGrpcService};
|
||||
|
||||
pub(crate) async fn monitor_health(
|
||||
mut health_reporter: HealthReporter,
|
||||
@@ -16,21 +16,18 @@ pub(crate) async fn monitor_health(
|
||||
shutdown: CancellationToken,
|
||||
) {
|
||||
let generate_service = GenerateGrpcService::NAME;
|
||||
let control_service = ControlGrpcService::NAME;
|
||||
let status = ServingStatus::NotServing;
|
||||
let health_event_first = tokio::select! {
|
||||
result = engine_health.wait_for(|healthy| !*healthy) => {
|
||||
match result {
|
||||
Ok(_) => warn!(
|
||||
generate_service,
|
||||
overall_service = true,
|
||||
status = ?status,
|
||||
reason = "engine_unhealthy",
|
||||
"marking gRPC health services as not serving"
|
||||
),
|
||||
Err(error) => warn!(
|
||||
%error,
|
||||
generate_service,
|
||||
overall_service = true,
|
||||
status = ?status,
|
||||
reason = "health_channel_closed",
|
||||
"engine health channel closed; marking gRPC health services as not serving"
|
||||
@@ -40,8 +37,6 @@ pub(crate) async fn monitor_health(
|
||||
}
|
||||
_ = shutdown.cancelled() => {
|
||||
info!(
|
||||
generate_service,
|
||||
overall_service = true,
|
||||
status = ?status,
|
||||
reason = "server_shutdown",
|
||||
"server shutting down; marking gRPC health services as not serving"
|
||||
@@ -51,20 +46,20 @@ pub(crate) async fn monitor_health(
|
||||
};
|
||||
|
||||
health_reporter.set_not_serving::<GenerateGrpcService>().await;
|
||||
// Generate is currently the only engine-backed gRPC service, so overall
|
||||
// server health intentionally mirrors it.
|
||||
health_reporter.set_not_serving::<ControlGrpcService>().await;
|
||||
// Both gRPC services use the same engine client, so overall server health
|
||||
// mirrors their shared engine health.
|
||||
health_reporter.set_service_status("", status).await;
|
||||
|
||||
if health_event_first {
|
||||
shutdown.cancelled().await;
|
||||
info!(
|
||||
generate_service,
|
||||
overall_service = true,
|
||||
reason = "server_shutdown",
|
||||
"server shutting down; closing gRPC health watches"
|
||||
);
|
||||
}
|
||||
|
||||
health_reporter.clear_service_status(generate_service).await;
|
||||
health_reporter.clear_service_status(control_service).await;
|
||||
health_reporter.clear_service_status("").await;
|
||||
}
|
||||
|
||||
@@ -26,8 +26,10 @@ pub mod pb {
|
||||
}
|
||||
|
||||
pub(crate) use health::monitor_health;
|
||||
pub use pb::control_server::ControlServer;
|
||||
pub use pb::generate_server::GenerateServer;
|
||||
|
||||
pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
|
||||
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -44,6 +46,36 @@ impl GenerateServiceImpl {
|
||||
}
|
||||
}
|
||||
|
||||
/// gRPC control service backed by the shared application state.
|
||||
pub struct ControlServiceImpl {
|
||||
state: Arc<AppState>,
|
||||
}
|
||||
|
||||
impl ControlServiceImpl {
|
||||
pub fn new(state: Arc<AppState>) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl pb::control_server::Control for ControlServiceImpl {
|
||||
async fn abort(
|
||||
&self,
|
||||
request: Request<pb::AbortRequest>,
|
||||
) -> Result<Response<pb::AbortResponse>, Status> {
|
||||
let request_ids = request.into_inner().request_ids;
|
||||
if request_ids.is_empty() {
|
||||
return Ok(Response::new(pb::AbortResponse {}));
|
||||
}
|
||||
self.state
|
||||
.chat
|
||||
.abort(&request_ids)
|
||||
.await
|
||||
.map_err(|error| Status::internal(error.to_report_string()))?;
|
||||
Ok(Response::new(pb::AbortResponse {}))
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl pb::generate_server::Generate for GenerateServiceImpl {
|
||||
type GenerateStreamStream =
|
||||
|
||||
@@ -38,8 +38,9 @@ use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
use super::pb::control_client::ControlClient;
|
||||
use super::pb::generate_client::GenerateClient;
|
||||
use super::{GenerateServer, GenerateServiceImpl, pb};
|
||||
use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
|
||||
use crate::listener::{Listener, MaybeTlsListener};
|
||||
use crate::state::AppState;
|
||||
use crate::tls;
|
||||
@@ -153,10 +154,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
|
||||
dealer.recv().await.expect("recv engine message").into_vec()
|
||||
}
|
||||
|
||||
fn test_llm(client: EngineCoreClient) -> Llm {
|
||||
Llm::new(client).with_request_id_randomization(false)
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct FakeTextBackend;
|
||||
|
||||
@@ -206,6 +203,7 @@ async fn setup_grpc_service(
|
||||
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
||||
) -> (
|
||||
GenerateServer<GenerateServiceImpl>,
|
||||
ControlServer<ControlServiceImpl>,
|
||||
tokio::sync::watch::Receiver<bool>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
@@ -243,12 +241,13 @@ async fn setup_grpc_service(
|
||||
let engine_health = client.subscribe_health();
|
||||
|
||||
let chat = ChatLlm::from_shared_backend(
|
||||
test_llm(client),
|
||||
Llm::new(client),
|
||||
Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
|
||||
);
|
||||
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
|
||||
(
|
||||
GenerateServer::new(GenerateServiceImpl::new(state)),
|
||||
GenerateServer::new(GenerateServiceImpl::new(state.clone())),
|
||||
ControlServer::new(ControlServiceImpl::new(state)),
|
||||
engine_health,
|
||||
engine_task,
|
||||
)
|
||||
@@ -264,9 +263,11 @@ async fn grpc_test_server(
|
||||
tokio::task::JoinHandle<()>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let (generate_service, control_service, engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, output_specs).await;
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
svc,
|
||||
generate_service,
|
||||
control_service,
|
||||
engine_health,
|
||||
tokio_util::sync::CancellationToken::new(),
|
||||
)
|
||||
@@ -276,11 +277,13 @@ async fn grpc_test_server(
|
||||
|
||||
async fn start_grpc_test_server(
|
||||
generate_service: GenerateServer<GenerateServiceImpl>,
|
||||
control_service: ControlServer<ControlServiceImpl>,
|
||||
engine_health: tokio::sync::watch::Receiver<bool>,
|
||||
shutdown: tokio_util::sync::CancellationToken,
|
||||
) -> (Channel, tokio::task::JoinHandle<()>) {
|
||||
let (health_reporter, health_service) = health_reporter();
|
||||
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
|
||||
health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
@@ -289,6 +292,7 @@ async fn start_grpc_test_server(
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
let server = TonicServer::builder()
|
||||
.add_service(health_service)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||
let health_monitor =
|
||||
@@ -319,7 +323,8 @@ async fn grpc_tls_test_server(
|
||||
certs: &TestCerts,
|
||||
cert_reqs: i32,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let (generate_service, control_service, _engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, output_specs).await;
|
||||
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
|
||||
.expect("build grpc tls config");
|
||||
|
||||
@@ -329,7 +334,8 @@ async fn grpc_tls_test_server(
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc tls server");
|
||||
@@ -409,7 +415,7 @@ async fn grpc_server_with_keepalive(
|
||||
engine_id: impl Into<EngineId>,
|
||||
keepalive: Option<Duration>,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, _engine_health, engine_task) =
|
||||
let (generate_service, control_service, _engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, default_stream_output_specs()).await;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||
@@ -425,7 +431,8 @@ async fn grpc_server_with_keepalive(
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
builder
|
||||
.add_service(svc)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc server");
|
||||
@@ -1073,14 +1080,106 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
|
||||
server_task.abort();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn control_abort_resolves_external_id_and_empty_is_noop() {
|
||||
let (generate_service, control_service, engine_health, engine_task) =
|
||||
setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
generate_service,
|
||||
control_service,
|
||||
engine_health,
|
||||
tokio_util::sync::CancellationToken::new(),
|
||||
)
|
||||
.await;
|
||||
let mut generate_client = GenerateClient::new(channel.clone());
|
||||
let mut control_client = ControlClient::new(channel);
|
||||
let request_id = "test-abort-active";
|
||||
|
||||
let mut stream = generate_client
|
||||
.generate_stream(pb::GenerateRequest {
|
||||
request_id: request_id.to_string(),
|
||||
model: "test-model".to_string(),
|
||||
prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
|
||||
stopping: Some(pb::StoppingCriteria {
|
||||
max_new_tokens: 10,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect("start generation")
|
||||
.into_inner();
|
||||
|
||||
loop {
|
||||
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||
.await
|
||||
.expect("timed out waiting for active generation output")
|
||||
.expect("read active generation output")
|
||||
.expect("generation ended before producing output");
|
||||
if let Some(output) = response.outputs {
|
||||
assert!(
|
||||
output.finish_info.is_none(),
|
||||
"generation finished before abort behavior was exercised"
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest::default())
|
||||
.await
|
||||
.expect("empty abort should be a no-op");
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(100), stream.message())
|
||||
.await
|
||||
.is_err(),
|
||||
"empty abort unexpectedly ended the active generation"
|
||||
);
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest {
|
||||
request_ids: vec![
|
||||
request_id.to_string(),
|
||||
request_id.to_string(),
|
||||
"unknown".to_string(),
|
||||
],
|
||||
})
|
||||
.await
|
||||
.expect("abort active generation");
|
||||
|
||||
let finish_reason = loop {
|
||||
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||
.await
|
||||
.expect("timed out waiting for aborted generation")
|
||||
.expect("read aborted generation")
|
||||
.expect("generation ended without an aborted response");
|
||||
if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
|
||||
break finish_info.finish_reason;
|
||||
}
|
||||
};
|
||||
assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest {
|
||||
request_ids: vec![request_id.to_string()],
|
||||
})
|
||||
.await
|
||||
.expect("repeated abort should be idempotent");
|
||||
|
||||
engine_task.await.expect("mock engine task");
|
||||
server_task.abort();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
|
||||
let (generate_service, _connected_engine_health, _engine_task) =
|
||||
let (generate_service, control_service, _connected_engine_health, _engine_task) =
|
||||
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
|
||||
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
generate_service,
|
||||
control_service,
|
||||
engine_health,
|
||||
tokio_util::sync::CancellationToken::new(),
|
||||
)
|
||||
@@ -1088,7 +1187,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
|
||||
let mut health_client = HealthClient::new(channel);
|
||||
|
||||
let mut health_streams = Vec::new();
|
||||
for service in ["vllm.Generate", ""] {
|
||||
for service in ["vllm.Generate", "vllm.Control", ""] {
|
||||
let service_label = if service.is_empty() {
|
||||
"overall"
|
||||
} else {
|
||||
@@ -1143,14 +1242,19 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn grpc_health_watch_closes_on_graceful_shutdown() {
|
||||
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
|
||||
let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
|
||||
b"engine-grpc-health-shutdown",
|
||||
default_stream_output_specs(),
|
||||
)
|
||||
.await;
|
||||
let shutdown = tokio_util::sync::CancellationToken::new();
|
||||
let (channel, server_task) =
|
||||
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
generate_service,
|
||||
control_service,
|
||||
engine_health,
|
||||
shutdown.clone(),
|
||||
)
|
||||
.await;
|
||||
let mut health_client = HealthClient::new(channel);
|
||||
let mut stream = health_client
|
||||
.watch(HealthCheckRequest {
|
||||
|
||||
@@ -207,6 +207,9 @@ where
|
||||
let (health_reporter, health_service) = health_reporter();
|
||||
let engine_health = state.engine_core_client().subscribe_health();
|
||||
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
|
||||
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
|
||||
let control_service =
|
||||
grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
|
||||
let generate_service =
|
||||
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
|
||||
let svc = TonicServer::builder()
|
||||
@@ -214,6 +217,7 @@ where
|
||||
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
|
||||
.layer(middleware::request_runtime_layer(state.clone()))
|
||||
.add_service(health_service)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service);
|
||||
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
|
||||
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
|
||||
|
||||
+108
-72
@@ -38,32 +38,22 @@ trait_set! {
|
||||
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
|
||||
}
|
||||
|
||||
/// Raw text facade above [`Llm`].
|
||||
///
|
||||
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
|
||||
/// in, decoded text deltas and terminal metadata flow out.
|
||||
pub struct TextLlm {
|
||||
/// Generate-only client owned by this text facade.
|
||||
llm: Llm,
|
||||
/// Text request preparation shared by inference and render-only frontends.
|
||||
pub struct TextRequestProcessor {
|
||||
/// Tokenizer/model metadata backend responsible for prompt encode/decode
|
||||
/// and sampling hints.
|
||||
backend: DynTextBackend,
|
||||
/// Runtime context window size reported by the engine startup handshake.
|
||||
/// Render-only frontends supply the downstream engine's effective value.
|
||||
max_model_len: u32,
|
||||
/// Maximum number of top log probabilities accepted by this text facade.
|
||||
max_logprobs: i32,
|
||||
}
|
||||
|
||||
impl TextLlm {
|
||||
/// Create a new text-generation facade from a shared LLM client plus a text
|
||||
/// backend.
|
||||
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
|
||||
// The engine-reported value reflects the post-profiling, auto-fitted
|
||||
// KV cache limit used at runtime.
|
||||
let max_model_len = llm.engine_core_client().max_model_len();
|
||||
|
||||
impl TextRequestProcessor {
|
||||
/// Create a processor with the effective model context length.
|
||||
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
|
||||
Self {
|
||||
llm,
|
||||
backend,
|
||||
max_model_len,
|
||||
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
|
||||
@@ -78,61 +68,18 @@ impl TextLlm {
|
||||
self
|
||||
}
|
||||
|
||||
/// Return the backend model ID.
|
||||
pub fn model_id(&self) -> &str {
|
||||
self.backend.model_id()
|
||||
}
|
||||
|
||||
/// Expose the underlying engine-core client for low-level utility/admin
|
||||
/// calls.
|
||||
pub fn engine_core_client(&self) -> &EngineCoreClient {
|
||||
self.llm.engine_core_client()
|
||||
}
|
||||
|
||||
/// Return the tokenizer used by this text backend.
|
||||
/// Return the tokenizer used by this processor.
|
||||
pub fn tokenizer(&self) -> DynTokenizer {
|
||||
self.backend.tokenizer()
|
||||
}
|
||||
|
||||
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
|
||||
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
|
||||
pub fn tokenizer_vocab_size(&self) -> usize {
|
||||
self.backend.tokenizer_vocab_size()
|
||||
/// Return the effective model context length.
|
||||
pub fn max_model_len(&self) -> u32 {
|
||||
self.max_model_len
|
||||
}
|
||||
|
||||
/// Model vocabulary size from the model config, used to bound generated
|
||||
/// token IDs and logits-domain sampling controls.
|
||||
pub fn model_vocab_size(&self) -> usize {
|
||||
self.backend.model_vocab_size()
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and return the raw
|
||||
/// token stream.
|
||||
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
|
||||
let (_, raw_stream) = self.generate_inner(request).await?;
|
||||
Ok(raw_stream)
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and stream
|
||||
/// incrementally decoded text.
|
||||
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
|
||||
let (text_request, raw_stream) = self.generate_inner(request).await?;
|
||||
let tokenizer = self.backend.tokenizer();
|
||||
let decoded_stream = output::decoded_text_event_stream(
|
||||
text_request.request_id,
|
||||
tokenizer,
|
||||
raw_stream,
|
||||
text_request.decode_options,
|
||||
text_request.intermediate,
|
||||
);
|
||||
|
||||
Ok(decoded_stream)
|
||||
}
|
||||
|
||||
async fn generate_inner(
|
||||
&self,
|
||||
mut request: TextRequest,
|
||||
) -> Result<(TextRequest, GenerateOutputStream)> {
|
||||
/// Tokenize and lower one request without submitting it to an engine.
|
||||
pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
|
||||
request.validate()?;
|
||||
|
||||
if request.arrival_time.is_none() {
|
||||
@@ -146,7 +93,6 @@ impl TextLlm {
|
||||
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
|
||||
Prompt::TokenIds(token_ids) => token_ids,
|
||||
};
|
||||
|
||||
let sampling_hints = self.backend.sampling_hints()?;
|
||||
let sampling_limits = SamplingLimits {
|
||||
max_model_len: self.max_model_len,
|
||||
@@ -155,16 +101,106 @@ impl TextLlm {
|
||||
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
|
||||
};
|
||||
|
||||
let PreparedTextRequest {
|
||||
text_request,
|
||||
generate_request,
|
||||
} = lower_text_request(
|
||||
lower_text_request(
|
||||
request,
|
||||
prompt_token_ids,
|
||||
sampling_hints,
|
||||
sampling_limits,
|
||||
&*tokenizer,
|
||||
)?;
|
||||
tokenizer.as_ref(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw text facade above [`Llm`].
|
||||
///
|
||||
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
|
||||
/// in, decoded text deltas and terminal metadata flow out.
|
||||
pub struct TextLlm {
|
||||
/// Generate-only client owned by this text facade.
|
||||
llm: Llm,
|
||||
/// Shared engine-free request preparation.
|
||||
processor: TextRequestProcessor,
|
||||
}
|
||||
|
||||
impl TextLlm {
|
||||
/// Create a new text-generation facade from a shared LLM client plus a text
|
||||
/// backend.
|
||||
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
|
||||
// The engine-reported value reflects the post-profiling, auto-fitted
|
||||
// KV cache limit used at runtime.
|
||||
let max_model_len = llm.engine_core_client().max_model_len();
|
||||
|
||||
Self {
|
||||
llm,
|
||||
processor: TextRequestProcessor::new(backend, max_model_len),
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the maximum accepted logprobs count.
|
||||
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
|
||||
self.processor = self.processor.with_max_logprobs(max_logprobs);
|
||||
self
|
||||
}
|
||||
|
||||
/// Return the backend model ID.
|
||||
pub fn model_id(&self) -> &str {
|
||||
self.processor.backend.model_id()
|
||||
}
|
||||
|
||||
/// Expose the underlying engine-core client for low-level utility/admin
|
||||
/// calls.
|
||||
pub fn engine_core_client(&self) -> &EngineCoreClient {
|
||||
self.llm.engine_core_client()
|
||||
}
|
||||
|
||||
/// Return the tokenizer used by this text backend.
|
||||
pub fn tokenizer(&self) -> DynTokenizer {
|
||||
self.processor.tokenizer()
|
||||
}
|
||||
|
||||
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
|
||||
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
|
||||
pub fn tokenizer_vocab_size(&self) -> usize {
|
||||
self.processor.backend.tokenizer_vocab_size()
|
||||
}
|
||||
|
||||
/// Model vocabulary size from the model config, used to bound generated
|
||||
/// token IDs and logits-domain sampling controls.
|
||||
pub fn model_vocab_size(&self) -> usize {
|
||||
self.processor.backend.model_vocab_size()
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and return the raw
|
||||
/// token stream.
|
||||
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
|
||||
let (_, raw_stream) = self.generate_inner(request).await?;
|
||||
Ok(raw_stream)
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and stream
|
||||
/// incrementally decoded text.
|
||||
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
|
||||
let (text_request, raw_stream) = self.generate_inner(request).await?;
|
||||
let tokenizer = self.processor.tokenizer();
|
||||
let decoded_stream = output::decoded_text_event_stream(
|
||||
text_request.request_id,
|
||||
tokenizer,
|
||||
raw_stream,
|
||||
text_request.decode_options,
|
||||
text_request.intermediate,
|
||||
);
|
||||
|
||||
Ok(decoded_stream)
|
||||
}
|
||||
|
||||
async fn generate_inner(
|
||||
&self,
|
||||
request: TextRequest,
|
||||
) -> Result<(TextRequest, GenerateOutputStream)> {
|
||||
let PreparedTextRequest {
|
||||
text_request,
|
||||
generate_request,
|
||||
} = self.processor.prepare(request)?;
|
||||
|
||||
let raw_stream = self.llm.generate(generate_request).await?;
|
||||
Ok((text_request, raw_stream))
|
||||
|
||||
@@ -214,3 +214,65 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
|
||||
base_hash = CacheConfig().compute_hash()
|
||||
assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
|
||||
assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
|
||||
|
||||
|
||||
def test_envs_compile_factors_relocation_invariant(tmp_path):
|
||||
"""Relocating HOME or the XDG roots must not change the compile-cache
|
||||
env hash.
|
||||
|
||||
Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME,
|
||||
VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about
|
||||
compiled artifacts, only about where directories live. When they leak
|
||||
into compile_factors(), a cache produced under one HOME/XDG layout
|
||||
silently misses under another - which defeats copying or pre-baking a
|
||||
compile cache into a container image.
|
||||
"""
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
code = """
|
||||
import sys
|
||||
import logging
|
||||
logging.disable(logging.CRITICAL)
|
||||
from vllm import envs
|
||||
from vllm.config.utils import hash_factors
|
||||
print(hash_factors(envs.compile_factors()))
|
||||
"""
|
||||
|
||||
def hash_with(extra_env):
|
||||
env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"}
|
||||
# Drop explicit overrides so the derived defaults are what is
|
||||
# exercised, then apply the relocation under test.
|
||||
for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"):
|
||||
env.pop(key, None)
|
||||
env.update(extra_env)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
env=env,
|
||||
)
|
||||
return result.stdout.strip()
|
||||
|
||||
xdg_cache = tmp_path / "relocated-xdg-cache"
|
||||
xdg_config = tmp_path / "relocated-xdg-config"
|
||||
new_home = tmp_path / "relocated-home"
|
||||
for d in (xdg_cache, xdg_config, new_home):
|
||||
d.mkdir()
|
||||
|
||||
base = hash_with({})
|
||||
relocated_xdg = hash_with(
|
||||
{"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)}
|
||||
)
|
||||
relocated_home = hash_with({"HOME": str(new_home)})
|
||||
|
||||
assert relocated_xdg == base, (
|
||||
"XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache "
|
||||
"env hash - a location-only derived var is leaking into the key"
|
||||
)
|
||||
assert relocated_home == base, (
|
||||
"HOME relocation changed the compile-cache env hash - a "
|
||||
"location-only derived var is leaking into the key"
|
||||
)
|
||||
|
||||
@@ -43,6 +43,15 @@ def test_language_model_only_affects_model_hash():
|
||||
assert base_hash != lm_only_hash
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"])
|
||||
def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str):
|
||||
config = MultiModalConfig(
|
||||
media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}}
|
||||
)
|
||||
|
||||
assert config.use_gpu_video_backend()
|
||||
|
||||
|
||||
def test_mm_encoder_fp8_scale_path_requires_fp8():
|
||||
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
|
||||
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
|
||||
|
||||
@@ -254,8 +254,8 @@ def test_multiproc_executor_shutdown_cleanup():
|
||||
for worker in executor.workers:
|
||||
assert not worker.proc.is_alive(), "Worker processes should be terminated"
|
||||
|
||||
# Verify shutdown event is set
|
||||
assert executor.shutdown_event.is_set(), "Shutdown event should be set"
|
||||
# Verify shutdown flag is set
|
||||
assert executor.shutting_down, "Shutdown flag should be set"
|
||||
|
||||
# Multiple shutdowns should be safe (idempotent)
|
||||
executor.shutdown()
|
||||
@@ -292,7 +292,6 @@ def test_multiproc_executor_pipeline_parallel():
|
||||
"Max concurrent batches should follow the configured PP/async "
|
||||
"scheduling policy"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
@@ -338,6 +337,10 @@ def test_multiproc_executor_multi_node():
|
||||
- Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
|
||||
Total world_size = 4, nnodes = 2
|
||||
"""
|
||||
# Python 3.14+ changed default multiprocessing start method to 'forkserver'
|
||||
# which cannot pickle nested functions. Use 'fork' for this test.
|
||||
mp_ctx = multiprocessing.get_context("fork")
|
||||
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
port = s.getsockname()[1]
|
||||
@@ -405,12 +408,12 @@ def test_multiproc_executor_multi_node():
|
||||
executor.shutdown()
|
||||
|
||||
# Create a queue to collect results from both processes
|
||||
result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
|
||||
result_queue: multiprocessing.Queue[dict[str, int | bool]] = mp_ctx.Queue()
|
||||
|
||||
# Start both node processes
|
||||
processes = []
|
||||
for node_rank in range(2):
|
||||
p = multiprocessing.Process(
|
||||
p = mp_ctx.Process(
|
||||
target=run_node,
|
||||
args=(node_rank, result_queue, port),
|
||||
name=f"Node{node_rank}",
|
||||
|
||||
@@ -9,8 +9,6 @@ from pydantic import TypeAdapter, ValidationError
|
||||
from vllm import PoolingParams
|
||||
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
CohereEmbedContent,
|
||||
CohereEmbedInput,
|
||||
CohereEmbedRequest,
|
||||
EmbeddingBatchChatInputRequest,
|
||||
EmbeddingBatchChatRequest,
|
||||
@@ -19,7 +17,10 @@ from vllm.entrypoints.pooling.embed.protocol import (
|
||||
EmbeddingCompletionRequest,
|
||||
EmbeddingRequest,
|
||||
)
|
||||
from vllm.entrypoints.pooling.typing import PoolingServeContext
|
||||
from vllm.entrypoints.pooling.typing import (
|
||||
PoolingEngineInput,
|
||||
PoolingServeContext,
|
||||
)
|
||||
from vllm.outputs import PoolingOutput, PoolingRequestOutput
|
||||
|
||||
|
||||
@@ -410,6 +411,7 @@ class TestChunkedEmbeddingProcessing:
|
||||
def _make_handler(cls):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.model_config = cls._FakeModelConfig()
|
||||
handler.enable_chunked_processing = True
|
||||
return handler
|
||||
|
||||
@staticmethod
|
||||
@@ -421,15 +423,29 @@ class TestChunkedEmbeddingProcessing:
|
||||
}
|
||||
)
|
||||
assert isinstance(request, EmbeddingCompletionRequest)
|
||||
pooling_params = PoolingParams()
|
||||
return PoolingServeContext(
|
||||
request=request,
|
||||
pooling_params=PoolingParams(),
|
||||
pooling_params=pooling_params,
|
||||
model_name="test",
|
||||
request_id="embd-client-prompt-999-chunk-888",
|
||||
engine_inputs=[
|
||||
{"prompt_token_ids": [0, 1, 2, 3, 4]},
|
||||
{"prompt_token_ids": [10, 11]},
|
||||
PoolingEngineInput(
|
||||
prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},
|
||||
params=pooling_params,
|
||||
lora_requests=None,
|
||||
priorities=0,
|
||||
),
|
||||
PoolingEngineInput(
|
||||
prompts={"prompt_token_ids": [10, 11]},
|
||||
params=pooling_params,
|
||||
lora_requests=None,
|
||||
priorities=0,
|
||||
),
|
||||
],
|
||||
lora_request=None,
|
||||
priorities=0,
|
||||
prompt_extras=None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@@ -450,7 +466,7 @@ class TestChunkedEmbeddingProcessing:
|
||||
handler = self._make_handler()
|
||||
ctx = self._make_context()
|
||||
|
||||
handler._pre_process_chunked(ctx)
|
||||
handler.maybe_pre_process_chunked(ctx)
|
||||
|
||||
assert ctx.prompt_request_ids == [
|
||||
"embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
|
||||
@@ -488,227 +504,3 @@ class TestChunkedEmbeddingProcessing:
|
||||
ctx.final_res_batch[1].outputs.data,
|
||||
torch.tensor([9.0, 9.0]),
|
||||
)
|
||||
|
||||
|
||||
class TestPreProcessCohereOnline:
|
||||
"""Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
|
||||
|
||||
@staticmethod
|
||||
def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
|
||||
return PoolingServeContext(
|
||||
request=CohereEmbedRequest(model="test", **request_kwargs),
|
||||
pooling_params=PoolingParams(),
|
||||
model_name="test",
|
||||
request_id="embd-test",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _make_handler():
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler._validate_input_type = lambda _input_type: None
|
||||
return handler
|
||||
|
||||
def test_text_only_without_task_prefix_uses_completion_path(self):
|
||||
handler = self._make_handler()
|
||||
ctx = self._make_context(texts=["hello"])
|
||||
calls: list[tuple[str, object]] = []
|
||||
|
||||
def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
|
||||
calls.append(("completion", prompt_input))
|
||||
return ["completion"]
|
||||
|
||||
handler._get_task_instruction_prefix = lambda _input_type: None
|
||||
handler._has_chat_template = lambda: False
|
||||
handler._preprocess_cmpl_online = preprocess_cmpl_online
|
||||
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
|
||||
"text-only request should not require chat rendering"
|
||||
)
|
||||
|
||||
handler._pre_process_cohere_online(ctx)
|
||||
|
||||
assert ctx.engine_inputs == ["completion"]
|
||||
assert calls == [("completion", ["hello"])]
|
||||
|
||||
def test_text_only_falls_back_to_prefixed_completion_without_template(self):
|
||||
handler = self._make_handler()
|
||||
ctx = self._make_context(texts=["hello"], input_type="query")
|
||||
calls: list[tuple[str, object]] = []
|
||||
|
||||
def preprocess_cmpl(request, prompt_input, prompt_embeds):
|
||||
calls.append(("completion", prompt_input))
|
||||
return ["fallback"]
|
||||
|
||||
handler._get_task_instruction_prefix = lambda _input_type: "query: "
|
||||
handler._has_chat_template = lambda: False
|
||||
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
|
||||
"chat rendering should be skipped without a template"
|
||||
)
|
||||
handler._preprocess_cmpl_online = preprocess_cmpl
|
||||
|
||||
handler._pre_process_cohere_online(ctx)
|
||||
|
||||
assert ctx.engine_inputs == ["fallback"]
|
||||
assert calls == [("completion", ["query: hello"])]
|
||||
|
||||
def test_text_only_with_template_uses_chat_path(self):
|
||||
handler = self._make_handler()
|
||||
ctx = self._make_context(texts=["hello"], input_type="query")
|
||||
calls: list[tuple[str, object]] = []
|
||||
|
||||
def batch_render_chat(
|
||||
request,
|
||||
all_messages,
|
||||
truncate_prompt_tokens,
|
||||
truncation_side,
|
||||
):
|
||||
calls.append(
|
||||
(
|
||||
"chat",
|
||||
{
|
||||
"request": request,
|
||||
"all_messages": all_messages,
|
||||
"truncate_prompt_tokens": truncate_prompt_tokens,
|
||||
"truncation_side": truncation_side,
|
||||
},
|
||||
)
|
||||
)
|
||||
return ["chat"]
|
||||
|
||||
handler._get_task_instruction_prefix = lambda _input_type: "query: "
|
||||
handler._has_chat_template = lambda: True
|
||||
handler._batch_render_chat = batch_render_chat
|
||||
handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail(
|
||||
"completion path should be skipped when a template exists"
|
||||
)
|
||||
|
||||
handler._pre_process_cohere_online(ctx)
|
||||
|
||||
assert ctx.engine_inputs == ["chat"]
|
||||
assert calls == [
|
||||
(
|
||||
"chat",
|
||||
{
|
||||
"request": ctx.request,
|
||||
"all_messages": [
|
||||
handler._mixed_input_to_messages(
|
||||
CohereEmbedInput(
|
||||
content=[CohereEmbedContent(type="text", text="hello")]
|
||||
),
|
||||
task_prefix="query: ",
|
||||
)
|
||||
],
|
||||
"truncate_prompt_tokens": -1,
|
||||
"truncation_side": None,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
class TestPreProcessOpenAIEmbeddingChatOnline:
|
||||
"""Unit tests for OpenAI embedding chat preprocessing."""
|
||||
|
||||
class _FakeModelConfig:
|
||||
max_model_len = 128
|
||||
encoder_config: dict[str, object] = {}
|
||||
pooler_config = None
|
||||
multimodal_config = None
|
||||
is_encoder_decoder = False
|
||||
|
||||
class _FakeRenderer:
|
||||
tokenizer = object()
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def render_chat(
|
||||
self,
|
||||
all_messages,
|
||||
chat_params,
|
||||
tok_params,
|
||||
prompt_extras=None,
|
||||
):
|
||||
self.calls.append(
|
||||
{
|
||||
"all_messages": all_messages,
|
||||
"chat_params": chat_params,
|
||||
"tok_params": tok_params,
|
||||
"prompt_extras": prompt_extras,
|
||||
}
|
||||
)
|
||||
return all_messages, [
|
||||
{"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def _make_handler(cls, renderer):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.renderer = renderer
|
||||
handler.model_config = cls._FakeModelConfig()
|
||||
handler.chat_template = "template"
|
||||
handler.chat_template_content_format = "auto"
|
||||
handler.trust_request_chat_template = False
|
||||
handler.enable_chunked_processing = False
|
||||
return handler
|
||||
|
||||
@staticmethod
|
||||
def _make_context(
|
||||
request: (
|
||||
EmbeddingChatRequest
|
||||
| EmbeddingBatchChatRequest
|
||||
| EmbeddingChatInputRequest
|
||||
| EmbeddingBatchChatInputRequest
|
||||
),
|
||||
) -> PoolingServeContext[
|
||||
EmbeddingChatRequest
|
||||
| EmbeddingBatchChatRequest
|
||||
| EmbeddingChatInputRequest
|
||||
| EmbeddingBatchChatInputRequest
|
||||
]:
|
||||
return PoolingServeContext(
|
||||
request=request,
|
||||
pooling_params=PoolingParams(),
|
||||
model_name="test",
|
||||
request_id="embd-test",
|
||||
)
|
||||
|
||||
def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
|
||||
request = TypeAdapter(EmbeddingRequest).validate_python(
|
||||
{
|
||||
"model": "test",
|
||||
"input": [
|
||||
[{"role": "user", "content": "hello"}],
|
||||
[{"role": "user", "content": "goodbye"}],
|
||||
],
|
||||
"add_generation_prompt": True,
|
||||
"chat_template_kwargs": {"instruction": "Represent the query: "},
|
||||
"mm_processor_kwargs": {"max_pixels": 1},
|
||||
"cache_salt": "salt",
|
||||
}
|
||||
)
|
||||
assert isinstance(request, EmbeddingBatchChatInputRequest)
|
||||
|
||||
renderer = self._FakeRenderer()
|
||||
handler = self._make_handler(renderer)
|
||||
ctx = self._make_context(request)
|
||||
|
||||
handler.pre_process_online(ctx)
|
||||
|
||||
assert ctx.engine_inputs == [
|
||||
{"prompt_token_ids": [0]},
|
||||
{"prompt_token_ids": [1]},
|
||||
]
|
||||
assert len(renderer.calls) == 1
|
||||
|
||||
call = renderer.calls[0]
|
||||
assert call["all_messages"] == request.messages
|
||||
assert call["prompt_extras"] == {
|
||||
"mm_processor_kwargs": {"max_pixels": 1},
|
||||
"cache_salt": "salt",
|
||||
}
|
||||
|
||||
chat_template_kwargs = call["chat_params"].chat_template_kwargs
|
||||
assert chat_template_kwargs["instruction"] == "Represent the query: "
|
||||
assert chat_template_kwargs["add_generation_prompt"] is True
|
||||
assert chat_template_kwargs["continue_final_message"] is False
|
||||
assert "tools" not in chat_template_kwargs
|
||||
assert chat_template_kwargs["tokenize"] is False
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import weakref
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -10,10 +9,7 @@ import torch
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
|
||||
from vllm.entrypoints.pooling.scoring.typing import ScoringData
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.renderers import TokenizeParams
|
||||
|
||||
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
|
||||
PROMPT = "The chef prepared a delicious meal."
|
||||
@@ -145,45 +141,6 @@ def test_max_tokens_per_doc(llm: LLM):
|
||||
assert with_limit_tokens < no_limit_tokens
|
||||
|
||||
|
||||
def test_token_type_ids_follow_post_tokenization():
|
||||
processor = object.__new__(CrossEncoderIOProcessor)
|
||||
processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
|
||||
processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
|
||||
processor.model_config = None
|
||||
processor.get_score_prompt = lambda **_: (
|
||||
"",
|
||||
{
|
||||
"prompt_token_ids": list(range(32)),
|
||||
"token_type_ids": [0] * 16 + [1] * 16,
|
||||
},
|
||||
)
|
||||
|
||||
engine_inputs, pooling_params = processor._pre_process(
|
||||
ScoringData(data_1=["query"], data_2=["document"]),
|
||||
TokenizeParams(
|
||||
max_total_tokens=None,
|
||||
truncate_prompt_tokens=16,
|
||||
truncation_side="left",
|
||||
),
|
||||
PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
|
||||
)
|
||||
|
||||
assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
|
||||
assert pooling_params[0].extra_kwargs == {
|
||||
"cache_salt": "salt",
|
||||
"compressed_token_type_ids": 0,
|
||||
}
|
||||
|
||||
engine_inputs, pooling_params = processor._pre_process(
|
||||
ScoringData(data_1=["query"], data_2=["document"]),
|
||||
TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
|
||||
PoolingParams(task="classify"),
|
||||
)
|
||||
|
||||
assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
|
||||
assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
|
||||
|
||||
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
outputs = llm.score(
|
||||
|
||||
@@ -161,7 +161,9 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
|
||||
return serving
|
||||
|
||||
|
||||
def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
|
||||
def _make_pooling_ctx(
|
||||
model_name: str, serving: PoolingBaseServing
|
||||
) -> PoolingServeContext:
|
||||
mock_request = MagicMock()
|
||||
mock_request.model = model_name
|
||||
return PoolingServeContext(
|
||||
@@ -169,6 +171,9 @@ def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
|
||||
model_name=MODEL_NAME,
|
||||
request_id="test-id",
|
||||
pooling_params=PoolingParams(),
|
||||
lora_request=serving._maybe_get_adapters(mock_request),
|
||||
priorities=0,
|
||||
prompt_extras=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -176,9 +181,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
|
||||
"""LoRA adapter name must populate ctx.lora_request without raising."""
|
||||
lora_name = "bot-embed-lora"
|
||||
serving = _make_pooling_serving(lora_name)
|
||||
ctx = _make_pooling_ctx(lora_name)
|
||||
|
||||
ctx.lora_request = serving._maybe_get_adapters(ctx.request)
|
||||
ctx = _make_pooling_ctx(lora_name, serving)
|
||||
|
||||
assert ctx.lora_request is not None
|
||||
assert ctx.lora_request.lora_name == lora_name
|
||||
@@ -187,7 +190,6 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
|
||||
def test_pooling_maybe_get_adapters_unknown_model_raises():
|
||||
"""An unrecognised model name must still raise VLLMNotFoundError."""
|
||||
serving = _make_pooling_serving("some-lora")
|
||||
ctx = _make_pooling_ctx("unknown-model")
|
||||
|
||||
with pytest.raises(VLLMNotFoundError):
|
||||
serving._maybe_get_adapters(ctx.request)
|
||||
_make_pooling_ctx("unknown-model", serving)
|
||||
|
||||
@@ -6,6 +6,7 @@ max_concurrency: 100
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--max-num-batched-tokens 32768
|
||||
--safetensors-load-strategy prefetch
|
||||
--moe-backend flashinfer_cutlass
|
||||
--prefill-context-parallel-size 4
|
||||
|
||||
@@ -6,6 +6,7 @@ max_concurrency: 100
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--max-num-batched-tokens 32768
|
||||
--safetensors-load-strategy prefetch
|
||||
--moe-backend flashinfer_cutlass
|
||||
--tensor-parallel-size 2
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
model_name: "poolside/Laguna-XS.2-NVFP4"
|
||||
accuracy_threshold: 0.86
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
startup_max_wait_seconds: 1200
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--trust-remote-code
|
||||
@@ -2,4 +2,7 @@ model_name: "google/gemma-4-E4B-it-qat-mobile-ct"
|
||||
accuracy_threshold: 0.50
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 4096"
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--speculative-config '{"method":"mtp","model":"google/gemma-4-E4B-it-assistant","num_speculative_tokens":4}'
|
||||
|
||||
@@ -3,3 +3,4 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
|
||||
Qwen1.5-MoE-W4A16-CT.yaml
|
||||
DeepSeek-V2-Lite-Instruct-FP8.yaml
|
||||
Qwen3-30B-A3B-NVFP4.yaml
|
||||
Laguna-XS.2-NVFP4.yaml
|
||||
|
||||
@@ -546,8 +546,11 @@ def test_flash_attn_accepts_handled_fp8_variants(
|
||||
):
|
||||
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
|
||||
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
|
||||
import vllm.v1.attention.backends.flash_attn as fa_mod
|
||||
import vllm.v1.attention.backends.fa_utils as fa_utils_mod
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
|
||||
monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
|
||||
# The fp8 decision is made in fa_utils, using its own current_platform
|
||||
# binding, so patch is_xpu there (not on flash_attn's) to stay robust to
|
||||
# import order across earlier tests that patch vllm.platforms.current_platform.
|
||||
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
|
||||
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
|
||||
@@ -11,6 +11,9 @@ from vllm._custom_ops import (
|
||||
scaled_fp8_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.ops.triton_merge_attn_states import (
|
||||
mask_empty_context,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_merge_attn_states import (
|
||||
merge_attn_states as merge_attn_states_triton,
|
||||
)
|
||||
@@ -73,6 +76,59 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
|
||||
all_case_info: list[tuple] = []
|
||||
|
||||
|
||||
def test_mask_empty_context() -> None:
|
||||
query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
|
||||
query_start_loc = torch.cat(
|
||||
(torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
|
||||
).cuda()
|
||||
context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
|
||||
context_start_loc = torch.cat(
|
||||
(torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
|
||||
).cuda()
|
||||
num_heads, num_tokens, head_dim = 4, 165, 16
|
||||
lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
|
||||
# Empty-context rows carry undefined (possibly non-finite) attention output.
|
||||
output[33:164] = float("nan")
|
||||
|
||||
expected_lse = lse.clone()
|
||||
expected_lse[:, 33:164] = float("-inf")
|
||||
expected_output = output.clone()
|
||||
expected_output[33:164] = 0.0
|
||||
|
||||
mask_empty_context(lse, output, query_start_loc, context_start_loc)
|
||||
|
||||
torch.testing.assert_close(lse, expected_lse)
|
||||
torch.testing.assert_close(output, expected_output)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
|
||||
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
|
||||
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
|
||||
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
|
||||
scales must not surface as NaN in the merged output."""
|
||||
num_tokens, num_heads, head_size = 6, 8, 128
|
||||
prefix_output = torch.zeros(
|
||||
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
|
||||
)
|
||||
prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
suffix_output = torch.zeros(
|
||||
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
|
||||
)
|
||||
suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
|
||||
# Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
|
||||
# their outputs and set both LSEs to -inf).
|
||||
empty = slice(2, 4)
|
||||
prefix_lse[:, empty] = float("-inf")
|
||||
suffix_lse[:, empty] = float("-inf")
|
||||
|
||||
output = torch.empty_like(prefix_output)
|
||||
merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
|
||||
|
||||
assert not output.isnan().any()
|
||||
|
||||
|
||||
def generate_markdown_table():
|
||||
global all_case_info
|
||||
table_header = (
|
||||
|
||||
@@ -40,10 +40,11 @@ BATCH_SIZE = 128
|
||||
CONTEXT_LEN = 8192
|
||||
PAGE_SIZE = 1
|
||||
|
||||
# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
|
||||
# KV-cache -> fp8_e4m3 kv.
|
||||
EXPECTED_Q_DTYPE = torch.bfloat16
|
||||
EXPECTED_KV_DTYPE = torch.float8_e4m3fn
|
||||
# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
|
||||
# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
|
||||
# instead of hardcoding a literal (see #47276).
|
||||
EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
|
||||
EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
|
||||
# is excluded: it holds raw device pointers, never equal across allocations.
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_default_torch_device():
|
||||
"""Several kernel tests call torch.set_default_device without restoring
|
||||
it, which poisons subsequent tests in the same pytest run (e.g. CPU
|
||||
tensors silently created on CUDA). Restore the factory default after
|
||||
every test.
|
||||
"""
|
||||
yield
|
||||
torch.set_default_device(None)
|
||||
@@ -33,12 +33,10 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402
|
||||
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
|
||||
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
|
||||
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
|
||||
seq_lens = torch.randint(
|
||||
1,
|
||||
130,
|
||||
(num_seqs,),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
rng_cpu = torch.Generator("cpu").manual_seed(1234)
|
||||
rng = torch.Generator("cuda").manual_seed(2345)
|
||||
|
||||
seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu)
|
||||
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
|
||||
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
|
||||
total_tokens = int(cu_seqlens[-1].item())
|
||||
@@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
|
||||
head_k_dim,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
generator=rng,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
k = torch.randn_like(q, generator=rng)
|
||||
v = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
@@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
|
||||
head_v_dim,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
generator=rng,
|
||||
)
|
||||
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
|
||||
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
|
||||
a = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
|
||||
)
|
||||
b = torch.randn(
|
||||
1,
|
||||
total_tokens,
|
||||
num_v_heads,
|
||||
device="cuda",
|
||||
dtype=dtype,
|
||||
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
|
||||
)
|
||||
# Match upstream FLA GatedDeltaNet synthetic initialization:
|
||||
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
|
||||
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
|
||||
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(
|
||||
0, 16, generator=rng
|
||||
)
|
||||
A_log = torch.log(A)
|
||||
dt = torch.exp(
|
||||
torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
|
||||
torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng)
|
||||
* (math.log(0.1) - math.log(0.001))
|
||||
+ math.log(0.001)
|
||||
)
|
||||
@@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
|
||||
head_k_dim,
|
||||
device="cuda",
|
||||
dtype=state_dtype,
|
||||
generator=rng,
|
||||
)
|
||||
* 0.05
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@ import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
RoutingMethodType,
|
||||
get_routing_method_type,
|
||||
@@ -231,3 +232,119 @@ def test_dsv4_fast_topk(
|
||||
atol=2e-5,
|
||||
rtol=2e-5,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="This test is skipped on non-CUDA platform.",
|
||||
)
|
||||
@pytest.mark.parametrize("use_hash", [False, True])
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
@pytest.mark.parametrize("use_padding_mask", [False, True])
|
||||
@pytest.mark.parametrize("pad_with_nan", [False, True])
|
||||
@pytest.mark.parametrize("num_experts", [128, 256, 384])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
|
||||
def test_fused_topk_softplus_sqrt_padding(
|
||||
use_hash: bool,
|
||||
use_bias: bool,
|
||||
use_padding_mask: bool,
|
||||
pad_with_nan: bool,
|
||||
num_experts: int,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
"""Verify explicit padding and NaN-padded rows do not affect real rows."""
|
||||
torch.manual_seed(0)
|
||||
num_tokens = 8
|
||||
topk = 6
|
||||
indices_dtype = torch.int32
|
||||
|
||||
gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
|
||||
|
||||
padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
|
||||
padding_rows[1::2] = True
|
||||
if pad_with_nan:
|
||||
gating_output[padding_rows] = float("nan")
|
||||
is_padding = padding_rows if use_padding_mask else None
|
||||
|
||||
# A negative correction bias makes explicit pad rows look selectable unless
|
||||
# the kernel uses the is_padding guard.
|
||||
e_score_correction_bias = None
|
||||
if use_bias:
|
||||
e_score_correction_bias = (
|
||||
-torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0
|
||||
)
|
||||
|
||||
input_ids = None
|
||||
hash_indices_table = None
|
||||
if use_hash:
|
||||
vocab_size = 64
|
||||
hash_indices_table = torch.stack(
|
||||
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
|
||||
).to(device="cuda", dtype=indices_dtype)
|
||||
input_ids = torch.randint(
|
||||
0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda"
|
||||
)
|
||||
|
||||
topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda")
|
||||
topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda")
|
||||
token_expert_indices = torch.empty(
|
||||
num_tokens, topk, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
ops.topk_hash_softplus_sqrt(
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
token_expert_indices,
|
||||
gating_output,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
input_tokens=input_ids,
|
||||
hash_indices_table=hash_indices_table,
|
||||
is_padding=is_padding,
|
||||
)
|
||||
|
||||
if use_padding_mask:
|
||||
pad_ids = topk_ids[padding_rows]
|
||||
pad_weights = topk_weights[padding_rows]
|
||||
assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), (
|
||||
f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}"
|
||||
)
|
||||
assert (pad_weights == 0).all(), (
|
||||
"Explicit pad rows should have all-zero weights, "
|
||||
f"got {pad_weights.tolist()}"
|
||||
)
|
||||
|
||||
if pad_with_nan:
|
||||
nan_pad_weights = topk_weights[padding_rows]
|
||||
assert torch.isfinite(nan_pad_weights).all(), (
|
||||
f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}"
|
||||
)
|
||||
assert (nan_pad_weights == 0).all(), (
|
||||
"NaN-padded rows should have all-zero weights, "
|
||||
f"got {nan_pad_weights.tolist()}"
|
||||
)
|
||||
|
||||
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
|
||||
gating_output=gating_output,
|
||||
topk=topk,
|
||||
renormalize=True,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
input_ids=input_ids,
|
||||
hash_indices_table=hash_indices_table,
|
||||
)
|
||||
|
||||
rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda")
|
||||
if use_padding_mask or pad_with_nan:
|
||||
rows_to_compare = ~padding_rows
|
||||
|
||||
sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1)
|
||||
sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1)
|
||||
torch.testing.assert_close(
|
||||
sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0
|
||||
)
|
||||
|
||||
sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)
|
||||
sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)
|
||||
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
|
||||
|
||||
@@ -21,9 +21,9 @@ def test_gather_cache_oob():
|
||||
seq_starts causes the block_table offset to read out of bounds.
|
||||
"""
|
||||
|
||||
batch_size = 1
|
||||
block_size = 64
|
||||
entry_size = 128
|
||||
# The kernel only supports the MLA entry sizes.
|
||||
entry_size = 576
|
||||
|
||||
block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
|
||||
|
||||
@@ -34,6 +34,7 @@ def test_gather_cache_oob():
|
||||
|
||||
seq_len = 65
|
||||
cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
|
||||
token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
|
||||
|
||||
# src_cache: [num_blocks, block_size, entry_size]
|
||||
num_blocks = 5
|
||||
@@ -51,7 +52,8 @@ def test_gather_cache_oob():
|
||||
dst,
|
||||
block_table,
|
||||
cu_seq_lens,
|
||||
batch_size,
|
||||
token_to_seq,
|
||||
seq_len,
|
||||
"auto", # kv_cache_dtype
|
||||
scale,
|
||||
seq_starts,
|
||||
|
||||
@@ -13,6 +13,7 @@ These tests cover:
|
||||
"""
|
||||
|
||||
import math
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -27,6 +28,7 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
|
||||
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
|
||||
_launch_two_stage_sparse_attn_compressor,
|
||||
)
|
||||
from vllm.models.deepseek_v4.compressor import _get_c128_boundary
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
|
||||
@@ -58,6 +60,25 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
|
||||
return x_fp8, scales
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("starts", "query_start_loc", "expected"),
|
||||
[
|
||||
([0], [0, 127], False),
|
||||
([0], [0, 128], True),
|
||||
([127], [0, 1], True),
|
||||
([128], [0, 127], False),
|
||||
([1, 255], [0, 1, 2], True),
|
||||
(None, [0, 1], None),
|
||||
],
|
||||
)
|
||||
def test_get_c128_boundary(starts, query_start_loc, expected):
|
||||
metadata = SimpleNamespace(
|
||||
_num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
|
||||
query_start_loc_cpu=torch.tensor(query_start_loc),
|
||||
)
|
||||
assert _get_c128_boundary(metadata) is expected
|
||||
|
||||
|
||||
# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
|
||||
device = torch.device("cuda")
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
|
||||
model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
|
||||
)
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
|
||||
|
||||
|
||||
@@ -726,14 +726,19 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
|
||||
This catches stride/layout bugs that only manifest when the einsum
|
||||
kernel actually consumes the quantized activations.
|
||||
"""
|
||||
from deep_gemm.utils.math import ceil_div
|
||||
|
||||
from vllm.utils.deep_gemm import (
|
||||
fp8_einsum,
|
||||
is_deep_gemm_supported,
|
||||
per_block_cast_to_fp8,
|
||||
transform_sf_into_required_layout,
|
||||
)
|
||||
|
||||
if not is_deep_gemm_supported():
|
||||
pytest.skip("DeepGEMM not supported on this platform")
|
||||
|
||||
def ceil_div(a: int, b: int) -> int:
|
||||
return (a + b - 1) // b
|
||||
|
||||
heads_per_group = num_heads // n_groups
|
||||
d = heads_per_group * HEAD_DIM
|
||||
o_lora_rank = 1024
|
||||
@@ -809,8 +814,12 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
|
||||
# -- Checks --
|
||||
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
|
||||
# come from fp32 ordering and UE8M0 boundary shifts only.
|
||||
# Use relative diff (same metric as test_fp8_einsum.py).
|
||||
from deep_gemm.testing import calc_diff
|
||||
# Use relative diff (same metric as deep_gemm.testing.calc_diff).
|
||||
def calc_diff(x, y):
|
||||
x, y = x.double(), y.double()
|
||||
denominator = (x * x + y * y).sum()
|
||||
sim = 2 * (x * y).sum() / denominator
|
||||
return 1 - sim
|
||||
|
||||
z_diff = calc_diff(z_fused, z_ref)
|
||||
assert z_diff < 0.01, (
|
||||
|
||||
@@ -84,6 +84,22 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
|
||||
return roped
|
||||
|
||||
|
||||
def assert_fp8_cache_close(kv_cache, expected_kv_cache):
|
||||
"""Compare two e4m3 caches allowing 1 ulp.
|
||||
|
||||
On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
|
||||
reshape_and_cache_flash reference quantizes the bf16-materialized value, so
|
||||
rounding-boundary values may differ by one e4m3 code.
|
||||
"""
|
||||
byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
|
||||
got = kv_cache.view(torch.float8_e4m3fn).float()
|
||||
exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
|
||||
ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
|
||||
assert bool(ok.all()), (
|
||||
f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
|
||||
)
|
||||
|
||||
|
||||
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
|
||||
|
||||
|
||||
@@ -265,7 +281,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
|
||||
scale,
|
||||
scale,
|
||||
)
|
||||
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
|
||||
assert_fp8_cache_close(kv_cache, expected_kv_cache)
|
||||
else:
|
||||
for t in range(num_tokens):
|
||||
s = slot_mapping[t].item()
|
||||
@@ -383,7 +399,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
|
||||
scale,
|
||||
scale,
|
||||
)
|
||||
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
|
||||
assert_fp8_cache_close(kv_cache, expected_kv_cache)
|
||||
else:
|
||||
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
|
||||
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
|
||||
|
||||
@@ -41,11 +41,12 @@ def test_fused_recurrent_packed_decode_matches_reference(
|
||||
A_log = torch.randn((HV,), device=device, dtype=dtype)
|
||||
dt_bias = torch.randn((HV,), device=device, dtype=dtype)
|
||||
|
||||
# Continuous batching indices (include PAD_SLOT_ID=-1 cases).
|
||||
ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
|
||||
# Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
|
||||
# reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
|
||||
ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
|
||||
ssm_state_indices[-3:] = -1
|
||||
|
||||
state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
|
||||
state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
|
||||
state_ref = state0.clone()
|
||||
state_packed = state0.clone()
|
||||
|
||||
@@ -94,5 +95,8 @@ def test_fused_recurrent_packed_decode_matches_reference(
|
||||
|
||||
atol = 2e-2 if dtype != torch.float32 else 1e-4
|
||||
rtol = 1e-2 if dtype != torch.float32 else 1e-4
|
||||
torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
|
||||
# Output rows for PAD_SLOT_ID entries are never written (uninitialized in
|
||||
# both paths), so compare only the valid rows.
|
||||
valid = ssm_state_indices > 0
|
||||
torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
|
||||
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
|
||||
|
||||
@@ -58,10 +58,12 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
|
||||
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
|
||||
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
|
||||
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
|
||||
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
|
||||
# state indices start at 1.
|
||||
ssm_state = torch.rand(
|
||||
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
|
||||
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
|
||||
)
|
||||
state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
|
||||
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
|
||||
cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
|
||||
|
||||
beta = b.sigmoid()
|
||||
@@ -144,13 +146,14 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
|
||||
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
|
||||
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
|
||||
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
|
||||
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
|
||||
# state indices start at 1.
|
||||
ssm_state = torch.rand(
|
||||
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
|
||||
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
|
||||
)
|
||||
state_indices = torch.randperm(
|
||||
total_entries,
|
||||
dtype=torch.int32,
|
||||
)[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
|
||||
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
|
||||
:num_tokens
|
||||
].view(num_reqs, num_speculative_tokens + 1)
|
||||
num_accepted_tokens = torch.randint(
|
||||
1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
|
||||
)
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.model_loader.mtp_validation import (
|
||||
disable_mtp_completeness_check,
|
||||
is_mtp_completeness_check_enabled,
|
||||
)
|
||||
|
||||
|
||||
def test_disable_mtp_completeness_check_is_scoped():
|
||||
assert is_mtp_completeness_check_enabled()
|
||||
|
||||
with pytest.raises(RuntimeError), disable_mtp_completeness_check():
|
||||
assert not is_mtp_completeness_check_enabled()
|
||||
raise RuntimeError
|
||||
|
||||
assert is_mtp_completeness_check_enabled()
|
||||
@@ -6,15 +6,43 @@ import time
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm.config.multimodal as multimodal_config_module
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.multimodal.gpu_ipc_memory import (
|
||||
MultiModalGPUMemoryPool,
|
||||
get_mm_gpu_ipc_pool,
|
||||
maybe_init_mm_gpu_ipc_pool,
|
||||
reserve_mm_ipc_gpu_memory,
|
||||
set_mm_gpu_ipc_pool,
|
||||
)
|
||||
from vllm.multimodal.video import (
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
|
||||
PYNVVIDEOCODEC_MAX_RETAINED_DECODERS,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
|
||||
|
||||
def _mm_config(
|
||||
*,
|
||||
mm_ipc_gpu_memory_gb: float = 0,
|
||||
video_backend: str | None = None,
|
||||
) -> MultiModalConfig:
|
||||
video_kwargs = {} if video_backend is None else {"video_backend": video_backend}
|
||||
return MultiModalConfig(
|
||||
mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb,
|
||||
media_io_kwargs={"video": video_kwargs} if video_kwargs else {},
|
||||
)
|
||||
|
||||
|
||||
def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int:
|
||||
return api_process_count * (
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS
|
||||
+ PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES
|
||||
)
|
||||
|
||||
|
||||
def test_acquire_release_accounting():
|
||||
pool = MultiModalGPUMemoryPool(total_bytes=100)
|
||||
assert pool.available_bytes == 100
|
||||
@@ -143,3 +171,72 @@ def test_global_pool_splits_budget_across_api_processes():
|
||||
def test_global_pool_rejects_invalid_api_process_count():
|
||||
with pytest.raises(ValueError):
|
||||
maybe_init_mm_gpu_ipc_pool(2, api_process_count=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("video_backend", [None, "opencv"])
|
||||
def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
video_backend: str | None,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
mm_config = _mm_config(
|
||||
mm_ipc_gpu_memory_gb=0.25,
|
||||
video_backend=video_backend,
|
||||
)
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(GiB_bytes, mm_config) == int(0.75 * GiB_bytes)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
mm_config = _mm_config(
|
||||
mm_ipc_gpu_memory_gb=0.25,
|
||||
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
available_bytes = 4 * GiB_bytes
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(available_bytes, mm_config) == (
|
||||
available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget()
|
||||
)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
available_bytes = 4 * GiB_bytes
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(available_bytes, _mm_config()) == (
|
||||
available_bytes - _pynvvideocodec_decoder_budget()
|
||||
)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_scales_decoder_budget_by_api_servers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
multimodal_config_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
available_bytes = 8 * GiB_bytes
|
||||
|
||||
assert reserve_mm_ipc_gpu_memory(
|
||||
available_bytes,
|
||||
_mm_config(),
|
||||
api_process_count=3,
|
||||
) == available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
|
||||
|
||||
@@ -920,3 +920,45 @@ class TestDelegatingParserLargeDelta:
|
||||
assert output.tool_calls[0]["name"] == "get_weather"
|
||||
args = json.loads(output.tool_calls[0]["arguments"])
|
||||
assert args == {"location": "Berlin"}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_size",
|
||||
[1, 2, 3, 5, None],
|
||||
ids=lambda c: f"chunk={c}",
|
||||
)
|
||||
def test_eos_not_leaked_when_reasoning_never_ends(self, chunk_size):
|
||||
"""EOS must not leak into reasoning_content when the model never
|
||||
emits </think> (generation ends while still in REASONING state)."""
|
||||
eos_text = "<|end▁of▁sentence|>"
|
||||
eos_id = 128801
|
||||
vocab = {
|
||||
**_DSV4_FULL_VOCAB,
|
||||
eos_text: eos_id,
|
||||
}
|
||||
|
||||
reasoning_text = "Good morning! How can I help you today?"
|
||||
tokens: list[tuple[int, str]] = []
|
||||
tid = 100
|
||||
for word in reasoning_text.split(" "):
|
||||
prefix = " " if tokens else ""
|
||||
tokens.append((tid, prefix + word))
|
||||
tid += 1
|
||||
tokens.append((eos_id, eos_text))
|
||||
|
||||
tokenizer = MockTokenizer(vocab=vocab, tokens=tokens)
|
||||
parser = _DeepSeekV4Delegating(
|
||||
tokenizer,
|
||||
chat_template_kwargs={"thinking": True},
|
||||
)
|
||||
deltas = replay_streaming(
|
||||
parser,
|
||||
tokens,
|
||||
chunk_size=chunk_size,
|
||||
finished_on_last=True,
|
||||
)
|
||||
output = collect_output(deltas)
|
||||
|
||||
assert reasoning_text in output.reasoning
|
||||
assert eos_text not in output.reasoning
|
||||
assert output.content == ""
|
||||
assert output.tool_calls == []
|
||||
|
||||
@@ -1624,19 +1624,41 @@ class TestDropSpecialTokens:
|
||||
assert delta is not None
|
||||
assert "<bos>" in delta.reasoning
|
||||
|
||||
def test_drops_suppressed_with_skip_tool_parsing(self):
|
||||
"""When skip_tool_parsing is active, drop tokens are preserved
|
||||
as content so a later tool-call pass can see them."""
|
||||
def test_drops_applied_with_skip_tool_parsing(self):
|
||||
"""Drop tokens are always dropped, even with skip_tool_parsing.
|
||||
DROP_TERMINALs have no transitions by construction, so no parser
|
||||
pass can use them."""
|
||||
for initial_state in (ParserState.REASONING, ParserState.CONTENT):
|
||||
engine = _make_engine(
|
||||
vocab=_DROP_VOCAB,
|
||||
special_tokens=list(_DROP_VOCAB.keys()),
|
||||
)
|
||||
engine._engine.skip_tool_parsing = True
|
||||
engine._engine.reset(initial_state=initial_state)
|
||||
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
|
||||
delta = engine._events_to_delta(events)
|
||||
assert delta is not None
|
||||
output = (delta.reasoning or "") + (delta.content or "")
|
||||
assert "<bos>" not in output, f"<bos> leaked in state {initial_state}"
|
||||
|
||||
def test_transitions_unaffected_by_drop_in_reasoning_with_skip_tool_parsing(self):
|
||||
"""With skip_tool_parsing in REASONING state, drop tokens are
|
||||
removed but configured terminals still fire their transitions."""
|
||||
engine = _make_engine(
|
||||
vocab=_DROP_VOCAB,
|
||||
special_tokens=list(_DROP_VOCAB.keys()),
|
||||
)
|
||||
engine._engine.skip_tool_parsing = True
|
||||
engine._engine.reset()
|
||||
events = engine._engine.feed("hello<bos>world", [72, 204, 73])
|
||||
delta = engine._events_to_delta(events)
|
||||
assert delta is not None
|
||||
assert "<bos>" in delta.reasoning
|
||||
events = engine._engine.feed("thought<bos></think>answer", [72, 204, 201, 73])
|
||||
types = [e.type for e in events]
|
||||
assert EventType.REASONING_CHUNK in types
|
||||
assert EventType.REASONING_END in types
|
||||
assert EventType.TEXT_CHUNK in types
|
||||
reasoning_text = "".join(
|
||||
e.value for e in events if e.type == EventType.REASONING_CHUNK
|
||||
)
|
||||
assert "<bos>" not in reasoning_text
|
||||
|
||||
def test_drops_in_tool_args_state(self):
|
||||
"""Drop tokens in TOOL_ARGS state are silently discarded."""
|
||||
|
||||
@@ -50,7 +50,7 @@ def test_oot_registration_embedding(
|
||||
m.setenv("VLLM_PLUGINS", "register_dummy_model")
|
||||
prompts = ["Hello, my name is", "The text does not matter"]
|
||||
llm = LLM(
|
||||
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=2048
|
||||
model=dummy_gemma2_embedding_path, load_format="dummy", max_model_len=512
|
||||
)
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
|
||||
@@ -5,13 +5,17 @@
|
||||
Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
from vllm.model_executor.layers.fused_moe import RoutedExperts
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import (
|
||||
AutoGPTQConfig,
|
||||
AutoGPTQLinearMethod,
|
||||
AutoGPTQMoEMethod,
|
||||
)
|
||||
|
||||
PROMPT = "On the surface of Mars, we found"
|
||||
@@ -54,3 +58,78 @@ def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
|
||||
def test_auto_gptq_config_get_name():
|
||||
"""Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
|
||||
assert AutoGPTQConfig.get_name() == "auto_gptq"
|
||||
|
||||
|
||||
def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
|
||||
method = object.__new__(AutoGPTQMoEMethod)
|
||||
method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
|
||||
method.input_dtype = None
|
||||
method.experts_cls = None
|
||||
layer = torch.nn.Module()
|
||||
|
||||
method.create_weights(
|
||||
layer=layer,
|
||||
num_experts=2,
|
||||
hidden_size=8,
|
||||
intermediate_size_per_partition=4,
|
||||
params_dtype=torch.float16,
|
||||
intermediate_size_full=4,
|
||||
weight_loader=lambda *args, **kwargs: None,
|
||||
)
|
||||
|
||||
assert layer.w13_bias.shape == (2, 8)
|
||||
assert layer.w2_bias.shape == (2, 8)
|
||||
assert torch.count_nonzero(layer.w13_bias) == 0
|
||||
assert torch.count_nonzero(layer.w2_bias) == 0
|
||||
|
||||
|
||||
def test_routed_experts_loads_per_expert_biases():
|
||||
class Loader:
|
||||
quant_config = None
|
||||
quant_method = object()
|
||||
moe_config = SimpleNamespace(
|
||||
is_act_and_mul=True,
|
||||
tp_rank=0,
|
||||
moe_parallel_config=SimpleNamespace(tp_size=1),
|
||||
)
|
||||
_get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
|
||||
_narrow_expert_data_for_padding = staticmethod(
|
||||
RoutedExperts._narrow_expert_data_for_padding
|
||||
)
|
||||
_load_w13 = RoutedExperts._load_w13
|
||||
_loaded_expert_biases = set()
|
||||
|
||||
@staticmethod
|
||||
def _map_global_expert_id_to_local_expert_id(expert_id):
|
||||
return expert_id
|
||||
|
||||
loader = Loader()
|
||||
w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
|
||||
w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
|
||||
|
||||
for shard_id, loaded in (
|
||||
("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
|
||||
("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
|
||||
):
|
||||
assert RoutedExperts.weight_loader(
|
||||
loader,
|
||||
w13_bias,
|
||||
loaded,
|
||||
weight_name="model.layers.0.mlp.experts.w13_bias",
|
||||
shard_id=shard_id,
|
||||
expert_id=0,
|
||||
return_success=True,
|
||||
)
|
||||
|
||||
assert RoutedExperts.weight_loader(
|
||||
loader,
|
||||
w2_bias,
|
||||
torch.tensor([9.0, 10.0, 11.0, 12.0]),
|
||||
weight_name="model.layers.0.mlp.experts.w2_bias",
|
||||
shard_id="w2",
|
||||
expert_id=0,
|
||||
return_success=True,
|
||||
)
|
||||
assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
|
||||
assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
|
||||
assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}
|
||||
|
||||
@@ -5,45 +5,194 @@ from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
ActivationOrdering,
|
||||
QuantizationArgs,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
|
||||
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
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
WNA16MoEBackend,
|
||||
_backend_incompatibility_reason,
|
||||
_convert_moe_wna16_humming_tensors,
|
||||
convert_to_wna16_moe_kernel_format,
|
||||
map_wna16_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import moe_wna16
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
|
||||
from vllm.model_executor.layers.quantization.moe_wna16 import (
|
||||
MoeWNA16Config,
|
||||
MoeWNA16Method,
|
||||
)
|
||||
|
||||
|
||||
@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 test_map_wna16_backend_supports_triton():
|
||||
assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
|
||||
|
||||
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,
|
||||
@pytest.mark.parametrize(
|
||||
("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
|
||||
[
|
||||
(
|
||||
WNA16MoEBackend.TRITON,
|
||||
AutoAWQConfig(4, 128, True, False),
|
||||
True,
|
||||
False,
|
||||
"AutoAWQ weight layout",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.TRITON,
|
||||
AutoGPTQConfig(4, 128, True, True, False, {}, {}),
|
||||
False,
|
||||
False,
|
||||
"activation ordering",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.TRITON,
|
||||
QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
group_size=128,
|
||||
actorder=ActivationOrdering.GROUP,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"activation ordering",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.TRITON,
|
||||
AutoGPTQConfig(4, 128, False, True, False, {}, {}),
|
||||
False,
|
||||
True,
|
||||
"bias",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.MARLIN,
|
||||
MoeWNA16Config(
|
||||
linear_quant_method="gptq",
|
||||
weight_bits=4,
|
||||
group_size=128,
|
||||
has_zp=False,
|
||||
lm_head_quantized=False,
|
||||
modules_to_not_convert=None,
|
||||
full_config={},
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"MoeWNA16 checkpoint layout",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_wna16_oracle_rejects_incompatible_quant_structures(
|
||||
backend, quant_config, may_have_zp, may_have_bias, expected
|
||||
):
|
||||
reason = _backend_incompatibility_reason(
|
||||
backend=backend,
|
||||
quant_config=quant_config,
|
||||
may_have_zp=may_have_zp,
|
||||
may_have_bias=may_have_bias,
|
||||
)
|
||||
|
||||
assert reason is not None
|
||||
assert expected in reason
|
||||
|
||||
|
||||
def test_compressed_tensors_weights_are_transposed_for_triton():
|
||||
quant_config = QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
group_size=32,
|
||||
)
|
||||
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
|
||||
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
|
||||
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
|
||||
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
|
||||
|
||||
converted = convert_to_wna16_moe_kernel_format(
|
||||
backend=WNA16MoEBackend.TRITON,
|
||||
layer=torch.nn.Module(),
|
||||
quant_config=quant_config,
|
||||
input_dtype=None,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
)
|
||||
|
||||
assert converted is not None
|
||||
assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
|
||||
assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
|
||||
assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
|
||||
assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
|
||||
|
||||
|
||||
def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
|
||||
method = object.__new__(MoeWNA16Method)
|
||||
method.experts_cls = object
|
||||
method.wna16_backend = WNA16MoEBackend.HUMMING
|
||||
method.moe = object()
|
||||
quant_config = object()
|
||||
method.get_fused_moe_quant_config = lambda layer: quant_config
|
||||
layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))
|
||||
captured = {}
|
||||
kernel = object()
|
||||
|
||||
def fake_make_wna16_moe_kernel(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return kernel
|
||||
|
||||
monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)
|
||||
|
||||
method._setup_kernel(layer)
|
||||
|
||||
assert method.moe_kernel is kernel
|
||||
assert captured["backend"] == WNA16MoEBackend.HUMMING
|
||||
assert captured["layer"] is layer
|
||||
|
||||
|
||||
def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
|
||||
qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8)
|
||||
scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4)
|
||||
qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2)
|
||||
|
||||
converted = _convert_moe_wna16_humming_tensors(
|
||||
{"qweight": qweight, "scales": scales, "qzeros": qzeros},
|
||||
has_zero_point=True,
|
||||
)
|
||||
|
||||
assert torch.equal(converted["weight"], qweight.view(torch.int32))
|
||||
assert converted["weight"].shape == (1, 4, 2)
|
||||
assert torch.equal(converted["weight_scale"], scales)
|
||||
expected_qzeros = (
|
||||
qzeros.transpose(-1, -2)
|
||||
.contiguous()
|
||||
.view(torch.int32)
|
||||
.transpose(-1, -2)
|
||||
.contiguous()
|
||||
)
|
||||
assert torch.equal(converted["zero_point"], expected_qzeros)
|
||||
assert converted["zero_point"].shape == (1, 2, 2)
|
||||
|
||||
|
||||
def test_moe_wna16_uses_humming_quant_config(monkeypatch):
|
||||
from vllm.model_executor.layers.quantization.utils import humming_utils
|
||||
|
||||
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,
|
||||
method.wna16_backend = WNA16MoEBackend.HUMMING
|
||||
layer = object()
|
||||
quant_config = object()
|
||||
monkeypatch.setattr(
|
||||
humming_utils,
|
||||
"get_humming_moe_quant_config",
|
||||
lambda actual_layer: quant_config if actual_layer is layer else 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
|
||||
assert method.get_fused_moe_quant_config(layer) is quant_config
|
||||
|
||||
@@ -1,13 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
|
||||
from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
|
||||
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
|
||||
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
|
||||
align_trtllm_fp4_moe_hidden_dim_for_fi,
|
||||
)
|
||||
|
||||
|
||||
def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
|
||||
monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
|
||||
|
||||
num_experts = 3
|
||||
layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
|
||||
w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
|
||||
w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
|
||||
w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
|
||||
w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
|
||||
weight_scale = torch.ones(num_experts)
|
||||
|
||||
outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
|
||||
backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
|
||||
layer=layer,
|
||||
w13=w13,
|
||||
w13_scale=w13_scale,
|
||||
w13_scale_2=weight_scale,
|
||||
a13_scale=torch.tensor([1.0, 2.0, 3.0]),
|
||||
w2=w2,
|
||||
w2_scale=w2_scale,
|
||||
w2_scale_2=weight_scale,
|
||||
a2_scale=torch.tensor([4.0, 5.0, 6.0]),
|
||||
is_act_and_mul=False,
|
||||
)
|
||||
a13_scale, a2_scale = outputs[3], outputs[7]
|
||||
|
||||
torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
|
||||
torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
|
||||
distinct_values = torch.arange(num_experts, dtype=torch.float32)
|
||||
a13_scale.copy_(distinct_values)
|
||||
a2_scale.copy_(distinct_values)
|
||||
torch.testing.assert_close(a13_scale, distinct_values)
|
||||
torch.testing.assert_close(a2_scale, distinct_values)
|
||||
|
||||
|
||||
def test_align_trtllm_fp4_moe_hidden_dim_noop():
|
||||
w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
|
||||
w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
|
||||
|
||||
@@ -46,9 +46,12 @@ BACKENDS_TO_TEST = [
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
# Use the platform's preferred FP8 type so the stored cache matches what the
|
||||
# backends reinterpret at runtime. On ROCm gfx94x this is e4m3fnuz, not e4m3fn;
|
||||
# storing e4m3fn bytes there would be re-read as fnuz and produce NaNs.
|
||||
FP8_KV_CACHE_DTYPES = {
|
||||
"fp8": torch.float8_e4m3fn,
|
||||
"fp8_e4m3": torch.float8_e4m3fn,
|
||||
"fp8": current_platform.fp8_dtype(),
|
||||
"fp8_e4m3": current_platform.fp8_dtype(),
|
||||
}
|
||||
|
||||
# Remove flashinfer from the list if it's not available
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
|
||||
from vllm.v1.core.kv_cache_utils import (
|
||||
_get_kv_cache_config_packed,
|
||||
_get_kv_cache_groups_uniform_groups,
|
||||
get_kv_cache_config_from_groups,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
@@ -16,6 +17,7 @@ from vllm.v1.kv_cache_interface import (
|
||||
KVCacheGroupSpec,
|
||||
KVCacheTensor,
|
||||
MLAAttentionSpec,
|
||||
SlidingWindowMLASpec,
|
||||
SlidingWindowSpec,
|
||||
UniformTypeKVCacheSpecs,
|
||||
)
|
||||
@@ -109,7 +111,132 @@ def _page_sizes_by_layer(
|
||||
return page_sizes
|
||||
|
||||
|
||||
def _packing_by_layer(
|
||||
tensors: list[KVCacheTensor],
|
||||
) -> dict[str, tuple[int, int]]:
|
||||
return {
|
||||
layer_name: (tensor.offset, tensor.block_stride)
|
||||
for tensor in tensors
|
||||
for layer_name in tensor.shared_by
|
||||
}
|
||||
|
||||
|
||||
def _make_views(
|
||||
groups: list[KVCacheGroupSpec],
|
||||
num_blocks: int,
|
||||
tensors: list[KVCacheTensor],
|
||||
) -> dict[str, torch.Tensor]:
|
||||
page_sizes = _page_sizes_by_layer(groups)
|
||||
packing = _packing_by_layer(tensors)
|
||||
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
|
||||
return {
|
||||
layer_name: torch.as_strided(
|
||||
backing,
|
||||
size=(num_blocks, page_size),
|
||||
stride=(packing[layer_name][1], 1),
|
||||
storage_offset=packing[layer_name][0],
|
||||
)
|
||||
for layer_name, page_size in page_sizes.items()
|
||||
}
|
||||
|
||||
|
||||
def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
|
||||
specs = {
|
||||
f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
|
||||
for i, page_size in enumerate(page_sizes)
|
||||
}
|
||||
return KVCacheGroupSpec(
|
||||
layer_names=list(specs),
|
||||
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
|
||||
)
|
||||
|
||||
|
||||
class TestInterleavedPacking:
|
||||
def test_compact_cache_overlays_fp32_state_group(self):
|
||||
full_specs = {}
|
||||
state_specs = {}
|
||||
for i in range(2):
|
||||
full_specs[f"mla.{i}"] = MLAAttentionSpec(
|
||||
block_size=256,
|
||||
num_kv_heads=1,
|
||||
head_size=512,
|
||||
dtype=torch.uint8,
|
||||
page_size_padded=32768,
|
||||
indexes_kv_by_block_stride=True,
|
||||
compress_ratio=4,
|
||||
)
|
||||
full_specs[f"indexer.{i}"] = MLAAttentionSpec(
|
||||
block_size=256,
|
||||
num_kv_heads=1,
|
||||
head_size=68,
|
||||
dtype=torch.uint8,
|
||||
page_size_padded=4608,
|
||||
compress_ratio=4,
|
||||
)
|
||||
state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
|
||||
block_size=4,
|
||||
num_kv_heads=1,
|
||||
head_size=2048,
|
||||
dtype=torch.float32,
|
||||
sliding_window=8,
|
||||
indexes_kv_by_block_stride=True,
|
||||
)
|
||||
state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
|
||||
block_size=4,
|
||||
num_kv_heads=1,
|
||||
head_size=512,
|
||||
dtype=torch.float32,
|
||||
sliding_window=8,
|
||||
indexes_kv_by_block_stride=True,
|
||||
)
|
||||
|
||||
grouped_specs = [
|
||||
UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
|
||||
UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
|
||||
]
|
||||
groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
|
||||
|
||||
assert len(groups) == 2
|
||||
assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
|
||||
assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
|
||||
4352
|
||||
}
|
||||
assert {
|
||||
state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
|
||||
} == {8192}
|
||||
|
||||
full_group_bytes = 2 * (32768 + 4608)
|
||||
state_group_bytes = 2 * (32768 + 8192)
|
||||
bytes_per_block = max(full_group_bytes, state_group_bytes)
|
||||
num_blocks, tensors = _get_kv_cache_config_packed(
|
||||
_mock_vllm_config(), groups, bytes_per_block * 32
|
||||
)
|
||||
assert num_blocks == 32
|
||||
assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
|
||||
|
||||
packing = _packing_by_layer(tensors)
|
||||
assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
|
||||
assert packing["indexer.0"][0] == 32768
|
||||
assert packing["indexer_state.0"][0] == 32768
|
||||
|
||||
def test_deepseek_v4_pro_stride(self):
|
||||
groups = [
|
||||
_make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
|
||||
_make_page_group("c4_state", [32768, 8192] * 30),
|
||||
_make_page_group("c128_state", [32768] * 31),
|
||||
_make_page_group("swa.0", [32768] * 31),
|
||||
_make_page_group("swa.1", [32768] * 30),
|
||||
]
|
||||
expected_stride = 1_228_800
|
||||
|
||||
num_blocks, tensors = _get_kv_cache_config_packed(
|
||||
_mock_vllm_config(), groups, expected_stride * 32
|
||||
)
|
||||
|
||||
assert num_blocks == 32
|
||||
assert {tensor.block_stride for tensor in tensors} == {expected_stride}
|
||||
assert {tensor.size for tensor in tensors} == {expected_stride * 32}
|
||||
|
||||
def test_all_tensors_have_block_stride(self):
|
||||
_, tensors = _run()
|
||||
for t in tensors:
|
||||
@@ -122,9 +249,30 @@ class TestInterleavedPacking:
|
||||
assert sizes.pop() > 0
|
||||
|
||||
def test_offsets_within_one_block(self):
|
||||
_, tensors = _run()
|
||||
for t in tensors:
|
||||
assert t.offset < t.block_stride
|
||||
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
|
||||
_, tensors = _get_kv_cache_config_packed(
|
||||
_mock_vllm_config(), groups, 100 * 1024 * 1024
|
||||
)
|
||||
page_sizes = _page_sizes_by_layer(groups)
|
||||
packing = _packing_by_layer(tensors)
|
||||
for layer_name, page_size in page_sizes.items():
|
||||
offset, block_stride = packing[layer_name]
|
||||
assert offset + page_size <= block_stride
|
||||
|
||||
def test_layouts_are_disjoint_within_each_group(self):
|
||||
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
|
||||
_, tensors = _get_kv_cache_config_packed(
|
||||
_mock_vllm_config(), groups, 100 * 1024 * 1024
|
||||
)
|
||||
page_sizes = _page_sizes_by_layer(groups)
|
||||
packing = _packing_by_layer(tensors)
|
||||
|
||||
for group in groups:
|
||||
ranges = sorted(
|
||||
(packing[name][0], packing[name][0] + page_sizes[name])
|
||||
for name in group.layer_names
|
||||
)
|
||||
assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
|
||||
|
||||
def test_all_layers_accounted_for(self):
|
||||
n_c4, n_c128, n_swa = 5, 4, 7
|
||||
@@ -135,29 +283,29 @@ class TestInterleavedPacking:
|
||||
expected = n_c4 * 2 + n_c128 + n_swa
|
||||
assert len(all_names) == expected
|
||||
|
||||
def test_strided_views_are_independent(self):
|
||||
def test_group_owned_blocks_do_not_alias(self):
|
||||
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
|
||||
page_sizes = _page_sizes_by_layer(groups)
|
||||
num_blocks, tensors = _get_kv_cache_config_packed(
|
||||
_mock_vllm_config(), groups, 100 * 1024 * 1024
|
||||
)
|
||||
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
|
||||
views = []
|
||||
for t in tensors:
|
||||
page_size = page_sizes[t.shared_by[0]]
|
||||
v = torch.as_strided(
|
||||
backing,
|
||||
size=(num_blocks, page_size),
|
||||
stride=(t.block_stride, 1),
|
||||
storage_offset=t.offset,
|
||||
)
|
||||
views.append(v)
|
||||
views = _make_views(groups, num_blocks, tensors)
|
||||
|
||||
for i, v in enumerate(views):
|
||||
v.fill_(i + 1)
|
||||
expected = {}
|
||||
value = 1
|
||||
for block_id, group in enumerate(groups):
|
||||
for layer_name in group.layer_names:
|
||||
views[layer_name][block_id].fill_(value)
|
||||
expected[layer_name] = (block_id, value)
|
||||
value += 1
|
||||
|
||||
for i, v in enumerate(views):
|
||||
assert (v == i + 1).all(), f"View {i} was corrupted"
|
||||
for layer_name, (block_id, value) in expected.items():
|
||||
assert (views[layer_name][block_id] == value).all()
|
||||
|
||||
# Once the first group releases its block, another group may reuse it.
|
||||
for layer_name in groups[1].layer_names:
|
||||
views[layer_name][0].fill_(255)
|
||||
for layer_name in groups[1].layer_names:
|
||||
assert (views[layer_name][0] == 255).all()
|
||||
|
||||
def test_hma_attention_groups_keep_default_backing(self):
|
||||
full = _make_full_spec()
|
||||
|
||||
@@ -25,14 +25,17 @@ from vllm.multimodal.inputs import (
|
||||
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
|
||||
from vllm.utils.hashing import sha256
|
||||
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
|
||||
from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
|
||||
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
|
||||
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
MambaSpec,
|
||||
)
|
||||
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
@@ -1068,6 +1071,141 @@ def test_preempt_during_execution():
|
||||
assert requests[1].output_token_ids[0] == 42
|
||||
|
||||
|
||||
def test_prefix_cache_query_not_inflated_by_connector_defer():
|
||||
"""The GPU prefix-cache query is recorded at admission, so a request the
|
||||
connector defers several times is counted once, not once per retry."""
|
||||
num_defers_before_matching = 3
|
||||
scheduler = create_scheduler(
|
||||
enable_prefix_caching=True,
|
||||
use_kv_connector=mock_kv(
|
||||
matched_tokens=0,
|
||||
is_async=False,
|
||||
num_defers_before_matching=num_defers_before_matching,
|
||||
),
|
||||
)
|
||||
request = create_requests(num_requests=1, num_tokens=32, block_size=16)[0]
|
||||
scheduler.add_request(request)
|
||||
|
||||
# Each deferred step re-runs the lookup but records nothing.
|
||||
for _ in range(num_defers_before_matching):
|
||||
assert not scheduler.schedule().scheduled_new_reqs
|
||||
|
||||
output = scheduler.schedule()
|
||||
assert any(r.req_id == request.request_id for r in output.scheduled_new_reqs)
|
||||
|
||||
stats = scheduler.kv_cache_manager.prefix_cache_stats
|
||||
assert stats is not None
|
||||
assert stats.requests == 1
|
||||
assert stats.queries == request.num_tokens
|
||||
|
||||
|
||||
def test_preemption_re_records_prefix_cache_query():
|
||||
"""A preempted request re-enters the lookup on resume, so its recomputation
|
||||
is counted again into the preempted stats."""
|
||||
scheduler = create_scheduler(enable_prefix_caching=True)
|
||||
request = create_requests(num_requests=1)[0]
|
||||
scheduler.add_request(request)
|
||||
|
||||
scheduler.schedule()
|
||||
stats = scheduler.kv_cache_manager.prefix_cache_stats
|
||||
assert stats is not None
|
||||
assert (stats.requests, stats.preempted_requests) == (1, 0)
|
||||
|
||||
scheduler.running.remove(request)
|
||||
scheduler._preempt_request(request, 0.0)
|
||||
assert request.status == RequestStatus.PREEMPTED
|
||||
|
||||
scheduler.schedule()
|
||||
assert request.status == RequestStatus.RUNNING
|
||||
assert stats.preempted_requests == 1
|
||||
|
||||
|
||||
def test_prefix_cache_stats_not_recorded_when_caching_disabled():
|
||||
"""With prefix caching off there is no local lookup, so admitting a request
|
||||
records no phantom miss."""
|
||||
scheduler = create_scheduler(enable_prefix_caching=False)
|
||||
for request in create_requests(num_requests=2):
|
||||
scheduler.add_request(request)
|
||||
|
||||
scheduler.schedule()
|
||||
|
||||
stats = scheduler.kv_cache_manager.prefix_cache_stats
|
||||
assert stats is not None
|
||||
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
|
||||
|
||||
|
||||
def test_prefix_cache_stats_counted_once_for_retried_then_scheduled_request():
|
||||
"""A real cache hit rejected once by allocate_slots and admitted on the next
|
||||
step is counted exactly once, hits included."""
|
||||
block_size = 16
|
||||
scheduler = create_scheduler(
|
||||
enable_prefix_caching=True,
|
||||
enable_chunked_prefill=False,
|
||||
block_size=block_size,
|
||||
)
|
||||
|
||||
# Seed the cache so the next request with the same prompt hits it.
|
||||
seed = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=block_size * 2,
|
||||
max_tokens=2,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["seed"],
|
||||
)[0]
|
||||
scheduler.add_request(seed)
|
||||
_step_until_done(
|
||||
scheduler,
|
||||
scheduler.schedule(),
|
||||
ModelRunnerOutput(
|
||||
req_ids=["seed"],
|
||||
req_id_to_index={"seed": 0},
|
||||
sampled_token_ids=[[1000]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
),
|
||||
)
|
||||
|
||||
# The seeding step swapped in a fresh accumulator, so re-read it; the
|
||||
# retried request must be the only thing recorded from here on.
|
||||
stats = scheduler.kv_cache_manager.prefix_cache_stats
|
||||
assert stats is not None
|
||||
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
|
||||
|
||||
retried = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=block_size * 3,
|
||||
max_tokens=1,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["retried"],
|
||||
)[0]
|
||||
scheduler.add_request(retried)
|
||||
|
||||
# Reject the first allocation attempt, then delegate to the real one.
|
||||
orig_allocate_slots = scheduler.kv_cache_manager.allocate_slots
|
||||
allocate_results: list = []
|
||||
|
||||
def spy_allocate_slots(*args, **kwargs):
|
||||
result = None if not allocate_results else orig_allocate_slots(*args, **kwargs)
|
||||
allocate_results.append(result)
|
||||
return result
|
||||
|
||||
scheduler.kv_cache_manager.allocate_slots = spy_allocate_slots
|
||||
|
||||
assert not scheduler.schedule().scheduled_new_reqs
|
||||
assert (stats.requests, stats.queries, stats.hits) == (0, 0, 0)
|
||||
|
||||
assert "retried" in scheduler.schedule().num_scheduled_tokens
|
||||
assert allocate_results[0] is None and allocate_results[1] is not None
|
||||
assert (stats.requests, stats.queries, stats.hits) == (
|
||||
1,
|
||||
retried.num_tokens,
|
||||
block_size * 2,
|
||||
)
|
||||
|
||||
|
||||
def test_scheduler_reset_prefix_cache():
|
||||
scheduler = create_scheduler(enable_prefix_caching=True)
|
||||
requests = create_requests(num_requests=10)
|
||||
@@ -5329,3 +5467,197 @@ def test_async_load_reservation_prevents_wedge_e2e():
|
||||
assert b.status == RequestStatus.WAITING
|
||||
assert b.num_preemptions == 0
|
||||
assert b.request_id not in req_to_blocks
|
||||
|
||||
|
||||
def _create_hybrid_mamba_connector_scheduler(
|
||||
matched_tokens: int,
|
||||
block_size: int = 16,
|
||||
num_blocks: int = 100,
|
||||
) -> Scheduler:
|
||||
"""FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
|
||||
model_config = ModelConfig(
|
||||
model="facebook/opt-125m",
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
seed=42,
|
||||
skip_tokenizer_init=True,
|
||||
)
|
||||
vllm_config = VllmConfig(
|
||||
scheduler_config=SchedulerConfig(
|
||||
max_num_seqs=4,
|
||||
max_num_batched_tokens=8192,
|
||||
max_model_len=8192,
|
||||
enable_chunked_prefill=True,
|
||||
is_encoder_decoder=False,
|
||||
watermark=0.0,
|
||||
),
|
||||
model_config=model_config,
|
||||
cache_config=CacheConfig(
|
||||
block_size=block_size,
|
||||
enable_prefix_caching=True,
|
||||
mamba_cache_mode="all",
|
||||
),
|
||||
kv_transfer_config=KVTransferConfig(
|
||||
kv_connector="MockKVConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={
|
||||
"matched_tokens": matched_tokens,
|
||||
"is_async": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
vllm_config.cache_config.num_gpu_blocks = num_blocks
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["fa"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["mamba"],
|
||||
MambaSpec(
|
||||
block_size=block_size,
|
||||
shapes=((1, 1),),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="all",
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
register_all_kvcache_specs(vllm_config)
|
||||
return Scheduler(
|
||||
vllm_config=vllm_config,
|
||||
kv_cache_config=kv_cache_config,
|
||||
structured_output_manager=StructuredOutputManager(vllm_config),
|
||||
block_size=block_size,
|
||||
hash_block_size=block_size,
|
||||
log_stats=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"matched_tokens,expected_num_computed",
|
||||
[
|
||||
# No external hit: resume on the deepest locally-consistent boundary
|
||||
# (block 0's state survives for both groups).
|
||||
(0, 16),
|
||||
# One external block on top of the reconciled local boundary.
|
||||
(16, 32),
|
||||
],
|
||||
)
|
||||
def test_hybrid_per_group_hit_divergence_with_connector(
|
||||
matched_tokens: int, expected_num_computed: int
|
||||
):
|
||||
"""Per-group prefix hits can diverge for hybrid models with a connector
|
||||
(#46453): under block pressure the FA prefix tail is evicted while a
|
||||
deeper Mamba state block survives. The scheduler must not report the
|
||||
deeper hit as locally computed (evicted FA blocks are not resident ->
|
||||
engine crash / dirty KV); it falls back to the reconciled boundary that
|
||||
every group is consistent at.
|
||||
"""
|
||||
block_size = 16
|
||||
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
|
||||
manager = scheduler.kv_cache_manager
|
||||
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
|
||||
|
||||
# Seed a 4-block prefix so both groups cache all four boundaries
|
||||
# (mamba cache mode "all" caches every block's state densely).
|
||||
[fill] = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=4 * block_size,
|
||||
max_tokens=1,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["fill"],
|
||||
)
|
||||
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
|
||||
blocks = manager.allocate_slots(
|
||||
fill, fill.num_tokens, num_computed, computed_blocks
|
||||
)
|
||||
fa_ids = [b.block_id for b in blocks.blocks[0]]
|
||||
mamba_ids = [b.block_id for b in blocks.blocks[1]]
|
||||
manager.free(fill)
|
||||
|
||||
# Evict the FA tail and the middle mamba states; block 0 (both groups)
|
||||
# and the deep mamba state at block 3 survive.
|
||||
manager.block_pool.evict_blocks({fa_ids[2], fa_ids[3], mamba_ids[1], mamba_ids[2]})
|
||||
|
||||
# A replay of the prefix plus one extra block now sees diverged
|
||||
# per-group hits: FA stops at the evicted tail, while the mamba lookup
|
||||
# finds the deeper surviving state.
|
||||
[replay] = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=5 * block_size,
|
||||
max_tokens=1,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["replay"],
|
||||
)
|
||||
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
|
||||
replay.block_hashes, replay.num_tokens - 1
|
||||
)
|
||||
assert per_group_hits == (2 * block_size, 4 * block_size) # diverged
|
||||
|
||||
scheduler.add_request(replay)
|
||||
output = scheduler.schedule()
|
||||
num_scheduled = output.num_scheduled_tokens[replay.request_id]
|
||||
assert replay.num_tokens - num_scheduled == expected_num_computed
|
||||
|
||||
|
||||
def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
|
||||
"""The opposite divergence: the FA prefix survives deeper than the Mamba
|
||||
state and the connector supplies nothing (ext == 0). Reporting the deep FA
|
||||
hit as locally computed would resume with no valid Mamba state at that
|
||||
boundary (silent bad output). The scheduler must fall back to the
|
||||
convergent boundary that every group agrees on (block 0's surviving state).
|
||||
"""
|
||||
block_size = 16
|
||||
scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0)
|
||||
manager = scheduler.kv_cache_manager
|
||||
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
|
||||
|
||||
# Seed a 4-block prefix in both groups.
|
||||
[fill] = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=4 * block_size,
|
||||
max_tokens=1,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["fill"],
|
||||
)
|
||||
computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
|
||||
blocks = manager.allocate_slots(
|
||||
fill, fill.num_tokens, num_computed, computed_blocks
|
||||
)
|
||||
mamba_ids = [b.block_id for b in blocks.blocks[1]]
|
||||
manager.free(fill)
|
||||
|
||||
# Keep all FA blocks; evict every mamba state but block 0. FA reaches 4
|
||||
# blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba).
|
||||
manager.block_pool.evict_blocks({mamba_ids[1], mamba_ids[2], mamba_ids[3]})
|
||||
|
||||
[replay] = create_requests(
|
||||
num_requests=1,
|
||||
num_tokens=5 * block_size,
|
||||
max_tokens=1,
|
||||
same_prompt=True,
|
||||
block_size=block_size,
|
||||
req_ids=["replay"],
|
||||
)
|
||||
_, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
|
||||
replay.block_hashes, replay.num_tokens - 1
|
||||
)
|
||||
assert per_group_hits == (4 * block_size, 1 * block_size) # FA deeper
|
||||
|
||||
scheduler.add_request(replay)
|
||||
output = scheduler.schedule()
|
||||
num_scheduled = output.num_scheduled_tokens[replay.request_id]
|
||||
# Must resume at the convergent boundary (block 0), not the deep FA hit.
|
||||
assert replay.num_tokens - num_scheduled == block_size
|
||||
|
||||
@@ -38,8 +38,12 @@ from vllm.v1.structured_output import StructuredOutputManager
|
||||
EOS_TOKEN_ID = 50256
|
||||
|
||||
|
||||
def mock_kv(matched_tokens: int, is_async: bool):
|
||||
return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
|
||||
def mock_kv(matched_tokens: int, is_async: bool, num_defers_before_matching: int = 0):
|
||||
return MockKVConfig(
|
||||
matched_tokens=matched_tokens,
|
||||
is_async=is_async,
|
||||
num_defers_before_matching=num_defers_before_matching,
|
||||
)
|
||||
|
||||
|
||||
def create_scheduler(
|
||||
@@ -115,6 +119,9 @@ def create_scheduler(
|
||||
kv_connector_extra_config={
|
||||
"matched_tokens": use_kv_connector.matched_tokens,
|
||||
"is_async": use_kv_connector.is_async,
|
||||
"num_defers_before_matching": (
|
||||
use_kv_connector.num_defers_before_matching
|
||||
),
|
||||
},
|
||||
)
|
||||
elif isinstance(use_kv_connector, str):
|
||||
|
||||
@@ -6,7 +6,6 @@ Unit tests for the breakable cudagraph primitives.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
from contextlib import nullcontext
|
||||
from unittest.mock import patch
|
||||
@@ -14,7 +13,21 @@ from unittest.mock import patch
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Enable breakable cudagraphs for this module's tests only.
|
||||
|
||||
eager_break_during_capture reads the env at decoration time, which
|
||||
happens inside the test bodies, so a per-test fixture suffices.
|
||||
monkeypatch restores the env so other test files running in the same
|
||||
pytest process are unaffected (a module-level os.environ assignment
|
||||
used to leak into test_cudagraph_dispatch.py and break it).
|
||||
"""
|
||||
import vllm.envs as envs
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
|
||||
envs.disable_envs_cache()
|
||||
|
||||
|
||||
def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
|
||||
@@ -86,10 +99,19 @@ def cuda_capture_stream():
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
from vllm.utils.torch_utils import _current_stream_tls
|
||||
|
||||
prev_stream = getattr(_current_stream_tls, "value", None)
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
yield stream
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
# Exiting torch.cuda.stream() records the default stream in vllm's
|
||||
# patched set_stream cache. A later CUDAGraphWrapper capture in the same
|
||||
# process would then run on the default stream, which cannot capture
|
||||
# (this broke test_cudagraph_dispatch.py when run in-process after this
|
||||
# file). Restore the pre-fixture value so this module leaves no trace.
|
||||
_current_stream_tls.value = prev_stream
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
CUDAGraphMode,
|
||||
ParallelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.distributed.device_communicators import pynccl_allocator
|
||||
from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
|
||||
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_graph_pool_id():
|
||||
pynccl_allocator._graph_pool_id = None
|
||||
yield
|
||||
pynccl_allocator._graph_pool_id = None
|
||||
|
||||
|
||||
def _create_vllm_config() -> MagicMock:
|
||||
compilation_config = CompilationConfig(
|
||||
cudagraph_mode="FULL",
|
||||
cudagraph_capture_sizes=[4],
|
||||
)
|
||||
compilation_config.max_cudagraph_capture_size = 4
|
||||
compilation_config.post_init_cudagraph_sizes()
|
||||
|
||||
vllm_config = MagicMock(spec=VllmConfig)
|
||||
vllm_config.compilation_config = compilation_config
|
||||
vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=4)
|
||||
vllm_config.parallel_config = ParallelConfig()
|
||||
vllm_config.speculative_config = None
|
||||
vllm_config.num_speculative_tokens = 0
|
||||
return vllm_config
|
||||
|
||||
|
||||
def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
|
||||
"""FULL capture must set graph_pool_id before entering torch.cuda.graph().
|
||||
|
||||
NCCL symmetric memory checks this global during graph capture; without
|
||||
it, capture fails with:
|
||||
AssertionError: graph_pool_id is not set under graph capture
|
||||
"""
|
||||
graph_pool = object()
|
||||
monkeypatch.setattr(
|
||||
gpu_cudagraph_utils,
|
||||
"get_pp_group",
|
||||
lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
gpu_cudagraph_utils.current_platform,
|
||||
"get_global_graph_pool",
|
||||
lambda: graph_pool,
|
||||
)
|
||||
|
||||
manager = gpu_cudagraph_utils.CudaGraphManager(
|
||||
vllm_config=_create_vllm_config(),
|
||||
device=torch.device("cpu"),
|
||||
cudagraph_mode=CUDAGraphMode.FULL,
|
||||
decode_query_len=1,
|
||||
)
|
||||
|
||||
desc = BatchExecutionDescriptor(
|
||||
cg_mode=CUDAGraphMode.FULL,
|
||||
num_tokens=4,
|
||||
num_reqs=4,
|
||||
uniform_token_count=1,
|
||||
)
|
||||
manager._capture_descs[CUDAGraphMode.FULL] = [desc]
|
||||
|
||||
def create_forward_fn(desc, warmup):
|
||||
return lambda _mode: None
|
||||
|
||||
@contextmanager
|
||||
def fake_graph_capture(*args, **kwargs):
|
||||
yield SimpleNamespace(stream=MagicMock())
|
||||
|
||||
fake_offloader = MagicMock()
|
||||
|
||||
def cuda_graph_enter(*args, **kwargs):
|
||||
assert pynccl_allocator._graph_pool_id is graph_pool
|
||||
|
||||
mock_cuda_graph_ctx = MagicMock()
|
||||
mock_cuda_graph_ctx.__enter__ = cuda_graph_enter
|
||||
mock_cuda_graph_ctx.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
with (
|
||||
patch.object(gpu_cudagraph_utils, "graph_capture", fake_graph_capture),
|
||||
patch.object(gpu_cudagraph_utils, "get_offloader", lambda: fake_offloader),
|
||||
patch.object(gpu_cudagraph_utils.torch.cuda, "CUDAGraph"),
|
||||
patch.object(
|
||||
gpu_cudagraph_utils.torch.cuda,
|
||||
"graph",
|
||||
return_value=mock_cuda_graph_ctx,
|
||||
) as mock_cuda_graph,
|
||||
):
|
||||
manager.capture(create_forward_fn)
|
||||
|
||||
mock_cuda_graph.assert_called_once()
|
||||
@@ -108,6 +108,11 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
|
||||
mgr.token_budgets = sorted(budgets)
|
||||
mgr.max_batch_size = 16
|
||||
mgr.use_dp = False
|
||||
mgr.config = EncoderCudaGraphConfig(
|
||||
modalities=["image"],
|
||||
buffer_keys=[],
|
||||
out_hidden_size=32,
|
||||
)
|
||||
mgr.budget_graphs = {"default": {}}
|
||||
mgr.graph_pool = None
|
||||
mgr.graph_hits = 0
|
||||
|
||||
@@ -0,0 +1,249 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import multi_gpu_marks
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.metrics.reader import Metric
|
||||
|
||||
DEEPSEEK_MTP_MAIN_RANDOM = "luccafong/deepseek_mtp_main_random"
|
||||
DEEPSEEK_MTP_DRAFT_RANDOM = "luccafong/deepseek_mtp_draft_random"
|
||||
|
||||
PROMPT = "The capital of France is"
|
||||
MAX_TOKENS = 8
|
||||
MAX_MODEL_LEN = 2048
|
||||
GPU_MEM_UTIL = 0.85
|
||||
# Spec / no-spec greedy output should match, but exact match is not guaranteed
|
||||
# across parallelism layouts (reduction order can flip near-tie argmaxes), so
|
||||
# gate on a similarity ratio like the spec-decode E2E tests.
|
||||
MIN_MATCH_RATIO = 0.8
|
||||
|
||||
|
||||
def _token_match_ratio(a: tuple[int, ...], b: tuple[int, ...]) -> float:
|
||||
n = min(len(a), len(b))
|
||||
return sum(x == y for x, y in zip(a, b)) / n if n else 0.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InlineConfig:
|
||||
id: str
|
||||
tp_size: int = 1
|
||||
pp_size: int = 1
|
||||
enable_expert_parallel: bool = False
|
||||
enable_eplb: bool = False
|
||||
skip_reason: str | None = None
|
||||
|
||||
|
||||
INLINE_CONFIGS = [
|
||||
InlineConfig(id="tp2", tp_size=2),
|
||||
InlineConfig(id="ep2", tp_size=2, enable_expert_parallel=True),
|
||||
InlineConfig(
|
||||
id="ep2_eplb",
|
||||
tp_size=2,
|
||||
enable_expert_parallel=True,
|
||||
enable_eplb=True,
|
||||
),
|
||||
InlineConfig(
|
||||
id="pp2",
|
||||
pp_size=2,
|
||||
skip_reason=(
|
||||
"DeepSeek MTP pipeline-parallel support is in flight upstream; "
|
||||
"see https://github.com/vllm-project/vllm/pull/38104"
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _inline_kwargs(config: InlineConfig, *, with_spec: bool) -> dict[str, Any]:
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": DEEPSEEK_MTP_MAIN_RANDOM,
|
||||
"tensor_parallel_size": config.tp_size,
|
||||
"pipeline_parallel_size": config.pp_size,
|
||||
"enable_expert_parallel": config.enable_expert_parallel,
|
||||
"max_model_len": MAX_MODEL_LEN,
|
||||
"gpu_memory_utilization": GPU_MEM_UTIL,
|
||||
"enforce_eager": True,
|
||||
"trust_remote_code": True,
|
||||
"disable_log_stats": False,
|
||||
# MLA + batch invariance currently disables prefix caching at runtime
|
||||
# (see vllm/v1/attention/backends/mla/common.py); set explicitly so the
|
||||
# configuration is unambiguous.
|
||||
"enable_prefix_caching": False,
|
||||
}
|
||||
if config.enable_eplb:
|
||||
kwargs["enable_eplb"] = True
|
||||
# Rearrangement first fires after step_interval // 4 forward steps
|
||||
# (the step counter starts at 3/4 of step_interval, see
|
||||
# vllm/distributed/eplb/eplb_state.py). Keep these small so the EPLB
|
||||
# routine actually runs within the short MAX_TOKENS generation instead
|
||||
# of never triggering.
|
||||
kwargs["eplb_config"] = {
|
||||
"num_redundant_experts": config.tp_size,
|
||||
"window_size": 2,
|
||||
"step_interval": 4,
|
||||
"log_balancedness": False,
|
||||
}
|
||||
if with_spec:
|
||||
kwargs["speculative_config"] = {
|
||||
"method": "mtp",
|
||||
"model": DEEPSEEK_MTP_DRAFT_RANDOM,
|
||||
"num_speculative_tokens": 1,
|
||||
}
|
||||
return kwargs
|
||||
|
||||
|
||||
def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]:
|
||||
outputs = llm.generate(
|
||||
[PROMPT],
|
||||
SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True),
|
||||
)
|
||||
assert outputs and outputs[0].outputs, "expected one completion"
|
||||
return tuple(outputs[0].outputs[0].token_ids)
|
||||
|
||||
|
||||
def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
|
||||
name2metric = {m.name: m for m in metrics}
|
||||
counter = name2metric.get("vllm:spec_decode_num_drafts")
|
||||
assert counter is not None, (
|
||||
"spec_decode_num_drafts metric missing; check disable_log_stats=False"
|
||||
)
|
||||
return int(counter.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config",
|
||||
[
|
||||
pytest.param(c, id=c.id, marks=multi_gpu_marks(num_gpus=2))
|
||||
for c in INLINE_CONFIGS
|
||||
],
|
||||
)
|
||||
def test_deepseek_mtp_load_inline(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
config: InlineConfig,
|
||||
):
|
||||
"""MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy."""
|
||||
if config.skip_reason is not None:
|
||||
pytest.skip(config.skip_reason)
|
||||
|
||||
# Reduces run-to-run nondeterminism so spec / no-spec stay close;
|
||||
# see tests/v1/distributed/test_eagle_dp.py for the same pattern.
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
|
||||
spec_llm = LLM(**_inline_kwargs(config, with_spec=True))
|
||||
try:
|
||||
spec_tokens = _generate_token_ids_inline(spec_llm)
|
||||
n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics())
|
||||
finally:
|
||||
del spec_llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False))
|
||||
try:
|
||||
no_spec_tokens = _generate_token_ids_inline(no_spec_llm)
|
||||
finally:
|
||||
del no_spec_llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
# Non-vacuity: a silently-broken MTP drafter that falls back to
|
||||
# verifier-only decoding would still produce matching output below.
|
||||
assert n_drafts > 0, (
|
||||
f"MTP drafter never fired under {config.id}: vllm:spec_decode_num_drafts == 0"
|
||||
)
|
||||
match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
|
||||
print(f"\n{config.id}: spec/no-spec match_ratio={match_ratio:.3f}")
|
||||
assert match_ratio >= MIN_MATCH_RATIO, (
|
||||
f"Spec / no-spec output divergence under {config.id}: "
|
||||
f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
|
||||
f" spec_tokens = {spec_tokens}\n"
|
||||
f" no_spec_tokens= {no_spec_tokens}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"dp_size",
|
||||
[
|
||||
pytest.param(2, marks=pytest.mark.distributed(num_gpus=2), id="dp2"),
|
||||
pytest.param(1, id="dp1"),
|
||||
],
|
||||
)
|
||||
async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: int):
|
||||
"""MTP loads under DP (on and off) via AsyncLLM; spec matches no-spec greedy."""
|
||||
if torch.accelerator.device_count() < dp_size:
|
||||
pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs")
|
||||
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
|
||||
base_args = AsyncEngineArgs(
|
||||
model=DEEPSEEK_MTP_MAIN_RANDOM,
|
||||
tensor_parallel_size=1,
|
||||
data_parallel_size=dp_size,
|
||||
data_parallel_backend="mp",
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
gpu_memory_utilization=GPU_MEM_UTIL,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_prefix_caching=False,
|
||||
)
|
||||
spec_args = replace(
|
||||
base_args,
|
||||
speculative_config={
|
||||
"method": "mtp",
|
||||
"model": DEEPSEEK_MTP_DRAFT_RANDOM,
|
||||
"num_speculative_tokens": 1,
|
||||
},
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=MAX_TOKENS,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.FINAL_ONLY,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]:
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
engine = AsyncLLM.from_engine_args(args)
|
||||
stack.callback(engine.shutdown)
|
||||
async for out in engine.generate(
|
||||
request_id=request_id,
|
||||
prompt=PROMPT,
|
||||
sampling_params=sampling_params,
|
||||
):
|
||||
token_ids = tuple(out.outputs[0].token_ids)
|
||||
assert len(token_ids) == MAX_TOKENS
|
||||
return token_ids
|
||||
raise AssertionError("AsyncLLM produced no output")
|
||||
finally:
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
spec_tokens = await asyncio.wait_for(
|
||||
_generate(spec_args, f"deepseek-mtp-dp{dp_size}-spec"), timeout=600
|
||||
)
|
||||
no_spec_tokens = await asyncio.wait_for(
|
||||
_generate(base_args, f"deepseek-mtp-dp{dp_size}-no-spec"), timeout=600
|
||||
)
|
||||
|
||||
match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens)
|
||||
print(f"\ndp{dp_size}: spec/no-spec match_ratio={match_ratio:.3f}")
|
||||
assert match_ratio >= MIN_MATCH_RATIO, (
|
||||
f"Spec / no-spec output divergence under dp{dp_size}: "
|
||||
f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n"
|
||||
f" spec_tokens = {spec_tokens}\n"
|
||||
f" no_spec_tokens= {no_spec_tokens}"
|
||||
)
|
||||
@@ -204,13 +204,12 @@ class FakeNixlWrapper:
|
||||
def _make_fake_nixl_pkg():
|
||||
"""Context manager that creates a temporary package making
|
||||
`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
|
||||
Also creates rixl package for ROCm compatibility.
|
||||
Also creates the ROCm NIXL packages.
|
||||
|
||||
Automatically cleans up the temporary directory when done.
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as td:
|
||||
# Create both nixl and rixl packages for cross-platform compatibility
|
||||
for pkg_name in ["nixl", "rixl"]:
|
||||
for pkg_name in ["nixl", "nixl_rocm"]:
|
||||
pkg_root = os.path.join(td, pkg_name, "_api")
|
||||
os.makedirs(pkg_root, exist_ok=True)
|
||||
|
||||
@@ -3130,10 +3129,12 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
)
|
||||
|
||||
|
||||
@patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
|
||||
FakeNixlWrapper,
|
||||
)
|
||||
def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
|
||||
"""kv_role='kv_both' should emit a deprecation log warning."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.logger import _print_warning_once
|
||||
|
||||
_print_warning_once.cache_clear()
|
||||
@@ -3156,10 +3157,12 @@ def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
|
||||
assert "deprecated" in msg
|
||||
|
||||
|
||||
@patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
|
||||
FakeNixlWrapper,
|
||||
)
|
||||
def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init):
|
||||
"""kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning."""
|
||||
from unittest.mock import patch
|
||||
|
||||
for role in ("kv_consumer", "kv_producer"):
|
||||
vllm_config = create_vllm_config(kv_role=role)
|
||||
with patch(
|
||||
|
||||
+7
-7
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Verify that GPU memory is fully released after RixlConnector shutdown on ROCm.
|
||||
"""Verify that GPU memory is released after NixlConnector shutdown on ROCm.
|
||||
|
||||
Regression test for ROCm/ucx#33: UCX rocm_ipc transport permanently pinned
|
||||
GPU memory via hsa_amd_ipc_memory_create during ucp_mem_map, causing
|
||||
@@ -62,7 +62,7 @@ def _full_gpu_cleanup():
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
|
||||
def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
def test_gpu_memory_nixl_hma(model_name, sw_size):
|
||||
"""Track GPU memory through NixlConnector create/infer/shutdown cycle."""
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
@@ -84,7 +84,7 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
}
|
||||
|
||||
print("\n" + "=" * 90)
|
||||
print("GPU MEMORY -- RIXL NixlConnector HMA (ROCm)")
|
||||
print("GPU MEMORY -- NIXL NixlConnector HMA (ROCm)")
|
||||
print("=" * 90)
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -169,14 +169,14 @@ def test_gpu_memory_rixl_hma(model_name, sw_size):
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["google/gemma-3-1b-it"])
|
||||
def test_gpu_memory_no_rixl_baseline(model_name):
|
||||
def test_gpu_memory_no_nixl_baseline(model_name):
|
||||
"""Same workload without NixlConnector. Comparing driver-level memory
|
||||
between this and test_gpu_memory_rixl_hma isolates UCX/RIXL impact."""
|
||||
between this and test_gpu_memory_nixl_hma isolates UCX/NIXL impact."""
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
|
||||
|
||||
print("\n" + "=" * 90)
|
||||
print("CONTROL -- same model, no RIXL connector")
|
||||
print("CONTROL -- same model, no NIXL connector")
|
||||
print("=" * 90)
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
@@ -209,7 +209,7 @@ def test_gpu_memory_no_rixl_baseline(model_name):
|
||||
drv_base = snap0["drv_used_mb"]
|
||||
drv_leaked = snap_final["drv_used_mb"] - drv_base
|
||||
drv_peak = snap_peak["drv_used_mb"] - drv_base
|
||||
print(f"\n Driver leaked (no rixl): {drv_leaked:.0f} MB")
|
||||
print(f"\n Driver leaked (no NIXL): {drv_leaked:.0f} MB")
|
||||
print("=" * 90)
|
||||
|
||||
leak_pct = (drv_leaked / drv_peak * 100) if drv_peak > 0 else 0
|
||||
@@ -362,6 +362,7 @@ class TestExampleConnector(ExampleConnector):
|
||||
class MockKVConfig:
|
||||
matched_tokens: int = 0
|
||||
is_async: bool = False
|
||||
num_defers_before_matching: int = 0
|
||||
|
||||
|
||||
class MockKVConnectorMetadata(KVConnectorMetadata):
|
||||
@@ -384,6 +385,12 @@ class MockKVConnector(KVConnectorBase_V1):
|
||||
self.config = MockKVConfig(
|
||||
matched_tokens=extra_config["matched_tokens"],
|
||||
is_async=extra_config["is_async"],
|
||||
num_defers_before_matching=extra_config.get(
|
||||
"num_defers_before_matching", 0
|
||||
),
|
||||
)
|
||||
self._defers_left: defaultdict[str, int] = defaultdict(
|
||||
lambda: self.config.num_defers_before_matching
|
||||
)
|
||||
|
||||
def get_num_new_matched_tokens(
|
||||
@@ -391,6 +398,9 @@ class MockKVConnector(KVConnectorBase_V1):
|
||||
request: Request,
|
||||
num_computed_tokens: int,
|
||||
) -> tuple[int | None, bool]:
|
||||
if self._defers_left[request.request_id] > 0:
|
||||
self._defers_left[request.request_id] -= 1
|
||||
return (None, False)
|
||||
return (self.config.matched_tokens, self.config.is_async)
|
||||
|
||||
def update_state_after_alloc(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Config-only resolution of DFlash draft attention causality.
|
||||
"""Config-only DFlash behavior.
|
||||
|
||||
``dflash_has_any_non_causal`` decides pre-build whether the draft needs a
|
||||
non-causal-capable backend, so its branch table (explicit override, SWA-derived
|
||||
@@ -13,8 +13,12 @@ import pytest
|
||||
|
||||
from vllm.model_executor.models.qwen3_dflash import (
|
||||
_dflash_layer_causal,
|
||||
_get_dflash_fc_input_size,
|
||||
dflash_has_any_non_causal,
|
||||
)
|
||||
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
|
||||
get_eagle3_aux_layers_from_config,
|
||||
)
|
||||
|
||||
|
||||
def _config(num_hidden_layers, layer_types=None, causal_override=None):
|
||||
@@ -53,3 +57,35 @@ def test_dflash_layer_causal_is_per_layer():
|
||||
config = _config(2, layer_types=["sliding_attention", "full_attention"])
|
||||
assert _dflash_layer_causal(config, 0) is True
|
||||
assert _dflash_layer_causal(config, 1) is False
|
||||
|
||||
|
||||
def _vllm_config(**draft_config):
|
||||
config = SimpleNamespace(**draft_config)
|
||||
return SimpleNamespace(
|
||||
speculative_config=SimpleNamespace(
|
||||
draft_model_config=SimpleNamespace(hf_config=config)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_dflash_fc_uses_aux_layer_count():
|
||||
vllm_config = _vllm_config(
|
||||
num_hidden_layers=5,
|
||||
hidden_size=4096,
|
||||
target_hidden_size=None,
|
||||
target_layer_ids=[1, 17, 32],
|
||||
)
|
||||
|
||||
assert _get_dflash_fc_input_size(vllm_config) == 3 * 4096
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_name", ["dflash_config", "eagle_config"])
|
||||
def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name):
|
||||
layer_ids = [1, 17, 32]
|
||||
vllm_config = _vllm_config(
|
||||
**{config_name: {"layer_ids": layer_ids}},
|
||||
)
|
||||
|
||||
assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple(
|
||||
layer_ids
|
||||
)
|
||||
|
||||
@@ -16,7 +16,8 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
|
||||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
# Not cpu_test: InputBatch allocates pinned (UVA) memory, which requires a
|
||||
# CUDA device even though the batch tensors live on the CPU.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -28,6 +29,7 @@ def mock_model_runner_with_input_batch():
|
||||
runner.requests = {}
|
||||
runner.max_num_reqs = 10
|
||||
runner.max_model_len = 1024
|
||||
runner.late_interaction_runner = Mock()
|
||||
|
||||
# Create a real InputBatch for e2e testing
|
||||
runner.input_batch = InputBatch(
|
||||
|
||||
@@ -16,7 +16,8 @@ from vllm.v1.core.sched.output import (
|
||||
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
|
||||
from vllm.v1.worker.gpu.states import RequestState
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
# Not cpu_test: RequestState allocates pinned (UVA) memory, which requires a
|
||||
# CUDA device even though the request state itself lives on the CPU.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -31,13 +32,12 @@ def mock_model_runner_with_req_states():
|
||||
num_speculative_steps=0,
|
||||
vocab_size=32000,
|
||||
device=torch.device("cpu"),
|
||||
model_dtype=torch.float32,
|
||||
cache_draft_logits=False,
|
||||
)
|
||||
runner.encoder_cache = None
|
||||
runner.model_state = Mock()
|
||||
runner.block_tables = Mock()
|
||||
runner.lora_state = Mock()
|
||||
runner.pp_handler = None
|
||||
runner.sampler = None
|
||||
runner.prompt_logprobs_worker = None
|
||||
runner.is_last_pp_rank = False
|
||||
|
||||
@@ -53,6 +53,8 @@ def create_scheduler() -> Scheduler:
|
||||
vllm_config.model_config = MagicMock()
|
||||
vllm_config.model_config.skip_tokenizer_init = True
|
||||
vllm_config.model_config.is_multimodal_model = False
|
||||
vllm_config.model_config.is_encoder_decoder = False
|
||||
vllm_config.model_config.is_diffusion = False
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.enable_return_routed_experts = False
|
||||
vllm_config.cache_config = MagicMock()
|
||||
@@ -496,7 +498,9 @@ class TestStreamingScheduler(unittest.TestCase):
|
||||
eco_cycle2 = eco_dict_cycle2[session.client_index].outputs[0]
|
||||
assert eco_cycle2.finish_reason == FinishReason.STOP
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
assert session in scheduler.waiting
|
||||
# Sessions paused for streaming input are blocked-waiting, so they
|
||||
# live in the skipped_waiting queue rather than the main waiting queue.
|
||||
assert session in scheduler.skipped_waiting
|
||||
assert session._all_token_ids == [1, 2, 3, 10, STOP_TOKEN]
|
||||
|
||||
# CRITICAL ASSERTION: Cached prompt_token_ids STILL must not have changed
|
||||
|
||||
@@ -6,122 +6,13 @@ from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm.v1.worker.gpu_worker as gpu_worker_module
|
||||
from vllm.multimodal.video import (
|
||||
PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
|
||||
PYNVVIDEOCODEC_MAX_RETAINED_DECODERS,
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
from vllm.v1.worker import startup_plan
|
||||
from vllm.v1.worker.gpu_worker import Worker
|
||||
from vllm.v1.worker.startup_plan import (
|
||||
maybe_apply_startup_plan,
|
||||
maybe_save_startup_plan,
|
||||
)
|
||||
|
||||
|
||||
def _worker_with_mm_config(
|
||||
mm_config: SimpleNamespace,
|
||||
*,
|
||||
api_process_count: int = 1,
|
||||
) -> Worker:
|
||||
worker = object.__new__(Worker)
|
||||
worker.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
worker.parallel_config = SimpleNamespace(_api_process_count=api_process_count)
|
||||
return worker
|
||||
|
||||
|
||||
def _mm_config(
|
||||
*,
|
||||
mm_ipc_gpu_memory_gb: float = 0,
|
||||
video_backend: str | None = None,
|
||||
) -> SimpleNamespace:
|
||||
video_kwargs = {} if video_backend is None else {"video_backend": video_backend}
|
||||
return SimpleNamespace(
|
||||
mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb,
|
||||
media_io_kwargs={"video": video_kwargs} if video_kwargs else {},
|
||||
)
|
||||
|
||||
|
||||
def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int:
|
||||
return api_process_count * (
|
||||
PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES * PYNVVIDEOCODEC_MAX_RETAINED_DECODERS
|
||||
+ PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("video_backend", [None, "opencv"])
|
||||
def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
video_backend: str | None,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
gpu_worker_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
worker = _worker_with_mm_config(
|
||||
_mm_config(mm_ipc_gpu_memory_gb=0.25, video_backend=video_backend)
|
||||
)
|
||||
|
||||
assert worker._reserve_mm_ipc_gpu_memory(GiB_bytes) == int(0.75 * GiB_bytes)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_includes_pynvvideocodec_decoder_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
gpu_worker_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
"opencv",
|
||||
)
|
||||
worker = _worker_with_mm_config(
|
||||
_mm_config(
|
||||
mm_ipc_gpu_memory_gb=0.25,
|
||||
video_backend=PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
)
|
||||
available_bytes = 4 * GiB_bytes
|
||||
|
||||
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
|
||||
available_bytes - int(0.25 * GiB_bytes) - _pynvvideocodec_decoder_budget()
|
||||
)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_uses_env_video_backend(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
gpu_worker_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
worker = _worker_with_mm_config(_mm_config())
|
||||
available_bytes = 4 * GiB_bytes
|
||||
|
||||
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
|
||||
available_bytes - _pynvvideocodec_decoder_budget()
|
||||
)
|
||||
|
||||
|
||||
def test_reserve_mm_ipc_gpu_memory_scales_pynvvideocodec_budget_by_api_servers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
gpu_worker_module.envs,
|
||||
"VLLM_VIDEO_LOADER_BACKEND",
|
||||
PYNVVIDEOCODEC_VIDEO_BACKEND,
|
||||
)
|
||||
worker = _worker_with_mm_config(_mm_config(), api_process_count=3)
|
||||
available_bytes = 8 * GiB_bytes
|
||||
|
||||
assert worker._reserve_mm_ipc_gpu_memory(available_bytes) == (
|
||||
available_bytes - _pynvvideocodec_decoder_budget(api_process_count=3)
|
||||
)
|
||||
|
||||
|
||||
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
|
||||
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
|
||||
|
||||
|
||||
+1
-20
@@ -36,22 +36,6 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
|
||||
]
|
||||
|
||||
|
||||
def write_coverage_objects(extensions: list[RustExtension], output: Path) -> None:
|
||||
artifacts = []
|
||||
for extension in extensions:
|
||||
for target in sorted(set(extension.target.values())):
|
||||
target_path = ROOT_DIR.joinpath(*target.split("."))
|
||||
if extension.binding == Binding.Exec:
|
||||
matches = [target_path]
|
||||
else:
|
||||
matches = sorted(target_path.parent.glob(f"{target_path.name}*.so"))
|
||||
if len(matches) != 1 or not matches[0].is_file():
|
||||
raise RuntimeError(f"unable to locate Rust artifact for {target}")
|
||||
artifacts.append(matches[0].relative_to(ROOT_DIR).as_posix())
|
||||
|
||||
output.write_text("\n".join(artifacts) + "\n")
|
||||
|
||||
|
||||
def rust_py_extension_module_names() -> list[str]:
|
||||
module_names = []
|
||||
for extension in rust_extensions():
|
||||
@@ -68,15 +52,12 @@ def rust_py_extension_module_names() -> list[str]:
|
||||
def build_binary(build_rust_args: list[str]) -> None:
|
||||
os.chdir(ROOT_DIR)
|
||||
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
|
||||
extensions = rust_extensions(optional=False)
|
||||
setup(
|
||||
name="vllm-rust-frontend-build",
|
||||
packages=[],
|
||||
rust_extensions=extensions,
|
||||
rust_extensions=rust_extensions(optional=False),
|
||||
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
|
||||
)
|
||||
if output := os.getenv("VLLM_RUST_COVERAGE_OBJECTS"):
|
||||
write_coverage_objects(extensions, Path(output))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user