forked from Karylab-cklius/vllm
Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5ea7cac55b | ||
|
|
be3476447f | ||
|
|
f1473092c4 | ||
|
|
4dcbae8670 | ||
|
|
65dac3a770 | ||
|
|
0ba2500ef0 | ||
|
|
ef576befd2 | ||
|
|
35e4a36107 | ||
|
|
da5803d46e | ||
|
|
75ddfaf909 | ||
|
|
f36fe52add | ||
|
|
391d918d4d | ||
|
|
fca040885b | ||
|
|
a7fd4c7482 | ||
|
|
5131691063 | ||
|
|
80e00e5ac6 |
@@ -813,8 +813,8 @@ steps:
|
||||
|
||||
# Download artifacts from current build
|
||||
echo "Downloading artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
# buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
# buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
|
||||
# # Run upload script
|
||||
bash .buildkite/scripts/upload-rocm-wheels.sh
|
||||
|
||||
@@ -29,7 +29,6 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
|
||||
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
|
||||
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
|
||||
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
|
||||
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
|
||||
|
||||
log_section() {
|
||||
echo "--- $*"
|
||||
@@ -107,18 +106,6 @@ install_cargo_nextest() {
|
||||
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
|
||||
}
|
||||
|
||||
install_cargo_llvm_cov() {
|
||||
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
|
||||
local toolchain
|
||||
toolchain="$(rust_toolchain)"
|
||||
rustup component add --toolchain "$toolchain" llvm-tools-preview
|
||||
cargo binstall \
|
||||
--no-confirm \
|
||||
--force \
|
||||
--secure \
|
||||
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
|
||||
}
|
||||
|
||||
install_uv() {
|
||||
log_section "Installing uv ${UV_VERSION}"
|
||||
curl -L --proto '=https' --tlsv1.2 -sSf \
|
||||
@@ -189,41 +176,14 @@ run_tests() {
|
||||
setup_pyo3_python
|
||||
install_cargo_binstall
|
||||
install_cargo_nextest
|
||||
install_cargo_llvm_cov
|
||||
|
||||
log_section "Running cargo nextest with Rust coverage"
|
||||
mkdir -p artifacts
|
||||
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
|
||||
cargo llvm-cov clean \
|
||||
--manifest-path rust/Cargo.toml \
|
||||
--profraw-only
|
||||
|
||||
set +e
|
||||
cargo llvm-cov nextest \
|
||||
log_section "Running cargo nextest"
|
||||
cargo nextest run \
|
||||
--manifest-path rust/Cargo.toml \
|
||||
--workspace \
|
||||
--all-features \
|
||||
--locked \
|
||||
--no-fail-fast \
|
||||
--no-clean \
|
||||
--lcov \
|
||||
--output-path artifacts/rust-unit.lcov \
|
||||
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
|
||||
local coverage_rc=$?
|
||||
|
||||
local upload_rc=0
|
||||
if [[ $coverage_rc -eq 0 ]]; then
|
||||
# shellcheck source=.buildkite/scripts/rust-coverage.sh
|
||||
source .buildkite/scripts/rust-coverage.sh
|
||||
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
|
||||
upload_rc=$?
|
||||
fi
|
||||
set -e
|
||||
|
||||
if [[ $coverage_rc -ne 0 ]]; then
|
||||
return "$coverage_rc"
|
||||
fi
|
||||
return "$upload_rc"
|
||||
--no-fail-fast
|
||||
}
|
||||
|
||||
install_protoc
|
||||
|
||||
@@ -1,182 +0,0 @@
|
||||
#!/bin/sh
|
||||
|
||||
RUST_CODECOV_VERSION="v11.3.1"
|
||||
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
|
||||
|
||||
rust_coverage_repo_root() {
|
||||
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
|
||||
printf '%s\n' /vllm-workspace
|
||||
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
|
||||
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
|
||||
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
|
||||
else
|
||||
git rev-parse --show-toplevel
|
||||
fi
|
||||
}
|
||||
|
||||
rust_coverage_start() {
|
||||
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
|
||||
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
|
||||
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
|
||||
mkdir -p "$RUST_COVERAGE_DIR"
|
||||
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
|
||||
export LLVM_PROFILE_FILE
|
||||
trap rust_coverage_finalize 0
|
||||
}
|
||||
|
||||
rust_coverage_objects() {
|
||||
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
|
||||
python3 - "$rust_cov_objects_manifest" <<'PY'
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
for relative in Path(sys.argv[1]).read_text().splitlines():
|
||||
for entry in sys.path:
|
||||
path = Path(entry or ".").resolve() / relative
|
||||
if path.is_file():
|
||||
print(path)
|
||||
break
|
||||
else:
|
||||
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
|
||||
PY
|
||||
}
|
||||
|
||||
rust_coverage_collect() {
|
||||
rust_cov_collect_flag=${1:?coverage flag is required}
|
||||
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
|
||||
|
||||
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
|
||||
rust_cov_collect_primary=
|
||||
set --
|
||||
while IFS= read -r rust_cov_collect_object; do
|
||||
if [ -z "$rust_cov_collect_primary" ]; then
|
||||
rust_cov_collect_primary=$rust_cov_collect_object
|
||||
else
|
||||
set -- "$@" "--object=$rust_cov_collect_object"
|
||||
fi
|
||||
done <<EOF
|
||||
$rust_cov_collect_objects
|
||||
EOF
|
||||
|
||||
llvm-profdata merge \
|
||||
-sparse \
|
||||
"$RUST_COVERAGE_DIR"/*.profraw \
|
||||
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
|
||||
llvm-cov export \
|
||||
"$rust_cov_collect_primary" \
|
||||
"$@" \
|
||||
--format=lcov \
|
||||
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
|
||||
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
|
||||
> "$rust_cov_collect_lcov" || return 1
|
||||
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
|
||||
export RUST_COVERAGE_LCOV
|
||||
}
|
||||
|
||||
rust_coverage_upload() {
|
||||
rust_cov_upload_lcov=${1:?LCOV path is required}
|
||||
rust_cov_upload_flag=${2:?coverage flag is required}
|
||||
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
|
||||
|
||||
if [ "$(uname -m)" != "x86_64" ]; then
|
||||
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|
||||
|| return 1
|
||||
curl -fsSL \
|
||||
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
|
||||
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
|
||||
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
|
||||
| sha256sum -c - || return 1
|
||||
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
|
||||
|
||||
rust_cov_upload_slug="vllm-project/vllm"
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
|
||||
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
|
||||
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
|
||||
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
|
||||
case "$rust_cov_upload_slug" in
|
||||
*/*) ;;
|
||||
*) rust_cov_upload_slug="vllm-project/vllm" ;;
|
||||
esac
|
||||
fi
|
||||
|
||||
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
|
||||
if [ -z "${CODECOV_TOKEN:-}" ]; then
|
||||
# Codecov accepts tokenless public uploads on unprotected branch names.
|
||||
# A colon-separated prefix keeps feature-branch and fork uploads from
|
||||
# requiring a repository secret.
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
|
||||
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
|
||||
else
|
||||
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
|
||||
fi
|
||||
fi
|
||||
|
||||
set --
|
||||
set -- "$@" upload-process
|
||||
set -- "$@" --file "$rust_cov_upload_lcov"
|
||||
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
|
||||
# source-line fix scanning, which is unrelated to path mapping.
|
||||
set -- "$@" --disable-search --disable-file-fixes
|
||||
set -- "$@" --fail-on-error --git-service github
|
||||
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
|
||||
set -- "$@" --branch "$rust_cov_upload_branch"
|
||||
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
|
||||
set -- "$@" --slug "$rust_cov_upload_slug"
|
||||
set -- "$@" --flag "$rust_cov_upload_flag"
|
||||
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
|
||||
set -- "$@" --dir "$rust_cov_upload_repo_root"
|
||||
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
|
||||
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
|
||||
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
|
||||
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
|
||||
fi
|
||||
|
||||
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
|
||||
# E2E steps run from tests/, so execute from the repository root to resolve
|
||||
# codecov.yml and repository paths consistently.
|
||||
(
|
||||
cd "$rust_cov_upload_repo_root" || exit 1
|
||||
"$rust_cov_upload_codecov_dir/codecov" "$@"
|
||||
) >"$rust_cov_upload_log" 2>&1
|
||||
rust_cov_upload_rc=$?
|
||||
cat "$rust_cov_upload_log"
|
||||
# v11.3.1 can log API failures while returning zero even with
|
||||
# --fail-on-error. Preserve the strict CI contract explicitly.
|
||||
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
|
||||
echo "Codecov CLI reported an upload error" >&2
|
||||
rust_cov_upload_rc=1
|
||||
fi
|
||||
rm -rf "$rust_cov_upload_codecov_dir"
|
||||
return "$rust_cov_upload_rc"
|
||||
}
|
||||
|
||||
rust_coverage_finalize() {
|
||||
rust_cov_finalize_test_rc=$?
|
||||
trap - 0
|
||||
set +e
|
||||
|
||||
rust_coverage_collect "$RUST_COVERAGE_FLAG"
|
||||
rust_cov_finalize_collect_rc=$?
|
||||
|
||||
rust_cov_finalize_upload_rc=0
|
||||
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
|
||||
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
|
||||
rust_cov_finalize_upload_rc=$?
|
||||
fi
|
||||
|
||||
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
|
||||
|
||||
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
|
||||
exit "$rust_cov_finalize_test_rc"
|
||||
fi
|
||||
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
|
||||
exit "$rust_cov_finalize_collect_rc"
|
||||
fi
|
||||
exit "$rust_cov_finalize_upload_rc"
|
||||
}
|
||||
@@ -10,9 +10,7 @@ steps:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
@@ -36,9 +34,7 @@ steps:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
@@ -8,11 +8,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/benchmarks/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
@@ -28,7 +23,6 @@ steps:
|
||||
- tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
@@ -49,11 +43,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
@@ -65,7 +54,6 @@ steps:
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
@@ -84,16 +72,10 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
@@ -104,17 +86,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/tool_parsers/
|
||||
- tests/utils.py
|
||||
- tests/tool_use/
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
@@ -125,11 +101,6 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- build_rust.sh
|
||||
- tools/build_rust.py
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
@@ -140,7 +111,6 @@ steps:
|
||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||
- tests/v1/distributed/test_internal_lb_dp.py
|
||||
commands:
|
||||
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
|
||||
@@ -26,7 +26,5 @@ steps:
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
||||
- .buildkite/scripts/rust-coverage.sh
|
||||
- codecov.yml
|
||||
commands:
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
|
||||
|
||||
@@ -47,7 +47,6 @@
|
||||
|
||||
# Rust Frontend
|
||||
/rust/ @BugenZhao @njhill
|
||||
/rust/src/bench @esmeetu
|
||||
/build_rust.sh @BugenZhao @njhill
|
||||
/rust-toolchain.toml @BugenZhao @njhill
|
||||
/.buildkite/test_areas/rust* @BugenZhao @njhill
|
||||
|
||||
@@ -257,4 +257,3 @@ vllm/grpc/vllm_engine_pb2.pyi
|
||||
|
||||
# Ignore generated cpu headers
|
||||
csrc/cpu/cpu_attn_dispatch_generated.h
|
||||
rust-coverage-tools/
|
||||
|
||||
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
|
||||
- Tool calling and reasoning parsers
|
||||
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
|
||||
- Efficient multi-LoRA support for dense and MoE layers
|
||||
- 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.
|
||||
- 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.
|
||||
|
||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||
|
||||
|
||||
@@ -8,8 +8,6 @@
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
|
||||
CARGO_LLVM_COV_VERSION="0.8.7"
|
||||
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
|
||||
|
||||
# Read the required toolchain from rust-toolchain.toml.
|
||||
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
||||
@@ -32,39 +30,4 @@ else
|
||||
PROFILE_ARG="--release"
|
||||
fi
|
||||
|
||||
rm -rf "$COVERAGE_TOOLS_DIR"
|
||||
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
|
||||
|
||||
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
|
||||
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
|
||||
# Keep instrumented intermediates isolated when local builds switch modes.
|
||||
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
|
||||
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
|
||||
cargo +"$TOOLCHAIN" install \
|
||||
--locked \
|
||||
--version "$CARGO_LLVM_COV_VERSION" \
|
||||
cargo-llvm-cov
|
||||
|
||||
eval "$(
|
||||
cargo +"$TOOLCHAIN" llvm-cov show-env \
|
||||
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
|
||||
--sh
|
||||
)"
|
||||
|
||||
# Build scripts and proc macros can run during compilation. Their profiles
|
||||
# are unrelated to runtime coverage and would otherwise pollute the tree.
|
||||
export LLVM_PROFILE_FILE=/dev/null
|
||||
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
|
||||
fi
|
||||
|
||||
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
|
||||
|
||||
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
|
||||
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
|
||||
--print target-libdir)")/bin"
|
||||
|
||||
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
|
||||
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
|
||||
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
|
||||
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
|
||||
fi
|
||||
|
||||
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
|
||||
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
|
||||
else()
|
||||
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
|
||||
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
|
||||
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
|
||||
|
||||
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _qutlass_fc_root)
|
||||
@@ -125,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
CUDA_ARCHS "${QUTLASS_ARCHS}"
|
||||
)
|
||||
|
||||
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
|
||||
# Keep it as its own extension (registers torch.ops._qutlass_C).
|
||||
define_extension_target(
|
||||
_qutlass_C
|
||||
DESTINATION vllm
|
||||
@@ -137,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
WITH_SOABI)
|
||||
|
||||
target_compile_definitions(_qutlass_C PRIVATE
|
||||
QUTLASS_MINIMAL_BUILD=1
|
||||
QUTLASS_DISABLE_PYBIND=1
|
||||
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL
|
||||
USE_CUDA)
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
|
||||
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
-13
@@ -10,16 +10,3 @@ fixes:
|
||||
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
|
||||
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
|
||||
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
|
||||
# Map Rust sources built in the E2E image and on Buildkite agents.
|
||||
- "/workspace/rust/::rust/"
|
||||
- "/var/lib/buildkite-agent/.*/rust/::rust/"
|
||||
|
||||
flags:
|
||||
rust-unit:
|
||||
paths:
|
||||
- rust/
|
||||
carryforward: false
|
||||
rust-e2e:
|
||||
paths:
|
||||
- rust/
|
||||
carryforward: false
|
||||
|
||||
@@ -102,9 +102,7 @@ class TileGemm82 {
|
||||
kv_cache_t* __restrict__ curr_b = b_tile;
|
||||
|
||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||
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;
|
||||
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
|
||||
@@ -39,15 +39,11 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
// AWQ zeros: (size_k // group_size, size_n // 8)
|
||||
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
|
||||
int32_t group_size) {
|
||||
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
|
||||
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
|
||||
// same row (stride 1) instead of striding across rows (stride size_n/8).
|
||||
int col = blockIdx.y * 32 + threadIdx.x;
|
||||
if (col >= size_n / 8) return;
|
||||
(void)size_k;
|
||||
|
||||
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
|
||||
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
|
||||
int32_t val =
|
||||
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
|
||||
int32_t zero =
|
||||
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
|
||||
blockIdx.y];
|
||||
int32_t new_val = 0;
|
||||
|
||||
#pragma unroll
|
||||
@@ -62,7 +58,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
zero >>= 4;
|
||||
}
|
||||
|
||||
output[blockIdx.x * (size_n / 8) + col] = new_val;
|
||||
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
|
||||
}
|
||||
|
||||
torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
@@ -106,7 +102,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
"qweight.size(0) % qzeros.size(0) != 0");
|
||||
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
|
||||
|
||||
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
|
||||
dim3 blocks(size_k / 32, size_n / 8);
|
||||
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
|
||||
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
|
||||
|
||||
@@ -294,9 +294,6 @@ FROM base AS rust-build
|
||||
ARG BUILD_OS
|
||||
ARG USE_SCCACHE
|
||||
ARG SCCACHE_ENDPOINT
|
||||
# Temporary default for the initial CI validation. Set this back to 0 when
|
||||
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
|
||||
ARG VLLM_RUST_COVERAGE=1
|
||||
|
||||
# Install native tools needed only for Rust/protoc builds.
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
@@ -905,14 +902,6 @@ COPY ./vllm/collect_env.py .
|
||||
# note that this uses vllm installed by `pip`
|
||||
FROM vllm-base AS test
|
||||
|
||||
COPY --from=rust-build \
|
||||
/workspace/rust-coverage-tools/ \
|
||||
/opt/vllm-rust-coverage/
|
||||
|
||||
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
|
||||
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
|
||||
ENV LLVM_PROFILE_FILE=/dev/null
|
||||
|
||||
ADD . /vllm-workspace/
|
||||
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
@@ -46,9 +46,6 @@
|
||||
"TORCH_CUDA_ARCH_LIST": {
|
||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
|
||||
},
|
||||
"VLLM_RUST_COVERAGE": {
|
||||
"default": "1"
|
||||
},
|
||||
"MAX_JOBS": {
|
||||
"default": "2"
|
||||
},
|
||||
|
||||
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
|
||||
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
|
||||
@@ -75,7 +75,7 @@ vllm serve <model> \
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
|
||||
@@ -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).
|
||||
- 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.18.38308.1) 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.14.37833.4) release, to avoid potential compatibility issue.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
@@ -58,40 +58,7 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
--8<-- [start:pre-built-images]
|
||||
|
||||
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>
|
||||
```
|
||||
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).
|
||||
|
||||
--8<-- [end:pre-built-images]
|
||||
--8<-- [start:build-image-from-source]
|
||||
|
||||
@@ -65,15 +65,6 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
!!! 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.
|
||||
|
||||
=== "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"
|
||||
|
||||
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
|
||||
|
||||
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
|
||||
- `response_format`: Format of the response ("json", "text") (optional)
|
||||
- `temperature`: Sampling temperature between 0 and 1 (optional)
|
||||
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
|
||||
|
||||
**Response Format:**
|
||||
|
||||
|
||||
Generated
+11
-28
@@ -3446,6 +3446,7 @@ dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
@@ -4875,7 +4876,6 @@ dependencies = [
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4937,9 +4937,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic"
|
||||
version = "0.14.6"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
@@ -4966,9 +4966,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-build"
|
||||
version = "0.14.6"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
|
||||
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -4976,24 +4976,11 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-health"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
|
||||
dependencies = [
|
||||
"prost",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.6"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
|
||||
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"prost",
|
||||
@@ -5002,9 +4989,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost-build"
|
||||
version = "0.14.6"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
|
||||
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -5497,13 +5484,10 @@ dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tiktoken-rs 0.9.1",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
@@ -5746,7 +5730,6 @@ dependencies = [
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tonic",
|
||||
"tonic-health",
|
||||
"tonic-prost",
|
||||
"tonic-prost-build",
|
||||
"tower",
|
||||
@@ -6338,9 +6321,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
|
||||
|
||||
[[package]]
|
||||
name = "xgrammar-structural-tag"
|
||||
version = "0.2.0+xgrammar.0.2.4.dd729e7"
|
||||
version = "0.1.0+xgrammar.0.2.2.4d145cc"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
|
||||
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
|
||||
dependencies = [
|
||||
"auto_impl",
|
||||
"serde",
|
||||
|
||||
+5
-6
@@ -118,11 +118,10 @@ tokio = { version = "1.47.1", features = [
|
||||
tokio-openssl = "0.6"
|
||||
tokio-stream = "0.1"
|
||||
tokio-util = { version = "0.7.18", features = ["rt"] }
|
||||
tonic = "0.14.6"
|
||||
tonic-build = "0.14.6"
|
||||
tonic-health = "0.14.6"
|
||||
tonic-prost = "0.14.6"
|
||||
tonic-prost-build = "0.14.6"
|
||||
tonic = "0.14.5"
|
||||
tonic-build = "0.14.5"
|
||||
tonic-prost = "0.14.5"
|
||||
tonic-prost-build = "0.14.5"
|
||||
tool-parser = "1.2.0"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
|
||||
@@ -144,7 +143,7 @@ vllm-server = { path = "src/server" }
|
||||
vllm-text = { path = "src/text" }
|
||||
vllm-tokenizer = { path = "src/tokenizer" }
|
||||
winnow = { version = "1.0.2", features = ["simd"] }
|
||||
xgrammar-structural-tag = "0.2.0"
|
||||
xgrammar-structural-tag = "0.1.0"
|
||||
zeromq = { version = "0.6.0", default-features = false, features = [
|
||||
"tokio-runtime",
|
||||
"all-transport",
|
||||
|
||||
@@ -20,19 +20,16 @@ mimalloc.workspace = true
|
||||
rand.workspace = true
|
||||
rand_distr.workspace = true
|
||||
rayon.workspace = true
|
||||
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
|
||||
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
|
||||
rlimit.workspace = true
|
||||
rustc-hash.workspace = true
|
||||
serde = { workspace = true, features = ["rc"] }
|
||||
serde_json = { workspace = true, features = ["raw_value"] }
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
tiktoken-rs.workspace = true
|
||||
tokenizers.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
|
||||
@@ -158,10 +158,9 @@ impl PoolingBackend {
|
||||
// (mirrors Python async_request_vllm_rerank).
|
||||
if let Some(ref list) = input.prompt_list {
|
||||
if list.len() < 2 {
|
||||
tracing::warn!(
|
||||
backend = "vllm-rerank",
|
||||
inputs = list.len(),
|
||||
"rerank request has no documents"
|
||||
eprintln!(
|
||||
"WARNING: vllm-rerank request has no documents \
|
||||
(prompt_list needs [query, doc, ...])"
|
||||
);
|
||||
}
|
||||
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
|
||||
@@ -176,10 +175,10 @@ impl PoolingBackend {
|
||||
// Legacy path: text prompt as query, documents via --extra-body.
|
||||
let query = input.prompt.as_ref();
|
||||
if query.is_empty() && input.prompt_token_ids.is_some() {
|
||||
tracing::warn!(
|
||||
backend = "vllm-rerank",
|
||||
dataset = "random",
|
||||
"rerank request has an empty query; use the random-rerank dataset"
|
||||
eprintln!(
|
||||
"WARNING: vllm-rerank received empty query (random dataset uses \
|
||||
token IDs only). Use --dataset-name random-rerank for meaningful \
|
||||
rerank benchmarks."
|
||||
);
|
||||
}
|
||||
serde_json::json!({
|
||||
|
||||
+84
-114
@@ -6,7 +6,6 @@ use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
|
||||
@@ -73,12 +72,12 @@ pub fn pre_resolve_dns(
|
||||
v4.extend(v6);
|
||||
if !v4.is_empty() {
|
||||
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
|
||||
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
|
||||
println!("Pre-resolved {host} -> {ips:?}");
|
||||
builder = builder.resolve_to_addrs(host, &v4);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
|
||||
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -347,14 +346,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||
(m.clone(), config.model_name.clone())
|
||||
} else {
|
||||
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||
println!("Model not specified, fetching first model from server...");
|
||||
let (name, id) =
|
||||
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
|
||||
tracing::info!(
|
||||
model_name = name,
|
||||
model_id = id,
|
||||
"selected first model from server"
|
||||
);
|
||||
println!("First model name: {name}, first model id: {id}");
|
||||
(id, Some(name))
|
||||
};
|
||||
|
||||
@@ -363,10 +358,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
None
|
||||
} else {
|
||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||
println!("Loading tokenizer: {tid}");
|
||||
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).await?;
|
||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
||||
println!("Tokenizer loaded successfully.");
|
||||
Some(t)
|
||||
};
|
||||
let has_tokenizer = tokenizer.is_some();
|
||||
@@ -426,12 +421,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.num_prompts, config.random_batch_size, config.is_reranker,
|
||||
),
|
||||
};
|
||||
tracing::info!(
|
||||
dataset = ?config.dataset_name,
|
||||
prompts = config.num_prompts,
|
||||
description = %dataset_label,
|
||||
"generating benchmark dataset"
|
||||
);
|
||||
println!("Generating {dataset_label}...");
|
||||
let gen_start = Instant::now();
|
||||
|
||||
let mut input_requests = match config.dataset_name {
|
||||
@@ -482,7 +472,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let path = match config.dataset_path.as_deref() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -522,8 +512,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
None => {
|
||||
downloaded = crate::datasets::speed_bench::download_speed_bench(
|
||||
config.speed_bench_config,
|
||||
)
|
||||
.await?;
|
||||
)?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -554,8 +543,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.hf_subset.as_deref(),
|
||||
config.hf_split.as_deref(),
|
||||
config.num_prompts,
|
||||
)
|
||||
.await?;
|
||||
)?;
|
||||
crate::datasets::hf_dataset::load_hf_dataset(
|
||||
tok,
|
||||
&downloaded_path,
|
||||
@@ -620,19 +608,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
};
|
||||
|
||||
let gen_elapsed = gen_start.elapsed();
|
||||
tracing::info!(
|
||||
prompts = input_requests.len(),
|
||||
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||
"generated benchmark dataset"
|
||||
println!(
|
||||
"Generated {} prompts in {:.2}s",
|
||||
input_requests.len(),
|
||||
gen_elapsed.as_secs_f64()
|
||||
);
|
||||
|
||||
let filtered_count =
|
||||
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
|
||||
if filtered_count > 0 {
|
||||
tracing::info!(
|
||||
filtered_prompts = filtered_count,
|
||||
max_model_len = config.max_model_len.unwrap(),
|
||||
"filtered prompts above maximum model length"
|
||||
println!(
|
||||
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
|
||||
config.max_model_len.unwrap()
|
||||
);
|
||||
}
|
||||
if input_requests.is_empty() {
|
||||
@@ -683,7 +670,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
|
||||
// Ready check
|
||||
if config.ready_check_timeout_sec > 0 {
|
||||
tracing::info!("starting initial single-prompt test run");
|
||||
println!("Starting initial single prompt test run...");
|
||||
let test_output = wait_for_endpoint(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -698,7 +685,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
test_output.error
|
||||
)));
|
||||
}
|
||||
tracing::info!("initial single-prompt test run completed");
|
||||
println!("Initial test run completed.");
|
||||
}
|
||||
|
||||
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
|
||||
@@ -716,15 +703,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
DatasetName::Random | DatasetName::PrefixRepetition
|
||||
);
|
||||
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
|
||||
tracing::info!(
|
||||
reason = "prompt_token_ids",
|
||||
"skipping server tokenizer verification"
|
||||
);
|
||||
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
|
||||
}
|
||||
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
|
||||
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
|
||||
if is_tokenizer_verified(&cache_key) {
|
||||
tracing::info!(reason = "cached", "skipping server tokenizer verification");
|
||||
println!("Tokenizer verified in previous run (cached), skipping verification.");
|
||||
} else {
|
||||
let num_special =
|
||||
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
|
||||
@@ -739,17 +723,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
.await?
|
||||
{
|
||||
SampleVerifyOutcome::Passed => {
|
||||
tracing::info!("tokenizer sample verification passed");
|
||||
println!("Sample verification passed, skipping full verification.");
|
||||
mark_tokenizer_verified(&cache_key);
|
||||
}
|
||||
SampleVerifyOutcome::Skipped(reason) => {
|
||||
tracing::warn!(
|
||||
reason = %reason,
|
||||
"server tokenizer unavailable; skipping prompt verification"
|
||||
);
|
||||
println!("Server /tokenize unavailable ({reason}), skipping verification.");
|
||||
}
|
||||
SampleVerifyOutcome::Mismatch => {
|
||||
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
|
||||
println!("Sample verification found mismatch, running full verify+fix...");
|
||||
match verify_and_fix_prompt_lengths(
|
||||
&client,
|
||||
&config.base_url,
|
||||
@@ -761,16 +742,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
tracing::info!(
|
||||
prompts = input_requests.len(),
|
||||
"verified exact prompt token lengths"
|
||||
println!(
|
||||
"All {} prompts verified: exact token length match.",
|
||||
input_requests.len()
|
||||
);
|
||||
mark_tokenizer_verified(&cache_key);
|
||||
}
|
||||
Err(BenchError::TokenizeUnavailable(reason)) => {
|
||||
tracing::warn!(
|
||||
reason = %reason,
|
||||
"server tokenizer became unavailable; using client token counts"
|
||||
println!(
|
||||
"Server /tokenize became unavailable during verification \
|
||||
({reason}); proceeding with client-side token counts."
|
||||
);
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
@@ -782,7 +763,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
|
||||
// Warmup
|
||||
if config.num_warmups > 0 {
|
||||
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
|
||||
println!("Warming up with {} requests...", config.num_warmups);
|
||||
run_warmup(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -795,7 +776,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.disable_tqdm,
|
||||
)
|
||||
.await;
|
||||
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
|
||||
println!("Warmup run completed.");
|
||||
}
|
||||
|
||||
// Start profiler if requested (immediate mode — no batch threshold)
|
||||
@@ -833,22 +814,28 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let spec_decode_before =
|
||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||
if spec_decode_before.is_some() {
|
||||
tracing::info!("detected speculative decoding; collecting metrics");
|
||||
println!("Speculative decoding detected, will collect metrics.");
|
||||
}
|
||||
|
||||
// Main benchmark
|
||||
println!("Starting main benchmark run...");
|
||||
let distribution = if config.burstiness == 1.0 {
|
||||
"Poisson process"
|
||||
} else {
|
||||
"Gamma distribution"
|
||||
};
|
||||
tracing::info!(
|
||||
request_rate = config.request_rate,
|
||||
burstiness = config.burstiness,
|
||||
distribution,
|
||||
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
|
||||
prompts = config.num_prompts,
|
||||
"starting main benchmark run"
|
||||
println!(
|
||||
"Traffic request rate: {}",
|
||||
if config.request_rate.is_infinite() {
|
||||
"inf".to_string()
|
||||
} else {
|
||||
format!("{}", config.request_rate)
|
||||
}
|
||||
);
|
||||
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).
|
||||
@@ -860,11 +847,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()) {
|
||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||
tracing::info!(
|
||||
adapters = modules.len(),
|
||||
names = ?names,
|
||||
assignment = ?config.lora_assignment,
|
||||
"assigned LoRA adapters"
|
||||
println!(
|
||||
"LoRA adapters ({}): {:?} [assignment={:?}]",
|
||||
modules.len(),
|
||||
names,
|
||||
config.lora_assignment
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1138,7 +1125,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
if let Some((cancel_tx, task)) = profile_task {
|
||||
let _ = cancel_tx.send(());
|
||||
if let Err(e) = task.await {
|
||||
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||
eprintln!("WARNING: Profile background task failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1302,14 +1289,12 @@ pub(crate) async fn start_profiler_immediate(
|
||||
base_url: &str,
|
||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||
) {
|
||||
println!("Starting profiler...");
|
||||
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 {
|
||||
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
|
||||
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
|
||||
}
|
||||
Ok(true) => println!("Profiler started"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1319,14 +1304,12 @@ pub(crate) async fn stop_profiler_immediate(
|
||||
base_url: &str,
|
||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||
) {
|
||||
println!("Stopping profiler...");
|
||||
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 {
|
||||
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
|
||||
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
|
||||
}
|
||||
Ok(true) => println!("Profiler stopped"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1388,30 +1371,25 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
duration_secs: f64,
|
||||
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
|
||||
) {
|
||||
tracing::info!(
|
||||
threshold,
|
||||
duration_seconds = duration_secs,
|
||||
"waiting for profiler batch threshold"
|
||||
println!(
|
||||
"Waiting for batch size >= {threshold} before starting profiler \
|
||||
(will capture {duration_secs}s)..."
|
||||
);
|
||||
|
||||
loop {
|
||||
if let Some(running) = fetch_num_requests_running(client, base_url).await
|
||||
&& running >= threshold
|
||||
{
|
||||
tracing::info!(
|
||||
running_requests = running,
|
||||
threshold,
|
||||
"profiler batch threshold reached"
|
||||
);
|
||||
println!("Batch size {running} >= {threshold}, starting profiler...");
|
||||
break;
|
||||
}
|
||||
// Wait 500ms or until the benchmark signals cancellation
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
|
||||
_ = &mut cancel_rx => {
|
||||
tracing::warn!(
|
||||
threshold,
|
||||
"benchmark finished before profiler batch threshold; skipping profiling"
|
||||
eprintln!(
|
||||
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
|
||||
profiling skipped."
|
||||
);
|
||||
return;
|
||||
}
|
||||
@@ -1420,13 +1398,13 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
|
||||
let start_url = format!("{base_url}/start_profile");
|
||||
match send_profile_request(client, &start_url, extra_headers).await {
|
||||
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
|
||||
Ok(true) => println!("Profiler started"),
|
||||
Ok(false) => {
|
||||
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
|
||||
eprintln!("WARNING: Profiler start request returned non-success");
|
||||
return;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
|
||||
eprintln!("WARNING: Failed to start profiler: {e}");
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -1435,17 +1413,15 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
|
||||
_ = &mut cancel_rx => {
|
||||
tracing::info!("benchmark finished; stopping profiler early");
|
||||
println!("Benchmark finished, stopping profiler early...");
|
||||
}
|
||||
}
|
||||
|
||||
let stop_url = format!("{base_url}/stop_profile");
|
||||
match send_profile_request(client, &stop_url, extra_headers).await {
|
||||
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
|
||||
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
|
||||
}
|
||||
Ok(true) => println!("Profiler stopped after capturing"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1526,11 +1502,10 @@ async fn verify_and_fix_prompt_lengths(
|
||||
let excess = tokens.len().saturating_sub(expected_input_len);
|
||||
let compensate = if excess > 0 && last_excess == Some(excess) {
|
||||
if _iter == 1 {
|
||||
tracing::warn!(
|
||||
prompt_index = i,
|
||||
extra_tokens = excess,
|
||||
adjusted_target = expected_input_len.saturating_sub(excess),
|
||||
"server consistently adds prompt tokens; compensating verification target"
|
||||
eprintln!(
|
||||
"Prompt {i}: server consistently adds {excess} extra token(s) \
|
||||
(likely BOS), compensating target to {}.",
|
||||
expected_input_len.saturating_sub(excess),
|
||||
);
|
||||
}
|
||||
excess
|
||||
@@ -1588,10 +1563,7 @@ async fn verify_and_fix_prompt_lengths(
|
||||
|
||||
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
|
||||
if fc > 0 {
|
||||
tracing::info!(
|
||||
fixed_prompts = fc,
|
||||
"fixed prompt lengths using server tokenizer"
|
||||
);
|
||||
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -1846,7 +1818,7 @@ async fn sample_verify_prompts(
|
||||
let tokenize_url = format!("{base_url}/tokenize");
|
||||
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
|
||||
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
|
||||
println!("Sampling {sample_size} prompts for verification...");
|
||||
|
||||
for (i, request) in requests.iter().enumerate().take(sample_size) {
|
||||
let tokens = match server_tokenize(
|
||||
@@ -1869,11 +1841,9 @@ async fn sample_verify_prompts(
|
||||
|
||||
let expected = request.prompt_len + num_special;
|
||||
if tokens.len() != expected {
|
||||
tracing::warn!(
|
||||
prompt_index = i,
|
||||
expected_tokens = expected,
|
||||
actual_tokens = tokens.len(),
|
||||
"tokenizer verification sample mismatch"
|
||||
println!(
|
||||
"Prompt {i}: expected {expected} tokens, server returned {}",
|
||||
tokens.len()
|
||||
);
|
||||
return Ok(SampleVerifyOutcome::Mismatch);
|
||||
}
|
||||
|
||||
@@ -288,18 +288,10 @@ impl BenchConfig {
|
||||
}
|
||||
Some(other) => {
|
||||
// extra_body was not an object — just use sampling params
|
||||
let value_type = match &other {
|
||||
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"
|
||||
eprintln!(
|
||||
"Warning: --extra-body is not a JSON object, sampling params may be lost"
|
||||
);
|
||||
let _ = other;
|
||||
sampling_params
|
||||
}
|
||||
None => sampling_params,
|
||||
@@ -497,9 +489,9 @@ impl BenchConfig {
|
||||
_ => {}
|
||||
}
|
||||
if !args.skip_chat_template {
|
||||
tracing::warn!(
|
||||
dataset = "custom",
|
||||
"client-side chat template rendering is unsupported; sending prompts raw"
|
||||
eprintln!(
|
||||
"NOTE: client-side chat template rendering is not supported; custom \
|
||||
dataset prompts are sent raw (equivalent to --skip-chat-template)."
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -578,10 +570,9 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
if ignore_eos {
|
||||
tracing::warn!(
|
||||
ignore_eos,
|
||||
multi_turn = true,
|
||||
"output length limits may be ignored, causing unbounded context growth"
|
||||
eprintln!(
|
||||
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
|
||||
respect output length limits, causing unbounded context growth."
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -128,10 +128,8 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::SampleRequest;
|
||||
use super::progress::RowDownloadReporter;
|
||||
use crate::error::{BenchError, Result};
|
||||
use crate::tokenizer::TokenizerKind;
|
||||
|
||||
@@ -51,19 +50,18 @@ enum ColumnFormat {
|
||||
|
||||
/// Make a GET request with retry logic (3 retries with exponential backoff).
|
||||
/// Returns the parsed JSON response.
|
||||
async fn get_with_retry(
|
||||
client: &reqwest::Client,
|
||||
fn get_with_retry(
|
||||
client: &reqwest::blocking::Client,
|
||||
url: &str,
|
||||
label: &str,
|
||||
) -> Result<serde_json::Value> {
|
||||
let max_retries = 3;
|
||||
for attempt in 0..=max_retries {
|
||||
let resp = match client.get(url).send().await {
|
||||
let resp = match client.get(url).send() {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
|
||||
.await;
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
continue;
|
||||
}
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -82,7 +80,7 @@ async fn get_with_retry(
|
||||
}
|
||||
|
||||
if status.is_server_error() && attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -94,7 +92,6 @@ async fn get_with_retry(
|
||||
|
||||
let data: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
|
||||
return Ok(data);
|
||||
}
|
||||
@@ -108,7 +105,7 @@ async fn get_with_retry(
|
||||
/// 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
|
||||
/// or the dataset is exhausted.
|
||||
pub async fn download_hf_dataset(
|
||||
pub fn download_hf_dataset(
|
||||
dataset: &str,
|
||||
subset: Option<&str>,
|
||||
split: Option<&str>,
|
||||
@@ -118,7 +115,7 @@ pub async fn download_hf_dataset(
|
||||
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
|
||||
|
||||
let mut client_builder =
|
||||
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
|
||||
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
|
||||
|
||||
// Add HF_TOKEN auth header if available
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
@@ -141,7 +138,7 @@ pub async fn download_hf_dataset(
|
||||
// Call /info to discover available configs and splits
|
||||
let info_url =
|
||||
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
|
||||
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
|
||||
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
|
||||
|
||||
let dataset_info =
|
||||
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
|
||||
@@ -204,12 +201,7 @@ pub async fn download_hf_dataset(
|
||||
(resolved_config, resolved_split)
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
dataset,
|
||||
config = resolved_config,
|
||||
split = resolved_split,
|
||||
"resolved Hugging Face dataset"
|
||||
);
|
||||
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
|
||||
|
||||
// Check cache
|
||||
let dir = cache_dir();
|
||||
@@ -223,16 +215,11 @@ pub async fn download_hf_dataset(
|
||||
|
||||
if cache_path.exists() {
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
|
||||
println!("HF dataset cached: {path_str}");
|
||||
return Ok((path_str, resolved_config, resolved_split));
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
dataset,
|
||||
config = resolved_config,
|
||||
split = resolved_split,
|
||||
"downloading Hugging Face dataset"
|
||||
);
|
||||
println!("Downloading HF dataset '{dataset}' from datasets-server...");
|
||||
|
||||
let encoded_config: String =
|
||||
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
|
||||
@@ -242,7 +229,6 @@ pub async fn download_hf_dataset(
|
||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||
let mut offset = 0usize;
|
||||
let page_size = 100usize;
|
||||
let mut progress = RowDownloadReporter::new();
|
||||
|
||||
loop {
|
||||
let url = format!(
|
||||
@@ -254,7 +240,7 @@ pub async fn download_hf_dataset(
|
||||
&length={page_size}"
|
||||
);
|
||||
|
||||
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
|
||||
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
|
||||
|
||||
let rows = data["rows"]
|
||||
.as_array()
|
||||
@@ -274,14 +260,14 @@ pub async fn download_hf_dataset(
|
||||
offset += fetched;
|
||||
|
||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||
progress.update(offset, total);
|
||||
eprint!("\r Fetched {offset}/{total} rows...");
|
||||
|
||||
// Stop if we have enough rows or reached end of dataset
|
||||
if all_rows.len() >= num_rows_needed || fetched < page_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
progress.finish();
|
||||
eprintln!(); // newline after progress
|
||||
|
||||
if all_rows.is_empty() {
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -294,12 +280,7 @@ pub async fn download_hf_dataset(
|
||||
std::fs::write(&cache_path, &json_str)?;
|
||||
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(
|
||||
dataset,
|
||||
rows = all_rows.len(),
|
||||
path = %path_str,
|
||||
"saved Hugging Face dataset"
|
||||
);
|
||||
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
|
||||
Ok((path_str, resolved_config, resolved_split))
|
||||
}
|
||||
|
||||
@@ -502,31 +483,21 @@ pub fn load_hf_dataset(
|
||||
// Detect column format from first row
|
||||
let format = detect_column_format(&entries[0], text_column_override)?;
|
||||
|
||||
// Print detected format
|
||||
match &format {
|
||||
ColumnFormat::Chat(col) => {
|
||||
tracing::info!(
|
||||
format = "chat",
|
||||
column = col,
|
||||
"detected Hugging Face dataset format"
|
||||
);
|
||||
}
|
||||
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
|
||||
ColumnFormat::Text {
|
||||
prompt_col,
|
||||
output_col,
|
||||
} => {
|
||||
tracing::info!(
|
||||
format = "text",
|
||||
prompt_column = prompt_col,
|
||||
output_column = output_col.as_deref().unwrap_or("none"),
|
||||
"detected Hugging Face dataset format"
|
||||
);
|
||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
||||
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
|
||||
}
|
||||
ColumnFormat::Combined { cols, output_col } => {
|
||||
tracing::info!(
|
||||
format = "combined",
|
||||
prompt_columns = ?cols,
|
||||
output_column = output_col.as_deref().unwrap_or("none"),
|
||||
"detected Hugging Face dataset format"
|
||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
||||
println!(
|
||||
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
|
||||
cols
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -637,10 +608,9 @@ pub fn load_hf_dataset(
|
||||
if len == 0 { 128 } else { len }
|
||||
} else {
|
||||
if !warned_no_output {
|
||||
tracing::warn!(
|
||||
path = dataset_path,
|
||||
default_output_tokens = 128,
|
||||
"no dataset output column or --hf-output-len; using default output length"
|
||||
eprintln!(
|
||||
"WARNING: No output column detected and --hf-output-len not set. \
|
||||
Using default output length of 128 tokens."
|
||||
);
|
||||
warned_no_output = true;
|
||||
}
|
||||
@@ -648,10 +618,9 @@ pub fn load_hf_dataset(
|
||||
}
|
||||
} else {
|
||||
if !warned_no_output {
|
||||
tracing::warn!(
|
||||
path = dataset_path,
|
||||
default_output_tokens = 128,
|
||||
"no dataset output column or --hf-output-len; using default output length"
|
||||
eprintln!(
|
||||
"WARNING: No output column detected and --hf-output-len not set. \
|
||||
Using default output length of 128 tokens."
|
||||
);
|
||||
warned_no_output = true;
|
||||
}
|
||||
@@ -671,11 +640,9 @@ pub fn load_hf_dataset(
|
||||
// Oversample if needed
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
tracing::info!(
|
||||
dataset = "hf",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let original_len = samples.len();
|
||||
@@ -685,11 +652,9 @@ pub fn load_hf_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
tracing::info!(
|
||||
dataset = "hf",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
println!(
|
||||
"Oversampled HF dataset from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1037,10 +1002,8 @@ mod tests {
|
||||
|
||||
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
|
||||
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
|
||||
crate::tokenizer::TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
}
|
||||
|
||||
/// Write JSON data to a unique temp file and return the path string.
|
||||
|
||||
@@ -5,7 +5,6 @@ pub mod custom;
|
||||
pub mod hf_dataset;
|
||||
pub mod multi_turn;
|
||||
pub mod prefix_repetition;
|
||||
mod progress;
|
||||
pub mod random;
|
||||
pub mod random_mm;
|
||||
pub mod random_rerank;
|
||||
@@ -91,10 +90,9 @@ pub fn oversample_requests(
|
||||
return;
|
||||
}
|
||||
if no_oversample {
|
||||
tracing::info!(
|
||||
samples = requests.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
requests.len()
|
||||
);
|
||||
return;
|
||||
}
|
||||
@@ -105,10 +103,9 @@ pub fn oversample_requests(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
requests.push(req);
|
||||
}
|
||||
tracing::info!(
|
||||
original_samples = original_len,
|
||||
samples = requests.len(),
|
||||
"oversampled dataset"
|
||||
println!(
|
||||
"Oversampled requests from {original_len} to {} total samples.",
|
||||
requests.len()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -445,10 +445,9 @@ pub fn load_sharegpt_multi_turn(
|
||||
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
|
||||
conversations.push(conv);
|
||||
}
|
||||
tracing::info!(
|
||||
original_conversations = original_len,
|
||||
conversations = conversations.len(),
|
||||
"oversampled multi-turn conversations"
|
||||
println!(
|
||||
"Oversampled multi-turn conversations from {original_len} to {} total.",
|
||||
conversations.len()
|
||||
);
|
||||
}
|
||||
|
||||
@@ -526,12 +525,10 @@ mod tests {
|
||||
len
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_prefix_sharing_structure() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
fn test_prefix_sharing_structure() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 5,
|
||||
@@ -613,12 +610,10 @@ mod tests {
|
||||
println!("All prefix sharing checks passed!");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_per_turn_input_len_default_mode() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
fn test_per_turn_input_len_default_mode() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 4,
|
||||
@@ -655,12 +650,10 @@ mod tests {
|
||||
println!("per_turn_input_len default-mode checks passed!");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_variable_turns_range() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
fn test_variable_turns_range() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 50,
|
||||
@@ -691,12 +684,10 @@ mod tests {
|
||||
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_variable_turns_fixed() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
fn test_variable_turns_fixed() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 10,
|
||||
@@ -718,12 +709,10 @@ mod tests {
|
||||
println!("variable_turns_fixed checks passed!");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_per_turn_input_len_prefix_sharing() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
fn test_per_turn_input_len_prefix_sharing() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
|
||||
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
|
||||
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
|
||||
|
||||
@@ -41,13 +41,11 @@ pub fn generate_prefix_repetition_dataset(
|
||||
}
|
||||
let total = prompts_per_prefix * num_prefixes;
|
||||
if total != num_requests {
|
||||
tracing::info!(
|
||||
requested = num_requests,
|
||||
generated = total,
|
||||
prefixes = num_prefixes,
|
||||
prompts_per_prefix,
|
||||
dropped = num_requests - total,
|
||||
"adjusted prefix-repetition request count"
|
||||
println!(
|
||||
"prefix_repetition: generating {total} requests \
|
||||
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
|
||||
{} dropped to divide evenly)",
|
||||
num_requests - total
|
||||
);
|
||||
}
|
||||
|
||||
@@ -111,10 +109,8 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
|
||||
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Reports row download progress to an interactive progress bar, or through
|
||||
/// periodic tracing events when the progress bar is hidden on a non-TTY.
|
||||
pub(super) struct RowDownloadReporter {
|
||||
progress: ProgressBar,
|
||||
next_report: Instant,
|
||||
}
|
||||
|
||||
impl RowDownloadReporter {
|
||||
/// Creates a reporter that emits non-TTY updates every 10 seconds.
|
||||
pub fn new() -> Self {
|
||||
let progress = ProgressBar::new(0);
|
||||
progress.set_style(
|
||||
ProgressStyle::with_template(
|
||||
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
|
||||
)
|
||||
.unwrap()
|
||||
.progress_chars("#>-"),
|
||||
);
|
||||
Self {
|
||||
progress,
|
||||
next_report: Instant::now() + REPORT_INTERVAL,
|
||||
}
|
||||
}
|
||||
|
||||
/// Updates the current row count and reports progress when due.
|
||||
pub fn update(&mut self, rows: usize, total: u64) {
|
||||
let rows = rows as u64;
|
||||
let total = total.max(rows);
|
||||
self.progress.set_length(total);
|
||||
self.progress.set_position(rows);
|
||||
|
||||
if self.should_report(Instant::now()) {
|
||||
tracing::info!(rows, total, "fetching dataset rows");
|
||||
}
|
||||
}
|
||||
|
||||
/// Clears the interactive progress bar after the download completes.
|
||||
pub fn finish(self) {
|
||||
self.progress.finish_and_clear();
|
||||
}
|
||||
|
||||
fn should_report(&mut self, now: Instant) -> bool {
|
||||
if !self.progress.is_hidden() || now < self.next_report {
|
||||
return false;
|
||||
}
|
||||
self.next_report = now + REPORT_INTERVAL;
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn hidden_reporter_uses_ten_second_deadline() {
|
||||
let start = Instant::now();
|
||||
let mut reporter = RowDownloadReporter {
|
||||
progress: ProgressBar::hidden(),
|
||||
next_report: start + REPORT_INTERVAL,
|
||||
};
|
||||
|
||||
assert!(!reporter.should_report(start + Duration::from_secs(9)));
|
||||
assert!(reporter.should_report(start + Duration::from_secs(10)));
|
||||
assert!(!reporter.should_report(start + Duration::from_secs(19)));
|
||||
assert!(reporter.should_report(start + Duration::from_secs(20)));
|
||||
}
|
||||
}
|
||||
@@ -49,12 +49,9 @@ pub fn generate_random_dataset(
|
||||
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
|
||||
let (output_low, output_high) = range_ratio.output_bounds(output_len);
|
||||
if !range_ratio.is_fixed() {
|
||||
tracing::info!(
|
||||
input_low,
|
||||
input_high,
|
||||
output_low,
|
||||
output_high,
|
||||
"sampling random request lengths"
|
||||
println!(
|
||||
"Sampling input_len from [{input_low}, {input_high}] and \
|
||||
output_len from [{output_low}, {output_high}]"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -308,8 +305,7 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_generate_random_dataset_token_ids() {
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
10, // num_requests
|
||||
@@ -341,8 +337,7 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_generate_random_dataset_text() {
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
10, // num_requests
|
||||
@@ -376,8 +371,7 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_token_length_exact_local() {
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let target_len = 512;
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
@@ -411,11 +405,11 @@ mod tests {
|
||||
}
|
||||
|
||||
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_token_length_exact_tiktoken() {
|
||||
fn test_token_length_exact_tiktoken() {
|
||||
// Use Qwen2.5 which has a tiktoken-format tokenizer
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
||||
let tokenizer = match tokenizer {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
@@ -459,10 +453,10 @@ mod tests {
|
||||
|
||||
/// Test encode/decode roundtrip stability for tiktoken.
|
||||
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
|
||||
#[tokio::test]
|
||||
#[test]
|
||||
#[ignore]
|
||||
async fn test_tiktoken_roundtrip_stability() {
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||
fn test_tiktoken_roundtrip_stability() {
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
||||
let tokenizer = match tokenizer {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
|
||||
@@ -139,10 +139,8 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
}
|
||||
|
||||
fn fixed_ratio() -> RangeRatio {
|
||||
|
||||
@@ -22,21 +22,18 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
|
||||
|
||||
/// Download the default ShareGPT dataset from HuggingFace Hub.
|
||||
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
|
||||
pub async fn download_sharegpt_dataset() -> Result<String> {
|
||||
tracing::info!(
|
||||
repository = DEFAULT_SHAREGPT_REPO,
|
||||
file = DEFAULT_SHAREGPT_FILE,
|
||||
"downloading ShareGPT dataset"
|
||||
pub fn download_sharegpt_dataset() -> Result<String> {
|
||||
println!(
|
||||
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
|
||||
);
|
||||
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
|
||||
.map_err(BenchError::Config)?;
|
||||
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
|
||||
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
|
||||
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
|
||||
BenchError::Config(format!(
|
||||
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
|
||||
))
|
||||
})?;
|
||||
let path_str = path.to_string_lossy().to_string();
|
||||
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
|
||||
println!("ShareGPT dataset ready: {path_str}");
|
||||
Ok(path_str)
|
||||
}
|
||||
|
||||
@@ -138,11 +135,9 @@ pub fn load_sharegpt_dataset(
|
||||
// Oversample if dataset is smaller than requested
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
tracing::info!(
|
||||
dataset = "sharegpt",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let needed = num_requests - samples.len();
|
||||
@@ -152,11 +147,9 @@ pub fn load_sharegpt_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
tracing::info!(
|
||||
dataset = "sharegpt",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
println!(
|
||||
"Oversampled requests from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::SampleRequest;
|
||||
use super::progress::RowDownloadReporter;
|
||||
use crate::cli::SpeedBenchConfig;
|
||||
use crate::error::{BenchError, Result};
|
||||
use crate::tokenizer::TokenizerKind;
|
||||
@@ -26,7 +25,7 @@ fn cache_dir() -> std::path::PathBuf {
|
||||
|
||||
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
|
||||
/// Results are cached as JSON locally for subsequent runs.
|
||||
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let config_name = config.as_str();
|
||||
|
||||
let dir = cache_dir();
|
||||
@@ -36,13 +35,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
// Return cached file if it exists
|
||||
if cache_path.exists() {
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
|
||||
println!("SPEED-Bench ({config_name}) cached: {path_str}");
|
||||
return Ok(path_str);
|
||||
}
|
||||
|
||||
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
|
||||
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
|
||||
@@ -50,7 +49,6 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||
let mut offset = 0usize;
|
||||
let page_size = 100usize;
|
||||
let mut progress = RowDownloadReporter::new();
|
||||
|
||||
loop {
|
||||
let url = format!(
|
||||
@@ -66,14 +64,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let max_retries = 3;
|
||||
let mut data: Option<serde_json::Value> = None;
|
||||
for attempt in 0..=max_retries {
|
||||
let resp = match client.get(&url).send().await {
|
||||
let resp = match client.get(&url).send() {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(
|
||||
std::thread::sleep(std::time::Duration::from_secs(
|
||||
2 * (attempt as u64 + 1),
|
||||
))
|
||||
.await;
|
||||
));
|
||||
continue;
|
||||
}
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -83,7 +80,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
};
|
||||
|
||||
if resp.status().is_server_error() && attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -94,7 +91,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
)));
|
||||
}
|
||||
|
||||
data = Some(resp.json().await.map_err(|e| {
|
||||
data = Some(resp.json().map_err(|e| {
|
||||
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
|
||||
})?);
|
||||
break;
|
||||
@@ -119,14 +116,15 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let fetched = rows.len();
|
||||
offset += fetched;
|
||||
|
||||
// Print progress
|
||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||
progress.update(offset, total);
|
||||
eprint!("\r Fetched {offset}/{total} rows...");
|
||||
|
||||
if fetched < page_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
progress.finish();
|
||||
eprintln!(); // newline after progress
|
||||
|
||||
if all_rows.is_empty() {
|
||||
return Err(BenchError::Config(
|
||||
@@ -139,11 +137,9 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
std::fs::write(&cache_path, &json_str)?;
|
||||
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(
|
||||
config = config_name,
|
||||
rows = all_rows.len(),
|
||||
path = %path_str,
|
||||
"saved SPEED-Bench dataset"
|
||||
println!(
|
||||
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
|
||||
all_rows.len()
|
||||
);
|
||||
Ok(path_str)
|
||||
}
|
||||
@@ -267,11 +263,9 @@ pub fn load_speed_bench_dataset(
|
||||
// Oversample if needed
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let original_len = samples.len();
|
||||
@@ -281,11 +275,9 @@ pub fn load_speed_bench_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
println!(
|
||||
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -296,6 +288,7 @@ pub fn load_speed_bench_dataset(
|
||||
));
|
||||
}
|
||||
|
||||
// Print category distribution
|
||||
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
|
||||
for entry in &filtered[..filtered.len().min(samples.len())] {
|
||||
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
|
||||
@@ -304,7 +297,7 @@ pub fn load_speed_bench_dataset(
|
||||
let mut cats: Vec<_> = cat_counts.into_iter().collect();
|
||||
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
|
||||
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
|
||||
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
|
||||
println!("SPEED-Bench categories: {}", cat_str.join(", "));
|
||||
|
||||
Ok(samples)
|
||||
}
|
||||
|
||||
+35
-20
@@ -1,39 +1,54 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::path::PathBuf;
|
||||
//! 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 hf_hub::Repo;
|
||||
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
|
||||
pub struct HubRepo {
|
||||
repo: ApiRepo,
|
||||
repo: hf_hub::Repo,
|
||||
}
|
||||
|
||||
impl HubRepo {
|
||||
pub fn model(model_id: String) -> Result<Self, String> {
|
||||
Self::new(Repo::model(model_id))
|
||||
pub fn model(model_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::model(model_id),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dataset(repo_id: String) -> Result<Self, String> {
|
||||
Self::new(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));
|
||||
pub fn dataset(repo_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::dataset(repo_id),
|
||||
}
|
||||
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.
|
||||
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
|
||||
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
self.repo.get(filename).await.map_err(|e| format!("{e}"))
|
||||
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
let repo = self.repo.clone();
|
||||
let filename = filename.to_string();
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
|
||||
rt.block_on(async move {
|
||||
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
|
||||
})
|
||||
})
|
||||
.join()
|
||||
.map_err(|_| "HF Hub download thread panicked".to_string())?
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ pub fn prepare_process() {
|
||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||
&& new > 1024
|
||||
{
|
||||
tracing::info!(soft_limit = new, "raised open-file limit");
|
||||
eprintln!("Open-file limit: {new}");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,19 +19,7 @@ struct Cli {
|
||||
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<()> {
|
||||
init_tracing();
|
||||
|
||||
let cli = Cli::parse();
|
||||
vllm_bench::prepare_process();
|
||||
|
||||
|
||||
@@ -7,22 +7,6 @@ use crate::datasets::SampleRequest;
|
||||
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
|
||||
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.
|
||||
///
|
||||
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
|
||||
@@ -79,7 +63,14 @@ pub fn calculate_metrics(
|
||||
|
||||
let failed = outputs.len() - completed;
|
||||
|
||||
log_failed_requests(outputs);
|
||||
// Print failed request errors (capped to 10)
|
||||
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
|
||||
let mut max_output_tokens_per_s = 0.0_f64;
|
||||
@@ -304,7 +295,14 @@ pub fn calculate_embedding_metrics(
|
||||
|
||||
let failed = outputs.len() - completed;
|
||||
|
||||
log_failed_requests(outputs);
|
||||
// Print failed request errors (capped to 10)
|
||||
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
|
||||
let successful_outputs: Vec<&RequestFuncOutput> =
|
||||
|
||||
@@ -6,7 +6,6 @@ use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
|
||||
@@ -73,13 +72,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||
(m.clone(), config.model_name.clone())
|
||||
} else {
|
||||
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||
println!("Model not specified, fetching first model from server...");
|
||||
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
|
||||
tracing::info!(
|
||||
model_name = name,
|
||||
model_id = id,
|
||||
"selected first model from server"
|
||||
);
|
||||
println!("First model name: {name}, first model id: {id}");
|
||||
(id, Some(name))
|
||||
};
|
||||
|
||||
@@ -88,19 +83,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
None
|
||||
} else {
|
||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||
println!("Loading tokenizer: {tid}");
|
||||
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).await?;
|
||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
||||
println!("Tokenizer loaded successfully.");
|
||||
Some(t)
|
||||
};
|
||||
|
||||
// Generate/load conversations
|
||||
tracing::info!(
|
||||
dataset = ?config.dataset_name,
|
||||
conversations = config.num_prompts,
|
||||
"generating multi-turn conversations"
|
||||
);
|
||||
println!("Generating multi-turn conversations...");
|
||||
let gen_start = Instant::now();
|
||||
|
||||
let mut conversations = match config.dataset_name {
|
||||
@@ -140,7 +131,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let path = match config.dataset_path.as_deref() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -188,11 +179,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let (filtered_conversations, filtered_turns) =
|
||||
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
|
||||
if filtered_turns > 0 || filtered_conversations > 0 {
|
||||
tracing::info!(
|
||||
filtered_turns,
|
||||
filtered_conversations,
|
||||
max_model_len,
|
||||
"filtered conversations above maximum model length"
|
||||
println!(
|
||||
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
|
||||
);
|
||||
}
|
||||
if conversations.is_empty() {
|
||||
@@ -204,11 +192,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
|
||||
let gen_elapsed = gen_start.elapsed();
|
||||
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
|
||||
tracing::info!(
|
||||
conversations = conversations.len(),
|
||||
println!(
|
||||
"Generated {} conversations ({} total turns) in {:.2}s",
|
||||
conversations.len(),
|
||||
total_turns,
|
||||
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||
"generated multi-turn conversations"
|
||||
gen_elapsed.as_secs_f64()
|
||||
);
|
||||
|
||||
// Log prefix sharing info
|
||||
@@ -220,15 +208,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let conv_tokens =
|
||||
(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);
|
||||
tracing::info!(
|
||||
global_ratio = config.multi_turn_prefix_global_ratio,
|
||||
println!(
|
||||
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
|
||||
config.multi_turn_prefix_global_ratio * 100.0,
|
||||
global_tokens,
|
||||
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
|
||||
conversation_tokens = conv_tokens,
|
||||
config.multi_turn_prefix_conversation_ratio * 100.0,
|
||||
conv_tokens,
|
||||
(1.0 - config.multi_turn_prefix_global_ratio
|
||||
- config.multi_turn_prefix_conversation_ratio)
|
||||
* 100.0,
|
||||
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 {
|
||||
@@ -262,7 +253,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
tracing::info!("starting initial single-prompt test run");
|
||||
println!("Starting initial single prompt test run...");
|
||||
let test_output = crate::ready_checker::wait_for_endpoint(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -277,7 +268,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
test_output.error
|
||||
)));
|
||||
}
|
||||
tracing::info!("initial single-prompt test run completed");
|
||||
println!("Initial test run completed.");
|
||||
}
|
||||
|
||||
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
|
||||
@@ -292,10 +283,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
"min_tokens".to_string(),
|
||||
serde_json::json!(config.random_output_len),
|
||||
);
|
||||
tracing::info!(
|
||||
min_tokens = config.random_output_len,
|
||||
dataset = "random",
|
||||
"set minimum output tokens for multi-turn dataset"
|
||||
println!(
|
||||
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
|
||||
config.random_output_len
|
||||
);
|
||||
}
|
||||
Some(body)
|
||||
@@ -307,7 +297,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let spec_decode_before =
|
||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||
if spec_decode_before.is_some() {
|
||||
tracing::info!("detected speculative decoding; collecting metrics");
|
||||
println!("Speculative decoding detected, will collect metrics.");
|
||||
}
|
||||
|
||||
// Start profiler if requested (immediate mode — no batch threshold)
|
||||
@@ -340,13 +330,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
};
|
||||
|
||||
// Main benchmark
|
||||
tracing::info!(
|
||||
conversations = conversations.len(),
|
||||
total_turns,
|
||||
concurrency,
|
||||
inter_turn_delay_ms = config.multi_turn_delay_ms,
|
||||
"starting multi-turn benchmark"
|
||||
);
|
||||
println!("Starting multi-turn benchmark...");
|
||||
println!("Conversations: {}", conversations.len());
|
||||
println!("Concurrency: {concurrency}");
|
||||
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
|
||||
|
||||
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
|
||||
|
||||
@@ -377,12 +364,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
);
|
||||
if let Some(modules) = config.lora_modules.as_ref() {
|
||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||
tracing::info!(
|
||||
adapters = modules.len(),
|
||||
names = ?names,
|
||||
assignment = ?config.lora_assignment,
|
||||
scope = "conversation",
|
||||
"assigned LoRA adapters"
|
||||
println!(
|
||||
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
|
||||
modules.len(),
|
||||
names,
|
||||
config.lora_assignment
|
||||
);
|
||||
}
|
||||
|
||||
@@ -447,7 +433,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
match handle.await {
|
||||
Ok(output) => all_outputs.push(output),
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e.as_report(), "conversation task panicked");
|
||||
eprintln!("Conversation task panicked: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -467,7 +453,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
if let Some((cancel_tx, task)) = profile_task {
|
||||
let _ = cancel_tx.send(());
|
||||
if let Err(e) = task.await {
|
||||
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||
eprintln!("WARNING: Profile background task failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -603,7 +603,7 @@ fn add_metric_stats(
|
||||
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
|
||||
let content = serde_json::to_string(json)?;
|
||||
std::fs::write(file_path, content)?;
|
||||
tracing::info!(path = file_path, "saved benchmark results");
|
||||
println!("Results saved to {file_path}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
|
||||
file.write_all(b"\n")?;
|
||||
}
|
||||
file.write_all(content.as_bytes())?;
|
||||
tracing::info!(path = file_path, "appended benchmark results");
|
||||
println!("Results appended to {file_path}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -23,11 +23,7 @@ pub async fn wait_for_endpoint(
|
||||
let backend = get_backend(backend)?;
|
||||
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
|
||||
|
||||
tracing::info!(
|
||||
timeout_seconds,
|
||||
retry_interval,
|
||||
"waiting for endpoint readiness"
|
||||
);
|
||||
println!("Waiting for endpoint to become up in {timeout_seconds}s");
|
||||
|
||||
let pb = ProgressBar::new(timeout_seconds);
|
||||
pb.set_style(
|
||||
@@ -57,9 +53,7 @@ pub async fn wait_for_endpoint(
|
||||
Ok(output) => {
|
||||
let err = output.error.clone();
|
||||
let err_last_line = err.lines().last().unwrap_or(&err);
|
||||
pb.suspend(|| {
|
||||
tracing::warn!(error = err_last_line, "endpoint is not ready");
|
||||
});
|
||||
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
|
||||
last_error = err;
|
||||
}
|
||||
Err(e) => {
|
||||
|
||||
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
|
||||
.await
|
||||
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
|
||||
if resp.status().is_success() {
|
||||
tracing::info!(url = %url, "reset prefix cache");
|
||||
println!("Prefix cache reset successfully.");
|
||||
} else {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
|
||||
@@ -199,17 +199,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
|
||||
}
|
||||
};
|
||||
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
|
||||
tracing::info!(
|
||||
encoding,
|
||||
kind = "built-in-tiktoken",
|
||||
vocab_size,
|
||||
"loaded tokenizer"
|
||||
);
|
||||
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
|
||||
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
|
||||
}
|
||||
|
||||
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
|
||||
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
// Phase 1: If model_id is a local directory, look for tiktoken files there
|
||||
let local_dir = Path::new(model_id);
|
||||
if local_dir.is_dir() {
|
||||
@@ -217,7 +212,7 @@ pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
}
|
||||
|
||||
// Phase 2: Fall back to HuggingFace Hub download
|
||||
try_load_tiktoken_from_hf(model_id).await
|
||||
try_load_tiktoken_from_hf(model_id)
|
||||
}
|
||||
|
||||
/// Common tiktoken model filenames to search for.
|
||||
@@ -252,28 +247,25 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
|
||||
}
|
||||
|
||||
/// Load a tiktoken tokenizer from a HuggingFace model repo.
|
||||
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
||||
|
||||
let mut model_path = None;
|
||||
for filename in TIKTOKEN_MODEL_FILENAMES {
|
||||
if let Ok(path) = repo.get(filename).await {
|
||||
model_path = Some(path);
|
||||
break;
|
||||
}
|
||||
}
|
||||
let model_path = model_path.ok_or_else(|| {
|
||||
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
||||
})?;
|
||||
let model_path = repo
|
||||
.get("tiktoken.model")
|
||||
.or_else(|_| repo.get("qwen.tiktoken"))
|
||||
.or_else(|_| repo.get("vocab.tiktoken"))
|
||||
.map_err(|_| {
|
||||
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
||||
})?;
|
||||
|
||||
let num_base_tokens = count_base_tokens(&model_path)?;
|
||||
|
||||
let config = match repo.get("tokenizer_config.json").await {
|
||||
let config = match repo.get("tokenizer_config.json") {
|
||||
Ok(config_path) => read_tokenizer_config(&config_path),
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
let pattern = extract_pat_str_from_repo(&repo).await;
|
||||
let pattern = extract_pat_str_from_repo(&repo);
|
||||
|
||||
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
|
||||
}
|
||||
@@ -317,16 +309,15 @@ fn build_tiktoken(
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
base_tokens = num_base_tokens,
|
||||
special_tokens = all_special_tokens.len(),
|
||||
pattern = if pattern.is_some() {
|
||||
println!(
|
||||
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
|
||||
num_base_tokens,
|
||||
all_special_tokens.len(),
|
||||
if pattern.is_some() {
|
||||
"custom"
|
||||
} else {
|
||||
"default"
|
||||
},
|
||||
"loading tiktoken model"
|
||||
);
|
||||
|
||||
TiktokenTokenizer::from_file(
|
||||
@@ -406,12 +397,9 @@ 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.
|
||||
/// Returns None if unavailable or unparsable.
|
||||
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
||||
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
||||
// Try common Python tokenizer filenames
|
||||
let py_path = match repo.get("tokenization_kimi.py").await {
|
||||
Ok(path) => path,
|
||||
Err(_) => repo.get("tokenizer.py").await.ok()?,
|
||||
};
|
||||
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
|
||||
|
||||
let source = std::fs::read_to_string(&py_path).ok()?;
|
||||
|
||||
@@ -450,9 +438,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
|
||||
|
||||
if !fragments.is_empty() {
|
||||
let pattern = fragments.join("|");
|
||||
tracing::debug!(
|
||||
fragments = fragments.len(),
|
||||
"extracted tiktoken pattern from Python source"
|
||||
println!(
|
||||
"Extracted pat_str from Python source: {} fragments",
|
||||
fragments.len()
|
||||
);
|
||||
return Some(pattern);
|
||||
}
|
||||
|
||||
@@ -2,10 +2,8 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::future::Future;
|
||||
use std::path::Path;
|
||||
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
use crate::error::{BenchError, Result};
|
||||
@@ -20,8 +18,7 @@ pub enum TokenizerKind {
|
||||
|
||||
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
|
||||
pub struct ServerTokenizer {
|
||||
client: reqwest::Client,
|
||||
runtime: tokio::runtime::Handle,
|
||||
client: reqwest::blocking::Client,
|
||||
tokenize_url: String,
|
||||
detokenize_url: String,
|
||||
model: String,
|
||||
@@ -30,8 +27,8 @@ pub struct ServerTokenizer {
|
||||
|
||||
impl ServerTokenizer {
|
||||
/// Create a new server tokenizer and verify connectivity.
|
||||
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
|
||||
let client = reqwest::Client::builder()
|
||||
pub fn new(base_url: &str, model: &str) -> Result<Self> {
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
|
||||
@@ -41,7 +38,6 @@ impl ServerTokenizer {
|
||||
|
||||
let st = Self {
|
||||
client,
|
||||
runtime: tokio::runtime::Handle::current(),
|
||||
tokenize_url,
|
||||
detokenize_url,
|
||||
model: model.to_string(),
|
||||
@@ -49,7 +45,7 @@ impl ServerTokenizer {
|
||||
};
|
||||
|
||||
// Probe the endpoint to verify it works and discover vocab size
|
||||
let test_tokens = st.encode_async("test").await?;
|
||||
let test_tokens = st.encode_inner("test")?;
|
||||
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
|
||||
let estimated_vocab = (max_id * 2).max(131072);
|
||||
|
||||
@@ -60,10 +56,6 @@ impl ServerTokenizer {
|
||||
}
|
||||
|
||||
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!({
|
||||
"model": self.model,
|
||||
"prompt": text,
|
||||
@@ -74,7 +66,6 @@ impl ServerTokenizer {
|
||||
.post(&self.tokenize_url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
@@ -84,7 +75,7 @@ impl ServerTokenizer {
|
||||
)));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
||||
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
|
||||
})?;
|
||||
|
||||
@@ -104,10 +95,6 @@ impl ServerTokenizer {
|
||||
}
|
||||
|
||||
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!({
|
||||
"model": self.model,
|
||||
"tokens": ids,
|
||||
@@ -118,7 +105,6 @@ impl ServerTokenizer {
|
||||
.post(&self.detokenize_url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
@@ -128,7 +114,7 @@ impl ServerTokenizer {
|
||||
)));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
||||
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
|
||||
})?;
|
||||
|
||||
@@ -137,26 +123,6 @@ impl ServerTokenizer {
|
||||
.map(|s| s.to_string())
|
||||
.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 ---
|
||||
@@ -226,7 +192,7 @@ impl TokenizerKind {
|
||||
/// 3. Server-side /tokenize + /detokenize endpoints
|
||||
///
|
||||
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
|
||||
pub async fn load_tokenizer(
|
||||
pub fn load_tokenizer(
|
||||
model_id: &str,
|
||||
_trust_remote_code: bool,
|
||||
server_info: Option<(&str, &str)>,
|
||||
@@ -246,48 +212,31 @@ pub async fn load_tokenizer(
|
||||
}
|
||||
|
||||
// 1. Try local HuggingFace tokenizer (tokenizer.json)
|
||||
match try_load_local(model_id).await {
|
||||
match try_load_local(model_id) {
|
||||
Ok(tok) => {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "local",
|
||||
vocab_size = tok.get_vocab_size(true),
|
||||
"loaded tokenizer"
|
||||
);
|
||||
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
|
||||
Ok(TokenizerKind::Local(Box::new(tok)))
|
||||
}
|
||||
Err(local_err) => {
|
||||
// 2. Try tiktoken format
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
error = %local_err.as_report(),
|
||||
"local tokenizer unavailable; trying tiktoken"
|
||||
);
|
||||
match crate::tiktoken::try_load_tiktoken(model_id).await {
|
||||
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
|
||||
match crate::tiktoken::try_load_tiktoken(model_id) {
|
||||
Ok(tok) => {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "tiktoken",
|
||||
vocab_size = tok.vocab_size(),
|
||||
"loaded tokenizer"
|
||||
);
|
||||
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
|
||||
Ok(TokenizerKind::Tiktoken(tok))
|
||||
}
|
||||
Err(tiktoken_err) => {
|
||||
// 3. Try server-side fallback
|
||||
if let Some((base_url, model)) = server_info {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
error = %tiktoken_err.as_report(),
|
||||
"tiktoken unavailable; trying server-side tokenization"
|
||||
println!(
|
||||
"Tiktoken also not available ({tiktoken_err}), \
|
||||
trying server-side tokenization..."
|
||||
);
|
||||
match ServerTokenizer::new(base_url, model).await {
|
||||
match ServerTokenizer::new(base_url, model) {
|
||||
Ok(srv) => {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "server",
|
||||
vocab_size = srv.cached_vocab_size,
|
||||
"loaded tokenizer"
|
||||
println!(
|
||||
"Tokenizer: Server (vocab_size≈{})",
|
||||
srv.cached_vocab_size
|
||||
);
|
||||
return Ok(TokenizerKind::Server(srv));
|
||||
}
|
||||
@@ -315,7 +264,7 @@ pub async fn load_tokenizer(
|
||||
}
|
||||
|
||||
/// Try loading tokenizer.json from local path or HuggingFace Hub.
|
||||
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
// 1. Try local directory with tokenizer.json
|
||||
let local_path = Path::new(model_id).join("tokenizer.json");
|
||||
if local_path.exists() {
|
||||
@@ -341,37 +290,11 @@ async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
}
|
||||
|
||||
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
||||
let tokenizer_path = repo
|
||||
.get("tokenizer.json")
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
|
||||
|
||||
Tokenizer::from_file(&tokenizer_path)
|
||||
.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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
|
||||
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
|
||||
};
|
||||
|
||||
apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
|
||||
apply_structural_tag_constraint(request, parser.structural_tag_model())?;
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
|
||||
@@ -7,8 +7,7 @@ use thiserror_ext::AsReport;
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::StructuralTagBuilder;
|
||||
use xgrammar_structural_tag::builders::StructuralTagOptions;
|
||||
use vllm_parser::tool::StructuralTagModel;
|
||||
use xgrammar_structural_tag::{
|
||||
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
|
||||
build_structural_tag,
|
||||
@@ -21,9 +20,9 @@ use crate::{Error, Result as ChatResult};
|
||||
/// support and the request's tool choice.
|
||||
pub(super) fn apply_structural_tag_constraint(
|
||||
request: &mut ChatRequest,
|
||||
builder: Option<&dyn StructuralTagBuilder>,
|
||||
model: Option<StructuralTagModel>,
|
||||
) -> ChatResult<()> {
|
||||
let Some(builder) = builder else {
|
||||
let Some(model) = model else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(tool_choice) = structural_tag_tool_choice(request) else {
|
||||
@@ -43,16 +42,11 @@ pub(super) fn apply_structural_tag_constraint(
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let structural_tag = build_structural_tag(
|
||||
builder,
|
||||
&tools,
|
||||
tool_choice,
|
||||
StructuralTagOptions::default().with_reasoning(false),
|
||||
)
|
||||
.and_then(|tag| tag.to_json_string())
|
||||
.map_err(|error| Error::StructuralTag {
|
||||
message: error.to_report_string(),
|
||||
})?;
|
||||
let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
|
||||
.and_then(|tag| tag.to_json_string())
|
||||
.map_err(|error| Error::StructuralTag {
|
||||
message: error.to_report_string(),
|
||||
})?;
|
||||
|
||||
// Overwrite any existing structured output settings with the structural tag constraint.
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
@@ -147,7 +141,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
@@ -160,7 +154,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
@@ -175,7 +169,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
@@ -190,7 +184,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
@@ -207,7 +201,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
@@ -227,7 +221,7 @@ mod tests {
|
||||
);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request).to_string();
|
||||
@@ -240,7 +234,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
@@ -255,7 +249,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
|
||||
@@ -26,21 +26,18 @@ const RESET: &str = "\x1b[0m";
|
||||
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
|
||||
format_description!("[month]-[day] [hour]:[minute]:[second]");
|
||||
|
||||
const PROCESS_LABEL: &str = "RustFrontend";
|
||||
|
||||
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
|
||||
pub(crate) fn init_tracing(process_label: &str) {
|
||||
pub(crate) fn init_tracing() {
|
||||
let filter = build_targets_filter(
|
||||
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
|
||||
env::var("RUST_LOG").ok().as_deref(),
|
||||
);
|
||||
let formatter = VllmEventFormatter::new(process_label);
|
||||
let formatter = VllmEventFormatter::new();
|
||||
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.event_format(formatter)
|
||||
.with_writer(std::io::stderr)
|
||||
.with_filter(filter),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
|
||||
.try_init();
|
||||
}
|
||||
|
||||
@@ -97,9 +94,9 @@ struct VllmEventFormatter {
|
||||
}
|
||||
|
||||
impl VllmEventFormatter {
|
||||
fn new(process_label: &str) -> Self {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
prefix: format!("({process_label} pid={})", process::id()),
|
||||
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
|
||||
timer: VllmLocalTimer::default(),
|
||||
}
|
||||
}
|
||||
@@ -294,13 +291,6 @@ fn map_python_log_level(level: &str) -> LevelFilter {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn formatter_prefix_uses_process_label() {
|
||||
let formatter = VllmEventFormatter::new("Bench");
|
||||
|
||||
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
|
||||
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
|
||||
|
||||
@@ -5,7 +5,6 @@ mod cli;
|
||||
mod logging;
|
||||
|
||||
use std::env;
|
||||
use std::ffi::OsStr;
|
||||
use std::process::ExitStatus;
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
@@ -83,14 +82,7 @@ fn shutdown_signal() -> CancellationToken {
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
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);
|
||||
|
||||
logging::init_tracing();
|
||||
let cli = Cli::parse();
|
||||
|
||||
let mut runtime = tokio::runtime::Builder::new_multi_thread();
|
||||
|
||||
@@ -460,12 +460,6 @@ impl EngineCoreClient {
|
||||
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.
|
||||
pub fn health_error(&self) -> Option<Arc<Error>> {
|
||||
self.inner.health_error()
|
||||
|
||||
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
|
||||
use parking_lot::Mutex;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::runtime::Handle;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, trace, warn};
|
||||
use vllm_metrics::METRICS;
|
||||
use zeromq::RouterSendHalf;
|
||||
@@ -36,7 +36,6 @@ pub(crate) struct ClientInner {
|
||||
request_reg: Mutex<RequestRegistry>,
|
||||
utility_reg: Mutex<UtilityRegistry>,
|
||||
health_error: ArcSwapOption<Error>,
|
||||
health_tx: watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl ClientInner {
|
||||
@@ -58,7 +57,6 @@ impl ClientInner {
|
||||
request_reg: Mutex::new(RequestRegistry::new(engines)),
|
||||
utility_reg: Mutex::new(UtilityRegistry::default()),
|
||||
health_error: ArcSwapOption::empty(),
|
||||
health_tx: watch::Sender::new(true),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,7 +169,6 @@ impl ClientInner {
|
||||
/// persistent health error.
|
||||
pub fn close_registries(&self, error: Arc<Error>) {
|
||||
let persistent_error = self.record_health_error(error);
|
||||
self.publish_unhealthy();
|
||||
let request_senders = self.request_reg.lock().close();
|
||||
let utility_senders = self.utility_reg.lock().close();
|
||||
|
||||
@@ -194,12 +191,6 @@ impl ClientInner {
|
||||
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
|
||||
/// waiting caller existed.
|
||||
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
|
||||
@@ -289,11 +280,6 @@ impl ClientInner {
|
||||
.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
|
||||
/// wrapping it for error returns when the client is already closed.
|
||||
fn closed_error(&self) -> Error {
|
||||
@@ -475,18 +461,13 @@ mod tests {
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn close_registries_records_first_health_error_only() {
|
||||
let inner = test_inner().await;
|
||||
let mut health = inner.subscribe_health();
|
||||
assert!(*health.borrow());
|
||||
|
||||
inner.close_registries(Arc::new(Error::EngineCoreDead));
|
||||
health.changed().await.expect("health sender remains open");
|
||||
assert!(!inner.is_healthy());
|
||||
assert!(!*health.borrow());
|
||||
assert!(matches!(
|
||||
inner.health_error().as_deref(),
|
||||
Some(Error::EngineCoreDead)
|
||||
));
|
||||
assert!(!*inner.subscribe_health().borrow());
|
||||
|
||||
inner.close_registries(Arc::new(client_closed!("shutdown")));
|
||||
assert!(matches!(
|
||||
|
||||
@@ -269,19 +269,6 @@ 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);
|
||||
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
|
||||
.data
|
||||
|
||||
@@ -303,49 +303,3 @@ 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]"
|
||||
);
|
||||
}
|
||||
|
||||
#[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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_parser::tool::{
|
||||
Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
|
||||
Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
|
||||
};
|
||||
use vllm_parser::unified::{
|
||||
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
|
||||
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
|
||||
self.inner.preserve_special_tokens()
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
self.inner.structural_tag_builder()
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
self.inner.structural_tag_model()
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.2 models.
|
||||
///
|
||||
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV32)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V4 models.
|
||||
///
|
||||
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV4)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
@@ -73,7 +73,7 @@ mod tests {
|
||||
|
||||
use super::DeepSeekV4ToolParser;
|
||||
use crate::tool::test_utils::{collect_stream, test_tools};
|
||||
use crate::tool::{ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -91,10 +91,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v4_exposes_structural_tag_builder() {
|
||||
fn deepseek_v4_exposes_structural_tag_model() {
|
||||
let parser = DeepSeekV4ToolParser::new(&test_tools());
|
||||
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
assert_eq!(
|
||||
parser.structural_tag_model(),
|
||||
Some(StructuralTagModel::DeepSeekV4)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
|
||||
///
|
||||
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekR1)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
|
||||
///
|
||||
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV31)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
|
||||
///
|
||||
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Glm47.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Glm47)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<tool_calls>";
|
||||
const TOOL_CALLS_END: &str = "</tool_calls>";
|
||||
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::HyV3.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::HyV3)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Hermes",
|
||||
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Hermes.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Hermes)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -12,9 +12,7 @@ use super::{
|
||||
argument_delta_event, tool_call_header_event,
|
||||
};
|
||||
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
|
||||
use crate::tool::{
|
||||
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
|
||||
};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum LlamaJsonMode {
|
||||
@@ -138,8 +136,8 @@ impl ToolParser for Llama3JsonToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Llama.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Llama)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Qwen XML",
|
||||
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Qwen3.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Qwen3)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
|
||||
|
||||
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
|
||||
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
|
||||
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Kimi.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Kimi)
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
|
||||
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
|
||||
const TOOL_CALL_START: &str = "<minimax:tool_call>";
|
||||
const TOOL_CALL_END: &str = "</minimax:tool_call>";
|
||||
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Minimax.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Minimax)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
|
||||
pub use seed_oss::SeedOssToolParser;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
|
||||
pub use xgrammar_structural_tag::Model as StructuralTagModel;
|
||||
|
||||
use crate::utils;
|
||||
|
||||
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
|
||||
false
|
||||
}
|
||||
|
||||
/// Return the xgrammar structural-tag builder used for strict tool calling.
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
/// Return the xgrammar structural-tag model used for strict tool calling.
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
None
|
||||
}
|
||||
|
||||
|
||||
@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
|
||||
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::Tool;
|
||||
|
||||
const TOOL_CALL_START: &str = "<tool_call>";
|
||||
const TOOL_CALL_END: &str = "</tool_call>";
|
||||
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Qwen3Coder)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
@@ -294,7 +294,7 @@ mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{Qwen3CoderToolParser, ToolParser};
|
||||
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
@@ -308,10 +308,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen_coder_exposes_structural_tag_builder() {
|
||||
fn qwen_coder_exposes_structural_tag_model() {
|
||||
let parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
assert_eq!(
|
||||
parser.structural_tag_model(),
|
||||
Some(StructuralTagModel::Qwen3Coder)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
|
||||
use crate::reasoning::ReasoningParser;
|
||||
use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Unified parser that composes existing reasoning and tool parsers.
|
||||
pub struct CombinedParser {
|
||||
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
|
||||
|| self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
|
||||
}
|
||||
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
@@ -269,7 +269,10 @@ mod tests {
|
||||
fn combined_parser_emits_tool_calls_from_visible_content() {
|
||||
let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
|
||||
let mut parser = CombinedParser::new(None, Some(tool));
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
assert!(matches!(
|
||||
parser.structural_tag_model(),
|
||||
Some(crate::tool::StructuralTagModel::Qwen3)
|
||||
));
|
||||
|
||||
let output = collect(
|
||||
&mut parser,
|
||||
|
||||
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use crate::reasoning::ReasoningError;
|
||||
use crate::tool::{
|
||||
StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
|
||||
StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
|
||||
};
|
||||
|
||||
/// Result alias for unified parser operations.
|
||||
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
|
||||
false
|
||||
}
|
||||
|
||||
/// Return the xgrammar structural-tag builder used for strict tool calling.
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
/// Return the xgrammar structural-tag model used for strict tool calling.
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
None
|
||||
}
|
||||
|
||||
|
||||
@@ -35,7 +35,6 @@ tokio-openssl.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tonic.workspace = true
|
||||
tonic-health.workspace = true
|
||||
tonic-prost.workspace = true
|
||||
tower.workspace = true
|
||||
tower-http.workspace = true
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
// 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,7 +4,6 @@
|
||||
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
|
||||
|
||||
mod convert;
|
||||
mod health;
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
@@ -25,11 +24,8 @@ pub mod pb {
|
||||
tonic::include_proto!("vllm");
|
||||
}
|
||||
|
||||
pub(crate) use health::monitor_health;
|
||||
pub use pb::generate_server::GenerateServer;
|
||||
|
||||
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
|
||||
@@ -16,10 +16,6 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_openssl::SslStream;
|
||||
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 vllm_chat::{
|
||||
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
|
||||
@@ -204,11 +200,7 @@ impl ChatRenderer for FakeTextBackend {
|
||||
async fn setup_grpc_service(
|
||||
engine_id: impl Into<EngineId>,
|
||||
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
||||
) -> (
|
||||
GenerateServer<GenerateServiceImpl>,
|
||||
tokio::sync::watch::Receiver<bool>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
|
||||
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = engine_id.into();
|
||||
@@ -240,7 +232,6 @@ async fn setup_grpc_service(
|
||||
)
|
||||
.await
|
||||
.expect("connect client");
|
||||
let engine_health = client.subscribe_health();
|
||||
|
||||
let chat = ChatLlm::from_shared_backend(
|
||||
test_llm(client),
|
||||
@@ -249,7 +240,6 @@ async fn setup_grpc_service(
|
||||
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
|
||||
(
|
||||
GenerateServer::new(GenerateServiceImpl::new(state)),
|
||||
engine_health,
|
||||
engine_task,
|
||||
)
|
||||
}
|
||||
@@ -264,51 +254,25 @@ async fn grpc_test_server(
|
||||
tokio::task::JoinHandle<()>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
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 (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
let server = TonicServer::builder()
|
||||
.add_service(health_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||
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");
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc server");
|
||||
});
|
||||
|
||||
let channel = Endpoint::from_shared(format!("http://{addr}"))
|
||||
.expect("grpc endpoint")
|
||||
.connect()
|
||||
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
|
||||
.await
|
||||
.expect("connect grpc channel");
|
||||
.expect("connect grpc client");
|
||||
|
||||
(channel, server_task)
|
||||
(grpc_client, server_task, engine_task)
|
||||
}
|
||||
|
||||
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
|
||||
@@ -319,7 +283,7 @@ async fn grpc_tls_test_server(
|
||||
certs: &TestCerts,
|
||||
cert_reqs: i32,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
|
||||
.expect("build grpc tls config");
|
||||
|
||||
@@ -409,8 +373,7 @@ async fn grpc_server_with_keepalive(
|
||||
engine_id: impl Into<EngineId>,
|
||||
keepalive: Option<Duration>,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, _engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, default_stream_output_specs()).await;
|
||||
let (svc, 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 addr = listener.local_addr().expect("local addr").to_string();
|
||||
@@ -1072,129 +1035,3 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
|
||||
|
||||
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");
|
||||
}
|
||||
|
||||
+14
-28
@@ -39,7 +39,6 @@ use tokio::net::TcpListener;
|
||||
use tokio::time::{Instant, sleep_until};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tonic::transport::Server as TonicServer;
|
||||
use tonic_health::server::health_reporter;
|
||||
use tower::ServiceExt as _;
|
||||
use tracing::{info, trace, warn};
|
||||
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
|
||||
@@ -204,19 +203,14 @@ where
|
||||
.map(tls::build_grpc_server_config)
|
||||
.transpose()
|
||||
.context("invalid gRPC TLS configuration")?;
|
||||
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 = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
|
||||
let svc = TonicServer::builder()
|
||||
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
|
||||
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
|
||||
.layer(middleware::request_runtime_layer(state.clone()))
|
||||
.add_service(health_service)
|
||||
.add_service(generate_service);
|
||||
.add_service(svc);
|
||||
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
|
||||
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
|
||||
Some((grpc_listener, svc, grpc_tls))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -300,8 +294,7 @@ where
|
||||
let server_shutdown = server_shutdown.clone();
|
||||
let force_shutdown = force_shutdown.clone();
|
||||
async move {
|
||||
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
|
||||
else {
|
||||
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
|
||||
// No gRPC configured: just wait for shutdown so we do not race the
|
||||
// join! by resolving early and tripping the cancellation token.
|
||||
shutdown.cancelled().await;
|
||||
@@ -311,26 +304,19 @@ where
|
||||
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
|
||||
None => MaybeTlsListener::plain(grpc_listener),
|
||||
};
|
||||
let server =
|
||||
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
|
||||
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
|
||||
|
||||
let server = async move {
|
||||
let result = tokio::select! {
|
||||
result = server => {
|
||||
result.context("gRPC server failed")
|
||||
}
|
||||
_ = force_shutdown.cancelled() => {
|
||||
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
|
||||
server_shutdown.cancel();
|
||||
result
|
||||
let result = tokio::select! {
|
||||
result = server => {
|
||||
result.context("gRPC server failed")
|
||||
}
|
||||
_ = force_shutdown.cancelled() => {
|
||||
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
|
||||
let (result, ()) = tokio::join!(server, health_monitor);
|
||||
server_shutdown.cancel();
|
||||
result
|
||||
}
|
||||
};
|
||||
|
||||
@@ -60,8 +60,6 @@ def _make_quick_allreduce(
|
||||
qar.use_fp16_kernels = use_fp16_kernels
|
||||
qar.qr_quant_level = QuickReduceRegime[quant_level]
|
||||
qar.qr_max_size = qr_max_size
|
||||
qar.qr_min_size = None
|
||||
qar.qr_quantization_min_size = None
|
||||
return qar
|
||||
|
||||
|
||||
@@ -513,21 +511,13 @@ def test_quick_reduce_regime_values():
|
||||
assert QuickReduceRegime.INT8.value == 1
|
||||
assert QuickReduceRegime.INT6.value == 2
|
||||
assert QuickReduceRegime.INT4.value == 3
|
||||
assert QuickReduceRegime.INT3.value == 4
|
||||
assert QuickReduceRegime.NONE.value == 5
|
||||
assert QuickReduceRegime.NONE.value == 4
|
||||
|
||||
|
||||
def test_quick_reduce_regime_names():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
|
||||
|
||||
assert set(QuickReduceRegime.__members__) == {
|
||||
"FP",
|
||||
"INT8",
|
||||
"INT6",
|
||||
"INT4",
|
||||
"INT3",
|
||||
"NONE",
|
||||
}
|
||||
assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
|
||||
@@ -703,7 +693,7 @@ def test_quick_allreduce_min_size_table():
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
|
||||
min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
|
||||
assert len(min_sizes) == 5
|
||||
assert len(min_sizes) == 4
|
||||
assert all(size > 0 for size in min_sizes)
|
||||
|
||||
|
||||
|
||||
@@ -1,880 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end tests for routed-expert capture on the monolithic MoE path.
|
||||
|
||||
These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
|
||||
expert IDs picked by FlashInfer's fused router-and-experts kernels (the
|
||||
"monolithic" path). When ``set_capture_fn`` is installed on
|
||||
a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
|
||||
should:
|
||||
|
||||
* allocate an int16 ``(num_tokens, top_k)`` buffer,
|
||||
* pass it to FlashInfer as ``routing_replay_out``,
|
||||
* invoke the callback after the kernel returns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
RoutingMethodType,
|
||||
fp8_w8a8_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
|
||||
TrtLlmBf16ExpertsMonolithic,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
|
||||
TrtLlmFp8ExpertsMonolithic,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
|
||||
TrtLlmNvFp4ExpertsMonolithic,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
try:
|
||||
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer not available", allow_module_level=True)
|
||||
|
||||
if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"Requires FlashInfer TRT-LLM fused MoE on CUDA",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
"TRT-LLM fused MoE kernels require SM100+",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _shuffle_bf16_weights_block_major_k(
|
||||
w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
|
||||
) -> torch.Tensor:
|
||||
"""Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
|
||||
``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
|
||||
row shuffle.
|
||||
"""
|
||||
from flashinfer import shuffle_matrix_a
|
||||
from flashinfer.fused_moe import convert_to_block_layout
|
||||
|
||||
num_experts = w.shape[0]
|
||||
shuffled = []
|
||||
for i in range(num_experts):
|
||||
t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
|
||||
shuffled.append(convert_to_block_layout(t, block_k))
|
||||
return torch.stack(shuffled).view(torch.bfloat16)
|
||||
|
||||
|
||||
def _make_bf16_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
routing_method: RoutingMethodType,
|
||||
device: torch.device,
|
||||
) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
|
||||
"""Construct the monolithic BF16 experts plus the BlockMajorK weights
|
||||
expected by ``trtllm_bf16_moe``.
|
||||
"""
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=routing_method,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
quant_dtype=None,
|
||||
per_act_token_quant=False,
|
||||
per_out_ch_quant=False,
|
||||
block_shape=None,
|
||||
)
|
||||
|
||||
experts = TrtLlmBf16ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
|
||||
gemm1 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
2 * intermediate_size,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
gemm2 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
w13 = _shuffle_bf16_weights_block_major_k(gemm1)
|
||||
w2 = _shuffle_bf16_weights_block_major_k(gemm2)
|
||||
return experts, w13, w2
|
||||
|
||||
|
||||
def _run_bf16_monolithic(
|
||||
experts: TrtLlmBf16ExpertsMonolithic,
|
||||
hidden_states: torch.Tensor,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
*,
|
||||
n_group: int | None = None,
|
||||
topk_group: int | None = None,
|
||||
routed_scaling_factor: float | None = None,
|
||||
e_score_correction_bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states,
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
a1q_scale=None,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=n_group,
|
||||
topk_group=topk_group,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
)
|
||||
|
||||
|
||||
_DSV3_NUM_EXPERTS = 32
|
||||
_DSV3_N_GROUP = 4
|
||||
_DSV3_TOPK_GROUP = 2
|
||||
|
||||
|
||||
def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
|
||||
return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""The capture callback should receive the int16 routed-expert IDs the
|
||||
kernel actually used, the values should be valid expert indices, and
|
||||
each token should pick ``top_k`` distinct experts."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
|
||||
def capture_fn(replay_out: torch.Tensor) -> None:
|
||||
captured.append(replay_out.clone())
|
||||
|
||||
assert experts.supports_routing_replay_capture()
|
||||
experts.set_capture_fn(capture_fn)
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
n_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize(
|
||||
"routing_method",
|
||||
[
|
||||
RoutingMethodType.Renormalize,
|
||||
RoutingMethodType.RenormalizeNaive,
|
||||
],
|
||||
)
|
||||
def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
|
||||
num_tokens: int,
|
||||
routing_method: RoutingMethodType,
|
||||
) -> None:
|
||||
"""Routing replay works for non-DeepSeekV3 routing methods too.
|
||||
FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = 8
|
||||
top_k = 2
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=routing_method,
|
||||
device=device,
|
||||
)
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
|
||||
"""With no callback installed the kernel should not see a
|
||||
``routing_replay_out`` tensor — verify the helper short-circuits."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=2,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
# No callback installed.
|
||||
buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
|
||||
assert buf is None
|
||||
|
||||
# Dispatch is also a no-op.
|
||||
experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
|
||||
"""FlashInfer's ``routing_replay_out`` is supported by all routing
|
||||
methods, so ``supports_routing_replay_capture`` should be True
|
||||
regardless of routing method."""
|
||||
device = torch.device("cuda:0")
|
||||
for routing_method in (
|
||||
RoutingMethodType.DeepSeekV3,
|
||||
RoutingMethodType.Renormalize,
|
||||
RoutingMethodType.RenormalizeNaive,
|
||||
):
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=2,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=routing_method,
|
||||
device=device,
|
||||
)
|
||||
assert experts.supports_routing_replay_capture() is True, (
|
||||
f"{routing_method!r} should support routing replay capture"
|
||||
)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
|
||||
"""When capture is installed, the allocated buffer is int16 and shaped
|
||||
``(num_tokens, experts_per_token)``."""
|
||||
device = torch.device("cuda:0")
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=4,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
experts.set_capture_fn(lambda r: None)
|
||||
buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
|
||||
assert buf is not None
|
||||
assert buf.dtype == torch.int16
|
||||
assert buf.shape[0] >= 11
|
||||
assert buf.shape[1] == 4
|
||||
assert buf.device.type == "cuda"
|
||||
|
||||
|
||||
def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
|
||||
"""End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
|
||||
on the monolithic experts and verify the captured rows land in the
|
||||
capturer's device buffer at the correct layer slot.
|
||||
|
||||
Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
|
||||
for the monolithic path: a single closure is installed on the monolithic
|
||||
``fused_experts`` (in addition to ``router.set_capture_fn`` on the
|
||||
non-monolithic path) and the capturer routes per-layer based on the
|
||||
closed-over ``layer_id``.
|
||||
"""
|
||||
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
|
||||
RoutedExpertsCapturer,
|
||||
)
|
||||
|
||||
torch.manual_seed(7)
|
||||
device = torch.device("cuda:0")
|
||||
num_tokens = 4
|
||||
top_k = 2
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
|
||||
num_layers = 3
|
||||
layer_id = 1
|
||||
capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
|
||||
capturer.dp_rank = 0
|
||||
capturer.tp_size = 1
|
||||
capturer.device_buffer = torch.full(
|
||||
(num_tokens + 4, num_layers, top_k),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def capture_fn(replay_out: torch.Tensor) -> None:
|
||||
capturer.capture(layer_id, replay_out)
|
||||
|
||||
experts.set_capture_fn(capture_fn)
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
# Patch get_forward_context to return a dp_metadata=None context so the
|
||||
# capturer takes the single-DP branch.
|
||||
import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
|
||||
|
||||
with patch.object(
|
||||
rec,
|
||||
"get_forward_context",
|
||||
return_value=SimpleNamespace(dp_metadata=None),
|
||||
):
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
n_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=routing_bias,
|
||||
)
|
||||
|
||||
captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
|
||||
# Valid expert IDs at this layer.
|
||||
assert (captured >= 0).all()
|
||||
assert (captured < num_experts).all()
|
||||
for t in range(num_tokens):
|
||||
unique = captured[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts at layer "
|
||||
f"{layer_id}, got {unique.numel()}"
|
||||
)
|
||||
|
||||
# Other layers / trailing token rows untouched.
|
||||
for other_layer in range(num_layers):
|
||||
if other_layer == layer_id:
|
||||
continue
|
||||
assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
|
||||
f"layer {other_layer} should be untouched, got writes"
|
||||
)
|
||||
assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
|
||||
"tail rows beyond num_tokens should remain sentinel"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fp8_block_scale_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
device: torch.device,
|
||||
) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
|
||||
"""Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
|
||||
code path with DSV3 routing.
|
||||
|
||||
Weights are shuffled into the BlockMajorK layout the kernel expects
|
||||
(same helper the vLLM weight loader uses for DeepSeek-FP8 models).
|
||||
"""
|
||||
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
|
||||
_shuffle_deepseek_fp8_moe_weights,
|
||||
)
|
||||
|
||||
block_k = 128
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
|
||||
# Random fp8 weights + ones-block scales (the kernel decoder only cares
|
||||
# that the per-block scales are present and finite for routing/replay).
|
||||
gemm1 = torch.randn(
|
||||
num_experts, 2 * intermediate_size, hidden_size, device=device
|
||||
).to(torch.float8_e4m3fn)
|
||||
gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
|
||||
|
||||
w1_scale = torch.ones(
|
||||
num_experts,
|
||||
2 * intermediate_size // block_k,
|
||||
hidden_size // block_k,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w2_scale = torch.ones(
|
||||
num_experts,
|
||||
hidden_size // block_k,
|
||||
intermediate_size // block_k,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
quant_config = fp8_w8a8_moe_quant_config(
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
block_shape=[block_k, block_k],
|
||||
per_act_token_quant=False,
|
||||
)
|
||||
|
||||
experts = TrtLlmFp8ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
return experts, w13_shuffled, w2_shuffled
|
||||
|
||||
|
||||
def _run_fp8_block_scale_monolithic(
|
||||
experts: TrtLlmFp8ExpertsMonolithic,
|
||||
hidden_states_fp8: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
routing_bias: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states_fp8,
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
# The block-scale apply path reads ``a1q_scale`` and transposes it
|
||||
# to ``(hidden_size/128, num_tokens)`` for the kernel call.
|
||||
a1q_scale=hidden_states_scale,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
e_score_correction_bias=routing_bias,
|
||||
routed_scaling_factor=1.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
|
||||
path, DSV3 routing) captures valid expert IDs."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
device=device,
|
||||
)
|
||||
assert experts.supports_routing_replay_capture()
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
# Per-token / per-block hidden scales (ones is fine for the routing
|
||||
# path; the GEMM output isn't being asserted on).
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
).to(torch.float8_e4m3fn)
|
||||
hidden_states_scale = torch.ones(
|
||||
num_tokens, hidden_size // 128, device=device, dtype=torch.float32
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_fp8_block_scale_monolithic(
|
||||
experts,
|
||||
hidden_states_fp8=hidden_states,
|
||||
hidden_states_scale=hidden_states_scale,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
routing_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_nvfp4_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
device: torch.device,
|
||||
) -> tuple[
|
||||
TrtLlmNvFp4ExpertsMonolithic,
|
||||
torch.Tensor, # w13 (packed nvfp4 uint8)
|
||||
torch.Tensor, # w13 block-scale (fp8)
|
||||
torch.Tensor, # w2 (packed nvfp4 uint8)
|
||||
torch.Tensor, # w2 block-scale (fp8)
|
||||
torch.Tensor, # input global scale (per-tensor float32)
|
||||
]:
|
||||
"""Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
|
||||
and DSV3 routing.
|
||||
|
||||
NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
|
||||
"global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
|
||||
``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
|
||||
* weights: uint8 (packed fp4) ``(E, M, K//2)``
|
||||
* weight scales: fp8 ``(E, M, K//16)``
|
||||
* hidden states: uint8 (packed fp4) ``(N, K//2)``
|
||||
* hidden state scales: fp8 ``(N, K//16)``
|
||||
"""
|
||||
from flashinfer import fp4_quantize
|
||||
|
||||
block_size = 16
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
|
||||
gemm1 = torch.randn(
|
||||
num_experts,
|
||||
2 * intermediate_size,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
gemm2 = torch.randn(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
# Per-tensor weight scaling factor (used to build ``g1_alphas`` /
|
||||
# ``g2_alphas`` below).
|
||||
w_global_scale = torch.tensor(1.0, device=device)
|
||||
# Per-tensor input scaling factor.
|
||||
a_global_scale = torch.tensor(1.0, device=device)
|
||||
|
||||
w13_q, w13_scale = fp4_quantize(
|
||||
gemm1,
|
||||
w_global_scale,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
|
||||
num_experts, 2 * intermediate_size, hidden_size // block_size
|
||||
)
|
||||
w2_q, w2_scale = fp4_quantize(
|
||||
gemm2,
|
||||
w_global_scale,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
|
||||
num_experts, hidden_size, intermediate_size // block_size
|
||||
)
|
||||
|
||||
# NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
|
||||
# g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
|
||||
g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
|
||||
a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
|
||||
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
quant_dtype="nvfp4",
|
||||
per_act_token_quant=False,
|
||||
per_out_ch_quant=False,
|
||||
block_shape=None,
|
||||
w1_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
g1_alphas=g_alphas,
|
||||
g2_alphas=g_alphas,
|
||||
a1_gscale=a_global_scale,
|
||||
a2_gscale=a2_gscale,
|
||||
)
|
||||
|
||||
experts = TrtLlmNvFp4ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
|
||||
|
||||
|
||||
def _run_nvfp4_monolithic(
|
||||
experts: TrtLlmNvFp4ExpertsMonolithic,
|
||||
hidden_states_q: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
routing_bias: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""The monolithic NVFP4 apply expects packed fp4 hidden states + the
|
||||
matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
|
||||
# Stash the weight tensors on the experts in the locations the apply()
|
||||
# implementation reads from (it pulls them from quant_config / scales
|
||||
# already; w1/w2 come in as args).
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states_q,
|
||||
w1=experts._w13_packed,
|
||||
w2=experts._w2_packed,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
a1q_scale=hidden_states_scale,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
e_score_correction_bias=routing_bias,
|
||||
routed_scaling_factor=1.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
|
||||
on the DSV3 routing path."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
from flashinfer import fp4_quantize
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
block_size = 16
|
||||
|
||||
experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
device=device,
|
||||
)
|
||||
# The apply() reads w1/w2 from its args, but we keep them on the experts
|
||||
# for convenience of the helper.
|
||||
experts._w13_packed = w13_q
|
||||
experts._w2_packed = w2_q
|
||||
|
||||
assert experts.supports_routing_replay_capture()
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
hidden_states_q, hidden_states_scale = fp4_quantize(
|
||||
hidden_states,
|
||||
a_gs,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
# The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
|
||||
# ``hidden_states_scale`` in its native (uint8 packed) form.
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_nvfp4_monolithic(
|
||||
experts,
|
||||
hidden_states_q=hidden_states_q,
|
||||
hidden_states_scale=hidden_states_scale,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
routing_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
@@ -66,6 +66,7 @@ def test_worker_apply_lora(qwen3_lora_files):
|
||||
runner_type="generate",
|
||||
max_num_batched_tokens=32,
|
||||
max_num_seqs=32,
|
||||
max_num_partial_prefills=32,
|
||||
),
|
||||
device_config=DeviceConfig(DEVICE_TYPE),
|
||||
cache_config=CacheConfig(
|
||||
|
||||
@@ -1,306 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import ast
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.warmup.jit_warmup import (
|
||||
VllmJitKernel,
|
||||
WarmupIntRange,
|
||||
get_ast_full_name,
|
||||
zip_inputs,
|
||||
)
|
||||
|
||||
|
||||
def _next_power_of_2(value: int) -> int:
|
||||
return 1 << max(0, value - 1).bit_length()
|
||||
|
||||
|
||||
def _round_up(value: int, *, multiple: int) -> int:
|
||||
return ((value + multiple - 1) // multiple) * multiple
|
||||
|
||||
|
||||
def _config(
|
||||
*,
|
||||
bias: int = 0,
|
||||
disabled: bool = False,
|
||||
name: str = "base",
|
||||
vectorized: bool = False,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
bias=bias,
|
||||
disabled=disabled,
|
||||
name=name,
|
||||
vectorized=vectorized,
|
||||
)
|
||||
|
||||
|
||||
class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
block_size: int
|
||||
work: int
|
||||
vector_width: int
|
||||
descriptor: tuple[object, ...]
|
||||
enabled: bool
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
cfg: Any,
|
||||
lanes: int = 1,
|
||||
mode: str = "default",
|
||||
debug: int = 0,
|
||||
) -> CompileKey:
|
||||
block_size = _next_power_of_2(tokens)
|
||||
work: int = block_size * lanes + cfg.bias
|
||||
return self.CompileKey(
|
||||
block_size=block_size,
|
||||
work=work,
|
||||
vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
|
||||
descriptor=(
|
||||
cfg.name,
|
||||
mode,
|
||||
-block_size,
|
||||
block_size % 3,
|
||||
block_size**2,
|
||||
),
|
||||
enabled=not cfg.disabled,
|
||||
)
|
||||
|
||||
def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
|
||||
return self._trace_dispatch(self.dispatch)(
|
||||
tokens=WarmupIntRange(1, max_tokens + 1),
|
||||
cfg=cfg,
|
||||
# This argument is intentionally unused by dispatch expressions.
|
||||
debug=WarmupIntRange(0, 100),
|
||||
)
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class RecordingToyKernel(ToyKernel):
|
||||
def __init__(self) -> None:
|
||||
self.compiled: list[ToyKernel.CompileKey] = []
|
||||
super().__init__()
|
||||
|
||||
def compile(self, compile_key: ToyKernel.CompileKey) -> None:
|
||||
self.compiled.append(compile_key)
|
||||
|
||||
|
||||
def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
|
||||
cfg = _config()
|
||||
|
||||
assert ToyKernel().get_warmup_keys(5, cfg) == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
|
||||
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
|
||||
ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
|
||||
]
|
||||
|
||||
|
||||
def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
|
||||
cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
|
||||
|
||||
assert ToyKernel().compile_key(
|
||||
{
|
||||
"tokens": 4,
|
||||
"cfg": cfg,
|
||||
"lanes": 2,
|
||||
}
|
||||
) == ToyKernel.CompileKey(
|
||||
block_size=4,
|
||||
work=11,
|
||||
vector_width=4,
|
||||
descriptor=("cfg", "default", -4, 1, 16),
|
||||
enabled=False,
|
||||
)
|
||||
|
||||
|
||||
def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
|
||||
cfg = _config(vectorized=True)
|
||||
|
||||
keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
|
||||
zip_inputs(
|
||||
dict(tokens=1, mode="small"),
|
||||
dict(tokens=4, mode="wide"),
|
||||
),
|
||||
cfg=cfg,
|
||||
lanes=(1, 2),
|
||||
)
|
||||
|
||||
assert keys == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
|
||||
ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
|
||||
]
|
||||
|
||||
|
||||
def test_zip_inputs_validates_input_rows() -> None:
|
||||
with pytest.raises(ValueError, match="requires at least one"):
|
||||
zip_inputs()
|
||||
with pytest.raises(ValueError, match="rows must be mappings"):
|
||||
zip_inputs(cast(Any, ("tokens", 1)))
|
||||
with pytest.raises(ValueError, match="at least one dispatch input name"):
|
||||
zip_inputs({})
|
||||
with pytest.raises(ValueError, match="dispatch input names must be strings"):
|
||||
zip_inputs(cast(Any, {1: 2}))
|
||||
with pytest.raises(ValueError, match="same dispatch input names"):
|
||||
zip_inputs({"tokens": 1}, {"mode": "small"})
|
||||
|
||||
|
||||
def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
|
||||
kernel = ToyKernel()
|
||||
|
||||
with pytest.raises(TypeError, match="zip_inputs"):
|
||||
kernel._trace_dispatch(kernel.dispatch)(
|
||||
cast(Any, {"tokens": 1}),
|
||||
cfg=_config(),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="specified more than once"):
|
||||
kernel._trace_dispatch(kernel.dispatch)(
|
||||
zip_inputs(dict(tokens=1, mode="small")),
|
||||
tokens=2,
|
||||
cfg=_config(),
|
||||
)
|
||||
|
||||
|
||||
def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
|
||||
class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
block_size: int,
|
||||
) -> CompileKey:
|
||||
return self.CompileKey(value=_round_up(tokens, multiple=block_size))
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
block_size: int,
|
||||
) -> CompileKey:
|
||||
return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
assert HelperKernel().compile_key(
|
||||
{
|
||||
"tokens": 5,
|
||||
"block_size": 4,
|
||||
}
|
||||
) == HelperKernel.CompileKey(value=8)
|
||||
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
|
||||
StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
|
||||
|
||||
|
||||
def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
|
||||
class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
if value > 0:
|
||||
value = 1
|
||||
return self.CompileKey(value=value)
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
return self.CompileKey(**{"value": value})
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="local assignments"):
|
||||
BranchKernel()
|
||||
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
|
||||
KwargsReturnKernel()
|
||||
|
||||
|
||||
def test_dispatch_reports_unsupported_expression_with_context() -> None:
|
||||
class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: object
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
return self.CompileKey(value={value})
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
UnsupportedKernel().compile_key({"value": 1})
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "Unsupported dispatch expression" in message
|
||||
assert "{value}" in message
|
||||
assert "Supported dispatch expressions" in message
|
||||
|
||||
|
||||
def test_warmup_compiles_all_returned_keys_in_order() -> None:
|
||||
kernel = RecordingToyKernel()
|
||||
cfg = _config()
|
||||
|
||||
kernel.warmup(3, cfg)
|
||||
|
||||
assert kernel.compiled == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
|
||||
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
|
||||
]
|
||||
|
||||
|
||||
def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
|
||||
dotted_expr = ast.parse("foo.bar.baz").body[0]
|
||||
call_expr = ast.parse("foo()").body[0]
|
||||
assert isinstance(dotted_expr, ast.Expr)
|
||||
assert isinstance(call_expr, ast.Expr)
|
||||
|
||||
assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
|
||||
assert get_ast_full_name(call_expr.value) is None
|
||||
@@ -66,12 +66,6 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
|
||||
)
|
||||
|
||||
|
||||
def _make_modular_routed_experts():
|
||||
return types.SimpleNamespace(
|
||||
quant_method=types.SimpleNamespace(is_monolithic=False),
|
||||
)
|
||||
|
||||
|
||||
def test_base_router_capture_pre_eplb_mapping():
|
||||
router = _make_router()
|
||||
captured = []
|
||||
@@ -128,8 +122,6 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
|
||||
def __init__(self):
|
||||
self.layer_id = 7
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
class DummyCapturer:
|
||||
def __init__(self):
|
||||
@@ -168,8 +160,6 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
|
||||
def __init__(self):
|
||||
self.layer_id = 11
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
class DummyCapturer:
|
||||
def __init__(self):
|
||||
@@ -207,8 +197,6 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
|
||||
def __init__(self, layer_id):
|
||||
self.layer_id = layer_id
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
target_module = DummyFusedMoE(layer_id=7)
|
||||
draft_module = DummyFusedMoE(layer_id=0)
|
||||
@@ -234,49 +222,6 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
|
||||
assert draft_module.router.capture_fn is None
|
||||
|
||||
|
||||
def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
|
||||
from vllm.v1.worker import gpu_model_runner as gmr
|
||||
|
||||
class DummyFusedMoE:
|
||||
def __init__(self):
|
||||
self.layer_id = 3
|
||||
self.router = _make_router()
|
||||
# Use a concrete monolithic expert and override its capability
|
||||
# instead of instantiating the abstract base class directly.
|
||||
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
|
||||
CPUExpertsFp8,
|
||||
)
|
||||
|
||||
fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
|
||||
self.routed_experts = types.SimpleNamespace(
|
||||
quant_method=types.SimpleNamespace(
|
||||
is_monolithic=True,
|
||||
moe_kernel=types.SimpleNamespace(
|
||||
impl=types.SimpleNamespace(fused_experts=fused_experts)
|
||||
),
|
||||
)
|
||||
)
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
self._quant_method.moe_kernel.impl.fused_experts = fused_experts
|
||||
fused_experts.supports_routing_replay_capture = lambda: False
|
||||
|
||||
class DummyCapturer:
|
||||
def capture(self, layer_id, topk_ids):
|
||||
pass
|
||||
|
||||
dummy_module = DummyFusedMoE()
|
||||
import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
|
||||
|
||||
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
|
||||
|
||||
dummy_self = types.SimpleNamespace(
|
||||
model=types.SimpleNamespace(modules=lambda: [dummy_module])
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="monolithic MoE kernel"):
|
||||
gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
|
||||
|
||||
|
||||
def test_routed_experts_capturer_single_dp_no_metadata():
|
||||
"""dp_metadata is None: capture writes the full topk_ids rows."""
|
||||
capturer = _capturer_with_buffer(dp_rank=0)
|
||||
|
||||
@@ -9,7 +9,6 @@ import torch.nn.functional as F
|
||||
from transformers import AutoModel
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_constants import MiB_bytes
|
||||
|
||||
from ....conftest import HfRunner
|
||||
from ....utils import VLLM_PATH
|
||||
@@ -107,12 +106,7 @@ def test_prm_models(
|
||||
if current_platform.is_cpu():
|
||||
pytest.skip("CPU only supports V1")
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
kv_cache_memory_bytes=64 * MiB_bytes,
|
||||
) as vllm_model:
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
|
||||
@@ -151,12 +145,7 @@ def test_prm_models_with_golden_outputs(
|
||||
if not FIXTURE_REWARD_RESULT.get(model):
|
||||
pytest.skip(f"No available golden outputs for {model}.")
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
kv_cache_memory_bytes=64 * MiB_bytes,
|
||||
) as vllm_model:
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
|
||||
|
||||
@@ -70,19 +70,7 @@ def _assert_video_outputs(processor, processed) -> None:
|
||||
merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
|
||||
expected_tokens = int(grid_thw.prod()) // merge_size**2
|
||||
video_token_id = processor.info.get_hf_config().video_token_id
|
||||
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
|
||||
)
|
||||
assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_images", [1, 2])
|
||||
|
||||
@@ -74,44 +74,6 @@ def test_cosmos3_new_checkpoint_weights_mapper():
|
||||
)
|
||||
|
||||
|
||||
def test_cosmos3_modelopt_quantizer_weights_mapper():
|
||||
"""ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
|
||||
(``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
|
||||
``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
|
||||
(it has no parameter for them) while keeping the latter."""
|
||||
from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
|
||||
|
||||
mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
|
||||
|
||||
# Native ModelOpt quantizer buffers are dropped.
|
||||
assert (
|
||||
mapper.apply_list(
|
||||
[
|
||||
"layers.0.self_attn.to_q.input_quantizer._amax",
|
||||
"layers.0.self_attn.to_q.weight_quantizer._amax",
|
||||
"layers.0.self_attn.to_q.weight_quantizer._scale",
|
||||
"layers.0.mlp.down_proj.output_quantizer._amax",
|
||||
]
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
# The FP8 scale sidecars vLLM actually consumes are kept and remapped.
|
||||
assert mapper.apply_list(
|
||||
[
|
||||
"layers.0.self_attn.to_q.weight",
|
||||
"layers.0.self_attn.to_q.weight_scale",
|
||||
"layers.0.self_attn.to_q.input_scale",
|
||||
"layers.0.mlp.down_proj.input_scale",
|
||||
]
|
||||
) == [
|
||||
"language_model.model.layers.0.self_attn.q_proj.weight",
|
||||
"language_model.model.layers.0.self_attn.q_proj.weight_scale",
|
||||
"language_model.model.layers.0.self_attn.q_proj.input_scale",
|
||||
"language_model.model.layers.0.mlp.down_proj.input_scale",
|
||||
]
|
||||
|
||||
|
||||
def test_cosmos3_edge_checkpoint_weights_mapper():
|
||||
from vllm.model_executor.models.cosmos3_edge import (
|
||||
Cosmos3EdgeForConditionalGeneration,
|
||||
@@ -170,7 +132,6 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
|
||||
"layers.0.self_attn.to_add_out.weight",
|
||||
"layers.0.self_attn.norm_added_q.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.mlp_moe_gen.up_proj.weight",
|
||||
"norm_moe_gen.weight",
|
||||
|
||||
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
|
||||
import pytest
|
||||
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
|
||||
from ..utils import check_logprobs_close
|
||||
|
||||
|
||||
@@ -70,7 +70,13 @@ def test_models(
|
||||
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
|
||||
pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
|
||||
|
||||
if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
|
||||
if not (
|
||||
current_platform.is_xpu()
|
||||
or (
|
||||
get_flash_attn_version() == 3
|
||||
and current_platform.is_device_capability_family(90)
|
||||
)
|
||||
):
|
||||
pytest.skip(
|
||||
f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
|
||||
)
|
||||
|
||||
@@ -6,7 +6,6 @@ import mimetypes
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from io import BytesIO
|
||||
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
||||
|
||||
import aiohttp
|
||||
@@ -112,34 +111,6 @@ async def test_fetch_image_base64(
|
||||
assert _image_equals(data_image_sync, data_image_async)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_image_keep_original_mode():
|
||||
"""media_io_kwargs can disable the default RGB conversion."""
|
||||
# RGBA image: opaque black pixel on a fully transparent background
|
||||
rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
|
||||
rgba_image.putpixel((2, 2), (0, 0, 0, 255))
|
||||
buffer = BytesIO()
|
||||
rgba_image.save(buffer, "PNG")
|
||||
data_url = (
|
||||
f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
|
||||
)
|
||||
|
||||
# Default behavior: RGBA is composited onto a white background
|
||||
default_image = MediaConnector().fetch_image(data_url)
|
||||
assert default_image.mode == "RGB"
|
||||
assert default_image.getpixel((0, 0)) == (255, 255, 255)
|
||||
assert default_image.getpixel((2, 2)) == (0, 0, 0)
|
||||
|
||||
# image_mode=None via media_io_kwargs: original mode is preserved
|
||||
connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
|
||||
image_sync = connector.fetch_image(data_url)
|
||||
image_async = await connector.fetch_image_async(data_url)
|
||||
for image in (image_sync, image_async):
|
||||
assert image.mode == "RGBA"
|
||||
assert image.getpixel((0, 0)) == (0, 0, 0, 0)
|
||||
assert image.getpixel((2, 2)) == (0, 0, 0, 255)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_fetch_image_local_files(image_url: str):
|
||||
|
||||
@@ -80,29 +80,6 @@ def test_image_media_io_rgba_custom_background(tmp_path):
|
||||
assert green_numpy[0][0][2] == 0 # B
|
||||
|
||||
|
||||
def test_image_media_io_no_mode_conversion(tmp_path):
|
||||
"""image_mode=None skips conversion and preserves the original mode."""
|
||||
# RGBA image: opaque black pixel on a fully transparent background
|
||||
rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
|
||||
rgba_image.putpixel((5, 5), (0, 0, 0, 255))
|
||||
test_image_path = tmp_path / "test_rgba.png"
|
||||
rgba_image.save(test_image_path)
|
||||
|
||||
# Default behavior: RGBA is composited onto a white background
|
||||
image_io_default = ImageMediaIO()
|
||||
converted_default = image_io_default.load_file(test_image_path)
|
||||
assert converted_default.media.mode == "RGB"
|
||||
assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
|
||||
assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
|
||||
|
||||
# image_mode=None: original mode and alpha channel are preserved
|
||||
image_io_keep = ImageMediaIO(image_mode=None)
|
||||
converted_keep = image_io_keep.load_file(test_image_path)
|
||||
assert converted_keep.media.mode == "RGBA"
|
||||
assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
|
||||
assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
|
||||
|
||||
|
||||
def test_image_media_io_rgba_background_color_validation():
|
||||
"""Test that invalid rgba_background_color values are properly rejected."""
|
||||
|
||||
|
||||
@@ -304,13 +304,6 @@ 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(
|
||||
"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):
|
||||
expected_weights = [
|
||||
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
|
||||
{"token_id": 70, "weight": 0.09808349609375, "token": " the"},
|
||||
{"token_id": 83, "weight": 0.08154296875, "token": " is"},
|
||||
{"token_id": 111, "weight": 0.11810302734375, "token": " of"},
|
||||
{"token_id": 4865, "weight": 0.1171875, "token": " What"},
|
||||
{"token_id": 9942, "weight": 0.292236328125, "token": " France"},
|
||||
{"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
|
||||
{"token_id": 70, "weight": 0.09808349609375, "token": "the"},
|
||||
{"token_id": 83, "weight": 0.08154296875, "token": "is"},
|
||||
{"token_id": 111, "weight": 0.11810302734375, "token": "of"},
|
||||
{"token_id": 4865, "weight": 0.1171875, "token": "What"},
|
||||
{"token_id": 9942, "weight": 0.292236328125, "token": "France"},
|
||||
{"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
|
||||
]
|
||||
expected_embed = {x["token_id"]: x for x in expected_weights}
|
||||
|
||||
|
||||
@@ -93,6 +93,9 @@ def test_online_quantization(
|
||||
use_rocm_aiter: bool,
|
||||
monkeypatch,
|
||||
) -> 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:
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
|
||||
@@ -102,15 +105,9 @@ def test_online_quantization(
|
||||
if force_marlin:
|
||||
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(
|
||||
"facebook/opt-125m",
|
||||
quantization="fp8",
|
||||
dtype=model_dtype,
|
||||
enforce_eager=True,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
) as llm:
|
||||
|
||||
@@ -165,38 +165,6 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
|
||||
assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
|
||||
|
||||
|
||||
def test_modelopt_mixed_precision_composes_gemma4_mappers():
|
||||
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
|
||||
from vllm.model_executor.models.gemma4_mm import (
|
||||
Gemma4ForConditionalGeneration,
|
||||
)
|
||||
|
||||
config = _mixed_precision_config(
|
||||
{
|
||||
"model.language_model.layers.0.experts": {
|
||||
"quant_algo": "NVFP4",
|
||||
"group_size": 16,
|
||||
},
|
||||
"model.language_model.layers.1.moe.experts.gate_up_proj": {
|
||||
"quant_algo": "NVFP4",
|
||||
"group_size": 16,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
config.apply_vllm_mapper(
|
||||
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
|
||||
)
|
||||
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
|
||||
|
||||
expected_prefix = "language_model.model.layers.0.moe.experts"
|
||||
assert set(config.quantized_layers) == {
|
||||
expected_prefix,
|
||||
"language_model.model.layers.1.moe.gate_up_proj",
|
||||
}
|
||||
assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
|
||||
|
||||
|
||||
def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.utils.extensible_tensor import ExtensibleTensor
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
||||
|
||||
|
||||
def test_extensible_tensor_grows_without_moving() -> None:
|
||||
buffer = ExtensibleTensor(4096, device="cuda")
|
||||
try:
|
||||
base_ptr = buffer.base_ptr
|
||||
first_view = buffer.resize_(1024)
|
||||
assert first_view.data_ptr() == base_ptr
|
||||
first_view.fill_(7)
|
||||
|
||||
second_view = buffer.resize_(2048)
|
||||
assert second_view.data_ptr() == base_ptr
|
||||
assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7))
|
||||
|
||||
second_view[1024:].fill_(3)
|
||||
assert torch.equal(buffer.tensor, second_view)
|
||||
|
||||
full_view = buffer.full_view()
|
||||
assert full_view.data_ptr() == base_ptr
|
||||
assert full_view.numel() == 4096
|
||||
finally:
|
||||
buffer.free()
|
||||
|
||||
|
||||
def test_extensible_tensor_rejects_shrink_and_overflow() -> None:
|
||||
buffer = ExtensibleTensor(1024, device="cuda")
|
||||
try:
|
||||
buffer.resize_(512)
|
||||
with pytest.raises(ValueError, match="grow-only"):
|
||||
buffer.resize_(256)
|
||||
with pytest.raises(ValueError, match="exceeds the segment capacity"):
|
||||
buffer.resize_(1025)
|
||||
finally:
|
||||
buffer.free()
|
||||
|
||||
|
||||
def test_segments_grow_in_lockstep_and_zero_new() -> None:
|
||||
"""Each segment's committed prefix grows in lockstep.
|
||||
|
||||
Data written to a segment's committed prefix survives a grow; the newly
|
||||
committed range of each segment is zeroed with `zero_new=True` while old
|
||||
bytes are preserved.
|
||||
"""
|
||||
et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
|
||||
try:
|
||||
assert et.num_segments == 2
|
||||
assert et.segment_capacity_bytes == 4096
|
||||
|
||||
et.resize_per_segment_(256, zero_new=True)
|
||||
assert et.bytes_per_segment == 256
|
||||
assert et.num_bytes == 512
|
||||
fv = et.full_view()
|
||||
assert fv.shape == (8192,)
|
||||
# Committed prefixes start zeroed.
|
||||
assert torch.count_nonzero(fv[:256]) == 0
|
||||
assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0
|
||||
|
||||
pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8)
|
||||
pattern_b = 255 - pattern_a
|
||||
fv[:256].copy_(pattern_a)
|
||||
fv[4096 : 4096 + 256].copy_(pattern_b)
|
||||
|
||||
et.resize_per_segment_(1024, zero_new=True)
|
||||
fv2 = et.full_view()
|
||||
assert fv2.data_ptr() == fv.data_ptr()
|
||||
# Old bytes of both segments preserved; freshly committed ranges zeroed.
|
||||
assert torch.equal(fv2[:256], pattern_a)
|
||||
assert torch.equal(fv2[4096 : 4096 + 256], pattern_b)
|
||||
assert torch.count_nonzero(fv2[256:1024]) == 0
|
||||
assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0
|
||||
finally:
|
||||
et.free()
|
||||
|
||||
|
||||
def test_segments_at_granularity_scale() -> None:
|
||||
"""Segments spanning multiple mapping granules commit correctly.
|
||||
|
||||
Uses a segment capacity that is not a multiple of the allocation
|
||||
granularity, so a granule straddles the segment boundary and is shared by
|
||||
the first commit of one segment and a later commit of the other -- it must
|
||||
be mapped exactly once.
|
||||
"""
|
||||
probe = ExtensibleTensor(max_num_bytes=1, device="cuda")
|
||||
granularity = probe.capacity_bytes
|
||||
probe.free()
|
||||
# Two segments of 1.5 granules each; the middle granule straddles the
|
||||
# boundary.
|
||||
max_num_bytes = 3 * granularity
|
||||
et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2)
|
||||
try:
|
||||
seg = et.segment_capacity_bytes
|
||||
assert seg == max_num_bytes // 2
|
||||
|
||||
step = granularity // 2
|
||||
et.resize_per_segment_(step, zero_new=True)
|
||||
fv = et.full_view()
|
||||
fv[:step].fill_(1)
|
||||
fv[seg : seg + step].fill_(2)
|
||||
|
||||
# Grow to the full segment capacity: previously mapped granules
|
||||
# (including the boundary-straddling one) are reused, new ones are
|
||||
# committed and zeroed.
|
||||
et.resize_per_segment_(seg, zero_new=True)
|
||||
fv2 = et.full_view()
|
||||
assert torch.all(fv2[:step] == 1)
|
||||
assert torch.all(fv2[seg : seg + step] == 2)
|
||||
assert torch.count_nonzero(fv2[step:seg]) == 0
|
||||
assert torch.count_nonzero(fv2[seg + step :]) == 0
|
||||
finally:
|
||||
et.free()
|
||||
|
||||
|
||||
def test_multi_segment_invalid_usage_raises() -> None:
|
||||
"""Prefix-view APIs and invalid segment configs raise for multi-segment
|
||||
buffers."""
|
||||
with pytest.raises(ValueError):
|
||||
ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3)
|
||||
|
||||
et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
|
||||
try:
|
||||
with pytest.raises(ValueError):
|
||||
_ = et.tensor
|
||||
with pytest.raises(ValueError):
|
||||
et.resize_(256)
|
||||
|
||||
et.resize_per_segment_(256)
|
||||
with pytest.raises(ValueError):
|
||||
et.resize_per_segment_(128) # shrink
|
||||
with pytest.raises(ValueError):
|
||||
et.resize_per_segment_(et.segment_capacity_bytes + 1) # over capacity
|
||||
finally:
|
||||
et.free()
|
||||
@@ -67,23 +67,6 @@ def test_memory_profiling():
|
||||
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
|
||||
assert abs(non_torch_ratio - 1) <= 0.05
|
||||
assert result.torch_peak_increase == 1024 * 1024 * 1024
|
||||
|
||||
expected_total_consumed = (256 + 512) * 1024 * 1024
|
||||
total_consumed_ratio = result.total_consumed / expected_total_consumed
|
||||
assert abs(total_consumed_ratio - 1) <= 0.05, (
|
||||
f"total_consumed={result.total_consumed}, "
|
||||
f"expected={expected_total_consumed}, "
|
||||
f"ratio={total_consumed_ratio}"
|
||||
)
|
||||
|
||||
expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
|
||||
non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
|
||||
assert abs(non_kv_ratio - 1) <= 0.05, (
|
||||
f"non_kv_cache_memory={result.non_kv_cache_memory}, "
|
||||
f"expected={expected_non_kv}, "
|
||||
f"ratio={non_kv_ratio}"
|
||||
)
|
||||
|
||||
del weights
|
||||
lib.cudaFree(handle1)
|
||||
lib.cudaFree(handle2)
|
||||
|
||||
@@ -21,7 +21,6 @@ from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
is_quantized_kv_cache,
|
||||
is_torch_equal_or_newer,
|
||||
set_random_seed,
|
||||
)
|
||||
@@ -46,11 +45,6 @@ BACKENDS_TO_TEST = [
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
FP8_KV_CACHE_DTYPES = {
|
||||
"fp8": torch.float8_e4m3fn,
|
||||
"fp8_e4m3": torch.float8_e4m3fn,
|
||||
}
|
||||
|
||||
# Remove flashinfer from the list if it's not available
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
@@ -116,7 +110,6 @@ def create_and_prepopulate_kv_cache(
|
||||
num_blocks: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
randomize_blocks: bool = True,
|
||||
kv_cache_dtype: str = "auto",
|
||||
) -> torch.Tensor:
|
||||
"""Create and prepopulate a KV cache with context data.
|
||||
|
||||
@@ -147,18 +140,8 @@ def create_and_prepopulate_kv_cache(
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
|
||||
# For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
|
||||
# the higher-precision context tensors quantizes them, mirroring runtime.
|
||||
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
|
||||
storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
|
||||
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
2 * head_size,
|
||||
dtype=storage_dtype,
|
||||
device=device,
|
||||
num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
|
||||
)
|
||||
kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
|
||||
|
||||
@@ -212,12 +195,7 @@ def create_and_prepopulate_kv_cache(
|
||||
] * block_size + token_inter_block_offsets.to(device)
|
||||
|
||||
# Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
|
||||
kv_cache = kv_cache.transpose(1, 2).contiguous()
|
||||
|
||||
if fp8_kv_cache:
|
||||
kv_cache = kv_cache.view(torch.uint8)
|
||||
|
||||
return kv_cache
|
||||
return kv_cache.transpose(1, 2).contiguous()
|
||||
|
||||
|
||||
class MockAttentionLayer:
|
||||
@@ -246,7 +224,6 @@ def run_attention_backend(
|
||||
kv_cache: torch.Tensor,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
sliding_window: int | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
) -> torch.Tensor:
|
||||
"""Run attention computation using the specified backend's AttentionImpl."""
|
||||
|
||||
@@ -314,16 +291,13 @@ def run_attention_backend(
|
||||
alibi_slopes=None,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
kv_cache_dtype="auto",
|
||||
)
|
||||
|
||||
# Create mock layer and output buffer
|
||||
mock_layer = MockAttentionLayer(device)
|
||||
output = torch.empty_like(query)
|
||||
|
||||
if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
|
||||
query = query.to(current_platform.fp8_dtype())
|
||||
|
||||
# Run forward pass
|
||||
# NOTE: The query, key, and value are already shaped correctly
|
||||
# in the calling test function.
|
||||
@@ -350,7 +324,6 @@ def _test_backend_correctness(
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: int = 1,
|
||||
kv_cache_dtype: str = "auto",
|
||||
):
|
||||
"""
|
||||
Test that all backends produce similar outputs to a reference implementation
|
||||
@@ -399,7 +372,6 @@ def _test_backend_correctness(
|
||||
num_gpu_blocks=8192,
|
||||
hf_config_override=hf_config_override,
|
||||
)
|
||||
vllm_config.cache_config.cache_dtype = kv_cache_dtype
|
||||
device = torch.device(f"{DEVICE_TYPE}:0")
|
||||
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
|
||||
@@ -420,13 +392,6 @@ def _test_backend_correctness(
|
||||
block_size = vllm_config.cache_config.block_size
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
|
||||
if fp8_kv_cache:
|
||||
query_fp8_dtype = current_platform.fp8_dtype()
|
||||
kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
|
||||
atol = max(atol, 6e-2)
|
||||
rtol = max(rtol, 1e-1)
|
||||
|
||||
# 2. Generate data and compute SDPA reference output
|
||||
all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
|
||||
all_sdpa_outputs = []
|
||||
@@ -442,17 +407,10 @@ def _test_backend_correctness(
|
||||
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
|
||||
if fp8_kv_cache:
|
||||
q_ref = q.to(query_fp8_dtype).to(dtype)
|
||||
k_ref = k_full.to(kv_fp8_dtype).to(dtype)
|
||||
v_ref = v_full.to(kv_fp8_dtype).to(dtype)
|
||||
else:
|
||||
q_ref, k_ref, v_ref = q, k_full, v_full
|
||||
|
||||
# SDPA expects (N, H, L, D), so unsqueeze batch and permute
|
||||
q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)
|
||||
q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if num_q_heads != num_kv_heads:
|
||||
assert num_q_heads % num_kv_heads == 0, (
|
||||
@@ -513,7 +471,6 @@ def _test_backend_correctness(
|
||||
num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
randomize_blocks=True,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# 4. Run vLLM backends and compare
|
||||
@@ -531,12 +488,6 @@ def _test_backend_correctness(
|
||||
else:
|
||||
backend_cls = None
|
||||
|
||||
if is_quantized_kv_cache(kv_cache_dtype) and (
|
||||
backend_cls is None
|
||||
or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
):
|
||||
continue
|
||||
|
||||
if backend_name == AttentionBackendEnum.FLASHINFER:
|
||||
set_kv_cache_layout("HND")
|
||||
reset_kv_cache_layout = True
|
||||
@@ -570,7 +521,6 @@ def _test_backend_correctness(
|
||||
kv_cache_for_backend,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
finally:
|
||||
if reset_kv_cache_layout:
|
||||
@@ -620,13 +570,8 @@ def _test_backend_correctness(
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
|
||||
def test_causal_backend_correctness(
|
||||
default_vllm_config,
|
||||
batch_spec_name: str,
|
||||
model: str,
|
||||
tensor_parallel_size: int,
|
||||
kv_cache_dtype: str,
|
||||
default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
|
||||
):
|
||||
"""Test backend's correctness with causal attention."""
|
||||
|
||||
@@ -667,7 +612,6 @@ def test_causal_backend_correctness(
|
||||
SMALL_BLOCK_BACKENDS,
|
||||
causal_mask_mod,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# Fast FlexAttention needs to run with block_size=128
|
||||
@@ -679,7 +623,6 @@ def test_causal_backend_correctness(
|
||||
causal_mask_mod,
|
||||
block_size=128,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -119,7 +119,6 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
|
||||
layer.kv_b_proj.quant_method = None
|
||||
layer.is_aiter_triton_fp4_bmm_enabled = False
|
||||
layer.is_aiter_triton_fp8_bmm_enabled = False
|
||||
layer.dcp_q_replicate = False
|
||||
layer.quant_config = None
|
||||
layer.layer_name = "test"
|
||||
|
||||
|
||||
@@ -33,7 +33,6 @@ from vllm.v1.kv_cache_interface import (
|
||||
EncoderOnlyAttentionSpec,
|
||||
FullAttentionSpec,
|
||||
MambaSpec,
|
||||
get_kv_quant_mode,
|
||||
)
|
||||
|
||||
|
||||
@@ -179,7 +178,6 @@ def create_standard_kv_cache_spec(
|
||||
head_size=vllm_config.model_config.get_head_size(),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
sliding_window=vllm_config.model_config.get_sliding_window(),
|
||||
kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -149,6 +149,30 @@ def test_has_cache_restores_from_freeable():
|
||||
assert manager.num_freeable_slots == 6
|
||||
|
||||
|
||||
def test_make_profiling_reservation():
|
||||
assert (
|
||||
EncoderCacheManager.make_profiling_reservation(
|
||||
cache_size=0,
|
||||
embed_size=8,
|
||||
dtype=torch.float16,
|
||||
device="cpu",
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
reservation = EncoderCacheManager.make_profiling_reservation(
|
||||
cache_size=7,
|
||||
embed_size=8,
|
||||
dtype=torch.float16,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
assert reservation is not None
|
||||
assert reservation.shape == (7, 8)
|
||||
assert reservation.dtype == torch.float16
|
||||
assert reservation.device.type == "cpu"
|
||||
|
||||
|
||||
def test_get_freed_mm_hashes_clears_freed_list():
|
||||
manager = EncoderCacheManager(cache_size=10)
|
||||
req1 = MockRequest("reqA", ["a"], [5])
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
from collections.abc import Callable
|
||||
@@ -1478,128 +1477,6 @@ def test_get_max_concurrency_for_kv_cache_config():
|
||||
assert num_tokens == max_concurrency_hybrid_model * max_model_len
|
||||
assert max_concurrency == max_concurrency_hybrid_model
|
||||
|
||||
# Unequal group sizes in the standard layout: each group's pages cost
|
||||
# whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
|
||||
# same as the equal-hybrid case above, regardless of the second group
|
||||
# holding only 2 layers.
|
||||
kv_cache_config_unequal_groups = KVCacheConfig(
|
||||
num_blocks=1153 * 3,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
|
||||
KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
|
||||
],
|
||||
)
|
||||
assert (
|
||||
get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_unequal_groups
|
||||
)
|
||||
== 3
|
||||
)
|
||||
|
||||
# UniformTypeKVCacheSpecs group (worker config shape): the aggregated
|
||||
# spec's memory/page ratio equals a single layer's page count, so the
|
||||
# group needs 1024 blocks and the request 1153 in total. The previous
|
||||
# formula normalized both groups' memory by the first group's page size,
|
||||
# reporting 3459/1057 = 3.27 here instead of 3 — and a different value
|
||||
# again for the scheduler-config shape below.
|
||||
uniform_full_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=full_attention_spec.block_size,
|
||||
kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
|
||||
)
|
||||
kv_cache_config_uniform_group = KVCacheConfig(
|
||||
num_blocks=1153 * 3,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
|
||||
KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
|
||||
],
|
||||
)
|
||||
assert (
|
||||
get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_uniform_group
|
||||
)
|
||||
== 3
|
||||
)
|
||||
|
||||
# Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
|
||||
# uniform-type group's spec with a representative per-layer spec.
|
||||
# Capacity must not change between the two shapes (the engine computes
|
||||
# on the scheduler config, the worker loop on the worker config).
|
||||
kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
|
||||
[copy.deepcopy(kv_cache_config_uniform_group)]
|
||||
)
|
||||
assert get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_scheduler_shape
|
||||
) == get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_uniform_group
|
||||
)
|
||||
|
||||
|
||||
def test_get_max_concurrency_packed_kv_cache_config():
|
||||
from vllm.v1.core.kv_cache_utils import (
|
||||
_get_kv_cache_config_packed,
|
||||
_use_packed_kv_cache_config,
|
||||
)
|
||||
|
||||
model_config = ModelConfig(
|
||||
"Qwen/Qwen1.5-7B",
|
||||
runner="generate",
|
||||
dtype="float16",
|
||||
max_model_len=16384,
|
||||
)
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_batched_tokens=1024,
|
||||
enable_chunked_prefill=True,
|
||||
max_model_len=model_config.max_model_len,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
async_scheduling=False,
|
||||
)
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
scheduler_config=scheduler_config,
|
||||
)
|
||||
|
||||
# All-UniformTypeKVCacheSpecs groups select the packed layout.
|
||||
mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
|
||||
swa_specs = {
|
||||
f"layer_{i}": SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=576,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
)
|
||||
for i in range(4, 6)
|
||||
}
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
list(mla_specs),
|
||||
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
list(swa_specs),
|
||||
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
|
||||
),
|
||||
]
|
||||
assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
|
||||
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
|
||||
vllm_config, kv_cache_groups, 2 * GiB_bytes
|
||||
)
|
||||
assert num_blocks > 0
|
||||
kv_cache_config_packed = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=kv_cache_tensors,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
# Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
|
||||
# the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
|
||||
# previous formula normalized by the first group's page size and gave
|
||||
# 1061 blocks per request instead of 1097.
|
||||
assert get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_packed
|
||||
) == num_blocks / (1024 + 73)
|
||||
|
||||
|
||||
def test_allocate_with_lookahead():
|
||||
"""Verify that lookahead tokens correctly affect block allocation"""
|
||||
|
||||
@@ -49,6 +49,18 @@ def test_prefix_caching_from_cli():
|
||||
args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
|
||||
|
||||
|
||||
def test_extensible_kv_cache_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
|
||||
args = parser.parse_args([])
|
||||
engine_args = EngineArgs.from_cli_args(args=args)
|
||||
assert not engine_args.enable_extensible_kv_cache
|
||||
|
||||
args = parser.parse_args(["--enable-extensible-kv-cache"])
|
||||
engine_args = EngineArgs.from_cli_args(args=args)
|
||||
assert engine_args.enable_extensible_kv_cache
|
||||
|
||||
|
||||
@pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
|
||||
def test_prefix_caching_xxhash_from_cli():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user