Compare commits

...
Author SHA1 Message Date
Bugen ZhaoandOpenAI Codex e1a763558c [CI] Discover Rust coverage artifacts from build metadata
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 06:44:16 +00:00
Bugen ZhaoandOpenAI Codex 84aeec9f22 [CI] Simplify Rust coverage reporting
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 08:18:14 +00:00
Bugen ZhaoandOpenAI Codex 82a770ddbd [CI] Simplify Rust coverage aggregation
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:58:45 +00:00
Bugen Zhao a09a9bace1 [CI] Disable redundant Codecov file fixes 2026-07-21 13:57:38 +00:00
Bugen Zhao 6c20d467a2 [CI] Run Codecov from repository root 2026-07-21 13:34:22 +00:00
Bugen ZhaoandOpenAI Codex cb59d0a351 [CI] Collect Rust coverage in Buildkite
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 13:01:56 +00:00
Bugen ZhaoandOpenAI Codex 0ab1bded36 [CI] Instrument Rust artifacts for coverage
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 12:22:25 +00:00
Umut PolatandGitHub 040cbf95cc [Misc] Use VLLMValidationError in chat completion tool and batch validators (#49214)
Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com>
2026-07-21 11:38:20 +00:00
SyaGitHubLi, Jiang <jiang1.li@intel.com>
5b3762a7f0 [Bugfix][CPU] Fix Clang OpenMP build on macOS (#49021)
Signed-off-by: markyangcc <mmdou3@163.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-21 09:58:52 +00:00
bastefaniakandGitHub 4d30c510ce [bugfix] Fix Cosmos3 Edge checkpoint weights filtering, video loading, prompt expansion (#49190)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
2026-07-21 17:18:36 +08:00
6700813f86 [3/N][KV-Cache Layout Refactor] Standardize Mamba cache; drop get_transfer_cache_regions (#44456)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 09:16:15 +00:00
Bugen ZhaoandGitHub eb44b3aaa4 [Rust][Benchmark] Use async HTTP clients (#49295)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 16:53:57 +08:00
Nicolò LucchesiandGitHub 7a98c7a392 [Misc] Remove old now unsupported max_num_partial_prefills and max_long_partial_prefills (#49244)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-07-21 08:52:52 +00:00
Lena OnyshchenkoandGitHub 0d9e60619b [Misc][Docs] Fix XPU compute-runtime driver link version mismatch (#49299)
Signed-off-by: oonyshch <xonyshch@gmail.com>
2026-07-21 08:45:41 +00:00
1134545b6f Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) (#49033)
Co-authored-by: vllm-agent <vllm-agent@users.noreply.github.com>
2026-07-21 09:36:45 +01:00
Miłosz GrunwaldGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Harry Mellor
3e0c887511 [Bugfix] Fix Ovis2_5 special tokens for transformers v5 (#47298)
Signed-off-by: mgrunwal <milosz.grunwald@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-21 08:09:20 +00:00
Stefan KaestleandGitHub adfbbc1005 Propagate Flash Attention cache configuration to Ray workers (#49177)
Signed-off-by: Stefan Kaestle <skaestle@nvidia.com>
2026-07-21 07:47:53 +00:00
Roy WangandGitHub adc98f04d0 [Misc] Add @esmeetu to codeowners for rust/src/bench (#49298)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-07-21 07:44:24 +00:00
8def3cdde2 [Bugfix] Propagate quant_config to LFM2 ShortConv projections (#48917)
Signed-off-by: Alex Yuan <alex.yuan@liquid.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 06:51:55 +00:00
616c9bd0f4 [Frontend] Support additional sampling parameters for translation API (#45839)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-21 05:53:56 +00:00
Bugen ZhaoandGitHub 8688a06d67 [Rust][Benchmark] Use tracing for logs (#48937)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 05:18:13 +00:00
Salt SatoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f25953cc59 [Bugfix][Rust Frontend] Handle zero-column logprobs payloads without panicking (#49113)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
2026-07-21 04:30:29 +00:00
d9aa35161d Update BGE-M3 token expectations for leading spaces (#49269)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-07-21 03:49:17 +00:00
6bcda970fd [CI][NIXL] Isolate concurrent engine internal ports (#49129)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 22:28:11 -05:00
Isotr0pyandGitHub ea0e9c8f2e [MRV2] Add encoder cache profiling implementation (#47985)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-20 20:18:14 -07:00
ChaunceyandGitHub 94ed0bf4e0 [Bugfix][KV Offloading] Handle queued request aborts without allocated KV blocks (#49146)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-07-21 11:16:26 +08:00
Connor CarpenterGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
1940c8441e [Rust Frontend][gRPC] Add engine-aware health reporting (#48992)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
2026-07-21 10:54:39 +08:00
Simon MoandGitHub 72d16aee15 [CI] Exercise FA3 FP8 attention on SM90 (#49231)
Signed-off-by: Simon Mo <simon@inferact.ai>
2026-07-21 10:26:55 +08:00
Kunshang JiandGitHub e78a0c8e59 [XPU][Doc] Update XPU docker image documents (#49148)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-21 10:11:27 +08:00
92 changed files with 1925 additions and 839 deletions
@@ -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
+182
View File
@@ -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"
}
+30
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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/
+1 -1
View File
@@ -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:
+37
View File
@@ -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
View File
@@ -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
+3 -1
View File
@@ -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) {
+11
View File
@@ -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
+3
View File
@@ -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]
+9
View File
@@ -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.
+27 -10
View File
@@ -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
View File
@@ -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"] }
+4 -1
View File
@@ -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
+8 -7
View File
@@ -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
View File
@@ -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);
} }
+18 -9
View File
@@ -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"
); );
} }
+4 -2
View File
@@ -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]
+74 -37
View File
@@ -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.
+9 -6
View File
@@ -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"
); );
} }
+29 -18
View File
@@ -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]
+77
View File
@@ -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)));
}
}
+18 -12
View File
@@ -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) => {
+4 -2
View File
@@ -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 {
+19 -12
View File
@@ -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"
); );
} }
} }
+30 -23
View File
@@ -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
View File
@@ -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())?
} }
} }
+1 -1
View File
@@ -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");
} }
} }
+12
View File
@@ -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();
+18 -16
View File
@@ -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> =
+53 -39
View File
@@ -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");
} }
} }
+2 -2
View File
@@ -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(())
} }
+8 -2
View File
@@ -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) => {
+1 -1
View File
@@ -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();
+36 -24
View File
@@ -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);
} }
+98 -21
View File
@@ -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);
}
}
+17 -7
View File
@@ -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"));
+9 -1
View File
@@ -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()
+20 -1
View File
@@ -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"
);
}
+1
View File
@@ -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
+70
View File
@@ -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
View File
@@ -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;
+175 -12
View File
@@ -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
View File
@@ -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
} }
}; };
-1
View File
@@ -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])
+1
View File
@@ -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",
+7
View File
@@ -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}
+6 -3
View File
@@ -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.
-38
View File
@@ -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
View File
@@ -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:
+4 -40
View File
@@ -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(
-19
View File
@@ -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",
) )
+8 -1
View File
@@ -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,
+1
View File
@@ -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",
) )
+1
View File
@@ -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",
) )
+3 -4
View File
@@ -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):
+28 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1
View File
@@ -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_",
+9 -15
View File
@@ -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,
+4 -6
View File
@@ -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)
+13 -19
View File
@@ -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,
+11 -87
View File
@@ -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],
+1 -1
View File
@@ -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:
+43
View File
@@ -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]]
+21
View File
@@ -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:
+1 -3
View File
@@ -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)
+7 -19
View File
@@ -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(
+9 -22
View File
@@ -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
+5 -2
View File
@@ -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(