forked from Karylab-cklius/vllm
Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e1a763558c | ||
|
|
84aeec9f22 | ||
|
|
82a770ddbd | ||
|
|
a09a9bace1 | ||
|
|
6c20d467a2 | ||
|
|
cb59d0a351 | ||
|
|
0ab1bded36 | ||
|
|
040cbf95cc | ||
|
|
5b3762a7f0 | ||
|
|
4d30c510ce | ||
|
|
6700813f86 | ||
|
|
eb44b3aaa4 | ||
|
|
7a98c7a392 | ||
|
|
0d9e60619b | ||
|
|
1134545b6f | ||
|
|
3e0c887511 | ||
|
|
adfbbc1005 | ||
|
|
adc98f04d0 | ||
|
|
8def3cdde2 | ||
|
|
616c9bd0f4 | ||
|
|
8688a06d67 | ||
|
|
f25953cc59 | ||
|
|
d9aa35161d | ||
|
|
6bcda970fd | ||
|
|
ea0e9c8f2e | ||
|
|
94ed0bf4e0 | ||
|
|
1940c8441e | ||
|
|
72d16aee15 | ||
|
|
e78a0c8e59 |
@@ -29,6 +29,7 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
|
|||||||
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
|
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
|
||||||
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
|
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
|
||||||
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
|
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
|
||||||
|
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
|
||||||
|
|
||||||
log_section() {
|
log_section() {
|
||||||
echo "--- $*"
|
echo "--- $*"
|
||||||
@@ -106,6 +107,18 @@ install_cargo_nextest() {
|
|||||||
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
|
"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() {
|
install_uv() {
|
||||||
log_section "Installing uv ${UV_VERSION}"
|
log_section "Installing uv ${UV_VERSION}"
|
||||||
curl -L --proto '=https' --tlsv1.2 -sSf \
|
curl -L --proto '=https' --tlsv1.2 -sSf \
|
||||||
@@ -176,14 +189,41 @@ run_tests() {
|
|||||||
setup_pyo3_python
|
setup_pyo3_python
|
||||||
install_cargo_binstall
|
install_cargo_binstall
|
||||||
install_cargo_nextest
|
install_cargo_nextest
|
||||||
|
install_cargo_llvm_cov
|
||||||
|
|
||||||
log_section "Running cargo nextest"
|
log_section "Running cargo nextest with Rust coverage"
|
||||||
cargo nextest run \
|
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 \
|
||||||
--manifest-path rust/Cargo.toml \
|
--manifest-path rust/Cargo.toml \
|
||||||
--workspace \
|
--workspace \
|
||||||
--all-features \
|
--all-features \
|
||||||
--locked \
|
--locked \
|
||||||
--no-fail-fast
|
--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"
|
||||||
}
|
}
|
||||||
|
|
||||||
install_protoc
|
install_protoc
|
||||||
|
|||||||
@@ -0,0 +1,182 @@
|
|||||||
|
#!/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"
|
||||||
|
}
|
||||||
@@ -8,6 +8,11 @@ steps:
|
|||||||
working_dir: "/vllm-workspace/tests"
|
working_dir: "/vllm-workspace/tests"
|
||||||
source_file_dependencies:
|
source_file_dependencies:
|
||||||
- rust/
|
- rust/
|
||||||
|
- build_rust.sh
|
||||||
|
- tools/build_rust.py
|
||||||
|
- rust-toolchain.toml
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
- vllm/benchmarks/
|
- vllm/benchmarks/
|
||||||
- vllm/entrypoints/openai/
|
- vllm/entrypoints/openai/
|
||||||
- vllm/entrypoints/serve/
|
- vllm/entrypoints/serve/
|
||||||
@@ -23,6 +28,7 @@ steps:
|
|||||||
- tests/entrypoints/openai/test_uds.py
|
- tests/entrypoints/openai/test_uds.py
|
||||||
- tests/v1/sample/test_logprobs_e2e.py
|
- tests/v1/sample/test_logprobs_e2e.py
|
||||||
commands:
|
commands:
|
||||||
|
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||||
- export VLLM_USE_RUST_FRONTEND=1
|
- export VLLM_USE_RUST_FRONTEND=1
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||||
@@ -43,6 +49,11 @@ steps:
|
|||||||
working_dir: "/vllm-workspace/tests"
|
working_dir: "/vllm-workspace/tests"
|
||||||
source_file_dependencies:
|
source_file_dependencies:
|
||||||
- rust/
|
- rust/
|
||||||
|
- build_rust.sh
|
||||||
|
- tools/build_rust.py
|
||||||
|
- rust-toolchain.toml
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
- vllm/entrypoints/openai/
|
- vllm/entrypoints/openai/
|
||||||
- vllm/entrypoints/serve/
|
- vllm/entrypoints/serve/
|
||||||
- vllm/v1/engine/
|
- vllm/v1/engine/
|
||||||
@@ -54,6 +65,7 @@ steps:
|
|||||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||||
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
||||||
commands:
|
commands:
|
||||||
|
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||||
- export VLLM_USE_RUST_FRONTEND=1
|
- export VLLM_USE_RUST_FRONTEND=1
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||||
@@ -72,10 +84,16 @@ steps:
|
|||||||
working_dir: "/vllm-workspace/tests"
|
working_dir: "/vllm-workspace/tests"
|
||||||
source_file_dependencies:
|
source_file_dependencies:
|
||||||
- rust/
|
- rust/
|
||||||
|
- build_rust.sh
|
||||||
|
- tools/build_rust.py
|
||||||
|
- rust-toolchain.toml
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
- vllm/entrypoints/openai/
|
- vllm/entrypoints/openai/
|
||||||
- tests/utils.py
|
- tests/utils.py
|
||||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||||
commands:
|
commands:
|
||||||
|
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||||
- export VLLM_USE_RUST_FRONTEND=1
|
- export VLLM_USE_RUST_FRONTEND=1
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||||
@@ -86,11 +104,17 @@ steps:
|
|||||||
working_dir: "/vllm-workspace/tests"
|
working_dir: "/vllm-workspace/tests"
|
||||||
source_file_dependencies:
|
source_file_dependencies:
|
||||||
- rust/
|
- rust/
|
||||||
|
- build_rust.sh
|
||||||
|
- tools/build_rust.py
|
||||||
|
- rust-toolchain.toml
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
- vllm/entrypoints/openai/
|
- vllm/entrypoints/openai/
|
||||||
- vllm/tool_parsers/
|
- vllm/tool_parsers/
|
||||||
- tests/utils.py
|
- tests/utils.py
|
||||||
- tests/tool_use/
|
- tests/tool_use/
|
||||||
commands:
|
commands:
|
||||||
|
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||||
- export VLLM_USE_RUST_FRONTEND=1
|
- export VLLM_USE_RUST_FRONTEND=1
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- 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"
|
- 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"
|
||||||
@@ -101,6 +125,11 @@ steps:
|
|||||||
working_dir: "/vllm-workspace/tests"
|
working_dir: "/vllm-workspace/tests"
|
||||||
source_file_dependencies:
|
source_file_dependencies:
|
||||||
- rust/
|
- rust/
|
||||||
|
- build_rust.sh
|
||||||
|
- tools/build_rust.py
|
||||||
|
- rust-toolchain.toml
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
- vllm/distributed/
|
- vllm/distributed/
|
||||||
- vllm/engine/
|
- vllm/engine/
|
||||||
- vllm/executor/
|
- vllm/executor/
|
||||||
@@ -111,6 +140,7 @@ steps:
|
|||||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||||
- tests/v1/distributed/test_internal_lb_dp.py
|
- tests/v1/distributed/test_internal_lb_dp.py
|
||||||
commands:
|
commands:
|
||||||
|
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||||
- export VLLM_USE_RUST_FRONTEND=1
|
- export VLLM_USE_RUST_FRONTEND=1
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||||
|
|||||||
@@ -26,5 +26,7 @@ steps:
|
|||||||
- rust-toolchain.toml
|
- rust-toolchain.toml
|
||||||
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
||||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
||||||
|
- .buildkite/scripts/rust-coverage.sh
|
||||||
|
- codecov.yml
|
||||||
commands:
|
commands:
|
||||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
|
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
|
||||||
|
|||||||
@@ -47,6 +47,7 @@
|
|||||||
|
|
||||||
# Rust Frontend
|
# Rust Frontend
|
||||||
/rust/ @BugenZhao @njhill
|
/rust/ @BugenZhao @njhill
|
||||||
|
/rust/src/bench @esmeetu
|
||||||
/build_rust.sh @BugenZhao @njhill
|
/build_rust.sh @BugenZhao @njhill
|
||||||
/rust-toolchain.toml @BugenZhao @njhill
|
/rust-toolchain.toml @BugenZhao @njhill
|
||||||
/.buildkite/test_areas/rust* @BugenZhao @njhill
|
/.buildkite/test_areas/rust* @BugenZhao @njhill
|
||||||
|
|||||||
@@ -257,3 +257,4 @@ vllm/grpc/vllm_engine_pb2.pyi
|
|||||||
|
|
||||||
# Ignore generated cpu headers
|
# Ignore generated cpu headers
|
||||||
csrc/cpu/cpu_attn_dispatch_generated.h
|
csrc/cpu/cpu_attn_dispatch_generated.h
|
||||||
|
rust-coverage-tools/
|
||||||
|
|||||||
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
|
|||||||
- Tool calling and reasoning parsers
|
- Tool calling and reasoning parsers
|
||||||
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
|
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
|
||||||
- Efficient multi-LoRA support for dense and MoE layers
|
- Efficient multi-LoRA support for dense and MoE layers
|
||||||
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||||
|
|
||||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||||
|
|
||||||
|
|||||||
@@ -8,6 +8,8 @@
|
|||||||
set -euo pipefail
|
set -euo pipefail
|
||||||
|
|
||||||
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
|
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.
|
# Read the required toolchain from rust-toolchain.toml.
|
||||||
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
||||||
@@ -30,4 +32,39 @@ else
|
|||||||
PROFILE_ARG="--release"
|
PROFILE_ARG="--release"
|
||||||
fi
|
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"
|
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
|
||||||
|
|||||||
+13
@@ -10,3 +10,16 @@ fixes:
|
|||||||
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
|
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
|
||||||
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
|
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
|
||||||
- "/usr/lib/python3.*/site-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
|
||||||
|
|||||||
@@ -102,7 +102,9 @@ class TileGemm82 {
|
|||||||
kv_cache_t* __restrict__ curr_b = b_tile;
|
kv_cache_t* __restrict__ curr_b = b_tile;
|
||||||
|
|
||||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||||
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
|
auto fp32_b_regs = load_b_pair_vec(curr_b);
|
||||||
|
auto fp32_b_0_reg = fp32_b_regs.first;
|
||||||
|
auto fp32_b_1_reg = fp32_b_regs.second;
|
||||||
|
|
||||||
float* __restrict__ curr_m_a = curr_a;
|
float* __restrict__ curr_m_a = curr_a;
|
||||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||||
|
|||||||
@@ -294,6 +294,9 @@ FROM base AS rust-build
|
|||||||
ARG BUILD_OS
|
ARG BUILD_OS
|
||||||
ARG USE_SCCACHE
|
ARG USE_SCCACHE
|
||||||
ARG SCCACHE_ENDPOINT
|
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.
|
# Install native tools needed only for Rust/protoc builds.
|
||||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||||
@@ -902,6 +905,14 @@ COPY ./vllm/collect_env.py .
|
|||||||
# note that this uses vllm installed by `pip`
|
# note that this uses vllm installed by `pip`
|
||||||
FROM vllm-base AS test
|
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/
|
ADD . /vllm-workspace/
|
||||||
|
|
||||||
ARG PYTHON_VERSION
|
ARG PYTHON_VERSION
|
||||||
|
|||||||
@@ -46,6 +46,9 @@
|
|||||||
"TORCH_CUDA_ARCH_LIST": {
|
"TORCH_CUDA_ARCH_LIST": {
|
||||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
|
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
|
||||||
},
|
},
|
||||||
|
"VLLM_RUST_COVERAGE": {
|
||||||
|
"default": "1"
|
||||||
|
},
|
||||||
"MAX_JOBS": {
|
"MAX_JOBS": {
|
||||||
"default": "2"
|
"default": "2"
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
|
|||||||
|
|
||||||
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
|
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
|
||||||
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
|
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
|
||||||
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
|
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
git clone https://github.com/vllm-project/vllm.git
|
git clone https://github.com/vllm-project/vllm.git
|
||||||
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
|
|||||||
--8<-- [end:build-wheel-from-source]
|
--8<-- [end:build-wheel-from-source]
|
||||||
--8<-- [start:pre-built-images]
|
--8<-- [start:pre-built-images]
|
||||||
|
|
||||||
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
|
vLLM offers official Docker images for deployment.
|
||||||
|
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||||
|
|
||||||
|
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
|
||||||
|
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
|
||||||
|
|
||||||
|
```bash
|
||||||
|
docker run --rm \
|
||||||
|
--network=host \
|
||||||
|
--device /dev/dri:/dev/dri \
|
||||||
|
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||||
|
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||||
|
--env "HF_TOKEN=$HF_TOKEN" \
|
||||||
|
--ipc=host \
|
||||||
|
--privileged \
|
||||||
|
vllm/vllm-openai-xpu:<tag> \
|
||||||
|
--model Qwen/Qwen3-0.6B
|
||||||
|
```
|
||||||
|
|
||||||
|
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
|
||||||
|
|
||||||
|
???+ console "Commands"
|
||||||
|
```bash
|
||||||
|
docker run --rm -it \
|
||||||
|
--network=host \
|
||||||
|
--device /dev/dri:/dev/dri \
|
||||||
|
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||||
|
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||||
|
--env "HF_TOKEN=$HF_TOKEN" \
|
||||||
|
--ipc=host \
|
||||||
|
--privileged \
|
||||||
|
--entrypoint /bin/bash \
|
||||||
|
vllm/vllm-openai-xpu:<tag>
|
||||||
|
```
|
||||||
|
|
||||||
--8<-- [end:pre-built-images]
|
--8<-- [end:pre-built-images]
|
||||||
--8<-- [start:build-image-from-source]
|
--8<-- [start:build-image-from-source]
|
||||||
|
|||||||
@@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform:
|
|||||||
!!! tip
|
!!! tip
|
||||||
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
|
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
|
||||||
|
|
||||||
|
=== "Intel GPU"
|
||||||
|
|
||||||
|
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
|
||||||
|
|
||||||
|
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||||
|
|
||||||
|
!!! tip
|
||||||
|
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
|
||||||
|
|
||||||
=== "Google TPU"
|
=== "Google TPU"
|
||||||
|
|
||||||
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
|
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
|
||||||
|
|||||||
Generated
+27
-10
@@ -3446,7 +3446,6 @@ dependencies = [
|
|||||||
"base64 0.22.1",
|
"base64 0.22.1",
|
||||||
"bytes",
|
"bytes",
|
||||||
"encoding_rs",
|
"encoding_rs",
|
||||||
"futures-channel",
|
|
||||||
"futures-core",
|
"futures-core",
|
||||||
"futures-util",
|
"futures-util",
|
||||||
"h2",
|
"h2",
|
||||||
@@ -4876,6 +4875,7 @@ dependencies = [
|
|||||||
"futures-core",
|
"futures-core",
|
||||||
"pin-project-lite",
|
"pin-project-lite",
|
||||||
"tokio",
|
"tokio",
|
||||||
|
"tokio-util",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -4937,9 +4937,9 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tonic"
|
name = "tonic"
|
||||||
version = "0.14.5"
|
version = "0.14.6"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
|
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"axum",
|
"axum",
|
||||||
@@ -4966,9 +4966,9 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tonic-build"
|
name = "tonic-build"
|
||||||
version = "0.14.5"
|
version = "0.14.6"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
|
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"prettyplease",
|
"prettyplease",
|
||||||
"proc-macro2",
|
"proc-macro2",
|
||||||
@@ -4977,10 +4977,23 @@ dependencies = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tonic-prost"
|
name = "tonic-health"
|
||||||
version = "0.14.5"
|
version = "0.14.6"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
|
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
|
||||||
|
dependencies = [
|
||||||
|
"prost",
|
||||||
|
"tokio",
|
||||||
|
"tokio-stream",
|
||||||
|
"tonic",
|
||||||
|
"tonic-prost",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "tonic-prost"
|
||||||
|
version = "0.14.6"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"bytes",
|
"bytes",
|
||||||
"prost",
|
"prost",
|
||||||
@@ -4989,9 +5002,9 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tonic-prost-build"
|
name = "tonic-prost-build"
|
||||||
version = "0.14.5"
|
version = "0.14.6"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
|
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"prettyplease",
|
"prettyplease",
|
||||||
"proc-macro2",
|
"proc-macro2",
|
||||||
@@ -5484,10 +5497,13 @@ dependencies = [
|
|||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
|
"thiserror-ext",
|
||||||
"tiktoken-rs 0.9.1",
|
"tiktoken-rs 0.9.1",
|
||||||
"tokenizers",
|
"tokenizers",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-stream",
|
"tokio-stream",
|
||||||
|
"tracing",
|
||||||
|
"tracing-subscriber",
|
||||||
"url",
|
"url",
|
||||||
"uuid",
|
"uuid",
|
||||||
]
|
]
|
||||||
@@ -5730,6 +5746,7 @@ dependencies = [
|
|||||||
"tokio-stream",
|
"tokio-stream",
|
||||||
"tokio-util",
|
"tokio-util",
|
||||||
"tonic",
|
"tonic",
|
||||||
|
"tonic-health",
|
||||||
"tonic-prost",
|
"tonic-prost",
|
||||||
"tonic-prost-build",
|
"tonic-prost-build",
|
||||||
"tower",
|
"tower",
|
||||||
|
|||||||
+5
-4
@@ -118,10 +118,11 @@ tokio = { version = "1.47.1", features = [
|
|||||||
tokio-openssl = "0.6"
|
tokio-openssl = "0.6"
|
||||||
tokio-stream = "0.1"
|
tokio-stream = "0.1"
|
||||||
tokio-util = { version = "0.7.18", features = ["rt"] }
|
tokio-util = { version = "0.7.18", features = ["rt"] }
|
||||||
tonic = "0.14.5"
|
tonic = "0.14.6"
|
||||||
tonic-build = "0.14.5"
|
tonic-build = "0.14.6"
|
||||||
tonic-prost = "0.14.5"
|
tonic-health = "0.14.6"
|
||||||
tonic-prost-build = "0.14.5"
|
tonic-prost = "0.14.6"
|
||||||
|
tonic-prost-build = "0.14.6"
|
||||||
tool-parser = "1.2.0"
|
tool-parser = "1.2.0"
|
||||||
tower = { version = "0.5.3", features = ["util"] }
|
tower = { version = "0.5.3", features = ["util"] }
|
||||||
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
|
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
|
||||||
|
|||||||
@@ -20,16 +20,19 @@ mimalloc.workspace = true
|
|||||||
rand.workspace = true
|
rand.workspace = true
|
||||||
rand_distr.workspace = true
|
rand_distr.workspace = true
|
||||||
rayon.workspace = true
|
rayon.workspace = true
|
||||||
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
|
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
|
||||||
rlimit.workspace = true
|
rlimit.workspace = true
|
||||||
rustc-hash.workspace = true
|
rustc-hash.workspace = true
|
||||||
serde = { workspace = true, features = ["rc"] }
|
serde = { workspace = true, features = ["rc"] }
|
||||||
serde_json = { workspace = true, features = ["raw_value"] }
|
serde_json = { workspace = true, features = ["raw_value"] }
|
||||||
thiserror.workspace = true
|
thiserror.workspace = true
|
||||||
|
thiserror-ext.workspace = true
|
||||||
tiktoken-rs.workspace = true
|
tiktoken-rs.workspace = true
|
||||||
tokenizers.workspace = true
|
tokenizers.workspace = true
|
||||||
tokio.workspace = true
|
tokio.workspace = true
|
||||||
tokio-stream.workspace = true
|
tokio-stream.workspace = true
|
||||||
|
tracing.workspace = true
|
||||||
|
tracing-subscriber.workspace = true
|
||||||
url.workspace = true
|
url.workspace = true
|
||||||
uuid.workspace = true
|
uuid.workspace = true
|
||||||
|
|
||||||
|
|||||||
@@ -158,9 +158,10 @@ impl PoolingBackend {
|
|||||||
// (mirrors Python async_request_vllm_rerank).
|
// (mirrors Python async_request_vllm_rerank).
|
||||||
if let Some(ref list) = input.prompt_list {
|
if let Some(ref list) = input.prompt_list {
|
||||||
if list.len() < 2 {
|
if list.len() < 2 {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"WARNING: vllm-rerank request has no documents \
|
backend = "vllm-rerank",
|
||||||
(prompt_list needs [query, doc, ...])"
|
inputs = list.len(),
|
||||||
|
"rerank request has no documents"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
|
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
|
||||||
@@ -175,10 +176,10 @@ impl PoolingBackend {
|
|||||||
// Legacy path: text prompt as query, documents via --extra-body.
|
// Legacy path: text prompt as query, documents via --extra-body.
|
||||||
let query = input.prompt.as_ref();
|
let query = input.prompt.as_ref();
|
||||||
if query.is_empty() && input.prompt_token_ids.is_some() {
|
if query.is_empty() && input.prompt_token_ids.is_some() {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"WARNING: vllm-rerank received empty query (random dataset uses \
|
backend = "vllm-rerank",
|
||||||
token IDs only). Use --dataset-name random-rerank for meaningful \
|
dataset = "random",
|
||||||
rerank benchmarks."
|
"rerank request has an empty query; use the random-rerank dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
serde_json::json!({
|
serde_json::json!({
|
||||||
|
|||||||
+114
-84
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
|||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
use indicatif::{ProgressBar, ProgressStyle};
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
|
use thiserror_ext::AsReport as _;
|
||||||
use tokio::sync::Semaphore;
|
use tokio::sync::Semaphore;
|
||||||
|
|
||||||
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
|
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
|
||||||
@@ -72,12 +73,12 @@ pub fn pre_resolve_dns(
|
|||||||
v4.extend(v6);
|
v4.extend(v6);
|
||||||
if !v4.is_empty() {
|
if !v4.is_empty() {
|
||||||
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
|
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
|
||||||
println!("Pre-resolved {host} -> {ips:?}");
|
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
|
||||||
builder = builder.resolve_to_addrs(host, &v4);
|
builder = builder.resolve_to_addrs(host, &v4);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
|
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||||
(m.clone(), config.model_name.clone())
|
(m.clone(), config.model_name.clone())
|
||||||
} else {
|
} else {
|
||||||
println!("Model not specified, fetching first model from server...");
|
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||||
let (name, id) =
|
let (name, id) =
|
||||||
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
|
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
|
||||||
println!("First model name: {name}, first model id: {id}");
|
tracing::info!(
|
||||||
|
model_name = name,
|
||||||
|
model_id = id,
|
||||||
|
"selected first model from server"
|
||||||
|
);
|
||||||
(id, Some(name))
|
(id, Some(name))
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -358,10 +363,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||||
println!("Loading tokenizer: {tid}");
|
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||||
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
||||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
let t =
|
||||||
println!("Tokenizer loaded successfully.");
|
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
|
||||||
Some(t)
|
Some(t)
|
||||||
};
|
};
|
||||||
let has_tokenizer = tokenizer.is_some();
|
let has_tokenizer = tokenizer.is_some();
|
||||||
@@ -421,7 +426,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
config.num_prompts, config.random_batch_size, config.is_reranker,
|
config.num_prompts, config.random_batch_size, config.is_reranker,
|
||||||
),
|
),
|
||||||
};
|
};
|
||||||
println!("Generating {dataset_label}...");
|
tracing::info!(
|
||||||
|
dataset = ?config.dataset_name,
|
||||||
|
prompts = config.num_prompts,
|
||||||
|
description = %dataset_label,
|
||||||
|
"generating benchmark dataset"
|
||||||
|
);
|
||||||
let gen_start = Instant::now();
|
let gen_start = Instant::now();
|
||||||
|
|
||||||
let mut input_requests = match config.dataset_name {
|
let mut input_requests = match config.dataset_name {
|
||||||
@@ -472,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
let path = match config.dataset_path.as_deref() {
|
let path = match config.dataset_path.as_deref() {
|
||||||
Some(p) => p,
|
Some(p) => p,
|
||||||
None => {
|
None => {
|
||||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||||
downloaded.as_str()
|
downloaded.as_str()
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -512,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
None => {
|
None => {
|
||||||
downloaded = crate::datasets::speed_bench::download_speed_bench(
|
downloaded = crate::datasets::speed_bench::download_speed_bench(
|
||||||
config.speed_bench_config,
|
config.speed_bench_config,
|
||||||
)?;
|
)
|
||||||
|
.await?;
|
||||||
downloaded.as_str()
|
downloaded.as_str()
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -543,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
config.hf_subset.as_deref(),
|
config.hf_subset.as_deref(),
|
||||||
config.hf_split.as_deref(),
|
config.hf_split.as_deref(),
|
||||||
config.num_prompts,
|
config.num_prompts,
|
||||||
)?;
|
)
|
||||||
|
.await?;
|
||||||
crate::datasets::hf_dataset::load_hf_dataset(
|
crate::datasets::hf_dataset::load_hf_dataset(
|
||||||
tok,
|
tok,
|
||||||
&downloaded_path,
|
&downloaded_path,
|
||||||
@@ -608,18 +620,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
};
|
};
|
||||||
|
|
||||||
let gen_elapsed = gen_start.elapsed();
|
let gen_elapsed = gen_start.elapsed();
|
||||||
println!(
|
tracing::info!(
|
||||||
"Generated {} prompts in {:.2}s",
|
prompts = input_requests.len(),
|
||||||
input_requests.len(),
|
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||||
gen_elapsed.as_secs_f64()
|
"generated benchmark dataset"
|
||||||
);
|
);
|
||||||
|
|
||||||
let filtered_count =
|
let filtered_count =
|
||||||
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
|
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
|
||||||
if filtered_count > 0 {
|
if filtered_count > 0 {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
|
filtered_prompts = filtered_count,
|
||||||
config.max_model_len.unwrap()
|
max_model_len = config.max_model_len.unwrap(),
|
||||||
|
"filtered prompts above maximum model length"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
if input_requests.is_empty() {
|
if input_requests.is_empty() {
|
||||||
@@ -670,7 +683,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
|
|
||||||
// Ready check
|
// Ready check
|
||||||
if config.ready_check_timeout_sec > 0 {
|
if config.ready_check_timeout_sec > 0 {
|
||||||
println!("Starting initial single prompt test run...");
|
tracing::info!("starting initial single-prompt test run");
|
||||||
let test_output = wait_for_endpoint(
|
let test_output = wait_for_endpoint(
|
||||||
config.backend,
|
config.backend,
|
||||||
&client,
|
&client,
|
||||||
@@ -685,7 +698,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
test_output.error
|
test_output.error
|
||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
println!("Initial test run completed.");
|
tracing::info!("initial single-prompt test run completed");
|
||||||
}
|
}
|
||||||
|
|
||||||
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
|
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
|
||||||
@@ -703,12 +716,15 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
DatasetName::Random | DatasetName::PrefixRepetition
|
DatasetName::Random | DatasetName::PrefixRepetition
|
||||||
);
|
);
|
||||||
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
|
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
|
||||||
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
|
tracing::info!(
|
||||||
|
reason = "prompt_token_ids",
|
||||||
|
"skipping server tokenizer verification"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
|
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
|
||||||
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
|
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
|
||||||
if is_tokenizer_verified(&cache_key) {
|
if is_tokenizer_verified(&cache_key) {
|
||||||
println!("Tokenizer verified in previous run (cached), skipping verification.");
|
tracing::info!(reason = "cached", "skipping server tokenizer verification");
|
||||||
} else {
|
} else {
|
||||||
let num_special =
|
let num_special =
|
||||||
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
|
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
|
||||||
@@ -723,14 +739,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
.await?
|
.await?
|
||||||
{
|
{
|
||||||
SampleVerifyOutcome::Passed => {
|
SampleVerifyOutcome::Passed => {
|
||||||
println!("Sample verification passed, skipping full verification.");
|
tracing::info!("tokenizer sample verification passed");
|
||||||
mark_tokenizer_verified(&cache_key);
|
mark_tokenizer_verified(&cache_key);
|
||||||
}
|
}
|
||||||
SampleVerifyOutcome::Skipped(reason) => {
|
SampleVerifyOutcome::Skipped(reason) => {
|
||||||
println!("Server /tokenize unavailable ({reason}), skipping verification.");
|
tracing::warn!(
|
||||||
|
reason = %reason,
|
||||||
|
"server tokenizer unavailable; skipping prompt verification"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
SampleVerifyOutcome::Mismatch => {
|
SampleVerifyOutcome::Mismatch => {
|
||||||
println!("Sample verification found mismatch, running full verify+fix...");
|
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
|
||||||
match verify_and_fix_prompt_lengths(
|
match verify_and_fix_prompt_lengths(
|
||||||
&client,
|
&client,
|
||||||
&config.base_url,
|
&config.base_url,
|
||||||
@@ -742,16 +761,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(()) => {
|
Ok(()) => {
|
||||||
println!(
|
tracing::info!(
|
||||||
"All {} prompts verified: exact token length match.",
|
prompts = input_requests.len(),
|
||||||
input_requests.len()
|
"verified exact prompt token lengths"
|
||||||
);
|
);
|
||||||
mark_tokenizer_verified(&cache_key);
|
mark_tokenizer_verified(&cache_key);
|
||||||
}
|
}
|
||||||
Err(BenchError::TokenizeUnavailable(reason)) => {
|
Err(BenchError::TokenizeUnavailable(reason)) => {
|
||||||
println!(
|
tracing::warn!(
|
||||||
"Server /tokenize became unavailable during verification \
|
reason = %reason,
|
||||||
({reason}); proceeding with client-side token counts."
|
"server tokenizer became unavailable; using client token counts"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Err(e) => return Err(e),
|
Err(e) => return Err(e),
|
||||||
@@ -763,7 +782,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
|
|
||||||
// Warmup
|
// Warmup
|
||||||
if config.num_warmups > 0 {
|
if config.num_warmups > 0 {
|
||||||
println!("Warming up with {} requests...", config.num_warmups);
|
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
|
||||||
run_warmup(
|
run_warmup(
|
||||||
config.backend,
|
config.backend,
|
||||||
&client,
|
&client,
|
||||||
@@ -776,7 +795,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
config.disable_tqdm,
|
config.disable_tqdm,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
println!("Warmup run completed.");
|
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
|
||||||
}
|
}
|
||||||
|
|
||||||
// Start profiler if requested (immediate mode — no batch threshold)
|
// Start profiler if requested (immediate mode — no batch threshold)
|
||||||
@@ -814,28 +833,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
let spec_decode_before =
|
let spec_decode_before =
|
||||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||||
if spec_decode_before.is_some() {
|
if spec_decode_before.is_some() {
|
||||||
println!("Speculative decoding detected, will collect metrics.");
|
tracing::info!("detected speculative decoding; collecting metrics");
|
||||||
}
|
}
|
||||||
|
|
||||||
// Main benchmark
|
// Main benchmark
|
||||||
println!("Starting main benchmark run...");
|
|
||||||
let distribution = if config.burstiness == 1.0 {
|
let distribution = if config.burstiness == 1.0 {
|
||||||
"Poisson process"
|
"Poisson process"
|
||||||
} else {
|
} else {
|
||||||
"Gamma distribution"
|
"Gamma distribution"
|
||||||
};
|
};
|
||||||
println!(
|
tracing::info!(
|
||||||
"Traffic request rate: {}",
|
request_rate = config.request_rate,
|
||||||
if config.request_rate.is_infinite() {
|
burstiness = config.burstiness,
|
||||||
"inf".to_string()
|
distribution,
|
||||||
} else {
|
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
|
||||||
format!("{}", config.request_rate)
|
prompts = config.num_prompts,
|
||||||
}
|
"starting main benchmark run"
|
||||||
);
|
|
||||||
println!("Burstiness factor: {} ({distribution})", config.burstiness);
|
|
||||||
println!(
|
|
||||||
"Maximum request concurrency: {}",
|
|
||||||
config.max_concurrency.unwrap_or(config.num_prompts)
|
|
||||||
);
|
);
|
||||||
|
|
||||||
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
|
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
|
||||||
@@ -847,11 +860,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
);
|
);
|
||||||
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
|
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
|
||||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||||
println!(
|
tracing::info!(
|
||||||
"LoRA adapters ({}): {:?} [assignment={:?}]",
|
adapters = modules.len(),
|
||||||
modules.len(),
|
names = ?names,
|
||||||
names,
|
assignment = ?config.lora_assignment,
|
||||||
config.lora_assignment
|
"assigned LoRA adapters"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1125,7 +1138,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
|||||||
if let Some((cancel_tx, task)) = profile_task {
|
if let Some((cancel_tx, task)) = profile_task {
|
||||||
let _ = cancel_tx.send(());
|
let _ = cancel_tx.send(());
|
||||||
if let Err(e) = task.await {
|
if let Err(e) = task.await {
|
||||||
eprintln!("WARNING: Profile background task failed: {e}");
|
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1289,12 +1302,14 @@ pub(crate) async fn start_profiler_immediate(
|
|||||||
base_url: &str,
|
base_url: &str,
|
||||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||||
) {
|
) {
|
||||||
println!("Starting profiler...");
|
|
||||||
let profile_url = format!("{base_url}/start_profile");
|
let profile_url = format!("{base_url}/start_profile");
|
||||||
|
tracing::info!(url = %profile_url, "starting profiler");
|
||||||
match send_profile_request(client, &profile_url, extra_headers).await {
|
match send_profile_request(client, &profile_url, extra_headers).await {
|
||||||
Ok(true) => println!("Profiler started"),
|
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
|
||||||
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
|
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
|
||||||
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
|
Err(e) => {
|
||||||
|
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1304,12 +1319,14 @@ pub(crate) async fn stop_profiler_immediate(
|
|||||||
base_url: &str,
|
base_url: &str,
|
||||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||||
) {
|
) {
|
||||||
println!("Stopping profiler...");
|
|
||||||
let profile_url = format!("{base_url}/stop_profile");
|
let profile_url = format!("{base_url}/stop_profile");
|
||||||
|
tracing::info!(url = %profile_url, "stopping profiler");
|
||||||
match send_profile_request(client, &profile_url, extra_headers).await {
|
match send_profile_request(client, &profile_url, extra_headers).await {
|
||||||
Ok(true) => println!("Profiler stopped"),
|
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
|
||||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
|
||||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
Err(e) => {
|
||||||
|
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1371,25 +1388,30 @@ pub(crate) async fn profile_on_batch_threshold(
|
|||||||
duration_secs: f64,
|
duration_secs: f64,
|
||||||
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
|
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
|
||||||
) {
|
) {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Waiting for batch size >= {threshold} before starting profiler \
|
threshold,
|
||||||
(will capture {duration_secs}s)..."
|
duration_seconds = duration_secs,
|
||||||
|
"waiting for profiler batch threshold"
|
||||||
);
|
);
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
if let Some(running) = fetch_num_requests_running(client, base_url).await
|
if let Some(running) = fetch_num_requests_running(client, base_url).await
|
||||||
&& running >= threshold
|
&& running >= threshold
|
||||||
{
|
{
|
||||||
println!("Batch size {running} >= {threshold}, starting profiler...");
|
tracing::info!(
|
||||||
|
running_requests = running,
|
||||||
|
threshold,
|
||||||
|
"profiler batch threshold reached"
|
||||||
|
);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
// Wait 500ms or until the benchmark signals cancellation
|
// Wait 500ms or until the benchmark signals cancellation
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
|
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
|
||||||
_ = &mut cancel_rx => {
|
_ = &mut cancel_rx => {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
|
threshold,
|
||||||
profiling skipped."
|
"benchmark finished before profiler batch threshold; skipping profiling"
|
||||||
);
|
);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -1398,13 +1420,13 @@ pub(crate) async fn profile_on_batch_threshold(
|
|||||||
|
|
||||||
let start_url = format!("{base_url}/start_profile");
|
let start_url = format!("{base_url}/start_profile");
|
||||||
match send_profile_request(client, &start_url, extra_headers).await {
|
match send_profile_request(client, &start_url, extra_headers).await {
|
||||||
Ok(true) => println!("Profiler started"),
|
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
|
||||||
Ok(false) => {
|
Ok(false) => {
|
||||||
eprintln!("WARNING: Profiler start request returned non-success");
|
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
eprintln!("WARNING: Failed to start profiler: {e}");
|
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1413,15 +1435,17 @@ pub(crate) async fn profile_on_batch_threshold(
|
|||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
|
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
|
||||||
_ = &mut cancel_rx => {
|
_ = &mut cancel_rx => {
|
||||||
println!("Benchmark finished, stopping profiler early...");
|
tracing::info!("benchmark finished; stopping profiler early");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let stop_url = format!("{base_url}/stop_profile");
|
let stop_url = format!("{base_url}/stop_profile");
|
||||||
match send_profile_request(client, &stop_url, extra_headers).await {
|
match send_profile_request(client, &stop_url, extra_headers).await {
|
||||||
Ok(true) => println!("Profiler stopped after capturing"),
|
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
|
||||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
|
||||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
Err(e) => {
|
||||||
|
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1502,10 +1526,11 @@ async fn verify_and_fix_prompt_lengths(
|
|||||||
let excess = tokens.len().saturating_sub(expected_input_len);
|
let excess = tokens.len().saturating_sub(expected_input_len);
|
||||||
let compensate = if excess > 0 && last_excess == Some(excess) {
|
let compensate = if excess > 0 && last_excess == Some(excess) {
|
||||||
if _iter == 1 {
|
if _iter == 1 {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"Prompt {i}: server consistently adds {excess} extra token(s) \
|
prompt_index = i,
|
||||||
(likely BOS), compensating target to {}.",
|
extra_tokens = excess,
|
||||||
expected_input_len.saturating_sub(excess),
|
adjusted_target = expected_input_len.saturating_sub(excess),
|
||||||
|
"server consistently adds prompt tokens; compensating verification target"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
excess
|
excess
|
||||||
@@ -1563,7 +1588,10 @@ async fn verify_and_fix_prompt_lengths(
|
|||||||
|
|
||||||
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
|
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
|
||||||
if fc > 0 {
|
if fc > 0 {
|
||||||
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
|
tracing::info!(
|
||||||
|
fixed_prompts = fc,
|
||||||
|
"fixed prompt lengths using server tokenizer"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
@@ -1818,7 +1846,7 @@ async fn sample_verify_prompts(
|
|||||||
let tokenize_url = format!("{base_url}/tokenize");
|
let tokenize_url = format!("{base_url}/tokenize");
|
||||||
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||||
|
|
||||||
println!("Sampling {sample_size} prompts for verification...");
|
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
|
||||||
|
|
||||||
for (i, request) in requests.iter().enumerate().take(sample_size) {
|
for (i, request) in requests.iter().enumerate().take(sample_size) {
|
||||||
let tokens = match server_tokenize(
|
let tokens = match server_tokenize(
|
||||||
@@ -1841,9 +1869,11 @@ async fn sample_verify_prompts(
|
|||||||
|
|
||||||
let expected = request.prompt_len + num_special;
|
let expected = request.prompt_len + num_special;
|
||||||
if tokens.len() != expected {
|
if tokens.len() != expected {
|
||||||
println!(
|
tracing::warn!(
|
||||||
"Prompt {i}: expected {expected} tokens, server returned {}",
|
prompt_index = i,
|
||||||
tokens.len()
|
expected_tokens = expected,
|
||||||
|
actual_tokens = tokens.len(),
|
||||||
|
"tokenizer verification sample mismatch"
|
||||||
);
|
);
|
||||||
return Ok(SampleVerifyOutcome::Mismatch);
|
return Ok(SampleVerifyOutcome::Mismatch);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -288,10 +288,18 @@ impl BenchConfig {
|
|||||||
}
|
}
|
||||||
Some(other) => {
|
Some(other) => {
|
||||||
// extra_body was not an object — just use sampling params
|
// extra_body was not an object — just use sampling params
|
||||||
eprintln!(
|
let value_type = match &other {
|
||||||
"Warning: --extra-body is not a JSON object, sampling params may be lost"
|
serde_json::Value::Null => "null",
|
||||||
|
serde_json::Value::Bool(_) => "boolean",
|
||||||
|
serde_json::Value::Number(_) => "number",
|
||||||
|
serde_json::Value::String(_) => "string",
|
||||||
|
serde_json::Value::Array(_) => "array",
|
||||||
|
serde_json::Value::Object(_) => unreachable!(),
|
||||||
|
};
|
||||||
|
tracing::warn!(
|
||||||
|
value_type,
|
||||||
|
"sampling parameters may be lost because --extra-body is not a JSON object"
|
||||||
);
|
);
|
||||||
let _ = other;
|
|
||||||
sampling_params
|
sampling_params
|
||||||
}
|
}
|
||||||
None => sampling_params,
|
None => sampling_params,
|
||||||
@@ -489,9 +497,9 @@ impl BenchConfig {
|
|||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
if !args.skip_chat_template {
|
if !args.skip_chat_template {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"NOTE: client-side chat template rendering is not supported; custom \
|
dataset = "custom",
|
||||||
dataset prompts are sent raw (equivalent to --skip-chat-template)."
|
"client-side chat template rendering is unsupported; sending prompts raw"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -570,9 +578,10 @@ impl BenchConfig {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if ignore_eos {
|
if ignore_eos {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
|
ignore_eos,
|
||||||
respect output length limits, causing unbounded context growth."
|
multi_turn = true,
|
||||||
|
"output length limits may be ignored, causing unbounded context growth"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -128,8 +128,10 @@ mod tests {
|
|||||||
|
|
||||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||||
fn test_tokenizer() -> TokenizerKind {
|
fn test_tokenizer() -> TokenizerKind {
|
||||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
TokenizerKind::Tiktoken(
|
||||||
.expect("gpt2 built-in tiktoken should always load without network")
|
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||||
|
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
|
|||||||
use rand::{Rng, SeedableRng};
|
use rand::{Rng, SeedableRng};
|
||||||
|
|
||||||
use super::SampleRequest;
|
use super::SampleRequest;
|
||||||
|
use super::progress::RowDownloadReporter;
|
||||||
use crate::error::{BenchError, Result};
|
use crate::error::{BenchError, Result};
|
||||||
use crate::tokenizer::TokenizerKind;
|
use crate::tokenizer::TokenizerKind;
|
||||||
|
|
||||||
@@ -50,18 +51,19 @@ enum ColumnFormat {
|
|||||||
|
|
||||||
/// Make a GET request with retry logic (3 retries with exponential backoff).
|
/// Make a GET request with retry logic (3 retries with exponential backoff).
|
||||||
/// Returns the parsed JSON response.
|
/// Returns the parsed JSON response.
|
||||||
fn get_with_retry(
|
async fn get_with_retry(
|
||||||
client: &reqwest::blocking::Client,
|
client: &reqwest::Client,
|
||||||
url: &str,
|
url: &str,
|
||||||
label: &str,
|
label: &str,
|
||||||
) -> Result<serde_json::Value> {
|
) -> Result<serde_json::Value> {
|
||||||
let max_retries = 3;
|
let max_retries = 3;
|
||||||
for attempt in 0..=max_retries {
|
for attempt in 0..=max_retries {
|
||||||
let resp = match client.get(url).send() {
|
let resp = match client.get(url).send().await {
|
||||||
Ok(r) => r,
|
Ok(r) => r,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
if attempt < max_retries {
|
if attempt < max_retries {
|
||||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
|
||||||
|
.await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
return Err(BenchError::Config(format!(
|
return Err(BenchError::Config(format!(
|
||||||
@@ -80,7 +82,7 @@ fn get_with_retry(
|
|||||||
}
|
}
|
||||||
|
|
||||||
if status.is_server_error() && attempt < max_retries {
|
if status.is_server_error() && attempt < max_retries {
|
||||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -92,6 +94,7 @@ fn get_with_retry(
|
|||||||
|
|
||||||
let data: serde_json::Value = resp
|
let data: serde_json::Value = resp
|
||||||
.json()
|
.json()
|
||||||
|
.await
|
||||||
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
|
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
|
||||||
return Ok(data);
|
return Ok(data);
|
||||||
}
|
}
|
||||||
@@ -105,7 +108,7 @@ fn get_with_retry(
|
|||||||
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
|
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
|
||||||
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
|
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
|
||||||
/// or the dataset is exhausted.
|
/// or the dataset is exhausted.
|
||||||
pub fn download_hf_dataset(
|
pub async fn download_hf_dataset(
|
||||||
dataset: &str,
|
dataset: &str,
|
||||||
subset: Option<&str>,
|
subset: Option<&str>,
|
||||||
split: Option<&str>,
|
split: Option<&str>,
|
||||||
@@ -115,7 +118,7 @@ pub fn download_hf_dataset(
|
|||||||
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
|
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
|
||||||
|
|
||||||
let mut client_builder =
|
let mut client_builder =
|
||||||
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
|
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
|
||||||
|
|
||||||
// Add HF_TOKEN auth header if available
|
// Add HF_TOKEN auth header if available
|
||||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||||
@@ -138,7 +141,7 @@ pub fn download_hf_dataset(
|
|||||||
// Call /info to discover available configs and splits
|
// Call /info to discover available configs and splits
|
||||||
let info_url =
|
let info_url =
|
||||||
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
|
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
|
||||||
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
|
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
|
||||||
|
|
||||||
let dataset_info =
|
let dataset_info =
|
||||||
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
|
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
|
||||||
@@ -201,7 +204,12 @@ pub fn download_hf_dataset(
|
|||||||
(resolved_config, resolved_split)
|
(resolved_config, resolved_split)
|
||||||
};
|
};
|
||||||
|
|
||||||
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
|
tracing::info!(
|
||||||
|
dataset,
|
||||||
|
config = resolved_config,
|
||||||
|
split = resolved_split,
|
||||||
|
"resolved Hugging Face dataset"
|
||||||
|
);
|
||||||
|
|
||||||
// Check cache
|
// Check cache
|
||||||
let dir = cache_dir();
|
let dir = cache_dir();
|
||||||
@@ -215,11 +223,16 @@ pub fn download_hf_dataset(
|
|||||||
|
|
||||||
if cache_path.exists() {
|
if cache_path.exists() {
|
||||||
let path_str = cache_path.to_string_lossy().to_string();
|
let path_str = cache_path.to_string_lossy().to_string();
|
||||||
println!("HF dataset cached: {path_str}");
|
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
|
||||||
return Ok((path_str, resolved_config, resolved_split));
|
return Ok((path_str, resolved_config, resolved_split));
|
||||||
}
|
}
|
||||||
|
|
||||||
println!("Downloading HF dataset '{dataset}' from datasets-server...");
|
tracing::info!(
|
||||||
|
dataset,
|
||||||
|
config = resolved_config,
|
||||||
|
split = resolved_split,
|
||||||
|
"downloading Hugging Face dataset"
|
||||||
|
);
|
||||||
|
|
||||||
let encoded_config: String =
|
let encoded_config: String =
|
||||||
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
|
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
|
||||||
@@ -229,6 +242,7 @@ pub fn download_hf_dataset(
|
|||||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||||
let mut offset = 0usize;
|
let mut offset = 0usize;
|
||||||
let page_size = 100usize;
|
let page_size = 100usize;
|
||||||
|
let mut progress = RowDownloadReporter::new();
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
let url = format!(
|
let url = format!(
|
||||||
@@ -240,7 +254,7 @@ pub fn download_hf_dataset(
|
|||||||
&length={page_size}"
|
&length={page_size}"
|
||||||
);
|
);
|
||||||
|
|
||||||
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
|
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
|
||||||
|
|
||||||
let rows = data["rows"]
|
let rows = data["rows"]
|
||||||
.as_array()
|
.as_array()
|
||||||
@@ -260,14 +274,14 @@ pub fn download_hf_dataset(
|
|||||||
offset += fetched;
|
offset += fetched;
|
||||||
|
|
||||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||||
eprint!("\r Fetched {offset}/{total} rows...");
|
progress.update(offset, total);
|
||||||
|
|
||||||
// Stop if we have enough rows or reached end of dataset
|
// Stop if we have enough rows or reached end of dataset
|
||||||
if all_rows.len() >= num_rows_needed || fetched < page_size {
|
if all_rows.len() >= num_rows_needed || fetched < page_size {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
eprintln!(); // newline after progress
|
progress.finish();
|
||||||
|
|
||||||
if all_rows.is_empty() {
|
if all_rows.is_empty() {
|
||||||
return Err(BenchError::Config(format!(
|
return Err(BenchError::Config(format!(
|
||||||
@@ -280,7 +294,12 @@ pub fn download_hf_dataset(
|
|||||||
std::fs::write(&cache_path, &json_str)?;
|
std::fs::write(&cache_path, &json_str)?;
|
||||||
|
|
||||||
let path_str = cache_path.to_string_lossy().to_string();
|
let path_str = cache_path.to_string_lossy().to_string();
|
||||||
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
|
tracing::info!(
|
||||||
|
dataset,
|
||||||
|
rows = all_rows.len(),
|
||||||
|
path = %path_str,
|
||||||
|
"saved Hugging Face dataset"
|
||||||
|
);
|
||||||
Ok((path_str, resolved_config, resolved_split))
|
Ok((path_str, resolved_config, resolved_split))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -483,21 +502,31 @@ pub fn load_hf_dataset(
|
|||||||
// Detect column format from first row
|
// Detect column format from first row
|
||||||
let format = detect_column_format(&entries[0], text_column_override)?;
|
let format = detect_column_format(&entries[0], text_column_override)?;
|
||||||
|
|
||||||
// Print detected format
|
|
||||||
match &format {
|
match &format {
|
||||||
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
|
ColumnFormat::Chat(col) => {
|
||||||
|
tracing::info!(
|
||||||
|
format = "chat",
|
||||||
|
column = col,
|
||||||
|
"detected Hugging Face dataset format"
|
||||||
|
);
|
||||||
|
}
|
||||||
ColumnFormat::Text {
|
ColumnFormat::Text {
|
||||||
prompt_col,
|
prompt_col,
|
||||||
output_col,
|
output_col,
|
||||||
} => {
|
} => {
|
||||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
tracing::info!(
|
||||||
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
|
format = "text",
|
||||||
|
prompt_column = prompt_col,
|
||||||
|
output_column = output_col.as_deref().unwrap_or("none"),
|
||||||
|
"detected Hugging Face dataset format"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
ColumnFormat::Combined { cols, output_col } => {
|
ColumnFormat::Combined { cols, output_col } => {
|
||||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
tracing::info!(
|
||||||
println!(
|
format = "combined",
|
||||||
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
|
prompt_columns = ?cols,
|
||||||
cols
|
output_column = output_col.as_deref().unwrap_or("none"),
|
||||||
|
"detected Hugging Face dataset format"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -608,9 +637,10 @@ pub fn load_hf_dataset(
|
|||||||
if len == 0 { 128 } else { len }
|
if len == 0 { 128 } else { len }
|
||||||
} else {
|
} else {
|
||||||
if !warned_no_output {
|
if !warned_no_output {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"WARNING: No output column detected and --hf-output-len not set. \
|
path = dataset_path,
|
||||||
Using default output length of 128 tokens."
|
default_output_tokens = 128,
|
||||||
|
"no dataset output column or --hf-output-len; using default output length"
|
||||||
);
|
);
|
||||||
warned_no_output = true;
|
warned_no_output = true;
|
||||||
}
|
}
|
||||||
@@ -618,9 +648,10 @@ pub fn load_hf_dataset(
|
|||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
if !warned_no_output {
|
if !warned_no_output {
|
||||||
eprintln!(
|
tracing::warn!(
|
||||||
"WARNING: No output column detected and --hf-output-len not set. \
|
path = dataset_path,
|
||||||
Using default output length of 128 tokens."
|
default_output_tokens = 128,
|
||||||
|
"no dataset output column or --hf-output-len; using default output length"
|
||||||
);
|
);
|
||||||
warned_no_output = true;
|
warned_no_output = true;
|
||||||
}
|
}
|
||||||
@@ -640,9 +671,11 @@ pub fn load_hf_dataset(
|
|||||||
// Oversample if needed
|
// Oversample if needed
|
||||||
if samples.len() < num_requests {
|
if samples.len() < num_requests {
|
||||||
if no_oversample {
|
if no_oversample {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
dataset = "hf",
|
||||||
samples.len()
|
samples = samples.len(),
|
||||||
|
requested = num_requests,
|
||||||
|
"skipping dataset oversampling"
|
||||||
);
|
);
|
||||||
} else if !samples.is_empty() {
|
} else if !samples.is_empty() {
|
||||||
let original_len = samples.len();
|
let original_len = samples.len();
|
||||||
@@ -652,9 +685,11 @@ pub fn load_hf_dataset(
|
|||||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||||
samples.push(req);
|
samples.push(req);
|
||||||
}
|
}
|
||||||
println!(
|
tracing::info!(
|
||||||
"Oversampled HF dataset from {original_len} to {} total samples.",
|
dataset = "hf",
|
||||||
samples.len()
|
original_samples = original_len,
|
||||||
|
samples = samples.len(),
|
||||||
|
"oversampled dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1002,8 +1037,10 @@ mod tests {
|
|||||||
|
|
||||||
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
|
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
|
||||||
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
|
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
|
||||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
crate::tokenizer::TokenizerKind::Tiktoken(
|
||||||
.expect("gpt2 built-in tiktoken should always load without network")
|
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||||
|
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Write JSON data to a unique temp file and return the path string.
|
/// Write JSON data to a unique temp file and return the path string.
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ pub mod custom;
|
|||||||
pub mod hf_dataset;
|
pub mod hf_dataset;
|
||||||
pub mod multi_turn;
|
pub mod multi_turn;
|
||||||
pub mod prefix_repetition;
|
pub mod prefix_repetition;
|
||||||
|
mod progress;
|
||||||
pub mod random;
|
pub mod random;
|
||||||
pub mod random_mm;
|
pub mod random_mm;
|
||||||
pub mod random_rerank;
|
pub mod random_rerank;
|
||||||
@@ -90,9 +91,10 @@ pub fn oversample_requests(
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
if no_oversample {
|
if no_oversample {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
samples = requests.len(),
|
||||||
requests.len()
|
requested = num_requests,
|
||||||
|
"skipping dataset oversampling"
|
||||||
);
|
);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -103,9 +105,10 @@ pub fn oversample_requests(
|
|||||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||||
requests.push(req);
|
requests.push(req);
|
||||||
}
|
}
|
||||||
println!(
|
tracing::info!(
|
||||||
"Oversampled requests from {original_len} to {} total samples.",
|
original_samples = original_len,
|
||||||
requests.len()
|
samples = requests.len(),
|
||||||
|
"oversampled dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
|
|||||||
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
|
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
|
||||||
conversations.push(conv);
|
conversations.push(conv);
|
||||||
}
|
}
|
||||||
println!(
|
tracing::info!(
|
||||||
"Oversampled multi-turn conversations from {original_len} to {} total.",
|
original_conversations = original_len,
|
||||||
conversations.len()
|
conversations = conversations.len(),
|
||||||
|
"oversampled multi-turn conversations"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -525,10 +526,12 @@ mod tests {
|
|||||||
len
|
len
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_prefix_sharing_structure() {
|
async fn test_prefix_sharing_structure() {
|
||||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let cfg = MultiTurnRandomConfig {
|
let cfg = MultiTurnRandomConfig {
|
||||||
num_conversations: 5,
|
num_conversations: 5,
|
||||||
@@ -610,10 +613,12 @@ mod tests {
|
|||||||
println!("All prefix sharing checks passed!");
|
println!("All prefix sharing checks passed!");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_per_turn_input_len_default_mode() {
|
async fn test_per_turn_input_len_default_mode() {
|
||||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let cfg = MultiTurnRandomConfig {
|
let cfg = MultiTurnRandomConfig {
|
||||||
num_conversations: 4,
|
num_conversations: 4,
|
||||||
@@ -650,10 +655,12 @@ mod tests {
|
|||||||
println!("per_turn_input_len default-mode checks passed!");
|
println!("per_turn_input_len default-mode checks passed!");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_variable_turns_range() {
|
async fn test_variable_turns_range() {
|
||||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let cfg = MultiTurnRandomConfig {
|
let cfg = MultiTurnRandomConfig {
|
||||||
num_conversations: 50,
|
num_conversations: 50,
|
||||||
@@ -684,10 +691,12 @@ mod tests {
|
|||||||
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
|
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_variable_turns_fixed() {
|
async fn test_variable_turns_fixed() {
|
||||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
let cfg = MultiTurnRandomConfig {
|
let cfg = MultiTurnRandomConfig {
|
||||||
num_conversations: 10,
|
num_conversations: 10,
|
||||||
@@ -709,10 +718,12 @@ mod tests {
|
|||||||
println!("variable_turns_fixed checks passed!");
|
println!("variable_turns_fixed checks passed!");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_per_turn_input_len_prefix_sharing() {
|
async fn test_per_turn_input_len_prefix_sharing() {
|
||||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
|
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
|
||||||
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
|
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
|
||||||
|
|||||||
@@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
|
|||||||
}
|
}
|
||||||
let total = prompts_per_prefix * num_prefixes;
|
let total = prompts_per_prefix * num_prefixes;
|
||||||
if total != num_requests {
|
if total != num_requests {
|
||||||
println!(
|
tracing::info!(
|
||||||
"prefix_repetition: generating {total} requests \
|
requested = num_requests,
|
||||||
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
|
generated = total,
|
||||||
{} dropped to divide evenly)",
|
prefixes = num_prefixes,
|
||||||
num_requests - total
|
prompts_per_prefix,
|
||||||
|
dropped = num_requests - total,
|
||||||
|
"adjusted prefix-repetition request count"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -109,8 +111,10 @@ mod tests {
|
|||||||
|
|
||||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||||
fn test_tokenizer() -> TokenizerKind {
|
fn test_tokenizer() -> TokenizerKind {
|
||||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
TokenizerKind::Tiktoken(
|
||||||
.expect("gpt2 built-in tiktoken should always load without network")
|
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||||
|
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -0,0 +1,77 @@
|
|||||||
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
|
|
||||||
|
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
|
||||||
|
|
||||||
|
/// Reports row download progress to an interactive progress bar, or through
|
||||||
|
/// periodic tracing events when the progress bar is hidden on a non-TTY.
|
||||||
|
pub(super) struct RowDownloadReporter {
|
||||||
|
progress: ProgressBar,
|
||||||
|
next_report: Instant,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RowDownloadReporter {
|
||||||
|
/// Creates a reporter that emits non-TTY updates every 10 seconds.
|
||||||
|
pub fn new() -> Self {
|
||||||
|
let progress = ProgressBar::new(0);
|
||||||
|
progress.set_style(
|
||||||
|
ProgressStyle::with_template(
|
||||||
|
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
.progress_chars("#>-"),
|
||||||
|
);
|
||||||
|
Self {
|
||||||
|
progress,
|
||||||
|
next_report: Instant::now() + REPORT_INTERVAL,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Updates the current row count and reports progress when due.
|
||||||
|
pub fn update(&mut self, rows: usize, total: u64) {
|
||||||
|
let rows = rows as u64;
|
||||||
|
let total = total.max(rows);
|
||||||
|
self.progress.set_length(total);
|
||||||
|
self.progress.set_position(rows);
|
||||||
|
|
||||||
|
if self.should_report(Instant::now()) {
|
||||||
|
tracing::info!(rows, total, "fetching dataset rows");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Clears the interactive progress bar after the download completes.
|
||||||
|
pub fn finish(self) {
|
||||||
|
self.progress.finish_and_clear();
|
||||||
|
}
|
||||||
|
|
||||||
|
fn should_report(&mut self, now: Instant) -> bool {
|
||||||
|
if !self.progress.is_hidden() || now < self.next_report {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
self.next_report = now + REPORT_INTERVAL;
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hidden_reporter_uses_ten_second_deadline() {
|
||||||
|
let start = Instant::now();
|
||||||
|
let mut reporter = RowDownloadReporter {
|
||||||
|
progress: ProgressBar::hidden(),
|
||||||
|
next_report: start + REPORT_INTERVAL,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert!(!reporter.should_report(start + Duration::from_secs(9)));
|
||||||
|
assert!(reporter.should_report(start + Duration::from_secs(10)));
|
||||||
|
assert!(!reporter.should_report(start + Duration::from_secs(19)));
|
||||||
|
assert!(reporter.should_report(start + Duration::from_secs(20)));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -49,9 +49,12 @@ pub fn generate_random_dataset(
|
|||||||
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
|
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
|
||||||
let (output_low, output_high) = range_ratio.output_bounds(output_len);
|
let (output_low, output_high) = range_ratio.output_bounds(output_len);
|
||||||
if !range_ratio.is_fixed() {
|
if !range_ratio.is_fixed() {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Sampling input_len from [{input_low}, {input_high}] and \
|
input_low,
|
||||||
output_len from [{output_low}, {output_high}]"
|
input_high,
|
||||||
|
output_low,
|
||||||
|
output_high,
|
||||||
|
"sampling random request lengths"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -305,7 +308,8 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_generate_random_dataset_token_ids() {
|
fn test_generate_random_dataset_token_ids() {
|
||||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
let tokenizer =
|
||||||
|
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||||
let requests = generate_random_dataset(
|
let requests = generate_random_dataset(
|
||||||
&tokenizer,
|
&tokenizer,
|
||||||
10, // num_requests
|
10, // num_requests
|
||||||
@@ -337,7 +341,8 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_generate_random_dataset_text() {
|
fn test_generate_random_dataset_text() {
|
||||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
let tokenizer =
|
||||||
|
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||||
let requests = generate_random_dataset(
|
let requests = generate_random_dataset(
|
||||||
&tokenizer,
|
&tokenizer,
|
||||||
10, // num_requests
|
10, // num_requests
|
||||||
@@ -371,7 +376,8 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_token_length_exact_local() {
|
fn test_token_length_exact_local() {
|
||||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
let tokenizer =
|
||||||
|
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||||
let target_len = 512;
|
let target_len = 512;
|
||||||
let requests = generate_random_dataset(
|
let requests = generate_random_dataset(
|
||||||
&tokenizer,
|
&tokenizer,
|
||||||
@@ -405,11 +411,11 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
|
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_token_length_exact_tiktoken() {
|
async fn test_token_length_exact_tiktoken() {
|
||||||
// Use Qwen2.5 which has a tiktoken-format tokenizer
|
// Use Qwen2.5 which has a tiktoken-format tokenizer
|
||||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||||
let tokenizer = match tokenizer {
|
let tokenizer = match tokenizer {
|
||||||
Ok(t) => t,
|
Ok(t) => t,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -453,10 +459,10 @@ mod tests {
|
|||||||
|
|
||||||
/// Test encode/decode roundtrip stability for tiktoken.
|
/// Test encode/decode roundtrip stability for tiktoken.
|
||||||
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
|
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
|
||||||
#[test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
fn test_tiktoken_roundtrip_stability() {
|
async fn test_tiktoken_roundtrip_stability() {
|
||||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||||
let tokenizer = match tokenizer {
|
let tokenizer = match tokenizer {
|
||||||
Ok(t) => t,
|
Ok(t) => t,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
|||||||
@@ -139,8 +139,10 @@ mod tests {
|
|||||||
|
|
||||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||||
fn test_tokenizer() -> TokenizerKind {
|
fn test_tokenizer() -> TokenizerKind {
|
||||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
TokenizerKind::Tiktoken(
|
||||||
.expect("gpt2 built-in tiktoken should always load without network")
|
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||||
|
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn fixed_ratio() -> RangeRatio {
|
fn fixed_ratio() -> RangeRatio {
|
||||||
|
|||||||
@@ -22,18 +22,21 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
|
|||||||
|
|
||||||
/// Download the default ShareGPT dataset from HuggingFace Hub.
|
/// Download the default ShareGPT dataset from HuggingFace Hub.
|
||||||
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
|
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
|
||||||
pub fn download_sharegpt_dataset() -> Result<String> {
|
pub async fn download_sharegpt_dataset() -> Result<String> {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
|
repository = DEFAULT_SHAREGPT_REPO,
|
||||||
|
file = DEFAULT_SHAREGPT_FILE,
|
||||||
|
"downloading ShareGPT dataset"
|
||||||
);
|
);
|
||||||
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
|
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
|
||||||
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
|
.map_err(BenchError::Config)?;
|
||||||
|
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
|
||||||
BenchError::Config(format!(
|
BenchError::Config(format!(
|
||||||
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
|
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
|
||||||
))
|
))
|
||||||
})?;
|
})?;
|
||||||
let path_str = path.to_string_lossy().to_string();
|
let path_str = path.to_string_lossy().to_string();
|
||||||
println!("ShareGPT dataset ready: {path_str}");
|
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
|
||||||
Ok(path_str)
|
Ok(path_str)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset(
|
|||||||
// Oversample if dataset is smaller than requested
|
// Oversample if dataset is smaller than requested
|
||||||
if samples.len() < num_requests {
|
if samples.len() < num_requests {
|
||||||
if no_oversample {
|
if no_oversample {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
dataset = "sharegpt",
|
||||||
samples.len()
|
samples = samples.len(),
|
||||||
|
requested = num_requests,
|
||||||
|
"skipping dataset oversampling"
|
||||||
);
|
);
|
||||||
} else if !samples.is_empty() {
|
} else if !samples.is_empty() {
|
||||||
let needed = num_requests - samples.len();
|
let needed = num_requests - samples.len();
|
||||||
@@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset(
|
|||||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||||
samples.push(req);
|
samples.push(req);
|
||||||
}
|
}
|
||||||
println!(
|
tracing::info!(
|
||||||
"Oversampled requests from {original_len} to {} total samples.",
|
dataset = "sharegpt",
|
||||||
samples.len()
|
original_samples = original_len,
|
||||||
|
samples = samples.len(),
|
||||||
|
"oversampled dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
|
|||||||
use rand::{Rng, SeedableRng};
|
use rand::{Rng, SeedableRng};
|
||||||
|
|
||||||
use super::SampleRequest;
|
use super::SampleRequest;
|
||||||
|
use super::progress::RowDownloadReporter;
|
||||||
use crate::cli::SpeedBenchConfig;
|
use crate::cli::SpeedBenchConfig;
|
||||||
use crate::error::{BenchError, Result};
|
use crate::error::{BenchError, Result};
|
||||||
use crate::tokenizer::TokenizerKind;
|
use crate::tokenizer::TokenizerKind;
|
||||||
@@ -25,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf {
|
|||||||
|
|
||||||
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
|
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
|
||||||
/// Results are cached as JSON locally for subsequent runs.
|
/// Results are cached as JSON locally for subsequent runs.
|
||||||
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||||
let config_name = config.as_str();
|
let config_name = config.as_str();
|
||||||
|
|
||||||
let dir = cache_dir();
|
let dir = cache_dir();
|
||||||
@@ -35,13 +36,13 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
// Return cached file if it exists
|
// Return cached file if it exists
|
||||||
if cache_path.exists() {
|
if cache_path.exists() {
|
||||||
let path_str = cache_path.to_string_lossy().to_string();
|
let path_str = cache_path.to_string_lossy().to_string();
|
||||||
println!("SPEED-Bench ({config_name}) cached: {path_str}");
|
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
|
||||||
return Ok(path_str);
|
return Ok(path_str);
|
||||||
}
|
}
|
||||||
|
|
||||||
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
|
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
|
||||||
|
|
||||||
let client = reqwest::blocking::Client::builder()
|
let client = reqwest::Client::builder()
|
||||||
.timeout(std::time::Duration::from_secs(120))
|
.timeout(std::time::Duration::from_secs(120))
|
||||||
.build()
|
.build()
|
||||||
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
|
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
|
||||||
@@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||||
let mut offset = 0usize;
|
let mut offset = 0usize;
|
||||||
let page_size = 100usize;
|
let page_size = 100usize;
|
||||||
|
let mut progress = RowDownloadReporter::new();
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
let url = format!(
|
let url = format!(
|
||||||
@@ -64,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
let max_retries = 3;
|
let max_retries = 3;
|
||||||
let mut data: Option<serde_json::Value> = None;
|
let mut data: Option<serde_json::Value> = None;
|
||||||
for attempt in 0..=max_retries {
|
for attempt in 0..=max_retries {
|
||||||
let resp = match client.get(&url).send() {
|
let resp = match client.get(&url).send().await {
|
||||||
Ok(r) => r,
|
Ok(r) => r,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
if attempt < max_retries {
|
if attempt < max_retries {
|
||||||
std::thread::sleep(std::time::Duration::from_secs(
|
tokio::time::sleep(std::time::Duration::from_secs(
|
||||||
2 * (attempt as u64 + 1),
|
2 * (attempt as u64 + 1),
|
||||||
));
|
))
|
||||||
|
.await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
return Err(BenchError::Config(format!(
|
return Err(BenchError::Config(format!(
|
||||||
@@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
};
|
};
|
||||||
|
|
||||||
if resp.status().is_server_error() && attempt < max_retries {
|
if resp.status().is_server_error() && attempt < max_retries {
|
||||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
data = Some(resp.json().map_err(|e| {
|
data = Some(resp.json().await.map_err(|e| {
|
||||||
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
|
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
|
||||||
})?);
|
})?);
|
||||||
break;
|
break;
|
||||||
@@ -116,15 +119,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
let fetched = rows.len();
|
let fetched = rows.len();
|
||||||
offset += fetched;
|
offset += fetched;
|
||||||
|
|
||||||
// Print progress
|
|
||||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||||
eprint!("\r Fetched {offset}/{total} rows...");
|
progress.update(offset, total);
|
||||||
|
|
||||||
if fetched < page_size {
|
if fetched < page_size {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
eprintln!(); // newline after progress
|
progress.finish();
|
||||||
|
|
||||||
if all_rows.is_empty() {
|
if all_rows.is_empty() {
|
||||||
return Err(BenchError::Config(
|
return Err(BenchError::Config(
|
||||||
@@ -137,9 +139,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
|||||||
std::fs::write(&cache_path, &json_str)?;
|
std::fs::write(&cache_path, &json_str)?;
|
||||||
|
|
||||||
let path_str = cache_path.to_string_lossy().to_string();
|
let path_str = cache_path.to_string_lossy().to_string();
|
||||||
println!(
|
tracing::info!(
|
||||||
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
|
config = config_name,
|
||||||
all_rows.len()
|
rows = all_rows.len(),
|
||||||
|
path = %path_str,
|
||||||
|
"saved SPEED-Bench dataset"
|
||||||
);
|
);
|
||||||
Ok(path_str)
|
Ok(path_str)
|
||||||
}
|
}
|
||||||
@@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset(
|
|||||||
// Oversample if needed
|
// Oversample if needed
|
||||||
if samples.len() < num_requests {
|
if samples.len() < num_requests {
|
||||||
if no_oversample {
|
if no_oversample {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
dataset = "speed-bench",
|
||||||
samples.len()
|
samples = samples.len(),
|
||||||
|
requested = num_requests,
|
||||||
|
"skipping dataset oversampling"
|
||||||
);
|
);
|
||||||
} else if !samples.is_empty() {
|
} else if !samples.is_empty() {
|
||||||
let original_len = samples.len();
|
let original_len = samples.len();
|
||||||
@@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset(
|
|||||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||||
samples.push(req);
|
samples.push(req);
|
||||||
}
|
}
|
||||||
println!(
|
tracing::info!(
|
||||||
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
|
dataset = "speed-bench",
|
||||||
samples.len()
|
original_samples = original_len,
|
||||||
|
samples = samples.len(),
|
||||||
|
"oversampled dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -288,7 +296,6 @@ pub fn load_speed_bench_dataset(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Print category distribution
|
|
||||||
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
|
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
|
||||||
for entry in &filtered[..filtered.len().min(samples.len())] {
|
for entry in &filtered[..filtered.len().min(samples.len())] {
|
||||||
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
|
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
|
||||||
@@ -297,7 +304,7 @@ pub fn load_speed_bench_dataset(
|
|||||||
let mut cats: Vec<_> = cat_counts.into_iter().collect();
|
let mut cats: Vec<_> = cat_counts.into_iter().collect();
|
||||||
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
|
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
|
||||||
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
|
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
|
||||||
println!("SPEED-Bench categories: {}", cat_str.join(", "));
|
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
|
||||||
|
|
||||||
Ok(samples)
|
Ok(samples)
|
||||||
}
|
}
|
||||||
|
|||||||
+20
-35
@@ -1,54 +1,39 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
//! Sync facade over the async `hf_hub` API.
|
|
||||||
//!
|
|
||||||
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
|
|
||||||
//! backend unconditionally pulls ureq's default rustls feature. So we use the
|
|
||||||
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
|
|
||||||
//! loaders, tokenizer fallback in rayon threads) by running each download on a
|
|
||||||
//! dedicated thread with its own single-threaded runtime.
|
|
||||||
|
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
use hf_hub::Repo;
|
||||||
|
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
|
||||||
|
|
||||||
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
|
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
|
||||||
pub struct HubRepo {
|
pub struct HubRepo {
|
||||||
repo: hf_hub::Repo,
|
repo: ApiRepo,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl HubRepo {
|
impl HubRepo {
|
||||||
pub fn model(model_id: String) -> Self {
|
pub fn model(model_id: String) -> Result<Self, String> {
|
||||||
Self {
|
Self::new(Repo::model(model_id))
|
||||||
repo: hf_hub::Repo::model(model_id),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn dataset(repo_id: String) -> Self {
|
pub fn dataset(repo_id: String) -> Result<Self, String> {
|
||||||
Self {
|
Self::new(Repo::dataset(repo_id))
|
||||||
repo: hf_hub::Repo::dataset(repo_id),
|
}
|
||||||
|
|
||||||
|
fn new(repo: Repo) -> Result<Self, String> {
|
||||||
|
let mut builder = ApiBuilder::from_env();
|
||||||
|
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||||
|
builder = builder.with_token(Some(token));
|
||||||
}
|
}
|
||||||
|
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||||
|
Ok(Self {
|
||||||
|
repo: api.repo(repo),
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Download (or fetch from cache) a single file from the repo.
|
/// Download (or fetch from cache) a single file from the repo.
|
||||||
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
|
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
|
||||||
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||||
let repo = self.repo.clone();
|
self.repo.get(filename).await.map_err(|e| format!("{e}"))
|
||||||
let filename = filename.to_string();
|
|
||||||
std::thread::spawn(move || {
|
|
||||||
let rt = tokio::runtime::Builder::new_current_thread()
|
|
||||||
.enable_all()
|
|
||||||
.build()
|
|
||||||
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
|
|
||||||
rt.block_on(async move {
|
|
||||||
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
|
|
||||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
|
||||||
builder = builder.with_token(Some(token));
|
|
||||||
}
|
|
||||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
|
||||||
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.join()
|
|
||||||
.map_err(|_| "HF Hub download thread panicked".to_string())?
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ pub fn prepare_process() {
|
|||||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||||
&& new > 1024
|
&& new > 1024
|
||||||
{
|
{
|
||||||
eprintln!("Open-file limit: {new}");
|
tracing::info!(soft_limit = new, "raised open-file limit");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -19,7 +19,19 @@ struct Cli {
|
|||||||
args: vllm_bench::BenchServeArgs,
|
args: vllm_bench::BenchServeArgs,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TODO: unify the tracing subscriber used by different binaries.
|
||||||
|
fn init_tracing() {
|
||||||
|
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||||
|
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||||
|
let _ = tracing_subscriber::fmt()
|
||||||
|
.with_env_filter(filter)
|
||||||
|
.with_writer(std::io::stderr)
|
||||||
|
.try_init();
|
||||||
|
}
|
||||||
|
|
||||||
fn main() -> anyhow::Result<()> {
|
fn main() -> anyhow::Result<()> {
|
||||||
|
init_tracing();
|
||||||
|
|
||||||
let cli = Cli::parse();
|
let cli = Cli::parse();
|
||||||
vllm_bench::prepare_process();
|
vllm_bench::prepare_process();
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,22 @@ use crate::datasets::SampleRequest;
|
|||||||
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
|
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
|
||||||
use crate::multi_turn::ConversationOutput;
|
use crate::multi_turn::ConversationOutput;
|
||||||
|
|
||||||
|
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
|
||||||
|
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
|
||||||
|
if failed_outputs.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::warn!(
|
||||||
|
failed_requests = failed_outputs.len(),
|
||||||
|
displayed_errors = failed_outputs.len().min(10),
|
||||||
|
"benchmark requests failed"
|
||||||
|
);
|
||||||
|
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
|
||||||
|
tracing::warn!(index, error = %output.error, "benchmark request failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Calculate benchmark metrics from request outputs.
|
/// Calculate benchmark metrics from request outputs.
|
||||||
///
|
///
|
||||||
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
|
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
|
||||||
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
|
|||||||
|
|
||||||
let failed = outputs.len() - completed;
|
let failed = outputs.len() - completed;
|
||||||
|
|
||||||
// Print failed request errors (capped to 10)
|
log_failed_requests(outputs);
|
||||||
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
|
|
||||||
if !failed_outputs.is_empty() {
|
|
||||||
eprintln!("Failed requests during benchmark run detected (capping to 10):");
|
|
||||||
for (i, err) in failed_outputs.iter().take(10).enumerate() {
|
|
||||||
eprintln!("Error {i}: {}", err.error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Calculate max output tokens per second and max concurrent requests
|
// Calculate max output tokens per second and max concurrent requests
|
||||||
let mut max_output_tokens_per_s = 0.0_f64;
|
let mut max_output_tokens_per_s = 0.0_f64;
|
||||||
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
|
|||||||
|
|
||||||
let failed = outputs.len() - completed;
|
let failed = outputs.len() - completed;
|
||||||
|
|
||||||
// Print failed request errors (capped to 10)
|
log_failed_requests(outputs);
|
||||||
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
|
|
||||||
if !failed_outputs.is_empty() {
|
|
||||||
eprintln!("Failed requests during benchmark run detected (capping to 10):");
|
|
||||||
for (i, err) in failed_outputs.iter().take(10).enumerate() {
|
|
||||||
eprintln!("Error {i}: {}", err.error);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Compute peak concurrent requests from start_time + latency windows
|
// Compute peak concurrent requests from start_time + latency windows
|
||||||
let successful_outputs: Vec<&RequestFuncOutput> =
|
let successful_outputs: Vec<&RequestFuncOutput> =
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
|||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
use indicatif::{ProgressBar, ProgressStyle};
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
|
use thiserror_ext::AsReport as _;
|
||||||
use tokio::sync::Semaphore;
|
use tokio::sync::Semaphore;
|
||||||
|
|
||||||
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
|
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
|
||||||
@@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||||
(m.clone(), config.model_name.clone())
|
(m.clone(), config.model_name.clone())
|
||||||
} else {
|
} else {
|
||||||
println!("Model not specified, fetching first model from server...");
|
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||||
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
|
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
|
||||||
println!("First model name: {name}, first model id: {id}");
|
tracing::info!(
|
||||||
|
model_name = name,
|
||||||
|
model_id = id,
|
||||||
|
"selected first model from server"
|
||||||
|
);
|
||||||
(id, Some(name))
|
(id, Some(name))
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -83,15 +88,19 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||||
println!("Loading tokenizer: {tid}");
|
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||||
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
||||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
let t =
|
||||||
println!("Tokenizer loaded successfully.");
|
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
|
||||||
Some(t)
|
Some(t)
|
||||||
};
|
};
|
||||||
|
|
||||||
// Generate/load conversations
|
// Generate/load conversations
|
||||||
println!("Generating multi-turn conversations...");
|
tracing::info!(
|
||||||
|
dataset = ?config.dataset_name,
|
||||||
|
conversations = config.num_prompts,
|
||||||
|
"generating multi-turn conversations"
|
||||||
|
);
|
||||||
let gen_start = Instant::now();
|
let gen_start = Instant::now();
|
||||||
|
|
||||||
let mut conversations = match config.dataset_name {
|
let mut conversations = match config.dataset_name {
|
||||||
@@ -131,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
let path = match config.dataset_path.as_deref() {
|
let path = match config.dataset_path.as_deref() {
|
||||||
Some(p) => p,
|
Some(p) => p,
|
||||||
None => {
|
None => {
|
||||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||||
downloaded.as_str()
|
downloaded.as_str()
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -179,8 +188,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
let (filtered_conversations, filtered_turns) =
|
let (filtered_conversations, filtered_turns) =
|
||||||
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
|
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
|
||||||
if filtered_turns > 0 || filtered_conversations > 0 {
|
if filtered_turns > 0 || filtered_conversations > 0 {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
|
filtered_turns,
|
||||||
|
filtered_conversations,
|
||||||
|
max_model_len,
|
||||||
|
"filtered conversations above maximum model length"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
if conversations.is_empty() {
|
if conversations.is_empty() {
|
||||||
@@ -192,11 +204,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
|
|
||||||
let gen_elapsed = gen_start.elapsed();
|
let gen_elapsed = gen_start.elapsed();
|
||||||
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
|
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
|
||||||
println!(
|
tracing::info!(
|
||||||
"Generated {} conversations ({} total turns) in {:.2}s",
|
conversations = conversations.len(),
|
||||||
conversations.len(),
|
|
||||||
total_turns,
|
total_turns,
|
||||||
gen_elapsed.as_secs_f64()
|
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||||
|
"generated multi-turn conversations"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Log prefix sharing info
|
// Log prefix sharing info
|
||||||
@@ -208,18 +220,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
let conv_tokens =
|
let conv_tokens =
|
||||||
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
|
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
|
||||||
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
|
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
|
||||||
println!(
|
tracing::info!(
|
||||||
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
|
global_ratio = config.multi_turn_prefix_global_ratio,
|
||||||
config.multi_turn_prefix_global_ratio * 100.0,
|
|
||||||
global_tokens,
|
global_tokens,
|
||||||
config.multi_turn_prefix_conversation_ratio * 100.0,
|
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
|
||||||
conv_tokens,
|
conversation_tokens = conv_tokens,
|
||||||
(1.0 - config.multi_turn_prefix_global_ratio
|
|
||||||
- config.multi_turn_prefix_conversation_ratio)
|
|
||||||
* 100.0,
|
|
||||||
unique_tokens,
|
unique_tokens,
|
||||||
|
history_accumulation = false,
|
||||||
|
"configured multi-turn prefix sharing"
|
||||||
);
|
);
|
||||||
println!("No history accumulation: each turn sends fixed-length prompt only.");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if config.dry_run {
|
if config.dry_run {
|
||||||
@@ -253,7 +262,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
..Default::default()
|
..Default::default()
|
||||||
};
|
};
|
||||||
|
|
||||||
println!("Starting initial single prompt test run...");
|
tracing::info!("starting initial single-prompt test run");
|
||||||
let test_output = crate::ready_checker::wait_for_endpoint(
|
let test_output = crate::ready_checker::wait_for_endpoint(
|
||||||
config.backend,
|
config.backend,
|
||||||
&client,
|
&client,
|
||||||
@@ -268,7 +277,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
test_output.error
|
test_output.error
|
||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
println!("Initial test run completed.");
|
tracing::info!("initial single-prompt test run completed");
|
||||||
}
|
}
|
||||||
|
|
||||||
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
|
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
|
||||||
@@ -283,9 +292,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
"min_tokens".to_string(),
|
"min_tokens".to_string(),
|
||||||
serde_json::json!(config.random_output_len),
|
serde_json::json!(config.random_output_len),
|
||||||
);
|
);
|
||||||
println!(
|
tracing::info!(
|
||||||
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
|
min_tokens = config.random_output_len,
|
||||||
config.random_output_len
|
dataset = "random",
|
||||||
|
"set minimum output tokens for multi-turn dataset"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Some(body)
|
Some(body)
|
||||||
@@ -297,7 +307,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
let spec_decode_before =
|
let spec_decode_before =
|
||||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||||
if spec_decode_before.is_some() {
|
if spec_decode_before.is_some() {
|
||||||
println!("Speculative decoding detected, will collect metrics.");
|
tracing::info!("detected speculative decoding; collecting metrics");
|
||||||
}
|
}
|
||||||
|
|
||||||
// Start profiler if requested (immediate mode — no batch threshold)
|
// Start profiler if requested (immediate mode — no batch threshold)
|
||||||
@@ -330,10 +340,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Main benchmark
|
// Main benchmark
|
||||||
println!("Starting multi-turn benchmark...");
|
tracing::info!(
|
||||||
println!("Conversations: {}", conversations.len());
|
conversations = conversations.len(),
|
||||||
println!("Concurrency: {concurrency}");
|
total_turns,
|
||||||
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
|
concurrency,
|
||||||
|
inter_turn_delay_ms = config.multi_turn_delay_ms,
|
||||||
|
"starting multi-turn benchmark"
|
||||||
|
);
|
||||||
|
|
||||||
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
|
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
|
||||||
|
|
||||||
@@ -364,11 +377,12 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
);
|
);
|
||||||
if let Some(modules) = config.lora_modules.as_ref() {
|
if let Some(modules) = config.lora_modules.as_ref() {
|
||||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||||
println!(
|
tracing::info!(
|
||||||
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
|
adapters = modules.len(),
|
||||||
modules.len(),
|
names = ?names,
|
||||||
names,
|
assignment = ?config.lora_assignment,
|
||||||
config.lora_assignment
|
scope = "conversation",
|
||||||
|
"assigned LoRA adapters"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
match handle.await {
|
match handle.await {
|
||||||
Ok(output) => all_outputs.push(output),
|
Ok(output) => all_outputs.push(output),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
eprintln!("Conversation task panicked: {e}");
|
tracing::error!(error = %e.as_report(), "conversation task panicked");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -453,7 +467,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
|||||||
if let Some((cancel_tx, task)) = profile_task {
|
if let Some((cancel_tx, task)) = profile_task {
|
||||||
let _ = cancel_tx.send(());
|
let _ = cancel_tx.send(());
|
||||||
if let Err(e) = task.await {
|
if let Err(e) = task.await {
|
||||||
eprintln!("WARNING: Profile background task failed: {e}");
|
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -603,7 +603,7 @@ fn add_metric_stats(
|
|||||||
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
|
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
|
||||||
let content = serde_json::to_string(json)?;
|
let content = serde_json::to_string(json)?;
|
||||||
std::fs::write(file_path, content)?;
|
std::fs::write(file_path, content)?;
|
||||||
println!("Results saved to {file_path}");
|
tracing::info!(path = file_path, "saved benchmark results");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
|
|||||||
file.write_all(b"\n")?;
|
file.write_all(b"\n")?;
|
||||||
}
|
}
|
||||||
file.write_all(content.as_bytes())?;
|
file.write_all(content.as_bytes())?;
|
||||||
println!("Results appended to {file_path}");
|
tracing::info!(path = file_path, "appended benchmark results");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
|
|||||||
let backend = get_backend(backend)?;
|
let backend = get_backend(backend)?;
|
||||||
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
|
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
|
||||||
|
|
||||||
println!("Waiting for endpoint to become up in {timeout_seconds}s");
|
tracing::info!(
|
||||||
|
timeout_seconds,
|
||||||
|
retry_interval,
|
||||||
|
"waiting for endpoint readiness"
|
||||||
|
);
|
||||||
|
|
||||||
let pb = ProgressBar::new(timeout_seconds);
|
let pb = ProgressBar::new(timeout_seconds);
|
||||||
pb.set_style(
|
pb.set_style(
|
||||||
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
|
|||||||
Ok(output) => {
|
Ok(output) => {
|
||||||
let err = output.error.clone();
|
let err = output.error.clone();
|
||||||
let err_last_line = err.lines().last().unwrap_or(&err);
|
let err_last_line = err.lines().last().unwrap_or(&err);
|
||||||
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
|
pb.suspend(|| {
|
||||||
|
tracing::warn!(error = err_last_line, "endpoint is not ready");
|
||||||
|
});
|
||||||
last_error = err;
|
last_error = err;
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
|
|||||||
.await
|
.await
|
||||||
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
|
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
|
||||||
if resp.status().is_success() {
|
if resp.status().is_success() {
|
||||||
println!("Prefix cache reset successfully.");
|
tracing::info!(url = %url, "reset prefix cache");
|
||||||
} else {
|
} else {
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let body = resp.text().await.unwrap_or_default();
|
let body = resp.text().await.unwrap_or_default();
|
||||||
|
|||||||
@@ -199,12 +199,17 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
|
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
|
||||||
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
|
tracing::info!(
|
||||||
|
encoding,
|
||||||
|
kind = "built-in-tiktoken",
|
||||||
|
vocab_size,
|
||||||
|
"loaded tokenizer"
|
||||||
|
);
|
||||||
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
|
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
|
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
|
||||||
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||||
// Phase 1: If model_id is a local directory, look for tiktoken files there
|
// Phase 1: If model_id is a local directory, look for tiktoken files there
|
||||||
let local_dir = Path::new(model_id);
|
let local_dir = Path::new(model_id);
|
||||||
if local_dir.is_dir() {
|
if local_dir.is_dir() {
|
||||||
@@ -212,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Phase 2: Fall back to HuggingFace Hub download
|
// Phase 2: Fall back to HuggingFace Hub download
|
||||||
try_load_tiktoken_from_hf(model_id)
|
try_load_tiktoken_from_hf(model_id).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Common tiktoken model filenames to search for.
|
/// Common tiktoken model filenames to search for.
|
||||||
@@ -247,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Load a tiktoken tokenizer from a HuggingFace model repo.
|
/// Load a tiktoken tokenizer from a HuggingFace model repo.
|
||||||
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||||
|
|
||||||
let model_path = repo
|
let mut model_path = None;
|
||||||
.get("tiktoken.model")
|
for filename in TIKTOKEN_MODEL_FILENAMES {
|
||||||
.or_else(|_| repo.get("qwen.tiktoken"))
|
if let Ok(path) = repo.get(filename).await {
|
||||||
.or_else(|_| repo.get("vocab.tiktoken"))
|
model_path = Some(path);
|
||||||
.map_err(|_| {
|
break;
|
||||||
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
}
|
||||||
})?;
|
}
|
||||||
|
let model_path = model_path.ok_or_else(|| {
|
||||||
|
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
||||||
|
})?;
|
||||||
|
|
||||||
let num_base_tokens = count_base_tokens(&model_path)?;
|
let num_base_tokens = count_base_tokens(&model_path)?;
|
||||||
|
|
||||||
let config = match repo.get("tokenizer_config.json") {
|
let config = match repo.get("tokenizer_config.json").await {
|
||||||
Ok(config_path) => read_tokenizer_config(&config_path),
|
Ok(config_path) => read_tokenizer_config(&config_path),
|
||||||
Err(_) => None,
|
Err(_) => None,
|
||||||
};
|
};
|
||||||
|
|
||||||
let pattern = extract_pat_str_from_repo(&repo);
|
let pattern = extract_pat_str_from_repo(&repo).await;
|
||||||
|
|
||||||
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
|
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
|
||||||
}
|
}
|
||||||
@@ -309,15 +317,16 @@ fn build_tiktoken(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
println!(
|
tracing::info!(
|
||||||
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
|
model = model_id,
|
||||||
num_base_tokens,
|
base_tokens = num_base_tokens,
|
||||||
all_special_tokens.len(),
|
special_tokens = all_special_tokens.len(),
|
||||||
if pattern.is_some() {
|
pattern = if pattern.is_some() {
|
||||||
"custom"
|
"custom"
|
||||||
} else {
|
} else {
|
||||||
"default"
|
"default"
|
||||||
},
|
},
|
||||||
|
"loading tiktoken model"
|
||||||
);
|
);
|
||||||
|
|
||||||
TiktokenTokenizer::from_file(
|
TiktokenTokenizer::from_file(
|
||||||
@@ -397,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
|
|||||||
|
|
||||||
/// Try to download the Python tokenizer source file and extract pat_str via regex.
|
/// Try to download the Python tokenizer source file and extract pat_str via regex.
|
||||||
/// Returns None if unavailable or unparsable.
|
/// Returns None if unavailable or unparsable.
|
||||||
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
||||||
// Try common Python tokenizer filenames
|
// Try common Python tokenizer filenames
|
||||||
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
|
let py_path = match repo.get("tokenization_kimi.py").await {
|
||||||
|
Ok(path) => path,
|
||||||
|
Err(_) => repo.get("tokenizer.py").await.ok()?,
|
||||||
|
};
|
||||||
|
|
||||||
let source = std::fs::read_to_string(&py_path).ok()?;
|
let source = std::fs::read_to_string(&py_path).ok()?;
|
||||||
|
|
||||||
@@ -438,9 +450,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
|
|||||||
|
|
||||||
if !fragments.is_empty() {
|
if !fragments.is_empty() {
|
||||||
let pattern = fragments.join("|");
|
let pattern = fragments.join("|");
|
||||||
println!(
|
tracing::debug!(
|
||||||
"Extracted pat_str from Python source: {} fragments",
|
fragments = fragments.len(),
|
||||||
fragments.len()
|
"extracted tiktoken pattern from Python source"
|
||||||
);
|
);
|
||||||
return Some(pattern);
|
return Some(pattern);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,8 +2,10 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
|
use std::future::Future;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
|
use thiserror_ext::AsReport as _;
|
||||||
use tokenizers::Tokenizer;
|
use tokenizers::Tokenizer;
|
||||||
|
|
||||||
use crate::error::{BenchError, Result};
|
use crate::error::{BenchError, Result};
|
||||||
@@ -18,7 +20,8 @@ pub enum TokenizerKind {
|
|||||||
|
|
||||||
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
|
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
|
||||||
pub struct ServerTokenizer {
|
pub struct ServerTokenizer {
|
||||||
client: reqwest::blocking::Client,
|
client: reqwest::Client,
|
||||||
|
runtime: tokio::runtime::Handle,
|
||||||
tokenize_url: String,
|
tokenize_url: String,
|
||||||
detokenize_url: String,
|
detokenize_url: String,
|
||||||
model: String,
|
model: String,
|
||||||
@@ -27,8 +30,8 @@ pub struct ServerTokenizer {
|
|||||||
|
|
||||||
impl ServerTokenizer {
|
impl ServerTokenizer {
|
||||||
/// Create a new server tokenizer and verify connectivity.
|
/// Create a new server tokenizer and verify connectivity.
|
||||||
pub fn new(base_url: &str, model: &str) -> Result<Self> {
|
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
|
||||||
let client = reqwest::blocking::Client::builder()
|
let client = reqwest::Client::builder()
|
||||||
.timeout(std::time::Duration::from_secs(30))
|
.timeout(std::time::Duration::from_secs(30))
|
||||||
.build()
|
.build()
|
||||||
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
|
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
|
||||||
@@ -38,6 +41,7 @@ impl ServerTokenizer {
|
|||||||
|
|
||||||
let st = Self {
|
let st = Self {
|
||||||
client,
|
client,
|
||||||
|
runtime: tokio::runtime::Handle::current(),
|
||||||
tokenize_url,
|
tokenize_url,
|
||||||
detokenize_url,
|
detokenize_url,
|
||||||
model: model.to_string(),
|
model: model.to_string(),
|
||||||
@@ -45,7 +49,7 @@ impl ServerTokenizer {
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Probe the endpoint to verify it works and discover vocab size
|
// Probe the endpoint to verify it works and discover vocab size
|
||||||
let test_tokens = st.encode_inner("test")?;
|
let test_tokens = st.encode_async("test").await?;
|
||||||
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
|
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
|
||||||
let estimated_vocab = (max_id * 2).max(131072);
|
let estimated_vocab = (max_id * 2).max(131072);
|
||||||
|
|
||||||
@@ -56,6 +60,10 @@ impl ServerTokenizer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
|
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
|
||||||
|
self.block_on(self.encode_async(text))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
|
||||||
let payload = serde_json::json!({
|
let payload = serde_json::json!({
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"prompt": text,
|
"prompt": text,
|
||||||
@@ -66,6 +74,7 @@ impl ServerTokenizer {
|
|||||||
.post(&self.tokenize_url)
|
.post(&self.tokenize_url)
|
||||||
.json(&payload)
|
.json(&payload)
|
||||||
.send()
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
|
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
|
||||||
|
|
||||||
if !resp.status().is_success() {
|
if !resp.status().is_success() {
|
||||||
@@ -75,7 +84,7 @@ impl ServerTokenizer {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||||
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
|
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
@@ -95,6 +104,10 @@ impl ServerTokenizer {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn decode_inner(&self, ids: &[u32]) -> Result<String> {
|
fn decode_inner(&self, ids: &[u32]) -> Result<String> {
|
||||||
|
self.block_on(self.decode_async(ids))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
|
||||||
let payload = serde_json::json!({
|
let payload = serde_json::json!({
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"tokens": ids,
|
"tokens": ids,
|
||||||
@@ -105,6 +118,7 @@ impl ServerTokenizer {
|
|||||||
.post(&self.detokenize_url)
|
.post(&self.detokenize_url)
|
||||||
.json(&payload)
|
.json(&payload)
|
||||||
.send()
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
|
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
|
||||||
|
|
||||||
if !resp.status().is_success() {
|
if !resp.status().is_success() {
|
||||||
@@ -114,7 +128,7 @@ impl ServerTokenizer {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
|
|
||||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||||
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
|
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
@@ -123,6 +137,26 @@ impl ServerTokenizer {
|
|||||||
.map(|s| s.to_string())
|
.map(|s| s.to_string())
|
||||||
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
|
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
|
||||||
|
if matches!(
|
||||||
|
self.runtime.runtime_flavor(),
|
||||||
|
tokio::runtime::RuntimeFlavor::CurrentThread
|
||||||
|
) {
|
||||||
|
return Err(BenchError::Tokenizer(
|
||||||
|
"Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
|
||||||
|
// Tokio workers must enter a blocking region before re-entering the runtime;
|
||||||
|
// Rayon workers can drive the future directly with the saved runtime handle.
|
||||||
|
if tokio::runtime::Handle::try_current().is_ok() {
|
||||||
|
tokio::task::block_in_place(|| self.runtime.block_on(future))
|
||||||
|
} else {
|
||||||
|
self.runtime.block_on(future)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- TokenizerKind methods ---
|
// --- TokenizerKind methods ---
|
||||||
@@ -192,7 +226,7 @@ impl TokenizerKind {
|
|||||||
/// 3. Server-side /tokenize + /detokenize endpoints
|
/// 3. Server-side /tokenize + /detokenize endpoints
|
||||||
///
|
///
|
||||||
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
|
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
|
||||||
pub fn load_tokenizer(
|
pub async fn load_tokenizer(
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
_trust_remote_code: bool,
|
_trust_remote_code: bool,
|
||||||
server_info: Option<(&str, &str)>,
|
server_info: Option<(&str, &str)>,
|
||||||
@@ -212,31 +246,48 @@ pub fn load_tokenizer(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 1. Try local HuggingFace tokenizer (tokenizer.json)
|
// 1. Try local HuggingFace tokenizer (tokenizer.json)
|
||||||
match try_load_local(model_id) {
|
match try_load_local(model_id).await {
|
||||||
Ok(tok) => {
|
Ok(tok) => {
|
||||||
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
|
tracing::info!(
|
||||||
|
model = model_id,
|
||||||
|
kind = "local",
|
||||||
|
vocab_size = tok.get_vocab_size(true),
|
||||||
|
"loaded tokenizer"
|
||||||
|
);
|
||||||
Ok(TokenizerKind::Local(Box::new(tok)))
|
Ok(TokenizerKind::Local(Box::new(tok)))
|
||||||
}
|
}
|
||||||
Err(local_err) => {
|
Err(local_err) => {
|
||||||
// 2. Try tiktoken format
|
// 2. Try tiktoken format
|
||||||
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
|
tracing::info!(
|
||||||
match crate::tiktoken::try_load_tiktoken(model_id) {
|
model = model_id,
|
||||||
|
error = %local_err.as_report(),
|
||||||
|
"local tokenizer unavailable; trying tiktoken"
|
||||||
|
);
|
||||||
|
match crate::tiktoken::try_load_tiktoken(model_id).await {
|
||||||
Ok(tok) => {
|
Ok(tok) => {
|
||||||
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
|
tracing::info!(
|
||||||
|
model = model_id,
|
||||||
|
kind = "tiktoken",
|
||||||
|
vocab_size = tok.vocab_size(),
|
||||||
|
"loaded tokenizer"
|
||||||
|
);
|
||||||
Ok(TokenizerKind::Tiktoken(tok))
|
Ok(TokenizerKind::Tiktoken(tok))
|
||||||
}
|
}
|
||||||
Err(tiktoken_err) => {
|
Err(tiktoken_err) => {
|
||||||
// 3. Try server-side fallback
|
// 3. Try server-side fallback
|
||||||
if let Some((base_url, model)) = server_info {
|
if let Some((base_url, model)) = server_info {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Tiktoken also not available ({tiktoken_err}), \
|
model = model_id,
|
||||||
trying server-side tokenization..."
|
error = %tiktoken_err.as_report(),
|
||||||
|
"tiktoken unavailable; trying server-side tokenization"
|
||||||
);
|
);
|
||||||
match ServerTokenizer::new(base_url, model) {
|
match ServerTokenizer::new(base_url, model).await {
|
||||||
Ok(srv) => {
|
Ok(srv) => {
|
||||||
println!(
|
tracing::info!(
|
||||||
"Tokenizer: Server (vocab_size≈{})",
|
model = model_id,
|
||||||
srv.cached_vocab_size
|
kind = "server",
|
||||||
|
vocab_size = srv.cached_vocab_size,
|
||||||
|
"loaded tokenizer"
|
||||||
);
|
);
|
||||||
return Ok(TokenizerKind::Server(srv));
|
return Ok(TokenizerKind::Server(srv));
|
||||||
}
|
}
|
||||||
@@ -264,7 +315,7 @@ pub fn load_tokenizer(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Try loading tokenizer.json from local path or HuggingFace Hub.
|
/// Try loading tokenizer.json from local path or HuggingFace Hub.
|
||||||
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||||
// 1. Try local directory with tokenizer.json
|
// 1. Try local directory with tokenizer.json
|
||||||
let local_path = Path::new(model_id).join("tokenizer.json");
|
let local_path = Path::new(model_id).join("tokenizer.json");
|
||||||
if local_path.exists() {
|
if local_path.exists() {
|
||||||
@@ -290,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
|
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
|
||||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||||
let tokenizer_path = repo
|
let tokenizer_path = repo
|
||||||
.get("tokenizer.json")
|
.get("tokenizer.json")
|
||||||
|
.await
|
||||||
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
|
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
|
||||||
|
|
||||||
Tokenizer::from_file(&tokenizer_path)
|
Tokenizer::from_file(&tokenizer_path)
|
||||||
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
|
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
|
async fn test_server_tokenizer_sync_bridge() {
|
||||||
|
let tokenizer = std::sync::Arc::new(ServerTokenizer {
|
||||||
|
client: reqwest::Client::new(),
|
||||||
|
runtime: tokio::runtime::Handle::current(),
|
||||||
|
tokenize_url: String::new(),
|
||||||
|
detokenize_url: String::new(),
|
||||||
|
model: String::new(),
|
||||||
|
cached_vocab_size: 0,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
|
||||||
|
|
||||||
|
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||||
|
rayon::spawn(move || {
|
||||||
|
let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
|
||||||
|
});
|
||||||
|
assert_eq!(rx.await.unwrap().unwrap(), 2);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -26,18 +26,21 @@ const RESET: &str = "\x1b[0m";
|
|||||||
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
|
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
|
||||||
format_description!("[month]-[day] [hour]:[minute]:[second]");
|
format_description!("[month]-[day] [hour]:[minute]:[second]");
|
||||||
|
|
||||||
const PROCESS_LABEL: &str = "RustFrontend";
|
|
||||||
|
|
||||||
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
|
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
|
||||||
pub(crate) fn init_tracing() {
|
pub(crate) fn init_tracing(process_label: &str) {
|
||||||
let filter = build_targets_filter(
|
let filter = build_targets_filter(
|
||||||
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
|
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
|
||||||
env::var("RUST_LOG").ok().as_deref(),
|
env::var("RUST_LOG").ok().as_deref(),
|
||||||
);
|
);
|
||||||
let formatter = VllmEventFormatter::new();
|
let formatter = VllmEventFormatter::new(process_label);
|
||||||
|
|
||||||
let _ = tracing_subscriber::registry()
|
let _ = tracing_subscriber::registry()
|
||||||
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
|
.with(
|
||||||
|
tracing_subscriber::fmt::layer()
|
||||||
|
.event_format(formatter)
|
||||||
|
.with_writer(std::io::stderr)
|
||||||
|
.with_filter(filter),
|
||||||
|
)
|
||||||
.try_init();
|
.try_init();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
|
|||||||
}
|
}
|
||||||
|
|
||||||
impl VllmEventFormatter {
|
impl VllmEventFormatter {
|
||||||
fn new() -> Self {
|
fn new(process_label: &str) -> Self {
|
||||||
Self {
|
Self {
|
||||||
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
|
prefix: format!("({process_label} pid={})", process::id()),
|
||||||
timer: VllmLocalTimer::default(),
|
timer: VllmLocalTimer::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -291,6 +294,13 @@ fn map_python_log_level(level: &str) -> LevelFilter {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn formatter_prefix_uses_process_label() {
|
||||||
|
let formatter = VllmEventFormatter::new("Bench");
|
||||||
|
|
||||||
|
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
|
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
|
||||||
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
|
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ mod cli;
|
|||||||
mod logging;
|
mod logging;
|
||||||
|
|
||||||
use std::env;
|
use std::env;
|
||||||
|
use std::ffi::OsStr;
|
||||||
use std::process::ExitStatus;
|
use std::process::ExitStatus;
|
||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn main() -> Result<()> {
|
fn main() -> Result<()> {
|
||||||
logging::init_tracing();
|
let process_label =
|
||||||
|
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
|
||||||
|
"bench" => "Bench",
|
||||||
|
"serve" | "frontend" => "RustFrontend",
|
||||||
|
_ => "Rust",
|
||||||
|
};
|
||||||
|
logging::init_tracing(process_label);
|
||||||
|
|
||||||
let cli = Cli::parse();
|
let cli = Cli::parse();
|
||||||
|
|
||||||
let mut runtime = tokio::runtime::Builder::new_multi_thread();
|
let mut runtime = tokio::runtime::Builder::new_multi_thread();
|
||||||
|
|||||||
@@ -460,6 +460,12 @@ impl EngineCoreClient {
|
|||||||
self.inner.is_healthy()
|
self.inner.is_healthy()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Subscribe to engine health changes. The current value is `true` while
|
||||||
|
/// the client is healthy and changes permanently to `false` on failure.
|
||||||
|
pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
|
||||||
|
self.inner.subscribe_health()
|
||||||
|
}
|
||||||
|
|
||||||
/// Return the first persistent health error observed by the client, if any.
|
/// Return the first persistent health error observed by the client, if any.
|
||||||
pub fn health_error(&self) -> Option<Arc<Error>> {
|
pub fn health_error(&self) -> Option<Arc<Error>> {
|
||||||
self.inner.health_error()
|
self.inner.health_error()
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
|
|||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
use thiserror_ext::AsReport as _;
|
use thiserror_ext::AsReport as _;
|
||||||
use tokio::runtime::Handle;
|
use tokio::runtime::Handle;
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::{mpsc, watch};
|
||||||
use tracing::{debug, info, trace, warn};
|
use tracing::{debug, info, trace, warn};
|
||||||
use vllm_metrics::METRICS;
|
use vllm_metrics::METRICS;
|
||||||
use zeromq::RouterSendHalf;
|
use zeromq::RouterSendHalf;
|
||||||
@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {
|
|||||||
request_reg: Mutex<RequestRegistry>,
|
request_reg: Mutex<RequestRegistry>,
|
||||||
utility_reg: Mutex<UtilityRegistry>,
|
utility_reg: Mutex<UtilityRegistry>,
|
||||||
health_error: ArcSwapOption<Error>,
|
health_error: ArcSwapOption<Error>,
|
||||||
|
health_tx: watch::Sender<bool>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ClientInner {
|
impl ClientInner {
|
||||||
@@ -57,6 +58,7 @@ impl ClientInner {
|
|||||||
request_reg: Mutex::new(RequestRegistry::new(engines)),
|
request_reg: Mutex::new(RequestRegistry::new(engines)),
|
||||||
utility_reg: Mutex::new(UtilityRegistry::default()),
|
utility_reg: Mutex::new(UtilityRegistry::default()),
|
||||||
health_error: ArcSwapOption::empty(),
|
health_error: ArcSwapOption::empty(),
|
||||||
|
health_tx: watch::Sender::new(true),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -169,6 +171,7 @@ impl ClientInner {
|
|||||||
/// persistent health error.
|
/// persistent health error.
|
||||||
pub fn close_registries(&self, error: Arc<Error>) {
|
pub fn close_registries(&self, error: Arc<Error>) {
|
||||||
let persistent_error = self.record_health_error(error);
|
let persistent_error = self.record_health_error(error);
|
||||||
|
self.publish_unhealthy();
|
||||||
let request_senders = self.request_reg.lock().close();
|
let request_senders = self.request_reg.lock().close();
|
||||||
let utility_senders = self.utility_reg.lock().close();
|
let utility_senders = self.utility_reg.lock().close();
|
||||||
|
|
||||||
@@ -191,6 +194,12 @@ impl ClientInner {
|
|||||||
self.health_error.load().is_none()
|
self.health_error.load().is_none()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Subscribe to engine health changes. The current value is `true` while
|
||||||
|
/// the client is healthy and changes permanently to `false` on failure.
|
||||||
|
pub fn subscribe_health(&self) -> watch::Receiver<bool> {
|
||||||
|
self.health_tx.subscribe()
|
||||||
|
}
|
||||||
|
|
||||||
/// Resolve one utility output to the waiting caller. Returns `true` if a
|
/// Resolve one utility output to the waiting caller. Returns `true` if a
|
||||||
/// waiting caller existed.
|
/// waiting caller existed.
|
||||||
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
|
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
|
||||||
@@ -280,6 +289,11 @@ impl ClientInner {
|
|||||||
.expect("health error must be recorded before registries close")
|
.expect("health error must be recorded before registries close")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Publish the sticky healthy-to-unhealthy transition.
|
||||||
|
fn publish_unhealthy(&self) {
|
||||||
|
self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
|
||||||
|
}
|
||||||
|
|
||||||
/// Assert there is a recorded health error and return a `Shared` variant
|
/// Assert there is a recorded health error and return a `Shared` variant
|
||||||
/// wrapping it for error returns when the client is already closed.
|
/// wrapping it for error returns when the client is already closed.
|
||||||
fn closed_error(&self) -> Error {
|
fn closed_error(&self) -> Error {
|
||||||
@@ -461,13 +475,18 @@ mod tests {
|
|||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||||
async fn close_registries_records_first_health_error_only() {
|
async fn close_registries_records_first_health_error_only() {
|
||||||
let inner = test_inner().await;
|
let inner = test_inner().await;
|
||||||
|
let mut health = inner.subscribe_health();
|
||||||
|
assert!(*health.borrow());
|
||||||
|
|
||||||
inner.close_registries(Arc::new(Error::EngineCoreDead));
|
inner.close_registries(Arc::new(Error::EngineCoreDead));
|
||||||
|
health.changed().await.expect("health sender remains open");
|
||||||
assert!(!inner.is_healthy());
|
assert!(!inner.is_healthy());
|
||||||
|
assert!(!*health.borrow());
|
||||||
assert!(matches!(
|
assert!(matches!(
|
||||||
inner.health_error().as_deref(),
|
inner.health_error().as_deref(),
|
||||||
Some(Error::EngineCoreDead)
|
Some(Error::EngineCoreDead)
|
||||||
));
|
));
|
||||||
|
assert!(!*inner.subscribe_health().borrow());
|
||||||
|
|
||||||
inner.close_registries(Arc::new(client_closed!("shutdown")));
|
inner.close_registries(Arc::new(client_closed!("shutdown")));
|
||||||
assert!(matches!(
|
assert!(matches!(
|
||||||
|
|||||||
@@ -269,6 +269,19 @@ impl WireLogprobs {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
|
||||||
|
if token_ids.rows == 0 {
|
||||||
|
return Ok(Logprobs {
|
||||||
|
positions: Vec::new(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if token_ids.cols == 0 {
|
||||||
|
bail_ext_value_decode!(
|
||||||
|
"{field_prefix}: zero-column logprobs payload with {} rows",
|
||||||
|
token_ids.rows
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let mut positions = Vec::with_capacity(token_ids.rows);
|
let mut positions = Vec::with_capacity(token_ids.rows);
|
||||||
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
|
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
|
||||||
.data
|
.data
|
||||||
|
|||||||
@@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() {
|
|||||||
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
|
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn decodes_zero_row_logprobs_as_empty() {
|
||||||
|
for shape in [[0usize, 0], [0, 3]] {
|
||||||
|
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
|
||||||
|
None,
|
||||||
|
Some(Value::Array(vec![
|
||||||
|
ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
|
||||||
|
ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
|
||||||
|
ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
|
||||||
|
Value::Nil,
|
||||||
|
])),
|
||||||
|
)))];
|
||||||
|
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||||
|
let logprobs = decoded.outputs[0]
|
||||||
|
.new_prompt_logprobs_tensors
|
||||||
|
.clone()
|
||||||
|
.unwrap()
|
||||||
|
.into_direct()
|
||||||
|
.unwrap();
|
||||||
|
assert!(logprobs.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn rejects_zero_column_logprobs_with_rows() {
|
||||||
|
let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
|
||||||
|
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
|
||||||
|
Some(Value::Array(vec![
|
||||||
|
ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
|
||||||
|
ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
|
||||||
|
ndarray_value("<i8", &[2], ranks),
|
||||||
|
Value::Nil,
|
||||||
|
])),
|
||||||
|
None,
|
||||||
|
)))];
|
||||||
|
|
||||||
|
let error = decode_engine_core_outputs(&frames).unwrap_err();
|
||||||
|
let crate::error::Error::ExtValueDecode { message } = &error else {
|
||||||
|
panic!("expected ExtValueDecode");
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
message,
|
||||||
|
"new_logprobs: zero-column logprobs payload with 2 rows"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ tokio-openssl.workspace = true
|
|||||||
tokio-stream.workspace = true
|
tokio-stream.workspace = true
|
||||||
tokio-util.workspace = true
|
tokio-util.workspace = true
|
||||||
tonic.workspace = true
|
tonic.workspace = true
|
||||||
|
tonic-health.workspace = true
|
||||||
tonic-prost.workspace = true
|
tonic-prost.workspace = true
|
||||||
tower.workspace = true
|
tower.workspace = true
|
||||||
tower-http.workspace = true
|
tower-http.workspace = true
|
||||||
|
|||||||
@@ -0,0 +1,70 @@
|
|||||||
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
|
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
use tokio::sync::watch;
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
|
use tonic::server::NamedService;
|
||||||
|
use tonic_health::ServingStatus;
|
||||||
|
use tonic_health::server::HealthReporter;
|
||||||
|
use tracing::{info, warn};
|
||||||
|
|
||||||
|
use super::GenerateGrpcService;
|
||||||
|
|
||||||
|
pub(crate) async fn monitor_health(
|
||||||
|
mut health_reporter: HealthReporter,
|
||||||
|
mut engine_health: watch::Receiver<bool>,
|
||||||
|
shutdown: CancellationToken,
|
||||||
|
) {
|
||||||
|
let generate_service = GenerateGrpcService::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"
|
||||||
|
),
|
||||||
|
}
|
||||||
|
true
|
||||||
|
}
|
||||||
|
_ = shutdown.cancelled() => {
|
||||||
|
info!(
|
||||||
|
generate_service,
|
||||||
|
overall_service = true,
|
||||||
|
status = ?status,
|
||||||
|
reason = "server_shutdown",
|
||||||
|
"server shutting down; marking gRPC health services as not serving"
|
||||||
|
);
|
||||||
|
false
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
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_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("").await;
|
||||||
|
}
|
||||||
@@ -4,6 +4,7 @@
|
|||||||
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
|
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
|
||||||
|
|
||||||
mod convert;
|
mod convert;
|
||||||
|
mod health;
|
||||||
|
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -24,8 +25,11 @@ pub mod pb {
|
|||||||
tonic::include_proto!("vllm");
|
tonic::include_proto!("vllm");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) use health::monitor_health;
|
||||||
pub use pb::generate_server::GenerateServer;
|
pub use pb::generate_server::GenerateServer;
|
||||||
|
|
||||||
|
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests;
|
mod tests;
|
||||||
|
|
||||||
|
|||||||
@@ -16,6 +16,10 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
|||||||
use tokio::net::TcpStream;
|
use tokio::net::TcpStream;
|
||||||
use tokio_openssl::SslStream;
|
use tokio_openssl::SslStream;
|
||||||
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
|
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
|
||||||
|
use tonic_health::pb::HealthCheckRequest;
|
||||||
|
use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
|
||||||
|
use tonic_health::pb::health_client::HealthClient;
|
||||||
|
use tonic_health::server::health_reporter;
|
||||||
use tower::service_fn;
|
use tower::service_fn;
|
||||||
use vllm_chat::{
|
use vllm_chat::{
|
||||||
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
|
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
|
||||||
@@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend {
|
|||||||
async fn setup_grpc_service(
|
async fn setup_grpc_service(
|
||||||
engine_id: impl Into<EngineId>,
|
engine_id: impl Into<EngineId>,
|
||||||
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
||||||
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
|
) -> (
|
||||||
|
GenerateServer<GenerateServiceImpl>,
|
||||||
|
tokio::sync::watch::Receiver<bool>,
|
||||||
|
MockEngineTask,
|
||||||
|
) {
|
||||||
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
||||||
let handshake_address = ipc.handshake_endpoint();
|
let handshake_address = ipc.handshake_endpoint();
|
||||||
let engine_id = engine_id.into();
|
let engine_id = engine_id.into();
|
||||||
@@ -232,6 +240,7 @@ async fn setup_grpc_service(
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.expect("connect client");
|
.expect("connect client");
|
||||||
|
let engine_health = client.subscribe_health();
|
||||||
|
|
||||||
let chat = ChatLlm::from_shared_backend(
|
let chat = ChatLlm::from_shared_backend(
|
||||||
test_llm(client),
|
test_llm(client),
|
||||||
@@ -240,6 +249,7 @@ async fn setup_grpc_service(
|
|||||||
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
|
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
|
||||||
(
|
(
|
||||||
GenerateServer::new(GenerateServiceImpl::new(state)),
|
GenerateServer::new(GenerateServiceImpl::new(state)),
|
||||||
|
engine_health,
|
||||||
engine_task,
|
engine_task,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -254,25 +264,51 @@ async fn grpc_test_server(
|
|||||||
tokio::task::JoinHandle<()>,
|
tokio::task::JoinHandle<()>,
|
||||||
MockEngineTask,
|
MockEngineTask,
|
||||||
) {
|
) {
|
||||||
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||||
|
let (channel, server_task) = start_grpc_test_server(
|
||||||
|
svc,
|
||||||
|
engine_health,
|
||||||
|
tokio_util::sync::CancellationToken::new(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
(GenerateClient::new(channel), server_task, engine_task)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn start_grpc_test_server(
|
||||||
|
generate_service: GenerateServer<GenerateServiceImpl>,
|
||||||
|
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;
|
||||||
|
|
||||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||||
let addr = listener.local_addr().expect("local addr");
|
let addr = listener.local_addr().expect("local addr");
|
||||||
|
|
||||||
let server_task = tokio::spawn(async move {
|
let server_task = tokio::spawn(async move {
|
||||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||||
TonicServer::builder()
|
let server = TonicServer::builder()
|
||||||
.add_service(svc)
|
.add_service(health_service)
|
||||||
.serve_with_incoming(incoming)
|
.add_service(generate_service)
|
||||||
.await
|
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||||
.expect("grpc server");
|
let health_monitor =
|
||||||
|
super::monitor_health(health_reporter, engine_health, shutdown.clone());
|
||||||
|
let server = async move {
|
||||||
|
let result = server.await;
|
||||||
|
shutdown.cancel();
|
||||||
|
result
|
||||||
|
};
|
||||||
|
let (server_result, ()) = tokio::join!(server, health_monitor);
|
||||||
|
server_result.expect("grpc server");
|
||||||
});
|
});
|
||||||
|
|
||||||
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
|
let channel = Endpoint::from_shared(format!("http://{addr}"))
|
||||||
|
.expect("grpc endpoint")
|
||||||
|
.connect()
|
||||||
.await
|
.await
|
||||||
.expect("connect grpc client");
|
.expect("connect grpc channel");
|
||||||
|
|
||||||
(grpc_client, server_task, engine_task)
|
(channel, server_task)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
|
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
|
||||||
@@ -283,7 +319,7 @@ async fn grpc_tls_test_server(
|
|||||||
certs: &TestCerts,
|
certs: &TestCerts,
|
||||||
cert_reqs: i32,
|
cert_reqs: i32,
|
||||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||||
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||||
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
|
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
|
||||||
.expect("build grpc tls config");
|
.expect("build grpc tls config");
|
||||||
|
|
||||||
@@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive(
|
|||||||
engine_id: impl Into<EngineId>,
|
engine_id: impl Into<EngineId>,
|
||||||
keepalive: Option<Duration>,
|
keepalive: Option<Duration>,
|
||||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||||
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
|
let (svc, _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");
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||||
let addr = listener.local_addr().expect("local addr").to_string();
|
let addr = listener.local_addr().expect("local addr").to_string();
|
||||||
@@ -1035,3 +1072,129 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
|
|||||||
|
|
||||||
server_task.abort();
|
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) =
|
||||||
|
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,
|
||||||
|
engine_health,
|
||||||
|
tokio_util::sync::CancellationToken::new(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let mut health_client = HealthClient::new(channel);
|
||||||
|
|
||||||
|
let mut health_streams = Vec::new();
|
||||||
|
for service in ["vllm.Generate", ""] {
|
||||||
|
let service_label = if service.is_empty() {
|
||||||
|
"overall"
|
||||||
|
} else {
|
||||||
|
service
|
||||||
|
};
|
||||||
|
let mut stream = health_client
|
||||||
|
.watch(HealthCheckRequest {
|
||||||
|
service: service.to_string(),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|error| {
|
||||||
|
panic!("failed to start health watch for {service_label}: {error}")
|
||||||
|
})
|
||||||
|
.into_inner();
|
||||||
|
let initial = stream
|
||||||
|
.message()
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|error| {
|
||||||
|
panic!("failed to read initial health status for {service_label}: {error}")
|
||||||
|
})
|
||||||
|
.unwrap_or_else(|| {
|
||||||
|
panic!("health watch for {service_label} ended before its initial status")
|
||||||
|
});
|
||||||
|
assert_eq!(
|
||||||
|
initial.status,
|
||||||
|
HealthServingStatus::Serving as i32,
|
||||||
|
"unexpected initial health status for {service_label}"
|
||||||
|
);
|
||||||
|
health_streams.push((service_label, stream));
|
||||||
|
}
|
||||||
|
|
||||||
|
engine_health_tx.send(false).expect("publish unhealthy engine state");
|
||||||
|
|
||||||
|
for (service_label, mut stream) in health_streams {
|
||||||
|
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))
|
||||||
|
.unwrap_or_else(|error| {
|
||||||
|
panic!("failed to read health update for {service_label}: {error}")
|
||||||
|
})
|
||||||
|
.unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));
|
||||||
|
assert_eq!(
|
||||||
|
update.status,
|
||||||
|
HealthServingStatus::NotServing as i32,
|
||||||
|
"unexpected health status for {service_label}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
server_task.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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(
|
||||||
|
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 mut health_client = HealthClient::new(channel);
|
||||||
|
let mut stream = health_client
|
||||||
|
.watch(HealthCheckRequest {
|
||||||
|
service: "vllm.Generate".to_string(),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("start health watch for vllm.Generate")
|
||||||
|
.into_inner();
|
||||||
|
|
||||||
|
let initial = stream
|
||||||
|
.message()
|
||||||
|
.await
|
||||||
|
.expect("read initial health status for vllm.Generate")
|
||||||
|
.expect("health watch ended before its initial status");
|
||||||
|
assert_eq!(
|
||||||
|
initial.status,
|
||||||
|
HealthServingStatus::Serving as i32,
|
||||||
|
"unexpected initial health status for vllm.Generate"
|
||||||
|
);
|
||||||
|
|
||||||
|
shutdown.cancel();
|
||||||
|
|
||||||
|
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for shutdown health update for vllm.Generate")
|
||||||
|
.expect("failed to read shutdown health update for vllm.Generate")
|
||||||
|
.expect("health watch ended before its shutdown update");
|
||||||
|
assert_eq!(
|
||||||
|
update.status,
|
||||||
|
HealthServingStatus::NotServing as i32,
|
||||||
|
"unexpected shutdown health status for vllm.Generate"
|
||||||
|
);
|
||||||
|
|
||||||
|
let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for vllm.Generate health watch to close")
|
||||||
|
.expect("failed while closing vllm.Generate health watch");
|
||||||
|
assert!(
|
||||||
|
stream_end.is_none(),
|
||||||
|
"vllm.Generate health watch remained open"
|
||||||
|
);
|
||||||
|
|
||||||
|
tokio::time::timeout(Duration::from_secs(2), server_task)
|
||||||
|
.await
|
||||||
|
.expect("timed out waiting for gRPC server shutdown")
|
||||||
|
.expect("gRPC server task failed");
|
||||||
|
}
|
||||||
|
|||||||
+28
-14
@@ -39,6 +39,7 @@ use tokio::net::TcpListener;
|
|||||||
use tokio::time::{Instant, sleep_until};
|
use tokio::time::{Instant, sleep_until};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
use tonic::transport::Server as TonicServer;
|
use tonic::transport::Server as TonicServer;
|
||||||
|
use tonic_health::server::health_reporter;
|
||||||
use tower::ServiceExt as _;
|
use tower::ServiceExt as _;
|
||||||
use tracing::{info, trace, warn};
|
use tracing::{info, trace, warn};
|
||||||
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
|
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
|
||||||
@@ -203,14 +204,19 @@ where
|
|||||||
.map(tls::build_grpc_server_config)
|
.map(tls::build_grpc_server_config)
|
||||||
.transpose()
|
.transpose()
|
||||||
.context("invalid gRPC TLS configuration")?;
|
.context("invalid gRPC TLS configuration")?;
|
||||||
let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
|
let (health_reporter, health_service) = health_reporter();
|
||||||
|
let engine_health = state.engine_core_client().subscribe_health();
|
||||||
|
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
|
||||||
|
let generate_service =
|
||||||
|
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
|
||||||
let svc = TonicServer::builder()
|
let svc = TonicServer::builder()
|
||||||
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
|
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
|
||||||
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
|
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
|
||||||
.layer(middleware::request_runtime_layer(state.clone()))
|
.layer(middleware::request_runtime_layer(state.clone()))
|
||||||
.add_service(svc);
|
.add_service(health_service)
|
||||||
|
.add_service(generate_service);
|
||||||
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
|
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
|
||||||
Some((grpc_listener, svc, grpc_tls))
|
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
@@ -294,7 +300,8 @@ where
|
|||||||
let server_shutdown = server_shutdown.clone();
|
let server_shutdown = server_shutdown.clone();
|
||||||
let force_shutdown = force_shutdown.clone();
|
let force_shutdown = force_shutdown.clone();
|
||||||
async move {
|
async move {
|
||||||
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
|
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
|
||||||
|
else {
|
||||||
// No gRPC configured: just wait for shutdown so we do not race the
|
// No gRPC configured: just wait for shutdown so we do not race the
|
||||||
// join! by resolving early and tripping the cancellation token.
|
// join! by resolving early and tripping the cancellation token.
|
||||||
shutdown.cancelled().await;
|
shutdown.cancelled().await;
|
||||||
@@ -304,19 +311,26 @@ where
|
|||||||
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
|
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
|
||||||
None => MaybeTlsListener::plain(grpc_listener),
|
None => MaybeTlsListener::plain(grpc_listener),
|
||||||
};
|
};
|
||||||
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
|
let server =
|
||||||
|
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||||
|
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
|
||||||
|
|
||||||
let result = tokio::select! {
|
let server = async move {
|
||||||
result = server => {
|
let result = tokio::select! {
|
||||||
result.context("gRPC server failed")
|
result = server => {
|
||||||
}
|
result.context("gRPC server failed")
|
||||||
_ = force_shutdown.cancelled() => {
|
}
|
||||||
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
_ = force_shutdown.cancelled() => {
|
||||||
Ok(())
|
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
||||||
}
|
Ok(())
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
server_shutdown.cancel();
|
||||||
|
result
|
||||||
};
|
};
|
||||||
|
|
||||||
server_shutdown.cancel();
|
let (result, ()) = tokio::join!(server, health_monitor);
|
||||||
result
|
result
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -66,7 +66,6 @@ def test_worker_apply_lora(qwen3_lora_files):
|
|||||||
runner_type="generate",
|
runner_type="generate",
|
||||||
max_num_batched_tokens=32,
|
max_num_batched_tokens=32,
|
||||||
max_num_seqs=32,
|
max_num_seqs=32,
|
||||||
max_num_partial_prefills=32,
|
|
||||||
),
|
),
|
||||||
device_config=DeviceConfig(DEVICE_TYPE),
|
device_config=DeviceConfig(DEVICE_TYPE),
|
||||||
cache_config=CacheConfig(
|
cache_config=CacheConfig(
|
||||||
|
|||||||
@@ -70,7 +70,19 @@ def _assert_video_outputs(processor, processed) -> None:
|
|||||||
merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
|
merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
|
||||||
expected_tokens = int(grid_thw.prod()) // merge_size**2
|
expected_tokens = int(grid_thw.prod()) // merge_size**2
|
||||||
video_token_id = processor.info.get_hf_config().video_token_id
|
video_token_id = processor.info.get_hf_config().video_token_id
|
||||||
assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
|
prompt_token_ids = processed["prompt_token_ids"]
|
||||||
|
assert prompt_token_ids.count(video_token_id) == expected_tokens
|
||||||
|
|
||||||
|
hf_processor = processor.info.get_hf_processor()
|
||||||
|
expected_frame_wrappers = int(grid_thw[:, 0].sum())
|
||||||
|
assert (
|
||||||
|
prompt_token_ids.count(hf_processor.vision_start_token_id)
|
||||||
|
== expected_frame_wrappers
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
prompt_token_ids.count(hf_processor.vision_end_token_id)
|
||||||
|
== expected_frame_wrappers
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("num_images", [1, 2])
|
@pytest.mark.parametrize("num_images", [1, 2])
|
||||||
|
|||||||
@@ -170,6 +170,7 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
|
|||||||
"layers.0.self_attn.to_add_out.weight",
|
"layers.0.self_attn.to_add_out.weight",
|
||||||
"layers.0.self_attn.norm_added_q.weight",
|
"layers.0.self_attn.norm_added_q.weight",
|
||||||
"layers.0.self_attn.norm_added_k.weight",
|
"layers.0.self_attn.norm_added_k.weight",
|
||||||
|
"layers.0.self_attn.k_norm_und_for_gen.weight",
|
||||||
"layers.0.self_attn.q_proj_moe_gen.weight",
|
"layers.0.self_attn.q_proj_moe_gen.weight",
|
||||||
"layers.0.mlp_moe_gen.up_proj.weight",
|
"layers.0.mlp_moe_gen.up_proj.weight",
|
||||||
"norm_moe_gen.weight",
|
"norm_moe_gen.weight",
|
||||||
|
|||||||
@@ -304,6 +304,13 @@ def test_pynvvideocodec_decoder_slot_retains_simple_decoder():
|
|||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|
||||||
|
|
||||||
|
def test_cosmos3_edge_uses_qwen3_vl_video_backend():
|
||||||
|
backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor")
|
||||||
|
|
||||||
|
assert backend == "qwen3_vl"
|
||||||
|
assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"model_repo, expected_loader_cls, hf_sample_kwargs",
|
"model_repo, expected_loader_cls, hf_sample_kwargs",
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -48,12 +48,12 @@ def _check_dense_embedding(data, index=0):
|
|||||||
def _check_sparse_embedding(data, check_tokens=False):
|
def _check_sparse_embedding(data, check_tokens=False):
|
||||||
expected_weights = [
|
expected_weights = [
|
||||||
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
|
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
|
||||||
{"token_id": 70, "weight": 0.09808349609375, "token": "the"},
|
{"token_id": 70, "weight": 0.09808349609375, "token": " the"},
|
||||||
{"token_id": 83, "weight": 0.08154296875, "token": "is"},
|
{"token_id": 83, "weight": 0.08154296875, "token": " is"},
|
||||||
{"token_id": 111, "weight": 0.11810302734375, "token": "of"},
|
{"token_id": 111, "weight": 0.11810302734375, "token": " of"},
|
||||||
{"token_id": 4865, "weight": 0.1171875, "token": "What"},
|
{"token_id": 4865, "weight": 0.1171875, "token": " What"},
|
||||||
{"token_id": 9942, "weight": 0.292236328125, "token": "France"},
|
{"token_id": 9942, "weight": 0.292236328125, "token": " France"},
|
||||||
{"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
|
{"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
|
||||||
]
|
]
|
||||||
expected_embed = {x["token_id"]: x for x in expected_weights}
|
expected_embed = {x["token_id"]: x for x in expected_weights}
|
||||||
|
|
||||||
|
|||||||
@@ -93,9 +93,6 @@ def test_online_quantization(
|
|||||||
use_rocm_aiter: bool,
|
use_rocm_aiter: bool,
|
||||||
monkeypatch,
|
monkeypatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
|
|
||||||
pytest.skip("FA3 currently rejects FP8 KV cache output dtype on SM90")
|
|
||||||
|
|
||||||
if use_rocm_aiter:
|
if use_rocm_aiter:
|
||||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
|
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||||
|
|
||||||
@@ -105,9 +102,15 @@ def test_online_quantization(
|
|||||||
if force_marlin:
|
if force_marlin:
|
||||||
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
|
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
|
||||||
|
|
||||||
|
model_dtype = "auto"
|
||||||
|
if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
|
||||||
|
# FA3 requires BF16 output when the query input is FP8.
|
||||||
|
model_dtype = "bfloat16"
|
||||||
|
|
||||||
with vllm_runner(
|
with vllm_runner(
|
||||||
"facebook/opt-125m",
|
"facebook/opt-125m",
|
||||||
quantization="fp8",
|
quantization="fp8",
|
||||||
|
dtype=model_dtype,
|
||||||
enforce_eager=True,
|
enforce_eager=True,
|
||||||
kv_cache_dtype=kv_cache_dtype,
|
kv_cache_dtype=kv_cache_dtype,
|
||||||
) as llm:
|
) as llm:
|
||||||
|
|||||||
@@ -83,6 +83,11 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
|
|||||||
ENFORCE_EAGER=${ENFORCE_EAGER:-1}
|
ENFORCE_EAGER=${ENFORCE_EAGER:-1}
|
||||||
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
|
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
|
||||||
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
|
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
|
||||||
|
# Pin concurrent prefiller and non-DP decoder engines to separate internal
|
||||||
|
# port windows. DP decoder ranks retain their existing internal port selection.
|
||||||
|
PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000}
|
||||||
|
DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000}
|
||||||
|
INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100}
|
||||||
|
|
||||||
# Resolve the repository root from the script location instead of `.git`.
|
# Resolve the repository root from the script location instead of `.git`.
|
||||||
# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
|
# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
|
||||||
@@ -154,12 +159,14 @@ run_tests_for_model() {
|
|||||||
PORT=$((8100 + i))
|
PORT=$((8100 + i))
|
||||||
# Calculate side channel port. Avoid clash with with TP workers.
|
# Calculate side channel port. Avoid clash with with TP workers.
|
||||||
SIDE_CHANNEL_PORT=$((5559 + i))
|
SIDE_CHANNEL_PORT=$((5559 + i))
|
||||||
|
INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
|
||||||
|
|
||||||
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
|
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
|
||||||
|
|
||||||
# Build the command with or without model-specific args
|
# Build the command with or without model-specific args
|
||||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||||
VLLM_KV_CACHE_LAYOUT='HND' \
|
VLLM_KV_CACHE_LAYOUT='HND' \
|
||||||
|
VLLM_PORT=$INTERNAL_PORT \
|
||||||
UCX_NET_DEVICES=all \
|
UCX_NET_DEVICES=all \
|
||||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||||
vllm serve $model_name \
|
vllm serve $model_name \
|
||||||
@@ -208,12 +215,18 @@ run_tests_for_model() {
|
|||||||
PORT=$((8200 + i))
|
PORT=$((8200 + i))
|
||||||
# Calculate side channel port
|
# Calculate side channel port
|
||||||
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
|
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
|
||||||
|
INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
|
||||||
|
DECODER_INTERNAL_PORT_ENV=
|
||||||
|
if [[ -z "${DP_EP:-}" ]]; then
|
||||||
|
DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT"
|
||||||
|
fi
|
||||||
|
|
||||||
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
|
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
|
||||||
|
|
||||||
# Build the command with or without model-specific args
|
# Build the command with or without model-specific args
|
||||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||||
VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
|
VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
|
||||||
|
$DECODER_INTERNAL_PORT_ENV \
|
||||||
UCX_NET_DEVICES=all \
|
UCX_NET_DEVICES=all \
|
||||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||||
vllm serve $model_name \
|
vllm serve $model_name \
|
||||||
|
|||||||
@@ -12,6 +12,9 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
|
|||||||
)
|
)
|
||||||
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
|
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
|
||||||
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
|
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
|
||||||
|
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
|
||||||
|
OffloadingConnectorMetadata,
|
||||||
|
)
|
||||||
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
|
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
|
||||||
OffloadingConnectorStats,
|
OffloadingConnectorStats,
|
||||||
_ConnectorMetricName,
|
_ConnectorMetricName,
|
||||||
@@ -109,6 +112,40 @@ def test_last_block_offloaded_at_request_finish(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||||
|
def test_abort_queued_request_does_not_build_store_job(
|
||||||
|
request_runner, async_scheduling: bool
|
||||||
|
):
|
||||||
|
"""Aborting a never-scheduled request must not store unallocated KV."""
|
||||||
|
block_size = 4
|
||||||
|
runner = request_runner(
|
||||||
|
block_size=block_size,
|
||||||
|
num_gpu_blocks=8,
|
||||||
|
async_scheduling=async_scheduling,
|
||||||
|
)
|
||||||
|
|
||||||
|
runner.new_request(token_ids=[0] * (block_size * 4))
|
||||||
|
runner.scheduler.schedule()
|
||||||
|
|
||||||
|
runner.new_request(token_ids=[1] * (block_size * 4))
|
||||||
|
queued_req_id = str(runner.req_id)
|
||||||
|
assert any(
|
||||||
|
request.request_id == queued_req_id for request in runner.scheduler.waiting
|
||||||
|
)
|
||||||
|
|
||||||
|
runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
|
||||||
|
req_status = runner.connector_scheduler._req_status[queued_req_id]
|
||||||
|
assert all(group_state.offload_keys for group_state in req_status.group_states)
|
||||||
|
assert all(not group_state.block_ids for group_state in req_status.group_states)
|
||||||
|
|
||||||
|
scheduler_output = runner.scheduler.schedule()
|
||||||
|
|
||||||
|
metadata = scheduler_output.kv_connector_metadata
|
||||||
|
assert isinstance(metadata, OffloadingConnectorMetadata)
|
||||||
|
assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
|
||||||
|
assert queued_req_id not in runner.connector_scheduler._req_status
|
||||||
|
|
||||||
|
|
||||||
def test_scheduler_reports_lookup_sync_delay(request_runner):
|
def test_scheduler_reports_lookup_sync_delay(request_runner):
|
||||||
runner = request_runner(
|
runner = request_runner(
|
||||||
block_size=4,
|
block_size=4,
|
||||||
|
|||||||
@@ -125,8 +125,7 @@ def test_register_kv_caches(backend):
|
|||||||
own dedicated tensors.
|
own dedicated tensors.
|
||||||
|
|
||||||
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
|
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
|
||||||
kv_caches, which automatically applies
|
the raw per-layer kv_caches registered by the connector.
|
||||||
_update_hybrid_attention_mamba_layout for hybrid models.
|
|
||||||
|
|
||||||
Verifies that the canonicalized CanonicalKVCaches has the correct
|
Verifies that the canonicalized CanonicalKVCaches has the correct
|
||||||
block tensors, tensor_idx references, and page sizes across all groups.
|
block tensors, tensor_idx references, and page sizes across all groups.
|
||||||
|
|||||||
@@ -5,13 +5,11 @@ import torch
|
|||||||
from torch import Generator
|
from torch import Generator
|
||||||
|
|
||||||
from tests.utils import large_gpu_mark
|
from tests.utils import large_gpu_mark
|
||||||
from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size
|
|
||||||
from vllm.platforms import current_platform
|
from vllm.platforms import current_platform
|
||||||
from vllm.triton_utils import HAS_TRITON
|
from vllm.triton_utils import HAS_TRITON
|
||||||
from vllm.utils.torch_utils import set_random_seed
|
from vllm.utils.torch_utils import set_random_seed
|
||||||
from vllm.v1.sample.ops.topk_topp_sampler import (
|
from vllm.v1.sample.ops.topk_topp_sampler import (
|
||||||
apply_top_k_top_p_pytorch,
|
apply_top_k_top_p_pytorch,
|
||||||
flashinfer_sample,
|
|
||||||
random_sample,
|
random_sample,
|
||||||
)
|
)
|
||||||
from vllm.v1.sample.sampler import Sampler
|
from vllm.v1.sample.sampler import Sampler
|
||||||
@@ -1045,39 +1043,3 @@ class TestFlashInferDistributionMatch:
|
|||||||
f"{label}: distribution differs from theoretical: "
|
f"{label}: distribution differs from theoretical: "
|
||||||
f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
|
f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.skipif(
|
|
||||||
not FLASHINFER_TOPK_TOPP_SUPPORTED,
|
|
||||||
reason="FlashInfer top-k/top-p sampler is not available on this platform.",
|
|
||||||
)
|
|
||||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
|
|
||||||
@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)])
|
|
||||||
def test_flashinfer_sample_padded_vocab(
|
|
||||||
dtype: torch.dtype, k: int | None, p: float | None
|
|
||||||
):
|
|
||||||
"""flashinfer_sample must accept the logits the sampler actually hands it.
|
|
||||||
|
|
||||||
compute_logits slices the padding off the vocab, so for a vocab that isn't a
|
|
||||||
multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model
|
|
||||||
dtype, while FlashInfer requires contiguous fp32.
|
|
||||||
"""
|
|
||||||
torch.set_default_device(DEVICE_TYPE)
|
|
||||||
batch_size = 8
|
|
||||||
org_vocab_size = 50272
|
|
||||||
padded_vocab_size = pad_vocab_size(org_vocab_size)
|
|
||||||
assert padded_vocab_size != org_vocab_size
|
|
||||||
|
|
||||||
logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[
|
|
||||||
..., :org_vocab_size
|
|
||||||
]
|
|
||||||
# A single row stays contiguous despite the padded stride, hence batch_size > 1.
|
|
||||||
assert not logits.is_contiguous()
|
|
||||||
|
|
||||||
token_ids = flashinfer_sample(
|
|
||||||
logits,
|
|
||||||
torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None,
|
|
||||||
torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None,
|
|
||||||
)
|
|
||||||
assert token_ids.shape == (batch_size,)
|
|
||||||
assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size))
|
|
||||||
|
|||||||
+20
-1
@@ -36,6 +36,22 @@ 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]:
|
def rust_py_extension_module_names() -> list[str]:
|
||||||
module_names = []
|
module_names = []
|
||||||
for extension in rust_extensions():
|
for extension in rust_extensions():
|
||||||
@@ -52,12 +68,15 @@ def rust_py_extension_module_names() -> list[str]:
|
|||||||
def build_binary(build_rust_args: list[str]) -> None:
|
def build_binary(build_rust_args: list[str]) -> None:
|
||||||
os.chdir(ROOT_DIR)
|
os.chdir(ROOT_DIR)
|
||||||
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
|
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
|
||||||
|
extensions = rust_extensions(optional=False)
|
||||||
setup(
|
setup(
|
||||||
name="vllm-rust-frontend-build",
|
name="vllm-rust-frontend-build",
|
||||||
packages=[],
|
packages=[],
|
||||||
rust_extensions=rust_extensions(optional=False),
|
rust_extensions=extensions,
|
||||||
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
|
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:
|
def main() -> None:
|
||||||
|
|||||||
@@ -67,16 +67,6 @@ class SchedulerConfig:
|
|||||||
In real usage, this should be set in `EngineArgs.create_engine_config`.
|
In real usage, this should be set in `EngineArgs.create_engine_config`.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
max_num_partial_prefills: int = Field(default=1, ge=1)
|
|
||||||
"""For chunked prefill, the maximum number of sequences that can be
|
|
||||||
partially prefilled concurrently."""
|
|
||||||
|
|
||||||
max_long_partial_prefills: int = Field(default=1, ge=1)
|
|
||||||
"""For chunked prefill, the maximum number of prompts longer than
|
|
||||||
long_prefill_token_threshold that will be prefilled concurrently. Setting
|
|
||||||
this less than max_num_partial_prefills will allow shorter prompts to jump
|
|
||||||
the queue in front of longer prompts in some cases, improving latency."""
|
|
||||||
|
|
||||||
long_prefill_token_threshold: int = Field(default=0, ge=0)
|
long_prefill_token_threshold: int = Field(default=0, ge=0)
|
||||||
"""For chunked prefill, a request is considered long if the prompt is
|
"""For chunked prefill, a request is considered long if the prompt is
|
||||||
longer than this number of tokens. 0 disables the cap (default)."""
|
longer than this number of tokens. 0 disables the cap (default)."""
|
||||||
@@ -254,19 +244,6 @@ class SchedulerConfig:
|
|||||||
self.max_num_batched_tokens,
|
self.max_num_batched_tokens,
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.max_num_partial_prefills > 1:
|
|
||||||
if self.long_prefill_token_threshold == 0:
|
|
||||||
self.long_prefill_token_threshold = int(max_model_len * 0.04)
|
|
||||||
|
|
||||||
logger.info(
|
|
||||||
"Concurrent partial prefills enabled with "
|
|
||||||
"max_num_partial_prefills=%d, max_long_partial_prefills=%d, "
|
|
||||||
"long_prefill_token_threshold=%d",
|
|
||||||
self.max_num_partial_prefills,
|
|
||||||
self.max_long_partial_prefills,
|
|
||||||
self.long_prefill_token_threshold,
|
|
||||||
)
|
|
||||||
|
|
||||||
self.verify_max_model_len(max_model_len)
|
self.verify_max_model_len(max_model_len)
|
||||||
|
|
||||||
def verify_max_model_len(self, max_model_len: int) -> Self:
|
def verify_max_model_len(self, max_model_len: int) -> Self:
|
||||||
@@ -298,24 +275,11 @@ class SchedulerConfig:
|
|||||||
self.max_num_seqs * max_model_len,
|
self.max_num_seqs * max_model_len,
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.max_num_partial_prefills > 1:
|
if self.long_prefill_token_threshold > max_model_len:
|
||||||
if not self.enable_chunked_prefill:
|
|
||||||
raise ValueError(
|
|
||||||
"Chunked prefill must be enabled to set "
|
|
||||||
"max_num_partial_prefills > 1."
|
|
||||||
)
|
|
||||||
|
|
||||||
if self.long_prefill_token_threshold > max_model_len:
|
|
||||||
raise ValueError(
|
|
||||||
"long_prefill_token_threshold "
|
|
||||||
f"({self.long_prefill_token_threshold}) cannot be greater "
|
|
||||||
f"than the max_model_len ({max_model_len})."
|
|
||||||
)
|
|
||||||
|
|
||||||
if self.max_long_partial_prefills > self.max_num_partial_prefills:
|
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"{self.max_long_partial_prefills=} must be less than or equal to "
|
"long_prefill_token_threshold "
|
||||||
f"{self.max_num_partial_prefills=}."
|
f"({self.long_prefill_token_threshold}) cannot be greater "
|
||||||
|
f"than the max_model_len ({max_model_len})."
|
||||||
)
|
)
|
||||||
|
|
||||||
return self
|
return self
|
||||||
|
|||||||
@@ -21,12 +21,10 @@ from vllm.logger import init_logger
|
|||||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||||
from vllm.platforms import current_platform
|
from vllm.platforms import current_platform
|
||||||
from vllm.v1.attention.backend import AttentionBackend
|
from vllm.v1.attention.backend import AttentionBackend
|
||||||
from vllm.v1.kv_cache_interface import MambaSpec
|
|
||||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
|
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
|
||||||
from vllm.v1.kv_cache_interface import KVCacheSpec
|
|
||||||
|
|
||||||
logger = init_logger(__name__)
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
@@ -594,32 +592,6 @@ class TransferTopology:
|
|||||||
abs_ratio = -tp_ratio
|
abs_ratio = -tp_ratio
|
||||||
return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
|
return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
|
||||||
|
|
||||||
def get_transfer_cache_regions(
|
|
||||||
self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
|
|
||||||
) -> list[torch.Tensor] | torch.Tensor:
|
|
||||||
"""Return the cache tensor(s) to register as NIXL memory regions,
|
|
||||||
also accounting for hybrid SSM models specificities.
|
|
||||||
"""
|
|
||||||
if isinstance(layer_spec, MambaSpec):
|
|
||||||
# Register the whole kv cache shared tensor, including
|
|
||||||
# SSM/Conv.
|
|
||||||
conv, ssm = cache
|
|
||||||
return [conv]
|
|
||||||
|
|
||||||
# Check may be hacky but it's matching
|
|
||||||
# `_update_hybrid_attention_mamba_layout`.
|
|
||||||
if self.is_mamba and cache.shape[0] == 2:
|
|
||||||
# When MAMBA is present, all backends are blocks first, so
|
|
||||||
# that blocks can be shared between attention layers and mamba
|
|
||||||
# layers. Runner already adjusted strides for FlashAttn-like
|
|
||||||
# backends so its num_blocks first.
|
|
||||||
# Swap [2<>num_blocks] dims for hybrid SSM layout.
|
|
||||||
cache = cache.transpose(0, 1)
|
|
||||||
|
|
||||||
# K and V are packed into one tensor (content dim), so each layer
|
|
||||||
# registers as a single region.
|
|
||||||
return [cache]
|
|
||||||
|
|
||||||
def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
|
def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
|
||||||
"""One-line summary of transfer config for logging."""
|
"""One-line summary of transfer config for logging."""
|
||||||
info = self._engines[(remote_engine_id, remote_pp_rank)]
|
info = self._engines[(remote_engine_id, remote_pp_rank)]
|
||||||
|
|||||||
@@ -1678,9 +1678,9 @@ class MooncakeConnectorWorker:
|
|||||||
conv, _ = cache_or_caches
|
conv, _ = cache_or_caches
|
||||||
cache_list = [conv]
|
cache_list = [conv]
|
||||||
else:
|
else:
|
||||||
cache_list = self.transfer_topo.get_transfer_cache_regions(
|
# K and V are packed into one blocks-first tensor per layer,
|
||||||
cache_or_caches, layer_spec
|
# so each layer registers as a single region.
|
||||||
)
|
cache_list = [cache_or_caches]
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"registering layer %s with %d cache tensor(s)",
|
"registering layer %s with %d cache tensor(s)",
|
||||||
|
|||||||
@@ -1092,7 +1092,7 @@ class NixlBaseConnectorWorker:
|
|||||||
# to better exploit the memory layout (ie num_blocks is the first dim).
|
# to better exploit the memory layout (ie num_blocks is the first dim).
|
||||||
tensor_size_bytes = None
|
tensor_size_bytes = None
|
||||||
|
|
||||||
for layer_name, cache_or_caches in xfer_buffers.items():
|
for layer_name, cache in xfer_buffers.items():
|
||||||
# NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
|
# NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
|
||||||
# that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
|
# that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
|
||||||
# However, physical page_size may differ when kernel requires a specific
|
# However, physical page_size may differ when kernel requires a specific
|
||||||
@@ -1109,9 +1109,6 @@ class NixlBaseConnectorWorker:
|
|||||||
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
|
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
|
||||||
# MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
|
# MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
|
||||||
layer_spec = layer_spec.kv_cache_specs[layer_name]
|
layer_spec = layer_spec.kv_cache_specs[layer_name]
|
||||||
cache_list = self.transfer_topo.get_transfer_cache_regions(
|
|
||||||
cache_or_caches, layer_spec
|
|
||||||
)
|
|
||||||
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
|
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
|
||||||
# the page_size assuming constant `self._logical_num_blocks`.
|
# the page_size assuming constant `self._logical_num_blocks`.
|
||||||
physical_page_size = (
|
physical_page_size = (
|
||||||
@@ -1120,8 +1117,6 @@ class NixlBaseConnectorWorker:
|
|||||||
else layer_spec.page_size_bytes
|
else layer_spec.page_size_bytes
|
||||||
// self._physical_blocks_per_logical_kv_block
|
// self._physical_blocks_per_logical_kv_block
|
||||||
)
|
)
|
||||||
# For when registering multiple tensors eg K/V in separate regions.
|
|
||||||
physical_page_size = physical_page_size // len(cache_list)
|
|
||||||
if self.transfer_topo._cross_layers_blocks:
|
if self.transfer_topo._cross_layers_blocks:
|
||||||
# When cross-layers blocks are used, multiply by number of layers
|
# When cross-layers blocks are used, multiply by number of layers
|
||||||
physical_page_size = physical_page_size * len(
|
physical_page_size = physical_page_size * len(
|
||||||
@@ -1136,66 +1131,61 @@ class NixlBaseConnectorWorker:
|
|||||||
# [`num_blocks` * `page_size`]
|
# [`num_blocks` * `page_size`]
|
||||||
curr_tensor_size_bytes = num_blocks * physical_page_size
|
curr_tensor_size_bytes = num_blocks * physical_page_size
|
||||||
|
|
||||||
# TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
|
base_addr = cache.data_ptr()
|
||||||
# registering a single tensor for both K/V and splitting logically like FI.
|
if base_addr in seen_base_addresses:
|
||||||
for cache in cache_list:
|
# NOTE (NickLucche) HMA employs memory pooling to share tensors
|
||||||
base_addr = cache.data_ptr()
|
# across groups. This results in skipping all tensors but the ones
|
||||||
if base_addr in seen_base_addresses:
|
# pointed to by group0. Also, generally we will have more blocks
|
||||||
# NOTE (NickLucche) HMA employs memory pooling to share tensors
|
# per tensor but fewer regions.
|
||||||
# across groups. This results in skipping all tensors but the ones
|
logger.debug("Skipping %s because it's already seen", layer_name)
|
||||||
# pointed to by group0. Also, generally we will have more blocks
|
continue
|
||||||
# per tensor but fewer regions.
|
logger.debug(
|
||||||
logger.debug("Skipping %s because it's already seen", layer_name)
|
"Registering layer %s with cache shape: %s", layer_name, cache.shape
|
||||||
continue
|
)
|
||||||
logger.debug(
|
seen_base_addresses.append(base_addr)
|
||||||
"Registering layer %s with cache shape: %s", layer_name, cache.shape
|
# Only record non-Mamba page sizes.
|
||||||
|
if isinstance(layer_spec, MambaSpec):
|
||||||
|
self.block_len_per_layer.append(
|
||||||
|
physical_page_size // self._physical_blocks_per_logical_kv_block
|
||||||
)
|
)
|
||||||
seen_base_addresses.append(base_addr)
|
else:
|
||||||
# Only record non-Mamba page sizes.
|
self.block_len_per_layer.append(physical_page_size)
|
||||||
if isinstance(layer_spec, MambaSpec):
|
is_mla_region = isinstance(
|
||||||
self.block_len_per_layer.append(
|
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
|
||||||
physical_page_size // self._physical_blocks_per_logical_kv_block
|
)
|
||||||
)
|
self._region_is_mla.append(is_mla_region)
|
||||||
else:
|
|
||||||
self.block_len_per_layer.append(physical_page_size)
|
if not is_mla_region:
|
||||||
is_mla_region = isinstance(
|
if tensor_size_bytes is None:
|
||||||
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
|
tensor_size_bytes = curr_tensor_size_bytes
|
||||||
|
assert tensor_size_bytes == curr_tensor_size_bytes, (
|
||||||
|
"All non-MLA kv cache tensors must have the same size"
|
||||||
)
|
)
|
||||||
self._region_is_mla.append(is_mla_region)
|
|
||||||
|
|
||||||
if not is_mla_region:
|
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
|
||||||
if tensor_size_bytes is None:
|
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
|
||||||
tensor_size_bytes = curr_tensor_size_bytes
|
if (
|
||||||
assert tensor_size_bytes == curr_tensor_size_bytes, (
|
self._physical_blocks_per_logical_kv_block == 1
|
||||||
"All non-MLA kv cache tensors must have the same size"
|
and cache.shape[0] != num_blocks
|
||||||
)
|
):
|
||||||
|
raise AssertionError(
|
||||||
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
|
"All kv cache tensors must have the same number of "
|
||||||
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
|
f"blocks; layer={layer_name}, "
|
||||||
if (
|
f"expected_num_blocks={num_blocks}, "
|
||||||
self._physical_blocks_per_logical_kv_block == 1
|
f"cache_shape={tuple(cache.shape)}, "
|
||||||
and cache.shape[0] != num_blocks
|
f"cache_stride={tuple(cache.stride())}, "
|
||||||
):
|
f"layer_spec={type(layer_spec).__name__}, "
|
||||||
raise AssertionError(
|
f"backend={self.backend_name}, "
|
||||||
"All kv cache tensors must have the same number of "
|
"all_backends="
|
||||||
f"blocks; layer={layer_name}, "
|
f"{[backend.get_name() for backend in self.attn_backends]}, "
|
||||||
f"expected_num_blocks={num_blocks}, "
|
f"kv_cache_layout={self.kv_cache_layout}"
|
||||||
f"cache_shape={tuple(cache.shape)}, "
|
|
||||||
f"cache_stride={tuple(cache.stride())}, "
|
|
||||||
f"layer_spec={type(layer_spec).__name__}, "
|
|
||||||
f"backend={self.backend_name}, "
|
|
||||||
"all_backends="
|
|
||||||
f"{[backend.get_name() for backend in self.attn_backends]}, "
|
|
||||||
f"kv_cache_layout={self.kv_cache_layout}"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Need to make sure the device ID is non-negative for NIXL,
|
|
||||||
# Torch uses -1 to indicate CPU tensors.
|
|
||||||
self.device_id = max(cache.get_device(), 0)
|
|
||||||
caches_data.append(
|
|
||||||
(base_addr, curr_tensor_size_bytes, self.device_id, "")
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Need to make sure the device ID is non-negative for NIXL,
|
||||||
|
# Torch uses -1 to indicate CPU tensors.
|
||||||
|
self.device_id = max(cache.get_device(), 0)
|
||||||
|
caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Different block lengths collected: %s", set(self.block_len_per_layer)
|
"Different block lengths collected: %s", set(self.block_len_per_layer)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -326,9 +326,12 @@ class RequestOffloadState:
|
|||||||
group_state.block_ids.extend(new_blocks)
|
group_state.block_ids.extend(new_blocks)
|
||||||
|
|
||||||
def storable_chunks(
|
def storable_chunks(
|
||||||
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
|
self,
|
||||||
|
group_config: "GroupOffloadConfig",
|
||||||
|
group_state: RequestGroupState,
|
||||||
|
num_offloadable_tokens: int,
|
||||||
) -> int:
|
) -> int:
|
||||||
"""Number of leading offloaded chunks eligible for store.
|
"""Number of allocated leading offloaded chunks eligible for store.
|
||||||
|
|
||||||
For eagle/MTP groups the volatile trailing chunk of the offloadable
|
For eagle/MTP groups the volatile trailing chunk of the offloadable
|
||||||
range is excluded while decoding: the draft-layer KV of the last
|
range is excluded while decoding: the draft-layer KV of the last
|
||||||
@@ -345,7 +348,10 @@ class RequestOffloadState:
|
|||||||
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
|
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
|
||||||
if group_config.is_eagle_group and is_decoding:
|
if group_config.is_eagle_group and is_decoding:
|
||||||
num_chunks = max(0, num_chunks - 1)
|
num_chunks = max(0, num_chunks - 1)
|
||||||
return num_chunks
|
num_allocated_chunks = (
|
||||||
|
len(group_state.block_ids) // self.config.blocks_per_chunk
|
||||||
|
)
|
||||||
|
return min(num_chunks, num_allocated_chunks)
|
||||||
|
|
||||||
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
|
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
|
||||||
# max(): at the prefill->decode transition of a chunk-aligned prompt,
|
# max(): at the prefill->decode transition of a chunk-aligned prompt,
|
||||||
@@ -356,7 +362,7 @@ class RequestOffloadState:
|
|||||||
):
|
):
|
||||||
group_state.next_stored_chunk_idx = max(
|
group_state.next_stored_chunk_idx = max(
|
||||||
group_state.next_stored_chunk_idx,
|
group_state.next_stored_chunk_idx,
|
||||||
self.storable_chunks(group_config, num_offloadable_tokens),
|
self.storable_chunks(group_config, group_state, num_offloadable_tokens),
|
||||||
)
|
)
|
||||||
|
|
||||||
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
|
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
|
||||||
@@ -991,7 +997,7 @@ class OffloadingConnectorScheduler:
|
|||||||
self.config.kv_group_configs, req_status.group_states
|
self.config.kv_group_configs, req_status.group_states
|
||||||
):
|
):
|
||||||
num_chunks = req_status.storable_chunks(
|
num_chunks = req_status.storable_chunks(
|
||||||
group_config, num_offloadable_tokens
|
group_config, group_state, num_offloadable_tokens
|
||||||
)
|
)
|
||||||
|
|
||||||
start_chunk_idx = group_state.next_stored_chunk_idx
|
start_chunk_idx = group_state.next_stored_chunk_idx
|
||||||
@@ -1068,7 +1074,7 @@ class OffloadingConnectorScheduler:
|
|||||||
group_config.sliding_window_size_in_chunks is not None
|
group_config.sliding_window_size_in_chunks is not None
|
||||||
)
|
)
|
||||||
num_chunks = req_status.storable_chunks(
|
num_chunks = req_status.storable_chunks(
|
||||||
group_config, num_offloadable_tokens
|
group_config, group_state, num_offloadable_tokens
|
||||||
)
|
)
|
||||||
start_chunk_idx = group_state.next_stored_chunk_idx
|
start_chunk_idx = group_state.next_stored_chunk_idx
|
||||||
block_ids = group_state.block_ids
|
block_ids = group_state.block_ids
|
||||||
|
|||||||
@@ -56,9 +56,7 @@ class OffloadingConnectorWorker:
|
|||||||
def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
|
def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
|
||||||
self.worker = self.spec.get_worker(kv_caches)
|
self.worker = self.spec.get_worker(kv_caches)
|
||||||
|
|
||||||
def register_kv_caches(
|
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
|
||||||
self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
|
|
||||||
):
|
|
||||||
kv_cache_config = self.kv_cache_config
|
kv_cache_config = self.kv_cache_config
|
||||||
num_blocks = kv_cache_config.num_blocks
|
num_blocks = kv_cache_config.num_blocks
|
||||||
|
|
||||||
@@ -120,24 +118,13 @@ class OffloadingConnectorWorker:
|
|||||||
)
|
)
|
||||||
|
|
||||||
elif isinstance(layer_kv_cache_spec, MambaSpec):
|
elif isinstance(layer_kv_cache_spec, MambaSpec):
|
||||||
state_tensors = kv_caches[layer_name]
|
layer_kv_cache = kv_caches[layer_name]
|
||||||
assert isinstance(state_tensors, list)
|
assert layer_kv_cache.dtype == torch.int8
|
||||||
|
tensors_per_block[layer_name] = (
|
||||||
# re-construct the raw (num_blocks, page_size) tensor
|
layer_kv_cache.view(
|
||||||
# from the first state tensor
|
num_blocks, layer_kv_cache_spec.page_size_bytes
|
||||||
assert len(state_tensors) > 0
|
),
|
||||||
first_state_tensor = state_tensors[0]
|
|
||||||
assert first_state_tensor.storage_offset() == 0
|
|
||||||
tensor = (
|
|
||||||
torch.tensor(
|
|
||||||
[],
|
|
||||||
dtype=torch.int8,
|
|
||||||
device=first_state_tensor.device,
|
|
||||||
)
|
|
||||||
.set_(first_state_tensor.untyped_storage())
|
|
||||||
.view((num_blocks, layer_kv_cache_spec.page_size_bytes))
|
|
||||||
)
|
)
|
||||||
tensors_per_block[layer_name] = (tensor,)
|
|
||||||
|
|
||||||
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
|
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
|
||||||
unpadded_page_size_bytes[layer_name] = replace(
|
unpadded_page_size_bytes[layer_name] = replace(
|
||||||
|
|||||||
@@ -527,8 +527,6 @@ class EngineArgs:
|
|||||||
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
|
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
|
||||||
max_num_batched_tokens: int | None = None
|
max_num_batched_tokens: int | None = None
|
||||||
max_num_scheduled_tokens: int | None = None
|
max_num_scheduled_tokens: int | None = None
|
||||||
max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
|
|
||||||
max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills
|
|
||||||
long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
|
long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
|
||||||
max_num_seqs: int | None = None
|
max_num_seqs: int | None = None
|
||||||
max_logprobs: int = ModelConfig.max_logprobs
|
max_logprobs: int = ModelConfig.max_logprobs
|
||||||
@@ -1440,13 +1438,6 @@ class EngineArgs:
|
|||||||
"default": None,
|
"default": None,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
scheduler_group.add_argument(
|
|
||||||
"--max-num-partial-prefills", **scheduler_kwargs["max_num_partial_prefills"]
|
|
||||||
)
|
|
||||||
scheduler_group.add_argument(
|
|
||||||
"--max-long-partial-prefills",
|
|
||||||
**scheduler_kwargs["max_long_partial_prefills"],
|
|
||||||
)
|
|
||||||
scheduler_group.add_argument(
|
scheduler_group.add_argument(
|
||||||
"--long-prefill-token-threshold",
|
"--long-prefill-token-threshold",
|
||||||
**scheduler_kwargs["long_prefill_token_threshold"],
|
**scheduler_kwargs["long_prefill_token_threshold"],
|
||||||
@@ -2190,8 +2181,6 @@ class EngineArgs:
|
|||||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||||
policy=self.scheduling_policy,
|
policy=self.scheduling_policy,
|
||||||
scheduler_cls=self.scheduler_cls,
|
scheduler_cls=self.scheduler_cls,
|
||||||
max_num_partial_prefills=self.max_num_partial_prefills,
|
|
||||||
max_long_partial_prefills=self.max_long_partial_prefills,
|
|
||||||
long_prefill_token_threshold=self.long_prefill_token_threshold,
|
long_prefill_token_threshold=self.long_prefill_token_threshold,
|
||||||
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
|
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
|
||||||
watermark=self.watermark,
|
watermark=self.watermark,
|
||||||
@@ -2402,14 +2391,6 @@ class EngineArgs:
|
|||||||
|
|
||||||
def _check_feature_supported(self):
|
def _check_feature_supported(self):
|
||||||
"""Raise an error if the feature is not supported."""
|
"""Raise an error if the feature is not supported."""
|
||||||
# No Concurrent Partial Prefills so far.
|
|
||||||
if (
|
|
||||||
self.max_num_partial_prefills != SchedulerConfig.max_num_partial_prefills
|
|
||||||
or self.max_long_partial_prefills
|
|
||||||
!= SchedulerConfig.max_long_partial_prefills
|
|
||||||
):
|
|
||||||
_raise_unsupported_error(feature_name="Concurrent Partial Prefill")
|
|
||||||
|
|
||||||
if self.pipeline_parallel_size > 1:
|
if self.pipeline_parallel_size > 1:
|
||||||
supports_pp = getattr(
|
supports_pp = getattr(
|
||||||
self.distributed_executor_backend, "supports_pp", False
|
self.distributed_executor_backend, "supports_pp", False
|
||||||
|
|||||||
@@ -687,9 +687,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
|||||||
skip_special_tokens=self.skip_special_tokens,
|
skip_special_tokens=self.skip_special_tokens,
|
||||||
spaces_between_special_tokens=self.spaces_between_special_tokens,
|
spaces_between_special_tokens=self.spaces_between_special_tokens,
|
||||||
include_stop_str_in_output=self.include_stop_str_in_output,
|
include_stop_str_in_output=self.include_stop_str_in_output,
|
||||||
output_kind=RequestOutputKind.DELTA
|
output_kind=(
|
||||||
if self.stream
|
RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY
|
||||||
else RequestOutputKind.FINAL_ONLY,
|
),
|
||||||
structured_outputs=self.extract_structured_outputs(),
|
structured_outputs=self.extract_structured_outputs(),
|
||||||
logit_bias=self.logit_bias,
|
logit_bias=self.logit_bias,
|
||||||
bad_words=self.bad_words,
|
bad_words=self.bad_words,
|
||||||
@@ -849,9 +849,10 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
|||||||
|
|
||||||
# Reject empty tools array, matching OpenAI API behavior
|
# Reject empty tools array, matching OpenAI API behavior
|
||||||
if data.get("tools") == []:
|
if data.get("tools") == []:
|
||||||
raise ValueError(
|
raise VLLMValidationError(
|
||||||
"`tools` must not be an empty array. "
|
"`tools` must not be an empty array. "
|
||||||
"Either provide at least one tool or omit the field entirely."
|
"Either provide at least one tool or omit the field entirely.",
|
||||||
|
parameter="tools",
|
||||||
)
|
)
|
||||||
|
|
||||||
# if "tool_choice" is not specified but tools are provided,
|
# if "tool_choice" is not specified but tools are provided,
|
||||||
@@ -1075,13 +1076,15 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
|
|||||||
if isinstance(data, BatchChatCompletionRequest):
|
if isinstance(data, BatchChatCompletionRequest):
|
||||||
data = data.model_dump(exclude_unset=True)
|
data = data.model_dump(exclude_unset=True)
|
||||||
if data.get("use_beam_search"):
|
if data.get("use_beam_search"):
|
||||||
raise ValueError(
|
raise VLLMValidationError(
|
||||||
"Batch chat completions do not support beam search. "
|
"Batch chat completions do not support beam search. "
|
||||||
"Please set `use_beam_search` to False."
|
"Please set `use_beam_search` to False.",
|
||||||
|
parameter="use_beam_search",
|
||||||
)
|
)
|
||||||
if data.get("logprob_token_ids") and not data.get("logprobs"):
|
if data.get("logprob_token_ids") and not data.get("logprobs"):
|
||||||
raise ValueError(
|
raise VLLMValidationError(
|
||||||
"when using `logprob_token_ids`, `logprobs` must be set to true."
|
"when using `logprob_token_ids`, `logprobs` must be set to true.",
|
||||||
|
parameter="logprob_token_ids",
|
||||||
)
|
)
|
||||||
response_format = data.get("response_format")
|
response_format = data.get("response_format")
|
||||||
rf_type = (
|
rf_type = (
|
||||||
@@ -1095,8 +1098,10 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
|
|||||||
validate_structured_outputs_structural_tag(structured_outputs)
|
validate_structured_outputs_structural_tag(structured_outputs)
|
||||||
n = data.get("n", 1)
|
n = data.get("n", 1)
|
||||||
if n is not None and n != 1:
|
if n is not None and n != 1:
|
||||||
raise ValueError(
|
raise VLLMValidationError(
|
||||||
"Batch chat completions do not support `n > 1`. Please set `n` to 1."
|
"Batch chat completions do not support `n > 1`. Please set `n` to 1.",
|
||||||
|
parameter="n",
|
||||||
|
value=n,
|
||||||
)
|
)
|
||||||
return data
|
return data
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
import json
|
||||||
import time
|
import time
|
||||||
from typing import TYPE_CHECKING, Literal, TypeAlias
|
from typing import TYPE_CHECKING, Literal, TypeAlias
|
||||||
|
|
||||||
@@ -78,7 +79,6 @@ class TranslationRequest(OpenAIBaseModel):
|
|||||||
`verbose_json`, or `vtt`.
|
`verbose_json`, or `vtt`.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
# TODO support additional sampling parameters
|
|
||||||
# --8<-- [start:translation-sampling-params]
|
# --8<-- [start:translation-sampling-params]
|
||||||
use_beam_search: bool = False
|
use_beam_search: bool = False
|
||||||
"""Whether or not beam search should be used."""
|
"""Whether or not beam search should be used."""
|
||||||
@@ -103,6 +103,28 @@ class TranslationRequest(OpenAIBaseModel):
|
|||||||
will use [log probability](https://en.wikipedia.org/wiki/Log_probability)
|
will use [log probability](https://en.wikipedia.org/wiki/Log_probability)
|
||||||
to automatically increase the temperature until certain thresholds are hit.
|
to automatically increase the temperature until certain thresholds are hit.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
top_p: float | None = None
|
||||||
|
"""Enables nucleus (top-p) sampling, where tokens are selected from the
|
||||||
|
smallest possible set whose cumulative probability exceeds `p`.
|
||||||
|
"""
|
||||||
|
|
||||||
|
top_k: int | None = None
|
||||||
|
"""Limits sampling to the `k` most probable tokens at each step."""
|
||||||
|
|
||||||
|
min_p: float | None = None
|
||||||
|
"""Filters out tokens with a probability lower than `min_p`, ensuring a
|
||||||
|
minimum likelihood threshold during sampling.
|
||||||
|
"""
|
||||||
|
|
||||||
|
frequency_penalty: float | None = 0.0
|
||||||
|
"""The frequency penalty to use for sampling."""
|
||||||
|
|
||||||
|
repetition_penalty: float | None = None
|
||||||
|
"""The repetition penalty to use for sampling."""
|
||||||
|
|
||||||
|
presence_penalty: float | None = 0.0
|
||||||
|
"""The presence penalty to use for sampling."""
|
||||||
# --8<-- [end:translation-sampling-params]
|
# --8<-- [end:translation-sampling-params]
|
||||||
|
|
||||||
# --8<-- [start:translation-extra-params]
|
# --8<-- [start:translation-extra-params]
|
||||||
@@ -139,11 +161,23 @@ class TranslationRequest(OpenAIBaseModel):
|
|||||||
|
|
||||||
max_completion_tokens: int | None = None
|
max_completion_tokens: int | None = None
|
||||||
"""The maximum number of tokens to generate."""
|
"""The maximum number of tokens to generate."""
|
||||||
|
|
||||||
|
vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(
|
||||||
|
default=None,
|
||||||
|
description=(
|
||||||
|
"Additional request parameters with (list of) string or "
|
||||||
|
"numeric values, used by custom extensions."
|
||||||
|
),
|
||||||
|
)
|
||||||
# --8<-- [end:translation-extra-params]
|
# --8<-- [end:translation-extra-params]
|
||||||
|
|
||||||
# Default sampling parameters for translation requests.
|
# Default sampling parameters for translation requests.
|
||||||
_DEFAULT_SAMPLING_PARAMS: dict = {
|
_DEFAULT_SAMPLING_PARAMS: dict = {
|
||||||
|
"repetition_penalty": 1.0,
|
||||||
"temperature": 0,
|
"temperature": 0,
|
||||||
|
"top_p": 1.0,
|
||||||
|
"top_k": 0,
|
||||||
|
"min_p": 0.0,
|
||||||
}
|
}
|
||||||
|
|
||||||
def build_stt_params(
|
def build_stt_params(
|
||||||
@@ -199,14 +233,38 @@ class TranslationRequest(OpenAIBaseModel):
|
|||||||
temperature = default_sampling_params.get(
|
temperature = default_sampling_params.get(
|
||||||
"temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
|
"temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
|
||||||
)
|
)
|
||||||
|
if (top_p := self.top_p) is None:
|
||||||
|
top_p = default_sampling_params.get(
|
||||||
|
"top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]
|
||||||
|
)
|
||||||
|
if (top_k := self.top_k) is None:
|
||||||
|
top_k = default_sampling_params.get(
|
||||||
|
"top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"]
|
||||||
|
)
|
||||||
|
if (min_p := self.min_p) is None:
|
||||||
|
min_p = default_sampling_params.get(
|
||||||
|
"min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"]
|
||||||
|
)
|
||||||
|
if (repetition_penalty := self.repetition_penalty) is None:
|
||||||
|
repetition_penalty = default_sampling_params.get(
|
||||||
|
"repetition_penalty",
|
||||||
|
self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"],
|
||||||
|
)
|
||||||
|
|
||||||
return SamplingParams.from_optional(
|
return SamplingParams.from_optional(
|
||||||
temperature=temperature,
|
temperature=temperature,
|
||||||
max_tokens=max_tokens,
|
max_tokens=max_tokens,
|
||||||
seed=self.seed,
|
seed=self.seed,
|
||||||
|
top_p=top_p,
|
||||||
|
top_k=top_k,
|
||||||
|
min_p=min_p,
|
||||||
|
frequency_penalty=self.frequency_penalty,
|
||||||
|
repetition_penalty=repetition_penalty,
|
||||||
|
presence_penalty=self.presence_penalty,
|
||||||
output_kind=RequestOutputKind.DELTA
|
output_kind=RequestOutputKind.DELTA
|
||||||
if self.stream
|
if self.stream
|
||||||
else RequestOutputKind.FINAL_ONLY,
|
else RequestOutputKind.FINAL_ONLY,
|
||||||
|
extra_args=self.vllm_xargs,
|
||||||
skip_clone=True, # Created fresh per request, safe to skip clone
|
skip_clone=True, # Created fresh per request, safe to skip clone
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -226,6 +284,16 @@ class TranslationRequest(OpenAIBaseModel):
|
|||||||
parameter=invalid_param,
|
parameter=invalid_param,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
xargs = data.get("vllm_xargs")
|
||||||
|
if isinstance(xargs, str):
|
||||||
|
try:
|
||||||
|
data["vllm_xargs"] = json.loads(xargs)
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
raise VLLMValidationError(
|
||||||
|
f"Failed to parse vllm_xargs. Must be valid JSON: {e}",
|
||||||
|
parameter="vllm_xargs",
|
||||||
|
) from e
|
||||||
|
|
||||||
return data
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -4,6 +4,8 @@
|
|||||||
|
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
from vllm.config import VllmConfig
|
from vllm.config import VllmConfig
|
||||||
from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
|
from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
|
||||||
from vllm.v1.kv_cache_interface import KVCacheSpec
|
from vllm.v1.kv_cache_interface import KVCacheSpec
|
||||||
@@ -21,6 +23,14 @@ class AttentionLayerBase(ABC):
|
|||||||
impl: "AttentionImpl"
|
impl: "AttentionImpl"
|
||||||
supports_dcp: bool = True
|
supports_dcp: bool = True
|
||||||
|
|
||||||
|
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
|
||||||
|
"""Bind the allocated KV cache tensor to this layer.
|
||||||
|
|
||||||
|
The default stores the cache view as-is; subclasses (e.g. Mamba)
|
||||||
|
override this to unpack the raw buffer into per-state views.
|
||||||
|
"""
|
||||||
|
self.kv_cache = kv_cache
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def get_attn_backend(self) -> type[AttentionBackend]:
|
def get_attn_backend(self) -> type[AttentionBackend]:
|
||||||
"""Get the attention backend class for this layer."""
|
"""Get the attention backend class for this layer."""
|
||||||
|
|||||||
@@ -2,11 +2,13 @@
|
|||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
from abc import abstractmethod
|
from abc import abstractmethod
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
|
from math import prod
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from vllm.config import VllmConfig
|
from vllm.config import VllmConfig
|
||||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||||
|
from vllm.utils.torch_utils import get_dtype_size
|
||||||
from vllm.v1.attention.backend import AttentionBackend
|
from vllm.v1.attention.backend import AttentionBackend
|
||||||
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
||||||
from vllm.v1.attention.selector import get_mamba_attn_backend
|
from vllm.v1.attention.selector import get_mamba_attn_backend
|
||||||
@@ -24,6 +26,22 @@ class MambaBase(AttentionLayerBase):
|
|||||||
kv_cache: tuple[torch.Tensor, ...]
|
kv_cache: tuple[torch.Tensor, ...]
|
||||||
supports_dcp: bool = False
|
supports_dcp: bool = False
|
||||||
|
|
||||||
|
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
|
||||||
|
"""Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.
|
||||||
|
|
||||||
|
Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
|
||||||
|
packed contiguously; slice them out and reinterpret per dtype/shape.
|
||||||
|
"""
|
||||||
|
pages = kv_cache.squeeze(dim=(1, 2))
|
||||||
|
states: list[torch.Tensor] = []
|
||||||
|
offset = 0
|
||||||
|
for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):
|
||||||
|
nbytes = prod(shape) * get_dtype_size(dtype)
|
||||||
|
state = pages[:, offset : offset + nbytes].view(dtype)
|
||||||
|
states.append(state.view(-1, *shape))
|
||||||
|
offset += nbytes
|
||||||
|
self.kv_cache = tuple(states)
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
|
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
|
|||||||
causal_conv1d_fn,
|
causal_conv1d_fn,
|
||||||
causal_conv1d_update,
|
causal_conv1d_update,
|
||||||
)
|
)
|
||||||
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||||
from vllm.platforms import current_platform
|
from vllm.platforms import current_platform
|
||||||
from vllm.utils.torch_utils import direct_register_custom_op
|
from vllm.utils.torch_utils import direct_register_custom_op
|
||||||
from vllm.v1.attention.backend import AttentionMetadata
|
from vllm.v1.attention.backend import AttentionMetadata
|
||||||
@@ -42,6 +43,7 @@ class ShortConv(MambaBase, CustomOp):
|
|||||||
layer_idx: int,
|
layer_idx: int,
|
||||||
model_config: ModelConfig | None = None,
|
model_config: ModelConfig | None = None,
|
||||||
cache_config: CacheConfig | None = None,
|
cache_config: CacheConfig | None = None,
|
||||||
|
quant_config: QuantizationConfig | None = None,
|
||||||
prefix: str = "",
|
prefix: str = "",
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
@@ -67,12 +69,14 @@ class ShortConv(MambaBase, CustomOp):
|
|||||||
input_size=dim,
|
input_size=dim,
|
||||||
output_sizes=[dim] * 3,
|
output_sizes=[dim] * 3,
|
||||||
bias=self.bias,
|
bias=self.bias,
|
||||||
|
quant_config=quant_config,
|
||||||
prefix=f"{prefix}.in_proj",
|
prefix=f"{prefix}.in_proj",
|
||||||
)
|
)
|
||||||
self.out_proj = RowParallelLinear(
|
self.out_proj = RowParallelLinear(
|
||||||
input_size=dim,
|
input_size=dim,
|
||||||
output_size=dim,
|
output_size=dim,
|
||||||
bias=self.bias,
|
bias=self.bias,
|
||||||
|
quant_config=quant_config,
|
||||||
prefix=f"{prefix}.out_proj",
|
prefix=f"{prefix}.out_proj",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -276,7 +276,13 @@ class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
|
|||||||
|
|
||||||
|
|
||||||
class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
|
class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
|
||||||
pass
|
@staticmethod
|
||||||
|
def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool:
|
||||||
|
# Cosmos renders each video as a full
|
||||||
|
# <|vision_start|><|video_pad|><|vision_end|> placeholder, and its
|
||||||
|
# reference processor replaces that entire triplet with timestamped,
|
||||||
|
# per-frame vision sequences.
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
|
class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
|
||||||
@@ -517,6 +523,7 @@ class Cosmos3EdgeForConditionalGeneration(
|
|||||||
},
|
},
|
||||||
orig_to_new_substr={
|
orig_to_new_substr={
|
||||||
"_moe_gen": None,
|
"_moe_gen": None,
|
||||||
|
"k_norm_und_for_gen": None,
|
||||||
".add_q_proj.": None,
|
".add_q_proj.": None,
|
||||||
".add_k_proj.": None,
|
".add_k_proj.": None,
|
||||||
".add_v_proj.": None,
|
".add_v_proj.": None,
|
||||||
|
|||||||
@@ -257,6 +257,7 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
|
|||||||
layer_idx=layer_idx,
|
layer_idx=layer_idx,
|
||||||
model_config=model_config,
|
model_config=model_config,
|
||||||
cache_config=cache_config,
|
cache_config=cache_config,
|
||||||
|
quant_config=quant_config,
|
||||||
prefix=f"{prefix}.conv",
|
prefix=f"{prefix}.conv",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -351,6 +351,7 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):
|
|||||||
layer_idx=layer_idx,
|
layer_idx=layer_idx,
|
||||||
model_config=model_config,
|
model_config=model_config,
|
||||||
cache_config=cache_config,
|
cache_config=cache_config,
|
||||||
|
quant_config=quant_config,
|
||||||
prefix=f"{prefix}.conv",
|
prefix=f"{prefix}.conv",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -402,12 +402,11 @@ class Ovis2_5MultiModalProcessor(BaseMultiModalProcessor[Ovis2_5ProcessingInfo])
|
|||||||
hf_processor_mm_kwargs: Mapping[str, object],
|
hf_processor_mm_kwargs: Mapping[str, object],
|
||||||
out_mm_kwargs: MultiModalKwargsItems,
|
out_mm_kwargs: MultiModalKwargsItems,
|
||||||
) -> list[PromptReplacement]:
|
) -> list[PromptReplacement]:
|
||||||
tokenizer = self.info.get_tokenizer()
|
hf_processor = self.info.get_hf_processor()
|
||||||
vocab = tokenizer.get_vocab()
|
|
||||||
|
|
||||||
placeholder = {
|
placeholder = {
|
||||||
"image": vocab[IMAGE_TOKEN],
|
"image": hf_processor.get_token_value("image_token"),
|
||||||
"video": vocab[VIDEO_TOKEN],
|
"video": hf_processor.get_token_value("video_token"),
|
||||||
}
|
}
|
||||||
|
|
||||||
def get_replacement_ovis(item_idx, modality: str):
|
def get_replacement_ovis(item_idx, modality: str):
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from collections.abc import Mapping
|
|||||||
|
|
||||||
from vllm.config import ModelConfig, VllmConfig
|
from vllm.config import ModelConfig, VllmConfig
|
||||||
from vllm.logger import init_logger
|
from vllm.logger import init_logger
|
||||||
|
from vllm.multimodal.inputs import MultiModalKwargsItem
|
||||||
from vllm.multimodal.processing import BaseMultiModalProcessor
|
from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||||
from vllm.multimodal.registry import MultiModalRegistry
|
from vllm.multimodal.registry import MultiModalRegistry
|
||||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||||
@@ -48,6 +49,8 @@ class MultiModalBudget:
|
|||||||
self,
|
self,
|
||||||
vllm_config: VllmConfig,
|
vllm_config: VllmConfig,
|
||||||
mm_registry: MultiModalRegistry,
|
mm_registry: MultiModalRegistry,
|
||||||
|
*,
|
||||||
|
enable_cache: bool = True,
|
||||||
) -> None:
|
) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
||||||
@@ -58,7 +61,11 @@ class MultiModalBudget:
|
|||||||
self.max_num_reqs = scheduler_config.max_num_seqs
|
self.max_num_reqs = scheduler_config.max_num_seqs
|
||||||
|
|
||||||
with set_default_torch_num_threads(): # Avoid hang during startup
|
with set_default_torch_num_threads(): # Avoid hang during startup
|
||||||
cache = mm_registry.processor_only_cache_from_config(vllm_config)
|
cache = (
|
||||||
|
mm_registry.processor_only_cache_from_config(vllm_config)
|
||||||
|
if enable_cache
|
||||||
|
else None
|
||||||
|
)
|
||||||
processor = mm_registry.create_processor(model_config, cache=cache)
|
processor = mm_registry.create_processor(model_config, cache=cache)
|
||||||
|
|
||||||
self.cache = cache
|
self.cache = cache
|
||||||
@@ -191,3 +198,23 @@ class MultiModalBudget:
|
|||||||
def reset_cache(self) -> None:
|
def reset_cache(self) -> None:
|
||||||
if self.cache is not None:
|
if self.cache is not None:
|
||||||
self.cache.clear_cache()
|
self.cache.clear_cache()
|
||||||
|
|
||||||
|
|
||||||
|
def get_dummy_encoder_profile_inputs(
|
||||||
|
mm_registry: MultiModalRegistry,
|
||||||
|
budget: MultiModalBudget,
|
||||||
|
) -> list[tuple[str, MultiModalKwargsItem]]:
|
||||||
|
if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
|
||||||
|
return []
|
||||||
|
|
||||||
|
modality = budget.get_modality_with_max_tokens()
|
||||||
|
max_items_per_batch = budget.mm_max_items_per_batch[modality]
|
||||||
|
dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
|
||||||
|
budget.model_config,
|
||||||
|
mm_counts={modality: 1},
|
||||||
|
processor=budget.processor,
|
||||||
|
)
|
||||||
|
dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
|
||||||
|
assert dummy_mm_item is not None, "Dummy item should be generated"
|
||||||
|
|
||||||
|
return [(modality, dummy_mm_item)] * max_items_per_batch
|
||||||
|
|||||||
@@ -1177,7 +1177,7 @@ class PyNvVideoCodecVideoBackend(VideoBackend):
|
|||||||
|
|
||||||
@VIDEO_LOADER_REGISTRY.register(
|
@VIDEO_LOADER_REGISTRY.register(
|
||||||
"qwen3_vl",
|
"qwen3_vl",
|
||||||
video_processor="Qwen3VLVideoProcessor",
|
video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
|
||||||
)
|
)
|
||||||
class Qwen3VLVideoBackend(VideoBackend):
|
class Qwen3VLVideoBackend(VideoBackend):
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ except json.JSONDecodeError:
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
DEFAULT_ENV_VAR_PREFIXES: set[str] = {
|
DEFAULT_ENV_VAR_PREFIXES: set[str] = {
|
||||||
"VLLM_",
|
"VLLM_",
|
||||||
|
"FLASH_ATTENTION_",
|
||||||
"LMCACHE_",
|
"LMCACHE_",
|
||||||
"NCCL_",
|
"NCCL_",
|
||||||
"UCX_",
|
"UCX_",
|
||||||
|
|||||||
@@ -78,7 +78,6 @@ class Ovis2_5Processor(ProcessorMixin):
|
|||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def extra_special_tokens(self):
|
def extra_special_tokens(self):
|
||||||
vocab = self.tokenizer.get_vocab()
|
|
||||||
required_tokens = {
|
required_tokens = {
|
||||||
"image_token": "<image>",
|
"image_token": "<image>",
|
||||||
"video_token": "<video>",
|
"video_token": "<video>",
|
||||||
@@ -90,21 +89,16 @@ class Ovis2_5Processor(ProcessorMixin):
|
|||||||
"image_pad": "<|image_pad|>",
|
"image_pad": "<|image_pad|>",
|
||||||
}
|
}
|
||||||
|
|
||||||
extra_special_tokens = {}
|
# The checkpoint defines both `additional_special_tokens` and
|
||||||
suggestion = (
|
# `extra_special_tokens`, with the latter empty. Transformers ignores
|
||||||
"please add '<image>', '<video>', '<ovis_visual_atom>', "
|
# the former because the latter is explicitly empty, so the tokens are
|
||||||
"'<ovis_image_start>', '<ovis_image_end>', '<ovis_video_start>', "
|
# missing from the vocab. Re-add them to restore the expected ids.
|
||||||
"'<ovis_video_end>' in 'additional_special_tokens' of "
|
self.tokenizer.add_tokens(list(required_tokens.values()), special_tokens=True)
|
||||||
"tokenizer_config.json, You can refer to "
|
|
||||||
"https://huggingface.co/AIDC-AI/Ovis2.6-30B-A3B/blob/main/tokenizer_config.json"
|
|
||||||
)
|
|
||||||
|
|
||||||
for key, token_name in required_tokens.items():
|
return {
|
||||||
if token_name not in vocab:
|
key: self.tokenizer.convert_tokens_to_ids(token_name)
|
||||||
raise ValueError(f"Can not find {token_name}, {suggestion}")
|
for key, token_name in required_tokens.items()
|
||||||
extra_special_tokens[key] = vocab[token_name]
|
}
|
||||||
|
|
||||||
return extra_special_tokens
|
|
||||||
|
|
||||||
def __call__(
|
def __call__(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -392,9 +392,8 @@ def apply_top_k_top_p_pytorch(
|
|||||||
logits_sort.masked_fill_(top_k_mask, -float("inf"))
|
logits_sort.masked_fill_(top_k_mask, -float("inf"))
|
||||||
|
|
||||||
if p is not None:
|
if p is not None:
|
||||||
# Apply top-p. The cumsum below runs over the whole vocab, so accumulating
|
# Apply top-p.
|
||||||
# in a low-precision dtype makes the nucleus undershoot p.
|
probs_sort = logits_sort.softmax(dim=-1)
|
||||||
probs_sort = logits_sort.softmax(dim=-1, dtype=torch.float32)
|
|
||||||
probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
|
probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
|
||||||
top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
|
top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
|
||||||
# at least one
|
# at least one
|
||||||
@@ -501,10 +500,9 @@ def flashinfer_sample(
|
|||||||
probs, k, deterministic=True
|
probs, k, deterministic=True
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Both top-k and top-p. FlashInfer requires contiguous fp32 logits; the
|
# Both top-k and top-p.
|
||||||
# branches above get that from softmax().
|
|
||||||
next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
|
next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
|
||||||
logits.float().contiguous(), k, p, deterministic=True
|
logits, k, p, deterministic=True
|
||||||
)
|
)
|
||||||
|
|
||||||
return next_token_ids.view(-1)
|
return next_token_ids.view(-1)
|
||||||
|
|||||||
@@ -131,7 +131,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs < VOCAB_SIZE
|
mask_n = offs < VOCAB_SIZE
|
||||||
logits_blk0 = tl.load(
|
logits_blk0 = tl.load(
|
||||||
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
# Exclude -inf values (e.g. from grammar bitmasks) from
|
# Exclude -inf values (e.g. from grammar bitmasks) from
|
||||||
# statistics to avoid NaN in pivot computation.
|
# statistics to avoid NaN in pivot computation.
|
||||||
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
||||||
@@ -164,7 +164,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs_n < VOCAB_SIZE
|
mask_n = offs_n < VOCAB_SIZE
|
||||||
logits_blk = tl.load(
|
logits_blk = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
|
|
||||||
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
||||||
# Exclude -inf from min to keep binary search bounds
|
# Exclude -inf from min to keep binary search bounds
|
||||||
@@ -305,7 +305,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs_n < VOCAB_SIZE
|
mask_n = offs_n < VOCAB_SIZE
|
||||||
logits_blk2 = tl.load(
|
logits_blk2 = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
|
|
||||||
above_0 = logits_blk2 > k_pivot_0
|
above_0 = logits_blk2 > k_pivot_0
|
||||||
above_1 = logits_blk2 > k_pivot_1
|
above_1 = logits_blk2 > k_pivot_1
|
||||||
@@ -457,7 +457,7 @@ def _topk_topp_kernel(
|
|||||||
LOGITS_ROW + offs_n,
|
LOGITS_ROW + offs_n,
|
||||||
mask=mask_n,
|
mask=mask_n,
|
||||||
other=-float("inf"),
|
other=-float("inf"),
|
||||||
).to(tl.float32)
|
)
|
||||||
|
|
||||||
outlier_mask = (probs_blk > min_logit) & mask_n
|
outlier_mask = (probs_blk > min_logit) & mask_n
|
||||||
|
|
||||||
@@ -600,7 +600,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs < VOCAB_SIZE
|
mask_n = offs < VOCAB_SIZE
|
||||||
logits_blk0 = tl.load(
|
logits_blk0 = tl.load(
|
||||||
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
# Exclude -inf values (e.g. from grammar bitmasks) from
|
# Exclude -inf values (e.g. from grammar bitmasks) from
|
||||||
# statistics to avoid NaN in pivot computation.
|
# statistics to avoid NaN in pivot computation.
|
||||||
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
||||||
@@ -626,7 +626,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs_n < VOCAB_SIZE
|
mask_n = offs_n < VOCAB_SIZE
|
||||||
logits_blk = tl.load(
|
logits_blk = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
||||||
# Exclude -inf from min to keep binary search bounds
|
# Exclude -inf from min to keep binary search bounds
|
||||||
# finite (avoids NaN pivots).
|
# finite (avoids NaN pivots).
|
||||||
@@ -660,7 +660,7 @@ def _topk_topp_kernel(
|
|||||||
|
|
||||||
probs_blk = tl.load(
|
probs_blk = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
probs_blk = tl.exp(probs_blk - max_sample)
|
probs_blk = tl.exp(probs_blk - max_sample)
|
||||||
probs_blk = probs_blk / sum_exp_logits
|
probs_blk = probs_blk / sum_exp_logits
|
||||||
|
|
||||||
@@ -754,7 +754,7 @@ def _topk_topp_kernel(
|
|||||||
|
|
||||||
probs_blk = tl.load(
|
probs_blk = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
probs_blk = tl.exp(probs_blk - max_sample)
|
probs_blk = tl.exp(probs_blk - max_sample)
|
||||||
probs_blk = probs_blk / sum_exp_logits
|
probs_blk = probs_blk / sum_exp_logits
|
||||||
tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)
|
tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)
|
||||||
@@ -835,7 +835,7 @@ def _topk_topp_kernel(
|
|||||||
mask_n = offs_n < VOCAB_SIZE
|
mask_n = offs_n < VOCAB_SIZE
|
||||||
logits_blk = tl.load(
|
logits_blk = tl.load(
|
||||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||||
).to(tl.float32)
|
)
|
||||||
keep_mask = (logits_blk > final_pivot) & mask_n
|
keep_mask = (logits_blk > final_pivot) & mask_n
|
||||||
|
|
||||||
# Duplicate logit handling
|
# Duplicate logit handling
|
||||||
@@ -878,7 +878,7 @@ def apply_top_k_top_p_triton(
|
|||||||
The masked logits tensor. It may or may not be modified in-place.
|
The masked logits tensor. It may or may not be modified in-place.
|
||||||
"""
|
"""
|
||||||
assert logits.ndim == 2
|
assert logits.ndim == 2
|
||||||
assert logits.dtype in (torch.float32, torch.bfloat16, torch.float16)
|
assert logits.dtype == torch.float32
|
||||||
batch_size, vocab_size = logits.shape
|
batch_size, vocab_size = logits.shape
|
||||||
topk_enabled = k is not None
|
topk_enabled = k is not None
|
||||||
topp_enabled = p is not None
|
topp_enabled = p is not None
|
||||||
@@ -911,9 +911,7 @@ def apply_top_k_top_p_triton(
|
|||||||
buffer = _TRITON_BUFFER_CACHE.get(buf_key)
|
buffer = _TRITON_BUFFER_CACHE.get(buf_key)
|
||||||
if buffer is None or buffer.shape[0] < NUM_PROGRAMS:
|
if buffer is None or buffer.shape[0] < NUM_PROGRAMS:
|
||||||
size = min(next_power_of_2(NUM_PROGRAMS), num_sm)
|
size = min(next_power_of_2(NUM_PROGRAMS), num_sm)
|
||||||
buffer = torch.empty(
|
buffer = logits.new_empty((size, vocab_size))
|
||||||
(size, vocab_size), dtype=torch.float32, device=logits.device
|
|
||||||
)
|
|
||||||
_TRITON_BUFFER_CACHE[buf_key] = buffer
|
_TRITON_BUFFER_CACHE[buf_key] = buffer
|
||||||
if buffer.shape[0] > NUM_PROGRAMS:
|
if buffer.shape[0] > NUM_PROGRAMS:
|
||||||
buffer = buffer[:NUM_PROGRAMS]
|
buffer = buffer[:NUM_PROGRAMS]
|
||||||
@@ -921,12 +919,8 @@ def apply_top_k_top_p_triton(
|
|||||||
# Cache lookup table entries on each device.
|
# Cache lookup table entries on each device.
|
||||||
tables = _TRITON_TABLE_CACHE.get(logits.device)
|
tables = _TRITON_TABLE_CACHE.get(logits.device)
|
||||||
if tables is None:
|
if tables is None:
|
||||||
normal_cdf_to_sigma_table = torch.tensor(
|
normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)
|
||||||
_NORMAL_CDF_TO_SIGMA_TABLE, dtype=torch.float32, device=logits.device
|
percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)
|
||||||
)
|
|
||||||
percentile_to_std_table = torch.tensor(
|
|
||||||
_PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=logits.device
|
|
||||||
)
|
|
||||||
_TRITON_TABLE_CACHE[logits.device] = (
|
_TRITON_TABLE_CACHE[logits.device] = (
|
||||||
normal_cdf_to_sigma_table,
|
normal_cdf_to_sigma_table,
|
||||||
percentile_to_std_table,
|
percentile_to_std_table,
|
||||||
|
|||||||
@@ -262,7 +262,7 @@ def _reshape_kv_cache(
|
|||||||
kv_cache_config: "KVCacheConfig | None" = None,
|
kv_cache_config: "KVCacheConfig | None" = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
kv_caches: dict[str, Any] = {}
|
kv_caches: dict[str, Any] = {}
|
||||||
has_attn, has_mamba = False, False
|
has_attn = False
|
||||||
|
|
||||||
layer_packing: dict[str, tuple[int, int]] = {}
|
layer_packing: dict[str, tuple[int, int]] = {}
|
||||||
if kv_cache_config is not None:
|
if kv_cache_config is not None:
|
||||||
@@ -340,38 +340,21 @@ def _reshape_kv_cache(
|
|||||||
)
|
)
|
||||||
|
|
||||||
elif isinstance(kv_cache_spec, MambaSpec):
|
elif isinstance(kv_cache_spec, MambaSpec):
|
||||||
has_mamba = True
|
page_size_bytes = kv_cache_spec.page_size_bytes
|
||||||
state_tensors = []
|
# Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
|
||||||
storage_offset_bytes = 0
|
# int8 page view per layer; the layer's bind_kv_cache unpacks
|
||||||
for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
|
# each block's bytes into its conv/ssm state views. Keeping
|
||||||
dtype_size = get_dtype_size(dtype)
|
# one tensor per layer lets the KV connector register it
|
||||||
num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size
|
# without special-casing Mamba.
|
||||||
target_shape = (num_blocks, *shape)
|
kv_caches[layer_name] = kv_raw_tensor[
|
||||||
stride = torch.empty(target_shape).stride()
|
: num_blocks * page_size_bytes
|
||||||
target_stride = (num_element_per_page, *stride[1:])
|
].view(num_blocks, 1, 1, page_size_bytes)
|
||||||
assert storage_offset_bytes % dtype_size == 0
|
|
||||||
tensor = torch.as_strided(
|
|
||||||
kv_raw_tensor.view(dtype),
|
|
||||||
size=target_shape,
|
|
||||||
stride=target_stride,
|
|
||||||
storage_offset=storage_offset_bytes // dtype_size,
|
|
||||||
)
|
|
||||||
state_tensors.append(tensor)
|
|
||||||
storage_offset_bytes += stride[0] * dtype_size
|
|
||||||
kv_caches[layer_name] = state_tensors
|
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
|
f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
|
||||||
)
|
)
|
||||||
|
|
||||||
if has_attn and has_mamba:
|
if has_attn and kv_cache_config is not None:
|
||||||
_update_hybrid_attention_layout(
|
|
||||||
attn_groups=attn_groups,
|
|
||||||
kv_caches=kv_caches,
|
|
||||||
kernel_block_sizes=kernel_block_sizes,
|
|
||||||
cache_dtype=cache_dtype,
|
|
||||||
)
|
|
||||||
elif has_attn and kv_cache_config is not None:
|
|
||||||
_align_mixed_attention_kv_cache_views(
|
_align_mixed_attention_kv_cache_views(
|
||||||
attn_groups=attn_groups,
|
attn_groups=attn_groups,
|
||||||
kv_caches=kv_caches,
|
kv_caches=kv_caches,
|
||||||
@@ -458,65 +441,6 @@ def _restride_blocks_first_kv_cache_to_kv_first_storage(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _update_hybrid_attention_layout(
|
|
||||||
attn_groups: Iterable[AttentionGroup],
|
|
||||||
kv_caches: dict[str, Any],
|
|
||||||
kernel_block_sizes: list[int],
|
|
||||||
cache_dtype: str,
|
|
||||||
) -> None:
|
|
||||||
for group in attn_groups:
|
|
||||||
if group.kv_cache_group_id >= len(kernel_block_sizes):
|
|
||||||
continue
|
|
||||||
|
|
||||||
kv_cache_spec = group.kv_cache_spec
|
|
||||||
if not isinstance(kv_cache_spec, AttentionSpec):
|
|
||||||
continue
|
|
||||||
# Mirror the per-layer dtype selection used when building the shape
|
|
||||||
# above. The block-dim index is dtype-independent for current backends
|
|
||||||
# (quantization only changes the last dim), so this is a no-op today,
|
|
||||||
# but it keeps both call sites consistent for skip layers.
|
|
||||||
layer_cache_dtype = (
|
|
||||||
"auto"
|
|
||||||
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
|
|
||||||
and not isinstance(kv_cache_spec, TQFullAttentionSpec)
|
|
||||||
else cache_dtype
|
|
||||||
)
|
|
||||||
block_dim = group.backend.get_kv_cache_block_dim(
|
|
||||||
kernel_block_sizes[group.kv_cache_group_id],
|
|
||||||
kv_cache_spec.num_kv_heads,
|
|
||||||
kv_cache_spec.head_size,
|
|
||||||
cache_dtype_str=layer_cache_dtype,
|
|
||||||
)
|
|
||||||
# if the first dim of the kvcache's layout is already num_blocks, continue
|
|
||||||
if block_dim == 0:
|
|
||||||
continue
|
|
||||||
|
|
||||||
assert block_dim == 1, (
|
|
||||||
"Expected the dim `num_blocks` at the second dim when updating"
|
|
||||||
" the kvcache's layout of full attention layer"
|
|
||||||
)
|
|
||||||
|
|
||||||
for layer_name in group.layer_names:
|
|
||||||
if layer_name not in kv_caches:
|
|
||||||
# Shared layer — will be aliased to its target after this pass.
|
|
||||||
continue
|
|
||||||
|
|
||||||
kv_cache = kv_caches[layer_name]
|
|
||||||
if kv_cache.shape[0] == 2:
|
|
||||||
assert kv_cache.shape[1] != 2, (
|
|
||||||
f"Cannot determine layout for tensor of shape {kv_cache.shape}"
|
|
||||||
)
|
|
||||||
hidden_size = kv_cache.shape[2:].numel()
|
|
||||||
kv_cache.as_strided_(
|
|
||||||
size=kv_cache.shape,
|
|
||||||
stride=(
|
|
||||||
hidden_size,
|
|
||||||
2 * hidden_size,
|
|
||||||
*kv_cache.stride()[2:],
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def init_kv_cache(
|
def init_kv_cache(
|
||||||
runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
|
runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
|
||||||
forward_context: dict[str, Any],
|
forward_context: dict[str, Any],
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ class EncoderCache:
|
|||||||
Clear the multi-modal cache that was used during profiling,
|
Clear the multi-modal cache that was used during profiling,
|
||||||
but no longer needed during inference.
|
but no longer needed during inference.
|
||||||
"""
|
"""
|
||||||
# TODO: Implement MM budget for encoder dummy run
|
# NOTE: v2 encoder cache profiling skips the multi-modal cache
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def reset_encoder_cache(self) -> None:
|
def reset_encoder_cache(self) -> None:
|
||||||
|
|||||||
@@ -3,7 +3,9 @@
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
|
from vllm.logger import init_logger
|
||||||
from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
|
from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
|
||||||
|
from vllm.multimodal.encoder_budget import MultiModalBudget
|
||||||
from vllm.multimodal.inputs import MultiModalKwargsItem
|
from vllm.multimodal.inputs import MultiModalKwargsItem
|
||||||
from vllm.multimodal.utils import (
|
from vllm.multimodal.utils import (
|
||||||
get_mm_features_in_window,
|
get_mm_features_in_window,
|
||||||
@@ -13,6 +15,8 @@ from vllm.multimodal.utils import (
|
|||||||
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
|
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
|
||||||
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
|
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
|
||||||
|
|
||||||
|
logger = init_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class EncoderRunner:
|
class EncoderRunner:
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -52,6 +56,45 @@ class EncoderRunner:
|
|||||||
|
|
||||||
return mm_hashes, mm_kwargs
|
return mm_hashes, mm_kwargs
|
||||||
|
|
||||||
|
@torch.inference_mode()
|
||||||
|
def profile_encoder_cache(
|
||||||
|
self,
|
||||||
|
dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
|
||||||
|
budget: MultiModalBudget,
|
||||||
|
) -> None:
|
||||||
|
"""Profile multimodal encoder and temporary encoder cache memory."""
|
||||||
|
if (encoder_budget := budget.get_encoder_budget()) <= 0:
|
||||||
|
return
|
||||||
|
|
||||||
|
if not budget.mm_max_toks_per_item:
|
||||||
|
logger.info(
|
||||||
|
"Skipping encoder profiling for embedding-only mode "
|
||||||
|
"(all modality limits=0 with enable_mm_embeds=True).",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
|
||||||
|
dummy_modality = dummy_mm_inputs[0][0]
|
||||||
|
max_mm_items_per_batch = len(dummy_mm_inputs)
|
||||||
|
|
||||||
|
logger.info_once(
|
||||||
|
"Encoder cache will be initialized with a budget of %s tokens, "
|
||||||
|
"and profiled with %s %s items of the maximum feature size.",
|
||||||
|
encoder_budget,
|
||||||
|
max_mm_items_per_batch,
|
||||||
|
dummy_modality,
|
||||||
|
)
|
||||||
|
|
||||||
|
dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
|
||||||
|
|
||||||
|
sanity_check_mm_encoder_outputs(
|
||||||
|
dummy_encoder_outputs,
|
||||||
|
expected_num_items=max_mm_items_per_batch,
|
||||||
|
)
|
||||||
|
self.encoder_cache.encoder_outputs.update(
|
||||||
|
(f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
|
||||||
|
)
|
||||||
|
|
||||||
@torch.inference_mode()
|
@torch.inference_mode()
|
||||||
def execute_mm_encoder(
|
def execute_mm_encoder(
|
||||||
self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
|
self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
|
||||||
|
|||||||
@@ -43,6 +43,10 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
|
|||||||
)
|
)
|
||||||
from vllm.model_executor.model_loader import get_model_loader
|
from vllm.model_executor.model_loader import get_model_loader
|
||||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||||
|
from vllm.multimodal.encoder_budget import (
|
||||||
|
MultiModalBudget,
|
||||||
|
get_dummy_encoder_profile_inputs,
|
||||||
|
)
|
||||||
from vllm.sequence import IntermediateTensors
|
from vllm.sequence import IntermediateTensors
|
||||||
from vllm.tasks import SupportedTask
|
from vllm.tasks import SupportedTask
|
||||||
from vllm.utils.math_utils import cdiv
|
from vllm.utils.math_utils import cdiv
|
||||||
@@ -671,6 +675,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
|||||||
|
|
||||||
@torch.inference_mode()
|
@torch.inference_mode()
|
||||||
def profile_run(self) -> None:
|
def profile_run(self) -> None:
|
||||||
|
if self.supports_mm_inputs and self.is_first_pp_rank:
|
||||||
|
mm_config = self.model_config.multimodal_config
|
||||||
|
if mm_config is not None and not mm_config.skip_mm_profiling:
|
||||||
|
mm_budget = MultiModalBudget(
|
||||||
|
self.vllm_config,
|
||||||
|
self.mm_registry,
|
||||||
|
enable_cache=False,
|
||||||
|
)
|
||||||
|
dummy_mm_inputs = get_dummy_encoder_profile_inputs(
|
||||||
|
self.mm_registry,
|
||||||
|
mm_budget,
|
||||||
|
)
|
||||||
|
self.model_state.encoder_runner.profile_encoder_cache(
|
||||||
|
dummy_mm_inputs, mm_budget
|
||||||
|
)
|
||||||
|
|
||||||
hidden_states, sample_hidden_states = self._dummy_run(
|
hidden_states, sample_hidden_states = self._dummy_run(
|
||||||
self.max_num_tokens, skip_attn=True, is_profile=True
|
self.max_num_tokens, skip_attn=True, is_profile=True
|
||||||
)
|
)
|
||||||
@@ -685,6 +705,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
|||||||
|
|
||||||
torch.accelerator.synchronize()
|
torch.accelerator.synchronize()
|
||||||
del hidden_states, sample_hidden_states
|
del hidden_states, sample_hidden_states
|
||||||
|
self.reset_encoder_cache()
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
def post_kv_cache_wake_up(self) -> None:
|
def post_kv_cache_wake_up(self) -> None:
|
||||||
|
|||||||
@@ -218,9 +218,7 @@ def _bias_kernel(
|
|||||||
mask=mask,
|
mask=mask,
|
||||||
)
|
)
|
||||||
bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
|
bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
|
||||||
logits = tl.load(
|
logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)
|
||||||
logits_ptr + token_idx * logits_stride + token_ids, mask=mask
|
|
||||||
).to(tl.float32)
|
|
||||||
logits += bias
|
logits += bias
|
||||||
tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)
|
tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)
|
||||||
|
|
||||||
|
|||||||
@@ -83,7 +83,11 @@ class Sampler:
|
|||||||
# that num_nans is computed before applying penalties and temperature.
|
# that num_nans is computed before applying penalties and temperature.
|
||||||
num_nans = get_num_nans(logits) if self.compute_nans else None
|
num_nans = get_num_nans(logits) if self.compute_nans else None
|
||||||
|
|
||||||
return_logprobs = self.returns_logprobs(idx_mapping_np)
|
max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
|
||||||
|
max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
|
||||||
|
idx_mapping_np
|
||||||
|
)
|
||||||
|
return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0
|
||||||
|
|
||||||
sampled, processed_logits = self.sample(
|
sampled, processed_logits = self.sample(
|
||||||
logits,
|
logits,
|
||||||
@@ -98,10 +102,6 @@ class Sampler:
|
|||||||
if return_logprobs:
|
if return_logprobs:
|
||||||
if self.logprobs_mode in ("processed_logprobs", "processed_logits"):
|
if self.logprobs_mode in ("processed_logprobs", "processed_logits"):
|
||||||
logits = processed_logits
|
logits = processed_logits
|
||||||
max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
|
|
||||||
max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
|
|
||||||
idx_mapping_np
|
|
||||||
)
|
|
||||||
expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
|
expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
|
||||||
cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
|
cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
|
||||||
num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
|
num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
|
||||||
@@ -142,13 +142,6 @@ class Sampler:
|
|||||||
)
|
)
|
||||||
return sampler_output
|
return sampler_output
|
||||||
|
|
||||||
def returns_logprobs(self, idx_mapping_np: np.ndarray) -> bool:
|
|
||||||
"""Whether any request in the batch produces logprobs this step."""
|
|
||||||
return (
|
|
||||||
self.sampling_states.max_num_logprobs(idx_mapping_np) != NO_LOGPROBS
|
|
||||||
or self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np) > 0
|
|
||||||
)
|
|
||||||
|
|
||||||
def apply_sampling_params(
|
def apply_sampling_params(
|
||||||
self,
|
self,
|
||||||
logits: torch.Tensor,
|
logits: torch.Tensor,
|
||||||
@@ -159,13 +152,8 @@ class Sampler:
|
|||||||
expanded_local_pos: torch.Tensor,
|
expanded_local_pos: torch.Tensor,
|
||||||
skip_top_k_top_p: bool = False,
|
skip_top_k_top_p: bool = False,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
# The ops below upcast to fp32 internally, so the input dtype is kept and
|
# Copy logits to a new FP32 tensor.
|
||||||
# mutated in place. Only raw_logprobs reads the unmodified logits
|
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
|
||||||
# afterward, so copy just for that case.
|
|
||||||
if self.logprobs_mode.startswith("raw_") and self.returns_logprobs(
|
|
||||||
idx_mapping_np
|
|
||||||
):
|
|
||||||
logits = logits.clone()
|
|
||||||
|
|
||||||
# Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
|
# Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
|
||||||
self.logit_bias_state.apply_logit_bias(
|
self.logit_bias_state.apply_logit_bias(
|
||||||
|
|||||||
@@ -126,7 +126,6 @@ from vllm.utils.torch_utils import (
|
|||||||
PIN_MEMORY,
|
PIN_MEMORY,
|
||||||
async_tensor_h2d,
|
async_tensor_h2d,
|
||||||
current_stream,
|
current_stream,
|
||||||
get_dtype_size,
|
|
||||||
is_quantized_kv_cache,
|
is_quantized_kv_cache,
|
||||||
kv_cache_dtype_str_to_dtype,
|
kv_cache_dtype_str_to_dtype,
|
||||||
)
|
)
|
||||||
@@ -7381,27 +7380,15 @@ class GPUModelRunner(
|
|||||||
elif isinstance(kv_cache_spec, MambaSpec):
|
elif isinstance(kv_cache_spec, MambaSpec):
|
||||||
has_mamba = True
|
has_mamba = True
|
||||||
raw_tensor = kv_cache_raw_tensors[layer_name]
|
raw_tensor = kv_cache_raw_tensors[layer_name]
|
||||||
state_tensors = []
|
page_size_bytes = kv_cache_spec.page_size_bytes
|
||||||
storage_offset_bytes = 0
|
# Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
|
||||||
for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
|
# int8 page view per layer; the layer's bind_kv_cache unpacks
|
||||||
dtype_size = get_dtype_size(dtype)
|
# each block's bytes into its conv/ssm state views. Keeping
|
||||||
num_element_per_page = (
|
# one tensor per layer lets the KV connector register it
|
||||||
kv_cache_spec.page_size_bytes // dtype_size
|
# without special-casing Mamba.
|
||||||
)
|
kv_caches[layer_name] = raw_tensor[
|
||||||
target_shape = (num_blocks, *shape)
|
: num_blocks * page_size_bytes
|
||||||
stride = torch.empty(target_shape).stride()
|
].view(num_blocks, 1, 1, page_size_bytes)
|
||||||
target_stride = (num_element_per_page, *stride[1:])
|
|
||||||
assert storage_offset_bytes % dtype_size == 0
|
|
||||||
tensor = torch.as_strided(
|
|
||||||
raw_tensor.view(dtype),
|
|
||||||
size=target_shape,
|
|
||||||
stride=target_stride,
|
|
||||||
storage_offset=storage_offset_bytes // dtype_size,
|
|
||||||
)
|
|
||||||
state_tensors.append(tensor)
|
|
||||||
storage_offset_bytes += stride[0] * dtype_size
|
|
||||||
|
|
||||||
kv_caches[layer_name] = state_tensors
|
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
|
|
||||||
|
|||||||
@@ -533,9 +533,12 @@ def bind_kv_cache(
|
|||||||
for layer_name in layer_names:
|
for layer_name in layer_names:
|
||||||
runner_kv_caches.append(kv_caches[layer_name])
|
runner_kv_caches.append(kv_caches[layer_name])
|
||||||
|
|
||||||
# Bind kv_caches to forward context
|
# Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
|
||||||
|
# its raw allocation into the per-layer view(s) it needs (e.g. Mamba
|
||||||
|
# splits conv/ssm), so the kv_caches dict can hold a single tensor per
|
||||||
|
# layer for the KV connector to register.
|
||||||
for layer_name, kv_cache in kv_caches.items():
|
for layer_name, kv_cache in kv_caches.items():
|
||||||
forward_context[layer_name].kv_cache = kv_cache
|
forward_context[layer_name].bind_kv_cache(kv_cache)
|
||||||
|
|
||||||
|
|
||||||
def copy_kv_cache_blocks_inplace(
|
def copy_kv_cache_blocks_inplace(
|
||||||
|
|||||||
Reference in New Issue
Block a user