forked from Karylab-cklius/vllm
Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a02155c787 | ||
|
|
319db65b68 | ||
|
|
83a7669827 | ||
|
|
401bed48ad | ||
|
|
d492d1e697 | ||
|
|
f37e113590 | ||
|
|
a10e369f06 | ||
|
|
aa3f2efe42 | ||
|
|
a32d1bff95 | ||
|
|
ee47a21fcd | ||
|
|
0f471a3088 |
@@ -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"
|
||||
}
|
||||
@@ -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/
|
||||
|
||||
+86
-42
@@ -227,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
cuda_archs_loose_intersection(CUDA_ARCHS
|
||||
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
|
||||
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
|
||||
|
||||
set(VLLM_COMPILED_CUDA_ARCHS)
|
||||
foreach(_ARCH ${CUDA_ARCHS})
|
||||
set(_COMPILED_ARCH "${_ARCH}")
|
||||
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
if(_ARCH MATCHES "^(10|11|12)\\.0$")
|
||||
set(_COMPILED_ARCH "${_ARCH}f")
|
||||
endif()
|
||||
elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
|
||||
if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
|
||||
set(_COMPILED_ARCH "${_ARCH}a")
|
||||
endif()
|
||||
endif()
|
||||
list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
|
||||
endforeach()
|
||||
list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
|
||||
else()
|
||||
#
|
||||
# For other GPU targets override the GPU architectures detected by cmake/torch
|
||||
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
#
|
||||
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
|
||||
#
|
||||
# Shared entry sources are part of the base extension source list, but some
|
||||
# optional kernel families below append source-local feature macros to them.
|
||||
set(VLLM_STABLE_TORCH_BINDINGS_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp")
|
||||
set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
|
||||
set(SCALED_MM_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
|
||||
set(NVFP4_QUANT_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
|
||||
set(NVFP4_SCALED_MM_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
|
||||
|
||||
set(VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp"
|
||||
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/activation_kernels.cu"
|
||||
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/sampler.cu"
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
|
||||
"csrc/libtorch_stable/cache_kernels.cu"
|
||||
"${CACHE_KERNELS_SRC}"
|
||||
"csrc/libtorch_stable/cache_kernels_fused.cu"
|
||||
"csrc/libtorch_stable/custom_all_reduce.cu"
|
||||
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
|
||||
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
|
||||
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
|
||||
"${SCALED_MM_ENTRY_SRC}"
|
||||
"${NVFP4_QUANT_ENTRY_SRC}"
|
||||
"${NVFP4_SCALED_MM_ENTRY_SRC}"
|
||||
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
|
||||
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
|
||||
DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")
|
||||
|
||||
#
|
||||
# Machete kernels
|
||||
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cooperative_topk.cu")
|
||||
"${COOPERATIVE_TOPK_SRC}")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
|
||||
SRCS "${COOPERATIVE_TOPK_SRC}"
|
||||
CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
|
||||
DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
|
||||
endif()
|
||||
|
||||
# Only build Marlin kernels if we are building for at least some compatible archs.
|
||||
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM90_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM90=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM120_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM120=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM100_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM100=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_C2X_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_C2X=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
|
||||
message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
|
||||
else()
|
||||
if (SCALED_MM_3X_ARCHS)
|
||||
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${CUTLASS_MOE_SM90_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
|
||||
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
|
||||
else()
|
||||
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
|
||||
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${CUTLASS_MOE_SM100_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
# The implementation is named sm100 historically, but it is built for the
|
||||
# SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
|
||||
# SM110. Keep the compile-time macro aligned with runtime dispatch.
|
||||
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
|
||||
else()
|
||||
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
|
||||
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
|
||||
#
|
||||
|
||||
# SM12x FP4 kernels. These share some generic NVFP4 quantization entry
|
||||
# sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
|
||||
# per-source flags, so shared files accumulate both SM12x and SM10x/11x
|
||||
# gencodes when both families are requested.
|
||||
# Shared FP4 implementation sources live next to the FP4 arch logic because
|
||||
# both SM12x and SM10x/11x append family-specific gencodes and feature macros
|
||||
# to them.
|
||||
set(FP4_SHARED_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
|
||||
set(FP4_SM120_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
${FP4_SHARED_SRCS}
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
)
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FP4_SM120_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM120_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
|
||||
DEFINITIONS ENABLE_NVFP4_SM120=1)
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
|
||||
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
|
||||
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
|
||||
set(FP4_SM100_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
${FP4_SHARED_SRCS}
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
|
||||
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
message(STATUS
|
||||
"Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
|
||||
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FP4_SM100_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM100_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_NVFP4_SM100=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
|
||||
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
SRCS "${CUTLASS_MLA_SRCS}"
|
||||
CUDA_ARCHS "${MLA_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
|
||||
# Add MLA-specific include directories only to MLA source files
|
||||
set_source_files_properties(${CUTLASS_MLA_SRCS}
|
||||
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
|
||||
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
VLLM_ENABLE_COOPERATIVE_TOPK=1)
|
||||
endif()
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
|
||||
endif()
|
||||
endmacro()
|
||||
|
||||
#
|
||||
# For a list of source files append preprocessor definitions to file-specific
|
||||
# compile options. Use this for optional kernel feature macros so toggling one
|
||||
# kernel family does not perturb the compile command for every source in the
|
||||
# extension target.
|
||||
#
|
||||
macro(set_compile_definitions_for_srcs)
|
||||
set(options)
|
||||
set(oneValueArgs)
|
||||
set(multiValueArgs SRCS DEFINITIONS)
|
||||
cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
|
||||
"${multiValueArgs}" ${ARGN})
|
||||
|
||||
foreach(_DEF ${arg_DEFINITIONS})
|
||||
set_property(
|
||||
SOURCE ${arg_SRCS}
|
||||
APPEND PROPERTY
|
||||
COMPILE_DEFINITIONS "${_DEF}"
|
||||
)
|
||||
endforeach()
|
||||
endmacro()
|
||||
|
||||
#
|
||||
# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
|
||||
# `<major>.<minor>[letter]` compute the "loose intersection" with the
|
||||
|
||||
-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) {
|
||||
|
||||
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
|
||||
torch::stable::Tensor const& perm);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
std::string get_compiled_cuda_archs();
|
||||
|
||||
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
|
||||
|
||||
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
|
||||
|
||||
#endif
|
||||
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
|
||||
torch::stable::Tensor const& a_tensors,
|
||||
torch::stable::Tensor const& b_tensors,
|
||||
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
|
||||
std::optional<torch::stable::Tensor> const& bias);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
|
||||
void get_cutlass_moe_mm_data_caller(
|
||||
const torch::stable::Tensor& topk_ids,
|
||||
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
|
||||
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
|
||||
// CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
|
||||
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
|
||||
// actual cutlass_moe_mm dispatch compiled into this file.
|
||||
// or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
|
||||
// SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
|
||||
// implementation. Only report archs that have an actual cutlass_moe_mm
|
||||
// dispatch compiled into this file.
|
||||
|
||||
#if defined CUDA_VERSION
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
|
||||
return CUDA_VERSION >= 12080;
|
||||
}
|
||||
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
|
||||
torch::stable::Tensor const& c_strides, bool per_act_token,
|
||||
bool per_out_ch) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
if (version_num >= 100 && version_num < 110) {
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
// Keep runtime dispatch aligned with the CMake arch list and support query:
|
||||
// CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
|
||||
if (version_num >= 100 && version_num < 120) {
|
||||
cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
|
||||
expert_offsets, problem_sizes, a_strides, b_strides,
|
||||
c_strides, per_act_token, per_out_ch);
|
||||
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
|
||||
// This function currently gets compiled only if we have a valid cutlass moe
|
||||
// mm to run it for.
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
|
||||
problem_sizes2, input_permutation,
|
||||
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
|
||||
torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
|
||||
const bool swap_ab) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
|
||||
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
|
||||
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
|
||||
// This function currently gets compiled only if we have a valid cutlass moe
|
||||
// mm to run it for.
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
|
||||
problem_sizes2, expert_num_tokens,
|
||||
|
||||
@@ -4,6 +4,10 @@
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
|
||||
#endif
|
||||
|
||||
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
|
||||
// Note: We register under namespace "_C" so ops are accessible as
|
||||
// torch.ops._C.<op_name> for compatibility with existing code.
|
||||
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.def("get_compiled_cuda_archs() -> str");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
|
||||
@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
|
||||
// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
|
||||
ops.impl("cutlass_scaled_mm_supports_fp8",
|
||||
TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
|
||||
ops.impl("cutlass_group_gemm_supported",
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
|
||||
uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
|
||||
```
|
||||
|
||||
!!! warning "CUDA architecture coverage"
|
||||
|
||||
Pre-built CUDA wheels are compiled for the CUDA architectures selected by
|
||||
vLLM's release and build pipelines. This list is intentionally smaller than
|
||||
every architecture CMake can build, because each additional architecture
|
||||
increases wheel size.
|
||||
|
||||
In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
|
||||
To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
|
||||
architecture-specific Blackwell/Thor targets, such as `sm_103` or
|
||||
`sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
|
||||
family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
|
||||
cover the corresponding major-version GPU family.
|
||||
|
||||
If vLLM logs a warning that your visible CUDA device is not covered by the
|
||||
wheel's compiled CUDA architectures, or if you see a CUDA error such as
|
||||
`no kernel image is available for execution on the device`, install a CUDA
|
||||
13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
|
||||
that includes your GPU.
|
||||
|
||||
#### Install the latest code
|
||||
|
||||
LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc
|
||||
|
||||
If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.
|
||||
|
||||
## CUDA architecture not covered by the wheel
|
||||
|
||||
If vLLM logs a warning that the current wheel was built for a set of CUDA
|
||||
architectures but one of your visible CUDA devices is not covered, the installed
|
||||
wheel may not contain native CUDA kernels for that GPU. The same issue may also
|
||||
surface later as a CUDA runtime error such as `no kernel image is available for
|
||||
execution on the device`.
|
||||
|
||||
The architecture list in a pre-built wheel is determined by the vLLM release and
|
||||
build pipelines, then filtered by CMake before individual kernels choose their
|
||||
own per-kernel architectures. It is not the same as the full set of
|
||||
architectures that vLLM can build from source.
|
||||
|
||||
CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
|
||||
keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
|
||||
Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
|
||||
them from source. CUDA 13 wheels can use family-specific targets such as
|
||||
`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
|
||||
major-version GPU family.
|
||||
|
||||
If you see this warning or error, install a CUDA 13 wheel when possible.
|
||||
Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
|
||||
your GPU's compute capability. See the CUDA installation guide's
|
||||
[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
|
||||
[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
|
||||
sections for details.
|
||||
|
||||
## Python multiprocessing
|
||||
|
||||
### `RuntimeError` Exception
|
||||
|
||||
Generated
+9
-26
@@ -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",
|
||||
|
||||
+4
-5
@@ -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"] }
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
};
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
|
||||
stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm._C_stable_libtorch",
|
||||
stable_libtorch_module,
|
||||
)
|
||||
from vllm.platforms.cuda import CudaPlatformBase
|
||||
|
||||
return CudaPlatformBase
|
||||
|
||||
|
||||
def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
|
||||
assert cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.1a", DeviceCapability(12, 1)
|
||||
)
|
||||
assert not cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.0a", DeviceCapability(12, 1)
|
||||
)
|
||||
|
||||
assert cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.0f", DeviceCapability(12, 1)
|
||||
)
|
||||
assert not cuda_platform_base._compiled_arch_covers_device(
|
||||
"10.0f", DeviceCapability(12, 1)
|
||||
)
|
||||
|
||||
|
||||
def test_warn_if_device_arch_not_compiled(
|
||||
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
|
||||
) -> None:
|
||||
def device_count(cls: type[Any]) -> int:
|
||||
return 2
|
||||
|
||||
def get_device_capability(
|
||||
cls: type[Any], device_id: int = 0
|
||||
) -> DeviceCapability | None:
|
||||
capabilities = {
|
||||
0: DeviceCapability(12, 1),
|
||||
1: DeviceCapability(10, 3),
|
||||
}
|
||||
return capabilities[device_id]
|
||||
|
||||
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
|
||||
return f"GPU {device_id}"
|
||||
|
||||
warnings: list[tuple[str, str, str]] = []
|
||||
|
||||
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
|
||||
warnings.append((message, compiled_archs, devices))
|
||||
|
||||
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base,
|
||||
"get_device_capability",
|
||||
classmethod(get_device_capability),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base, "get_device_name", classmethod(get_device_name)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.torch.ops",
|
||||
SimpleNamespace(
|
||||
_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.logger",
|
||||
SimpleNamespace(warning_once=warning_once),
|
||||
)
|
||||
|
||||
cuda_platform_base._warn_if_device_arch_not_compiled()
|
||||
|
||||
assert len(warnings) == 1
|
||||
assert warnings[0][1] == "12.0f, 10.0a"
|
||||
assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
|
||||
|
||||
|
||||
def test_warn_if_device_arch_not_compiled_no_warning(
|
||||
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
|
||||
) -> None:
|
||||
def device_count(cls: type[Any]) -> int:
|
||||
return 1
|
||||
|
||||
def get_device_capability(
|
||||
cls: type[Any], device_id: int = 0
|
||||
) -> DeviceCapability | None:
|
||||
return DeviceCapability(10, 3)
|
||||
|
||||
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
|
||||
raise AssertionError("get_device_name should not be called for covered devices")
|
||||
|
||||
warnings: list[tuple[str, str, str]] = []
|
||||
|
||||
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
|
||||
warnings.append((message, compiled_archs, devices))
|
||||
|
||||
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base,
|
||||
"get_device_capability",
|
||||
classmethod(get_device_capability),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base, "get_device_name", classmethod(get_device_name)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.torch.ops",
|
||||
SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.logger",
|
||||
SimpleNamespace(warning_once=warning_once),
|
||||
)
|
||||
|
||||
cuda_platform_base._warn_if_device_arch_not_compiled()
|
||||
|
||||
assert warnings == []
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
|
||||
def test_cutlass_group_gemm_python_guard_allows_thor(monkeypatch):
|
||||
seen_capabilities: list[int] = []
|
||||
|
||||
def fake_cutlass_group_gemm_supported(capability: int) -> bool:
|
||||
seen_capabilities.append(capability)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(
|
||||
torch.ops,
|
||||
"_C",
|
||||
SimpleNamespace(cutlass_group_gemm_supported=fake_cutlass_group_gemm_supported),
|
||||
)
|
||||
|
||||
# CUDA 12 reports Thor as SM101 and CUDA 13 reports it as SM110. Both
|
||||
# should reach the C++ query for the SM10x/SM11x CUTLASS MoE kernel.
|
||||
assert ops.cutlass_group_gemm_supported(101)
|
||||
assert ops.cutlass_group_gemm_supported(110)
|
||||
|
||||
# SM120 uses separate kernels and must not be advertised by this path.
|
||||
assert not ops.cutlass_group_gemm_supported(120)
|
||||
assert seen_capabilities == [101, 110]
|
||||
@@ -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(
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -170,7 +170,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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -83,11 +83,6 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
|
||||
ENFORCE_EAGER=${ENFORCE_EAGER:-1}
|
||||
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
|
||||
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
|
||||
# Pin concurrent prefiller and non-DP decoder engines to separate internal
|
||||
# port windows. DP decoder ranks retain their existing internal port selection.
|
||||
PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000}
|
||||
DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000}
|
||||
INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100}
|
||||
|
||||
# Resolve the repository root from the script location instead of `.git`.
|
||||
# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
|
||||
@@ -159,14 +154,12 @@ run_tests_for_model() {
|
||||
PORT=$((8100 + i))
|
||||
# Calculate side channel port. Avoid clash with with TP workers.
|
||||
SIDE_CHANNEL_PORT=$((5559 + i))
|
||||
INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
|
||||
|
||||
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
|
||||
|
||||
# Build the command with or without model-specific args
|
||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT='HND' \
|
||||
VLLM_PORT=$INTERNAL_PORT \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
vllm serve $model_name \
|
||||
@@ -215,18 +208,12 @@ run_tests_for_model() {
|
||||
PORT=$((8200 + i))
|
||||
# Calculate side channel port
|
||||
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
|
||||
INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
|
||||
DECODER_INTERNAL_PORT_ENV=
|
||||
if [[ -z "${DP_EP:-}" ]]; then
|
||||
DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT"
|
||||
fi
|
||||
|
||||
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
|
||||
|
||||
# Build the command with or without model-specific args
|
||||
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
|
||||
$DECODER_INTERNAL_PORT_ENV \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
vllm serve $model_name \
|
||||
|
||||
@@ -12,9 +12,6 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
|
||||
)
|
||||
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
|
||||
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
|
||||
OffloadingConnectorMetadata,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
|
||||
OffloadingConnectorStats,
|
||||
_ConnectorMetricName,
|
||||
@@ -112,40 +109,6 @@ def test_last_block_offloaded_at_request_finish(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_scheduling", [True, False])
|
||||
def test_abort_queued_request_does_not_build_store_job(
|
||||
request_runner, async_scheduling: bool
|
||||
):
|
||||
"""Aborting a never-scheduled request must not store unallocated KV."""
|
||||
block_size = 4
|
||||
runner = request_runner(
|
||||
block_size=block_size,
|
||||
num_gpu_blocks=8,
|
||||
async_scheduling=async_scheduling,
|
||||
)
|
||||
|
||||
runner.new_request(token_ids=[0] * (block_size * 4))
|
||||
runner.scheduler.schedule()
|
||||
|
||||
runner.new_request(token_ids=[1] * (block_size * 4))
|
||||
queued_req_id = str(runner.req_id)
|
||||
assert any(
|
||||
request.request_id == queued_req_id for request in runner.scheduler.waiting
|
||||
)
|
||||
|
||||
runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
|
||||
req_status = runner.connector_scheduler._req_status[queued_req_id]
|
||||
assert all(group_state.offload_keys for group_state in req_status.group_states)
|
||||
assert all(not group_state.block_ids for group_state in req_status.group_states)
|
||||
|
||||
scheduler_output = runner.scheduler.schedule()
|
||||
|
||||
metadata = scheduler_output.kv_connector_metadata
|
||||
assert isinstance(metadata, OffloadingConnectorMetadata)
|
||||
assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
|
||||
assert queued_req_id not in runner.connector_scheduler._req_status
|
||||
|
||||
|
||||
def test_scheduler_reports_lookup_sync_delay(request_runner):
|
||||
runner = request_runner(
|
||||
block_size=4,
|
||||
|
||||
@@ -125,7 +125,8 @@ def test_register_kv_caches(backend):
|
||||
own dedicated tensors.
|
||||
|
||||
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
|
||||
the raw per-layer kv_caches registered by the connector.
|
||||
kv_caches, which automatically applies
|
||||
_update_hybrid_attention_mamba_layout for hybrid models.
|
||||
|
||||
Verifies that the canonicalized CanonicalKVCaches has the correct
|
||||
block tensors, tensor_idx references, and page sizes across all groups.
|
||||
|
||||
@@ -5,11 +5,13 @@ import torch
|
||||
from torch import Generator
|
||||
|
||||
from tests.utils import large_gpu_mark
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import (
|
||||
apply_top_k_top_p_pytorch,
|
||||
flashinfer_sample,
|
||||
random_sample,
|
||||
)
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
@@ -1043,3 +1045,39 @@ class TestFlashInferDistributionMatch:
|
||||
f"{label}: distribution differs from theoretical: "
|
||||
f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not FLASHINFER_TOPK_TOPP_SUPPORTED,
|
||||
reason="FlashInfer top-k/top-p sampler is not available on this platform.",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
|
||||
@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)])
|
||||
def test_flashinfer_sample_padded_vocab(
|
||||
dtype: torch.dtype, k: int | None, p: float | None
|
||||
):
|
||||
"""flashinfer_sample must accept the logits the sampler actually hands it.
|
||||
|
||||
compute_logits slices the padding off the vocab, so for a vocab that isn't a
|
||||
multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model
|
||||
dtype, while FlashInfer requires contiguous fp32.
|
||||
"""
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
batch_size = 8
|
||||
org_vocab_size = 50272
|
||||
padded_vocab_size = pad_vocab_size(org_vocab_size)
|
||||
assert padded_vocab_size != org_vocab_size
|
||||
|
||||
logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[
|
||||
..., :org_vocab_size
|
||||
]
|
||||
# A single row stays contiguous despite the padded stride, hence batch_size > 1.
|
||||
assert not logits.is_contiguous()
|
||||
|
||||
token_ids = flashinfer_sample(
|
||||
logits,
|
||||
torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None,
|
||||
torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None,
|
||||
)
|
||||
assert token_ids.shape == (batch_size,)
|
||||
assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size))
|
||||
|
||||
+1
-20
@@ -36,22 +36,6 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
|
||||
]
|
||||
|
||||
|
||||
def write_coverage_objects(extensions: list[RustExtension], output: Path) -> None:
|
||||
artifacts = []
|
||||
for extension in extensions:
|
||||
for target in sorted(set(extension.target.values())):
|
||||
target_path = ROOT_DIR.joinpath(*target.split("."))
|
||||
if extension.binding == Binding.Exec:
|
||||
matches = [target_path]
|
||||
else:
|
||||
matches = sorted(target_path.parent.glob(f"{target_path.name}*.so"))
|
||||
if len(matches) != 1 or not matches[0].is_file():
|
||||
raise RuntimeError(f"unable to locate Rust artifact for {target}")
|
||||
artifacts.append(matches[0].relative_to(ROOT_DIR).as_posix())
|
||||
|
||||
output.write_text("\n".join(artifacts) + "\n")
|
||||
|
||||
|
||||
def rust_py_extension_module_names() -> list[str]:
|
||||
module_names = []
|
||||
for extension in rust_extensions():
|
||||
@@ -68,15 +52,12 @@ def rust_py_extension_module_names() -> list[str]:
|
||||
def build_binary(build_rust_args: list[str]) -> None:
|
||||
os.chdir(ROOT_DIR)
|
||||
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
|
||||
extensions = rust_extensions(optional=False)
|
||||
setup(
|
||||
name="vllm-rust-frontend-build",
|
||||
packages=[],
|
||||
rust_extensions=extensions,
|
||||
rust_extensions=rust_extensions(optional=False),
|
||||
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
|
||||
)
|
||||
if output := os.getenv("VLLM_RUST_COVERAGE_OBJECTS"):
|
||||
write_coverage_objects(extensions, Path(output))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
+7
-3
@@ -835,12 +835,16 @@ def cutlass_scaled_mm_azp(
|
||||
|
||||
|
||||
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
|
||||
if cuda_device_capability < 90 or cuda_device_capability >= 110:
|
||||
# CUTLASS grouped FP8 MoE uses the same sm100-named implementation for
|
||||
# SM10x and SM11x. Thor reports SM101 with CUDA 12 and SM110 with CUDA 13,
|
||||
# so keep this Python guard aligned with the C++ support query/dispatch.
|
||||
if cuda_device_capability < 90 or cuda_device_capability >= 120:
|
||||
return False
|
||||
try:
|
||||
return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
|
||||
except AttributeError:
|
||||
# Return False on non-CUDA platforms where it is not available
|
||||
except (AttributeError, RuntimeError):
|
||||
# Return False on non-CUDA platforms where it is not available or not
|
||||
# implemented for the current build.
|
||||
return False
|
||||
|
||||
|
||||
|
||||
@@ -67,6 +67,16 @@ class SchedulerConfig:
|
||||
In real usage, this should be set in `EngineArgs.create_engine_config`.
|
||||
"""
|
||||
|
||||
max_num_partial_prefills: int = Field(default=1, ge=1)
|
||||
"""For chunked prefill, the maximum number of sequences that can be
|
||||
partially prefilled concurrently."""
|
||||
|
||||
max_long_partial_prefills: int = Field(default=1, ge=1)
|
||||
"""For chunked prefill, the maximum number of prompts longer than
|
||||
long_prefill_token_threshold that will be prefilled concurrently. Setting
|
||||
this less than max_num_partial_prefills will allow shorter prompts to jump
|
||||
the queue in front of longer prompts in some cases, improving latency."""
|
||||
|
||||
long_prefill_token_threshold: int = Field(default=0, ge=0)
|
||||
"""For chunked prefill, a request is considered long if the prompt is
|
||||
longer than this number of tokens. 0 disables the cap (default)."""
|
||||
@@ -244,6 +254,19 @@ class SchedulerConfig:
|
||||
self.max_num_batched_tokens,
|
||||
)
|
||||
|
||||
if self.max_num_partial_prefills > 1:
|
||||
if self.long_prefill_token_threshold == 0:
|
||||
self.long_prefill_token_threshold = int(max_model_len * 0.04)
|
||||
|
||||
logger.info(
|
||||
"Concurrent partial prefills enabled with "
|
||||
"max_num_partial_prefills=%d, max_long_partial_prefills=%d, "
|
||||
"long_prefill_token_threshold=%d",
|
||||
self.max_num_partial_prefills,
|
||||
self.max_long_partial_prefills,
|
||||
self.long_prefill_token_threshold,
|
||||
)
|
||||
|
||||
self.verify_max_model_len(max_model_len)
|
||||
|
||||
def verify_max_model_len(self, max_model_len: int) -> Self:
|
||||
@@ -275,11 +298,24 @@ class SchedulerConfig:
|
||||
self.max_num_seqs * max_model_len,
|
||||
)
|
||||
|
||||
if self.long_prefill_token_threshold > max_model_len:
|
||||
if self.max_num_partial_prefills > 1:
|
||||
if not self.enable_chunked_prefill:
|
||||
raise ValueError(
|
||||
"Chunked prefill must be enabled to set "
|
||||
"max_num_partial_prefills > 1."
|
||||
)
|
||||
|
||||
if self.long_prefill_token_threshold > max_model_len:
|
||||
raise ValueError(
|
||||
"long_prefill_token_threshold "
|
||||
f"({self.long_prefill_token_threshold}) cannot be greater "
|
||||
f"than the max_model_len ({max_model_len})."
|
||||
)
|
||||
|
||||
if self.max_long_partial_prefills > self.max_num_partial_prefills:
|
||||
raise ValueError(
|
||||
"long_prefill_token_threshold "
|
||||
f"({self.long_prefill_token_threshold}) cannot be greater "
|
||||
f"than the max_model_len ({max_model_len})."
|
||||
f"{self.max_long_partial_prefills=} must be less than or equal to "
|
||||
f"{self.max_num_partial_prefills=}."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
@@ -21,10 +21,12 @@ from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backend import AttentionBackend
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
|
||||
from vllm.v1.kv_cache_interface import KVCacheSpec
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -592,6 +594,32 @@ class TransferTopology:
|
||||
abs_ratio = -tp_ratio
|
||||
return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
|
||||
|
||||
def get_transfer_cache_regions(
|
||||
self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
|
||||
) -> list[torch.Tensor] | torch.Tensor:
|
||||
"""Return the cache tensor(s) to register as NIXL memory regions,
|
||||
also accounting for hybrid SSM models specificities.
|
||||
"""
|
||||
if isinstance(layer_spec, MambaSpec):
|
||||
# Register the whole kv cache shared tensor, including
|
||||
# SSM/Conv.
|
||||
conv, ssm = cache
|
||||
return [conv]
|
||||
|
||||
# Check may be hacky but it's matching
|
||||
# `_update_hybrid_attention_mamba_layout`.
|
||||
if self.is_mamba and cache.shape[0] == 2:
|
||||
# When MAMBA is present, all backends are blocks first, so
|
||||
# that blocks can be shared between attention layers and mamba
|
||||
# layers. Runner already adjusted strides for FlashAttn-like
|
||||
# backends so its num_blocks first.
|
||||
# Swap [2<>num_blocks] dims for hybrid SSM layout.
|
||||
cache = cache.transpose(0, 1)
|
||||
|
||||
# K and V are packed into one tensor (content dim), so each layer
|
||||
# registers as a single region.
|
||||
return [cache]
|
||||
|
||||
def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
|
||||
"""One-line summary of transfer config for logging."""
|
||||
info = self._engines[(remote_engine_id, remote_pp_rank)]
|
||||
|
||||
@@ -1678,9 +1678,9 @@ class MooncakeConnectorWorker:
|
||||
conv, _ = cache_or_caches
|
||||
cache_list = [conv]
|
||||
else:
|
||||
# K and V are packed into one blocks-first tensor per layer,
|
||||
# so each layer registers as a single region.
|
||||
cache_list = [cache_or_caches]
|
||||
cache_list = self.transfer_topo.get_transfer_cache_regions(
|
||||
cache_or_caches, layer_spec
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"registering layer %s with %d cache tensor(s)",
|
||||
|
||||
@@ -1092,7 +1092,7 @@ class NixlBaseConnectorWorker:
|
||||
# to better exploit the memory layout (ie num_blocks is the first dim).
|
||||
tensor_size_bytes = None
|
||||
|
||||
for layer_name, cache in xfer_buffers.items():
|
||||
for layer_name, cache_or_caches in xfer_buffers.items():
|
||||
# NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
|
||||
# that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
|
||||
# However, physical page_size may differ when kernel requires a specific
|
||||
@@ -1109,6 +1109,9 @@ class NixlBaseConnectorWorker:
|
||||
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
|
||||
# MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
|
||||
layer_spec = layer_spec.kv_cache_specs[layer_name]
|
||||
cache_list = self.transfer_topo.get_transfer_cache_regions(
|
||||
cache_or_caches, layer_spec
|
||||
)
|
||||
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
|
||||
# the page_size assuming constant `self._logical_num_blocks`.
|
||||
physical_page_size = (
|
||||
@@ -1117,6 +1120,8 @@ class NixlBaseConnectorWorker:
|
||||
else layer_spec.page_size_bytes
|
||||
// self._physical_blocks_per_logical_kv_block
|
||||
)
|
||||
# For when registering multiple tensors eg K/V in separate regions.
|
||||
physical_page_size = physical_page_size // len(cache_list)
|
||||
if self.transfer_topo._cross_layers_blocks:
|
||||
# When cross-layers blocks are used, multiply by number of layers
|
||||
physical_page_size = physical_page_size * len(
|
||||
@@ -1131,61 +1136,66 @@ class NixlBaseConnectorWorker:
|
||||
# [`num_blocks` * `page_size`]
|
||||
curr_tensor_size_bytes = num_blocks * physical_page_size
|
||||
|
||||
base_addr = cache.data_ptr()
|
||||
if base_addr in seen_base_addresses:
|
||||
# NOTE (NickLucche) HMA employs memory pooling to share tensors
|
||||
# across groups. This results in skipping all tensors but the ones
|
||||
# pointed to by group0. Also, generally we will have more blocks
|
||||
# per tensor but fewer regions.
|
||||
logger.debug("Skipping %s because it's already seen", layer_name)
|
||||
continue
|
||||
logger.debug(
|
||||
"Registering layer %s with cache shape: %s", layer_name, cache.shape
|
||||
)
|
||||
seen_base_addresses.append(base_addr)
|
||||
# Only record non-Mamba page sizes.
|
||||
if isinstance(layer_spec, MambaSpec):
|
||||
self.block_len_per_layer.append(
|
||||
physical_page_size // self._physical_blocks_per_logical_kv_block
|
||||
# TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
|
||||
# registering a single tensor for both K/V and splitting logically like FI.
|
||||
for cache in cache_list:
|
||||
base_addr = cache.data_ptr()
|
||||
if base_addr in seen_base_addresses:
|
||||
# NOTE (NickLucche) HMA employs memory pooling to share tensors
|
||||
# across groups. This results in skipping all tensors but the ones
|
||||
# pointed to by group0. Also, generally we will have more blocks
|
||||
# per tensor but fewer regions.
|
||||
logger.debug("Skipping %s because it's already seen", layer_name)
|
||||
continue
|
||||
logger.debug(
|
||||
"Registering layer %s with cache shape: %s", layer_name, cache.shape
|
||||
)
|
||||
else:
|
||||
self.block_len_per_layer.append(physical_page_size)
|
||||
is_mla_region = isinstance(
|
||||
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
|
||||
)
|
||||
self._region_is_mla.append(is_mla_region)
|
||||
|
||||
if not is_mla_region:
|
||||
if tensor_size_bytes is None:
|
||||
tensor_size_bytes = curr_tensor_size_bytes
|
||||
assert tensor_size_bytes == curr_tensor_size_bytes, (
|
||||
"All non-MLA kv cache tensors must have the same size"
|
||||
seen_base_addresses.append(base_addr)
|
||||
# Only record non-Mamba page sizes.
|
||||
if isinstance(layer_spec, MambaSpec):
|
||||
self.block_len_per_layer.append(
|
||||
physical_page_size // self._physical_blocks_per_logical_kv_block
|
||||
)
|
||||
else:
|
||||
self.block_len_per_layer.append(physical_page_size)
|
||||
is_mla_region = isinstance(
|
||||
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
|
||||
)
|
||||
self._region_is_mla.append(is_mla_region)
|
||||
|
||||
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
|
||||
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
|
||||
if (
|
||||
self._physical_blocks_per_logical_kv_block == 1
|
||||
and cache.shape[0] != num_blocks
|
||||
):
|
||||
raise AssertionError(
|
||||
"All kv cache tensors must have the same number of "
|
||||
f"blocks; layer={layer_name}, "
|
||||
f"expected_num_blocks={num_blocks}, "
|
||||
f"cache_shape={tuple(cache.shape)}, "
|
||||
f"cache_stride={tuple(cache.stride())}, "
|
||||
f"layer_spec={type(layer_spec).__name__}, "
|
||||
f"backend={self.backend_name}, "
|
||||
"all_backends="
|
||||
f"{[backend.get_name() for backend in self.attn_backends]}, "
|
||||
f"kv_cache_layout={self.kv_cache_layout}"
|
||||
if not is_mla_region:
|
||||
if tensor_size_bytes is None:
|
||||
tensor_size_bytes = curr_tensor_size_bytes
|
||||
assert tensor_size_bytes == curr_tensor_size_bytes, (
|
||||
"All non-MLA kv cache tensors must have the same size"
|
||||
)
|
||||
|
||||
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
|
||||
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
|
||||
if (
|
||||
self._physical_blocks_per_logical_kv_block == 1
|
||||
and cache.shape[0] != num_blocks
|
||||
):
|
||||
raise AssertionError(
|
||||
"All kv cache tensors must have the same number of "
|
||||
f"blocks; layer={layer_name}, "
|
||||
f"expected_num_blocks={num_blocks}, "
|
||||
f"cache_shape={tuple(cache.shape)}, "
|
||||
f"cache_stride={tuple(cache.stride())}, "
|
||||
f"layer_spec={type(layer_spec).__name__}, "
|
||||
f"backend={self.backend_name}, "
|
||||
"all_backends="
|
||||
f"{[backend.get_name() for backend in self.attn_backends]}, "
|
||||
f"kv_cache_layout={self.kv_cache_layout}"
|
||||
)
|
||||
|
||||
# Need to make sure the device ID is non-negative for NIXL,
|
||||
# Torch uses -1 to indicate CPU tensors.
|
||||
self.device_id = max(cache.get_device(), 0)
|
||||
caches_data.append(
|
||||
(base_addr, curr_tensor_size_bytes, self.device_id, "")
|
||||
)
|
||||
|
||||
# Need to make sure the device ID is non-negative for NIXL,
|
||||
# Torch uses -1 to indicate CPU tensors.
|
||||
self.device_id = max(cache.get_device(), 0)
|
||||
caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))
|
||||
|
||||
logger.debug(
|
||||
"Different block lengths collected: %s", set(self.block_len_per_layer)
|
||||
)
|
||||
|
||||
@@ -326,12 +326,9 @@ class RequestOffloadState:
|
||||
group_state.block_ids.extend(new_blocks)
|
||||
|
||||
def storable_chunks(
|
||||
self,
|
||||
group_config: "GroupOffloadConfig",
|
||||
group_state: RequestGroupState,
|
||||
num_offloadable_tokens: int,
|
||||
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
|
||||
) -> int:
|
||||
"""Number of allocated leading offloaded chunks eligible for store.
|
||||
"""Number of leading offloaded chunks eligible for store.
|
||||
|
||||
For eagle/MTP groups the volatile trailing chunk of the offloadable
|
||||
range is excluded while decoding: the draft-layer KV of the last
|
||||
@@ -348,10 +345,7 @@ class RequestOffloadState:
|
||||
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
|
||||
if group_config.is_eagle_group and is_decoding:
|
||||
num_chunks = max(0, num_chunks - 1)
|
||||
num_allocated_chunks = (
|
||||
len(group_state.block_ids) // self.config.blocks_per_chunk
|
||||
)
|
||||
return min(num_chunks, num_allocated_chunks)
|
||||
return num_chunks
|
||||
|
||||
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
|
||||
# max(): at the prefill->decode transition of a chunk-aligned prompt,
|
||||
@@ -362,7 +356,7 @@ class RequestOffloadState:
|
||||
):
|
||||
group_state.next_stored_chunk_idx = max(
|
||||
group_state.next_stored_chunk_idx,
|
||||
self.storable_chunks(group_config, group_state, num_offloadable_tokens),
|
||||
self.storable_chunks(group_config, num_offloadable_tokens),
|
||||
)
|
||||
|
||||
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
|
||||
@@ -997,7 +991,7 @@ class OffloadingConnectorScheduler:
|
||||
self.config.kv_group_configs, req_status.group_states
|
||||
):
|
||||
num_chunks = req_status.storable_chunks(
|
||||
group_config, group_state, num_offloadable_tokens
|
||||
group_config, num_offloadable_tokens
|
||||
)
|
||||
|
||||
start_chunk_idx = group_state.next_stored_chunk_idx
|
||||
@@ -1074,7 +1068,7 @@ class OffloadingConnectorScheduler:
|
||||
group_config.sliding_window_size_in_chunks is not None
|
||||
)
|
||||
num_chunks = req_status.storable_chunks(
|
||||
group_config, group_state, num_offloadable_tokens
|
||||
group_config, num_offloadable_tokens
|
||||
)
|
||||
start_chunk_idx = group_state.next_stored_chunk_idx
|
||||
block_ids = group_state.block_ids
|
||||
|
||||
@@ -56,7 +56,9 @@ class OffloadingConnectorWorker:
|
||||
def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
|
||||
self.worker = self.spec.get_worker(kv_caches)
|
||||
|
||||
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
|
||||
def register_kv_caches(
|
||||
self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
|
||||
):
|
||||
kv_cache_config = self.kv_cache_config
|
||||
num_blocks = kv_cache_config.num_blocks
|
||||
|
||||
@@ -118,13 +120,24 @@ class OffloadingConnectorWorker:
|
||||
)
|
||||
|
||||
elif isinstance(layer_kv_cache_spec, MambaSpec):
|
||||
layer_kv_cache = kv_caches[layer_name]
|
||||
assert layer_kv_cache.dtype == torch.int8
|
||||
tensors_per_block[layer_name] = (
|
||||
layer_kv_cache.view(
|
||||
num_blocks, layer_kv_cache_spec.page_size_bytes
|
||||
),
|
||||
state_tensors = kv_caches[layer_name]
|
||||
assert isinstance(state_tensors, list)
|
||||
|
||||
# re-construct the raw (num_blocks, page_size) tensor
|
||||
# from the first state tensor
|
||||
assert len(state_tensors) > 0
|
||||
first_state_tensor = state_tensors[0]
|
||||
assert first_state_tensor.storage_offset() == 0
|
||||
tensor = (
|
||||
torch.tensor(
|
||||
[],
|
||||
dtype=torch.int8,
|
||||
device=first_state_tensor.device,
|
||||
)
|
||||
.set_(first_state_tensor.untyped_storage())
|
||||
.view((num_blocks, layer_kv_cache_spec.page_size_bytes))
|
||||
)
|
||||
tensors_per_block[layer_name] = (tensor,)
|
||||
|
||||
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
|
||||
unpadded_page_size_bytes[layer_name] = replace(
|
||||
|
||||
@@ -527,6 +527,8 @@ class EngineArgs:
|
||||
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
|
||||
max_num_batched_tokens: int | None = None
|
||||
max_num_scheduled_tokens: int | None = None
|
||||
max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
|
||||
max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills
|
||||
long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
|
||||
max_num_seqs: int | None = None
|
||||
max_logprobs: int = ModelConfig.max_logprobs
|
||||
@@ -1438,6 +1440,13 @@ class EngineArgs:
|
||||
"default": None,
|
||||
},
|
||||
)
|
||||
scheduler_group.add_argument(
|
||||
"--max-num-partial-prefills", **scheduler_kwargs["max_num_partial_prefills"]
|
||||
)
|
||||
scheduler_group.add_argument(
|
||||
"--max-long-partial-prefills",
|
||||
**scheduler_kwargs["max_long_partial_prefills"],
|
||||
)
|
||||
scheduler_group.add_argument(
|
||||
"--long-prefill-token-threshold",
|
||||
**scheduler_kwargs["long_prefill_token_threshold"],
|
||||
@@ -2181,6 +2190,8 @@ class EngineArgs:
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
policy=self.scheduling_policy,
|
||||
scheduler_cls=self.scheduler_cls,
|
||||
max_num_partial_prefills=self.max_num_partial_prefills,
|
||||
max_long_partial_prefills=self.max_long_partial_prefills,
|
||||
long_prefill_token_threshold=self.long_prefill_token_threshold,
|
||||
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
|
||||
watermark=self.watermark,
|
||||
@@ -2391,6 +2402,14 @@ class EngineArgs:
|
||||
|
||||
def _check_feature_supported(self):
|
||||
"""Raise an error if the feature is not supported."""
|
||||
# No Concurrent Partial Prefills so far.
|
||||
if (
|
||||
self.max_num_partial_prefills != SchedulerConfig.max_num_partial_prefills
|
||||
or self.max_long_partial_prefills
|
||||
!= SchedulerConfig.max_long_partial_prefills
|
||||
):
|
||||
_raise_unsupported_error(feature_name="Concurrent Partial Prefill")
|
||||
|
||||
if self.pipeline_parallel_size > 1:
|
||||
supports_pp = getattr(
|
||||
self.distributed_executor_backend, "supports_pp", False
|
||||
|
||||
@@ -687,9 +687,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
||||
skip_special_tokens=self.skip_special_tokens,
|
||||
spaces_between_special_tokens=self.spaces_between_special_tokens,
|
||||
include_stop_str_in_output=self.include_stop_str_in_output,
|
||||
output_kind=(
|
||||
RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY
|
||||
),
|
||||
output_kind=RequestOutputKind.DELTA
|
||||
if self.stream
|
||||
else RequestOutputKind.FINAL_ONLY,
|
||||
structured_outputs=self.extract_structured_outputs(),
|
||||
logit_bias=self.logit_bias,
|
||||
bad_words=self.bad_words,
|
||||
@@ -849,10 +849,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
||||
|
||||
# Reject empty tools array, matching OpenAI API behavior
|
||||
if data.get("tools") == []:
|
||||
raise VLLMValidationError(
|
||||
raise ValueError(
|
||||
"`tools` must not be an empty array. "
|
||||
"Either provide at least one tool or omit the field entirely.",
|
||||
parameter="tools",
|
||||
"Either provide at least one tool or omit the field entirely."
|
||||
)
|
||||
|
||||
# if "tool_choice" is not specified but tools are provided,
|
||||
@@ -1076,15 +1075,13 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
|
||||
if isinstance(data, BatchChatCompletionRequest):
|
||||
data = data.model_dump(exclude_unset=True)
|
||||
if data.get("use_beam_search"):
|
||||
raise VLLMValidationError(
|
||||
raise ValueError(
|
||||
"Batch chat completions do not support beam search. "
|
||||
"Please set `use_beam_search` to False.",
|
||||
parameter="use_beam_search",
|
||||
"Please set `use_beam_search` to False."
|
||||
)
|
||||
if data.get("logprob_token_ids") and not data.get("logprobs"):
|
||||
raise VLLMValidationError(
|
||||
"when using `logprob_token_ids`, `logprobs` must be set to true.",
|
||||
parameter="logprob_token_ids",
|
||||
raise ValueError(
|
||||
"when using `logprob_token_ids`, `logprobs` must be set to true."
|
||||
)
|
||||
response_format = data.get("response_format")
|
||||
rf_type = (
|
||||
@@ -1098,10 +1095,8 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
|
||||
validate_structured_outputs_structural_tag(structured_outputs)
|
||||
n = data.get("n", 1)
|
||||
if n is not None and n != 1:
|
||||
raise VLLMValidationError(
|
||||
"Batch chat completions do not support `n > 1`. Please set `n` to 1.",
|
||||
parameter="n",
|
||||
value=n,
|
||||
raise ValueError(
|
||||
"Batch chat completions do not support `n > 1`. Please set `n` to 1."
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Literal, TypeAlias
|
||||
|
||||
@@ -79,6 +78,7 @@ class TranslationRequest(OpenAIBaseModel):
|
||||
`verbose_json`, or `vtt`.
|
||||
"""
|
||||
|
||||
# TODO support additional sampling parameters
|
||||
# --8<-- [start:translation-sampling-params]
|
||||
use_beam_search: bool = False
|
||||
"""Whether or not beam search should be used."""
|
||||
@@ -103,28 +103,6 @@ class TranslationRequest(OpenAIBaseModel):
|
||||
will use [log probability](https://en.wikipedia.org/wiki/Log_probability)
|
||||
to automatically increase the temperature until certain thresholds are hit.
|
||||
"""
|
||||
|
||||
top_p: float | None = None
|
||||
"""Enables nucleus (top-p) sampling, where tokens are selected from the
|
||||
smallest possible set whose cumulative probability exceeds `p`.
|
||||
"""
|
||||
|
||||
top_k: int | None = None
|
||||
"""Limits sampling to the `k` most probable tokens at each step."""
|
||||
|
||||
min_p: float | None = None
|
||||
"""Filters out tokens with a probability lower than `min_p`, ensuring a
|
||||
minimum likelihood threshold during sampling.
|
||||
"""
|
||||
|
||||
frequency_penalty: float | None = 0.0
|
||||
"""The frequency penalty to use for sampling."""
|
||||
|
||||
repetition_penalty: float | None = None
|
||||
"""The repetition penalty to use for sampling."""
|
||||
|
||||
presence_penalty: float | None = 0.0
|
||||
"""The presence penalty to use for sampling."""
|
||||
# --8<-- [end:translation-sampling-params]
|
||||
|
||||
# --8<-- [start:translation-extra-params]
|
||||
@@ -161,23 +139,11 @@ class TranslationRequest(OpenAIBaseModel):
|
||||
|
||||
max_completion_tokens: int | None = None
|
||||
"""The maximum number of tokens to generate."""
|
||||
|
||||
vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Additional request parameters with (list of) string or "
|
||||
"numeric values, used by custom extensions."
|
||||
),
|
||||
)
|
||||
# --8<-- [end:translation-extra-params]
|
||||
|
||||
# Default sampling parameters for translation requests.
|
||||
_DEFAULT_SAMPLING_PARAMS: dict = {
|
||||
"repetition_penalty": 1.0,
|
||||
"temperature": 0,
|
||||
"top_p": 1.0,
|
||||
"top_k": 0,
|
||||
"min_p": 0.0,
|
||||
}
|
||||
|
||||
def build_stt_params(
|
||||
@@ -233,38 +199,14 @@ class TranslationRequest(OpenAIBaseModel):
|
||||
temperature = default_sampling_params.get(
|
||||
"temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
|
||||
)
|
||||
if (top_p := self.top_p) is None:
|
||||
top_p = default_sampling_params.get(
|
||||
"top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]
|
||||
)
|
||||
if (top_k := self.top_k) is None:
|
||||
top_k = default_sampling_params.get(
|
||||
"top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"]
|
||||
)
|
||||
if (min_p := self.min_p) is None:
|
||||
min_p = default_sampling_params.get(
|
||||
"min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"]
|
||||
)
|
||||
if (repetition_penalty := self.repetition_penalty) is None:
|
||||
repetition_penalty = default_sampling_params.get(
|
||||
"repetition_penalty",
|
||||
self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"],
|
||||
)
|
||||
|
||||
return SamplingParams.from_optional(
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
seed=self.seed,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
min_p=min_p,
|
||||
frequency_penalty=self.frequency_penalty,
|
||||
repetition_penalty=repetition_penalty,
|
||||
presence_penalty=self.presence_penalty,
|
||||
output_kind=RequestOutputKind.DELTA
|
||||
if self.stream
|
||||
else RequestOutputKind.FINAL_ONLY,
|
||||
extra_args=self.vllm_xargs,
|
||||
skip_clone=True, # Created fresh per request, safe to skip clone
|
||||
)
|
||||
|
||||
@@ -284,16 +226,6 @@ class TranslationRequest(OpenAIBaseModel):
|
||||
parameter=invalid_param,
|
||||
)
|
||||
|
||||
xargs = data.get("vllm_xargs")
|
||||
if isinstance(xargs, str):
|
||||
try:
|
||||
data["vllm_xargs"] = json.loads(xargs)
|
||||
except json.JSONDecodeError as e:
|
||||
raise VLLMValidationError(
|
||||
f"Failed to parse vllm_xargs. Must be valid JSON: {e}",
|
||||
parameter="vllm_xargs",
|
||||
) from e
|
||||
|
||||
return data
|
||||
|
||||
|
||||
|
||||
@@ -4,8 +4,6 @@
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
|
||||
from vllm.v1.kv_cache_interface import KVCacheSpec
|
||||
@@ -23,14 +21,6 @@ class AttentionLayerBase(ABC):
|
||||
impl: "AttentionImpl"
|
||||
supports_dcp: bool = True
|
||||
|
||||
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
|
||||
"""Bind the allocated KV cache tensor to this layer.
|
||||
|
||||
The default stores the cache view as-is; subclasses (e.g. Mamba)
|
||||
override this to unpack the raw buffer into per-state views.
|
||||
"""
|
||||
self.kv_cache = kv_cache
|
||||
|
||||
@abstractmethod
|
||||
def get_attn_backend(self) -> type[AttentionBackend]:
|
||||
"""Get the attention backend class for this layer."""
|
||||
|
||||
@@ -2,13 +2,11 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Iterable
|
||||
from math import prod
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.utils.torch_utils import get_dtype_size
|
||||
from vllm.v1.attention.backend import AttentionBackend
|
||||
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
||||
from vllm.v1.attention.selector import get_mamba_attn_backend
|
||||
@@ -26,22 +24,6 @@ class MambaBase(AttentionLayerBase):
|
||||
kv_cache: tuple[torch.Tensor, ...]
|
||||
supports_dcp: bool = False
|
||||
|
||||
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
|
||||
"""Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.
|
||||
|
||||
Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
|
||||
packed contiguously; slice them out and reinterpret per dtype/shape.
|
||||
"""
|
||||
pages = kv_cache.squeeze(dim=(1, 2))
|
||||
states: list[torch.Tensor] = []
|
||||
offset = 0
|
||||
for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):
|
||||
nbytes = prod(shape) * get_dtype_size(dtype)
|
||||
state = pages[:, offset : offset + nbytes].view(dtype)
|
||||
states.append(state.view(-1, *shape))
|
||||
offset += nbytes
|
||||
self.kv_cache = tuple(states)
|
||||
|
||||
@abstractmethod
|
||||
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
|
||||
"""
|
||||
|
||||
@@ -23,7 +23,6 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
|
||||
causal_conv1d_fn,
|
||||
causal_conv1d_update,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
from vllm.v1.attention.backend import AttentionMetadata
|
||||
@@ -43,7 +42,6 @@ class ShortConv(MambaBase, CustomOp):
|
||||
layer_idx: int,
|
||||
model_config: ModelConfig | None = None,
|
||||
cache_config: CacheConfig | None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
@@ -69,14 +67,12 @@ class ShortConv(MambaBase, CustomOp):
|
||||
input_size=dim,
|
||||
output_sizes=[dim] * 3,
|
||||
bias=self.bias,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.in_proj",
|
||||
)
|
||||
self.out_proj = RowParallelLinear(
|
||||
input_size=dim,
|
||||
output_size=dim,
|
||||
bias=self.bias,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.out_proj",
|
||||
)
|
||||
|
||||
|
||||
@@ -276,13 +276,7 @@ class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
|
||||
|
||||
|
||||
class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
|
||||
@staticmethod
|
||||
def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool:
|
||||
# Cosmos renders each video as a full
|
||||
# <|vision_start|><|video_pad|><|vision_end|> placeholder, and its
|
||||
# reference processor replaces that entire triplet with timestamped,
|
||||
# per-frame vision sequences.
|
||||
return False
|
||||
pass
|
||||
|
||||
|
||||
class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
|
||||
@@ -523,7 +517,6 @@ class Cosmos3EdgeForConditionalGeneration(
|
||||
},
|
||||
orig_to_new_substr={
|
||||
"_moe_gen": None,
|
||||
"k_norm_und_for_gen": None,
|
||||
".add_q_proj.": None,
|
||||
".add_k_proj.": None,
|
||||
".add_v_proj.": None,
|
||||
|
||||
@@ -257,7 +257,6 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
|
||||
layer_idx=layer_idx,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.conv",
|
||||
)
|
||||
|
||||
|
||||
@@ -351,7 +351,6 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):
|
||||
layer_idx=layer_idx,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.conv",
|
||||
)
|
||||
|
||||
|
||||
@@ -402,11 +402,12 @@ class Ovis2_5MultiModalProcessor(BaseMultiModalProcessor[Ovis2_5ProcessingInfo])
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
out_mm_kwargs: MultiModalKwargsItems,
|
||||
) -> list[PromptReplacement]:
|
||||
hf_processor = self.info.get_hf_processor()
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
vocab = tokenizer.get_vocab()
|
||||
|
||||
placeholder = {
|
||||
"image": hf_processor.get_token_value("image_token"),
|
||||
"video": hf_processor.get_token_value("video_token"),
|
||||
"image": vocab[IMAGE_TOKEN],
|
||||
"video": vocab[VIDEO_TOKEN],
|
||||
}
|
||||
|
||||
def get_replacement_ovis(item_idx, modality: str):
|
||||
|
||||
@@ -4,7 +4,6 @@ from collections.abc import Mapping
|
||||
|
||||
from vllm.config import ModelConfig, VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.multimodal.inputs import MultiModalKwargsItem
|
||||
from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||
from vllm.multimodal.registry import MultiModalRegistry
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
@@ -49,8 +48,6 @@ class MultiModalBudget:
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
mm_registry: MultiModalRegistry,
|
||||
*,
|
||||
enable_cache: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
@@ -61,11 +58,7 @@ class MultiModalBudget:
|
||||
self.max_num_reqs = scheduler_config.max_num_seqs
|
||||
|
||||
with set_default_torch_num_threads(): # Avoid hang during startup
|
||||
cache = (
|
||||
mm_registry.processor_only_cache_from_config(vllm_config)
|
||||
if enable_cache
|
||||
else None
|
||||
)
|
||||
cache = mm_registry.processor_only_cache_from_config(vllm_config)
|
||||
processor = mm_registry.create_processor(model_config, cache=cache)
|
||||
|
||||
self.cache = cache
|
||||
@@ -198,23 +191,3 @@ class MultiModalBudget:
|
||||
def reset_cache(self) -> None:
|
||||
if self.cache is not None:
|
||||
self.cache.clear_cache()
|
||||
|
||||
|
||||
def get_dummy_encoder_profile_inputs(
|
||||
mm_registry: MultiModalRegistry,
|
||||
budget: MultiModalBudget,
|
||||
) -> list[tuple[str, MultiModalKwargsItem]]:
|
||||
if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
|
||||
return []
|
||||
|
||||
modality = budget.get_modality_with_max_tokens()
|
||||
max_items_per_batch = budget.mm_max_items_per_batch[modality]
|
||||
dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
|
||||
budget.model_config,
|
||||
mm_counts={modality: 1},
|
||||
processor=budget.processor,
|
||||
)
|
||||
dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
|
||||
assert dummy_mm_item is not None, "Dummy item should be generated"
|
||||
|
||||
return [(modality, dummy_mm_item)] * max_items_per_batch
|
||||
|
||||
@@ -1177,7 +1177,7 @@ class PyNvVideoCodecVideoBackend(VideoBackend):
|
||||
|
||||
@VIDEO_LOADER_REGISTRY.register(
|
||||
"qwen3_vl",
|
||||
video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
|
||||
video_processor="Qwen3VLVideoProcessor",
|
||||
)
|
||||
class Qwen3VLVideoBackend(VideoBackend):
|
||||
@classmethod
|
||||
|
||||
@@ -14,6 +14,7 @@ from datetime import timedelta
|
||||
from functools import cache, lru_cache, wraps
|
||||
from typing import TYPE_CHECKING, NamedTuple, TypeVar
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
from torch.distributed import PrefixStore, ProcessGroup
|
||||
from torch.distributed.distributed_c10d import is_nccl_available
|
||||
@@ -48,6 +49,8 @@ _R = TypeVar("_R")
|
||||
|
||||
pynvml = import_pynvml()
|
||||
|
||||
_COMPILED_CUDA_ARCH_RE = re.compile(r"^(?P<major>\d+)\.(?P<minor>\d+)(?P<kind>[af])?$")
|
||||
|
||||
# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
|
||||
# see https://github.com/huggingface/diffusers/issues/9704 for details
|
||||
torch.backends.cuda.enable_cudnn_sdp(False)
|
||||
@@ -282,6 +285,62 @@ class CudaPlatformBase(Platform):
|
||||
def log_warnings(cls):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _compiled_arch_covers_device(arch: str, capability: DeviceCapability) -> bool:
|
||||
match = _COMPILED_CUDA_ARCH_RE.match(arch)
|
||||
if match is None:
|
||||
return False
|
||||
|
||||
major = int(match.group("major"))
|
||||
minor = int(match.group("minor"))
|
||||
kind = match.group("kind")
|
||||
|
||||
if kind == "f":
|
||||
return major == capability.major
|
||||
return major == capability.major and minor == capability.minor
|
||||
|
||||
@classmethod
|
||||
def _warn_if_device_arch_not_compiled(cls) -> None:
|
||||
try:
|
||||
compiled_archs_str = torch.ops._C.get_compiled_cuda_archs()
|
||||
except (AttributeError, RuntimeError):
|
||||
return
|
||||
|
||||
compiled_archs = [
|
||||
arch.strip() for arch in compiled_archs_str.split(",") if arch.strip()
|
||||
]
|
||||
if not compiled_archs:
|
||||
return
|
||||
|
||||
unsupported_devices: list[str] = []
|
||||
for device_id in range(cls.device_count()):
|
||||
capability = cls.get_device_capability(device_id)
|
||||
if capability is None:
|
||||
continue
|
||||
if any(
|
||||
cls._compiled_arch_covers_device(arch, capability)
|
||||
for arch in compiled_archs
|
||||
):
|
||||
continue
|
||||
device_name = cls.get_device_name(device_id)
|
||||
unsupported_devices.append(
|
||||
f"{device_id}: {device_name} (compute capability "
|
||||
f"{capability.as_version_str()})"
|
||||
)
|
||||
|
||||
if unsupported_devices:
|
||||
logger.warning_once(
|
||||
"The current vLLM wheel was built for CUDA architectures %s, "
|
||||
"but the following visible CUDA device(s) are not covered: %s. "
|
||||
"This may cause CUDA kernel launch failures or force slower "
|
||||
"fallback paths. Note that CUDA 13 family-specific targets "
|
||||
"such as 10.0f and 12.0f cover the corresponding major-version "
|
||||
"GPU family, while architecture-specific targets such as 12.1a "
|
||||
"cover only that exact compute capability.",
|
||||
", ".join(compiled_archs),
|
||||
"; ".join(unsupported_devices),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def is_pin_memory_available(cls) -> bool:
|
||||
if in_wsl():
|
||||
@@ -949,6 +1008,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
|
||||
@classmethod
|
||||
@with_nvml_context
|
||||
def log_warnings(cls):
|
||||
cls._warn_if_device_arch_not_compiled()
|
||||
device_ids: int = pynvml.nvmlDeviceGetCount()
|
||||
if device_ids > 1:
|
||||
device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
|
||||
|
||||
@@ -35,7 +35,6 @@ except json.JSONDecodeError:
|
||||
# ---------------------------------------------------------------------------
|
||||
DEFAULT_ENV_VAR_PREFIXES: set[str] = {
|
||||
"VLLM_",
|
||||
"FLASH_ATTENTION_",
|
||||
"LMCACHE_",
|
||||
"NCCL_",
|
||||
"UCX_",
|
||||
|
||||
@@ -78,6 +78,7 @@ class Ovis2_5Processor(ProcessorMixin):
|
||||
|
||||
@cached_property
|
||||
def extra_special_tokens(self):
|
||||
vocab = self.tokenizer.get_vocab()
|
||||
required_tokens = {
|
||||
"image_token": "<image>",
|
||||
"video_token": "<video>",
|
||||
@@ -89,16 +90,21 @@ class Ovis2_5Processor(ProcessorMixin):
|
||||
"image_pad": "<|image_pad|>",
|
||||
}
|
||||
|
||||
# The checkpoint defines both `additional_special_tokens` and
|
||||
# `extra_special_tokens`, with the latter empty. Transformers ignores
|
||||
# the former because the latter is explicitly empty, so the tokens are
|
||||
# missing from the vocab. Re-add them to restore the expected ids.
|
||||
self.tokenizer.add_tokens(list(required_tokens.values()), special_tokens=True)
|
||||
extra_special_tokens = {}
|
||||
suggestion = (
|
||||
"please add '<image>', '<video>', '<ovis_visual_atom>', "
|
||||
"'<ovis_image_start>', '<ovis_image_end>', '<ovis_video_start>', "
|
||||
"'<ovis_video_end>' in 'additional_special_tokens' of "
|
||||
"tokenizer_config.json, You can refer to "
|
||||
"https://huggingface.co/AIDC-AI/Ovis2.6-30B-A3B/blob/main/tokenizer_config.json"
|
||||
)
|
||||
|
||||
return {
|
||||
key: self.tokenizer.convert_tokens_to_ids(token_name)
|
||||
for key, token_name in required_tokens.items()
|
||||
}
|
||||
for key, token_name in required_tokens.items():
|
||||
if token_name not in vocab:
|
||||
raise ValueError(f"Can not find {token_name}, {suggestion}")
|
||||
extra_special_tokens[key] = vocab[token_name]
|
||||
|
||||
return extra_special_tokens
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
|
||||
@@ -392,8 +392,9 @@ def apply_top_k_top_p_pytorch(
|
||||
logits_sort.masked_fill_(top_k_mask, -float("inf"))
|
||||
|
||||
if p is not None:
|
||||
# Apply top-p.
|
||||
probs_sort = logits_sort.softmax(dim=-1)
|
||||
# Apply top-p. The cumsum below runs over the whole vocab, so accumulating
|
||||
# in a low-precision dtype makes the nucleus undershoot p.
|
||||
probs_sort = logits_sort.softmax(dim=-1, dtype=torch.float32)
|
||||
probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
|
||||
top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
|
||||
# at least one
|
||||
@@ -500,9 +501,10 @@ def flashinfer_sample(
|
||||
probs, k, deterministic=True
|
||||
)
|
||||
else:
|
||||
# Both top-k and top-p.
|
||||
# Both top-k and top-p. FlashInfer requires contiguous fp32 logits; the
|
||||
# branches above get that from softmax().
|
||||
next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
|
||||
logits, k, p, deterministic=True
|
||||
logits.float().contiguous(), k, p, deterministic=True
|
||||
)
|
||||
|
||||
return next_token_ids.view(-1)
|
||||
|
||||
@@ -131,7 +131,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs < VOCAB_SIZE
|
||||
logits_blk0 = tl.load(
|
||||
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
# Exclude -inf values (e.g. from grammar bitmasks) from
|
||||
# statistics to avoid NaN in pivot computation.
|
||||
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
||||
@@ -164,7 +164,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs_n < VOCAB_SIZE
|
||||
logits_blk = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
|
||||
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
||||
# Exclude -inf from min to keep binary search bounds
|
||||
@@ -305,7 +305,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs_n < VOCAB_SIZE
|
||||
logits_blk2 = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
|
||||
above_0 = logits_blk2 > k_pivot_0
|
||||
above_1 = logits_blk2 > k_pivot_1
|
||||
@@ -457,7 +457,7 @@ def _topk_topp_kernel(
|
||||
LOGITS_ROW + offs_n,
|
||||
mask=mask_n,
|
||||
other=-float("inf"),
|
||||
)
|
||||
).to(tl.float32)
|
||||
|
||||
outlier_mask = (probs_blk > min_logit) & mask_n
|
||||
|
||||
@@ -600,7 +600,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs < VOCAB_SIZE
|
||||
logits_blk0 = tl.load(
|
||||
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
# Exclude -inf values (e.g. from grammar bitmasks) from
|
||||
# statistics to avoid NaN in pivot computation.
|
||||
finite_mask = (logits_blk0 > -float("inf")) & mask_n
|
||||
@@ -626,7 +626,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs_n < VOCAB_SIZE
|
||||
logits_blk = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
|
||||
# Exclude -inf from min to keep binary search bounds
|
||||
# finite (avoids NaN pivots).
|
||||
@@ -660,7 +660,7 @@ def _topk_topp_kernel(
|
||||
|
||||
probs_blk = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
probs_blk = tl.exp(probs_blk - max_sample)
|
||||
probs_blk = probs_blk / sum_exp_logits
|
||||
|
||||
@@ -754,7 +754,7 @@ def _topk_topp_kernel(
|
||||
|
||||
probs_blk = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
probs_blk = tl.exp(probs_blk - max_sample)
|
||||
probs_blk = probs_blk / sum_exp_logits
|
||||
tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)
|
||||
@@ -835,7 +835,7 @@ def _topk_topp_kernel(
|
||||
mask_n = offs_n < VOCAB_SIZE
|
||||
logits_blk = tl.load(
|
||||
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
|
||||
)
|
||||
).to(tl.float32)
|
||||
keep_mask = (logits_blk > final_pivot) & mask_n
|
||||
|
||||
# Duplicate logit handling
|
||||
@@ -878,7 +878,7 @@ def apply_top_k_top_p_triton(
|
||||
The masked logits tensor. It may or may not be modified in-place.
|
||||
"""
|
||||
assert logits.ndim == 2
|
||||
assert logits.dtype == torch.float32
|
||||
assert logits.dtype in (torch.float32, torch.bfloat16, torch.float16)
|
||||
batch_size, vocab_size = logits.shape
|
||||
topk_enabled = k is not None
|
||||
topp_enabled = p is not None
|
||||
@@ -911,7 +911,9 @@ def apply_top_k_top_p_triton(
|
||||
buffer = _TRITON_BUFFER_CACHE.get(buf_key)
|
||||
if buffer is None or buffer.shape[0] < NUM_PROGRAMS:
|
||||
size = min(next_power_of_2(NUM_PROGRAMS), num_sm)
|
||||
buffer = logits.new_empty((size, vocab_size))
|
||||
buffer = torch.empty(
|
||||
(size, vocab_size), dtype=torch.float32, device=logits.device
|
||||
)
|
||||
_TRITON_BUFFER_CACHE[buf_key] = buffer
|
||||
if buffer.shape[0] > NUM_PROGRAMS:
|
||||
buffer = buffer[:NUM_PROGRAMS]
|
||||
@@ -919,8 +921,12 @@ def apply_top_k_top_p_triton(
|
||||
# Cache lookup table entries on each device.
|
||||
tables = _TRITON_TABLE_CACHE.get(logits.device)
|
||||
if tables is None:
|
||||
normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)
|
||||
percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)
|
||||
normal_cdf_to_sigma_table = torch.tensor(
|
||||
_NORMAL_CDF_TO_SIGMA_TABLE, dtype=torch.float32, device=logits.device
|
||||
)
|
||||
percentile_to_std_table = torch.tensor(
|
||||
_PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=logits.device
|
||||
)
|
||||
_TRITON_TABLE_CACHE[logits.device] = (
|
||||
normal_cdf_to_sigma_table,
|
||||
percentile_to_std_table,
|
||||
|
||||
@@ -262,7 +262,7 @@ def _reshape_kv_cache(
|
||||
kv_cache_config: "KVCacheConfig | None" = None,
|
||||
) -> dict[str, Any]:
|
||||
kv_caches: dict[str, Any] = {}
|
||||
has_attn = False
|
||||
has_attn, has_mamba = False, False
|
||||
|
||||
layer_packing: dict[str, tuple[int, int]] = {}
|
||||
if kv_cache_config is not None:
|
||||
@@ -340,21 +340,38 @@ def _reshape_kv_cache(
|
||||
)
|
||||
|
||||
elif isinstance(kv_cache_spec, MambaSpec):
|
||||
page_size_bytes = kv_cache_spec.page_size_bytes
|
||||
# Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
|
||||
# int8 page view per layer; the layer's bind_kv_cache unpacks
|
||||
# each block's bytes into its conv/ssm state views. Keeping
|
||||
# one tensor per layer lets the KV connector register it
|
||||
# without special-casing Mamba.
|
||||
kv_caches[layer_name] = kv_raw_tensor[
|
||||
: num_blocks * page_size_bytes
|
||||
].view(num_blocks, 1, 1, page_size_bytes)
|
||||
has_mamba = True
|
||||
state_tensors = []
|
||||
storage_offset_bytes = 0
|
||||
for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
|
||||
dtype_size = get_dtype_size(dtype)
|
||||
num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size
|
||||
target_shape = (num_blocks, *shape)
|
||||
stride = torch.empty(target_shape).stride()
|
||||
target_stride = (num_element_per_page, *stride[1:])
|
||||
assert storage_offset_bytes % dtype_size == 0
|
||||
tensor = torch.as_strided(
|
||||
kv_raw_tensor.view(dtype),
|
||||
size=target_shape,
|
||||
stride=target_stride,
|
||||
storage_offset=storage_offset_bytes // dtype_size,
|
||||
)
|
||||
state_tensors.append(tensor)
|
||||
storage_offset_bytes += stride[0] * dtype_size
|
||||
kv_caches[layer_name] = state_tensors
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
|
||||
)
|
||||
|
||||
if has_attn and kv_cache_config is not None:
|
||||
if has_attn and has_mamba:
|
||||
_update_hybrid_attention_layout(
|
||||
attn_groups=attn_groups,
|
||||
kv_caches=kv_caches,
|
||||
kernel_block_sizes=kernel_block_sizes,
|
||||
cache_dtype=cache_dtype,
|
||||
)
|
||||
elif has_attn and kv_cache_config is not None:
|
||||
_align_mixed_attention_kv_cache_views(
|
||||
attn_groups=attn_groups,
|
||||
kv_caches=kv_caches,
|
||||
@@ -441,6 +458,65 @@ def _restride_blocks_first_kv_cache_to_kv_first_storage(
|
||||
)
|
||||
|
||||
|
||||
def _update_hybrid_attention_layout(
|
||||
attn_groups: Iterable[AttentionGroup],
|
||||
kv_caches: dict[str, Any],
|
||||
kernel_block_sizes: list[int],
|
||||
cache_dtype: str,
|
||||
) -> None:
|
||||
for group in attn_groups:
|
||||
if group.kv_cache_group_id >= len(kernel_block_sizes):
|
||||
continue
|
||||
|
||||
kv_cache_spec = group.kv_cache_spec
|
||||
if not isinstance(kv_cache_spec, AttentionSpec):
|
||||
continue
|
||||
# Mirror the per-layer dtype selection used when building the shape
|
||||
# above. The block-dim index is dtype-independent for current backends
|
||||
# (quantization only changes the last dim), so this is a no-op today,
|
||||
# but it keeps both call sites consistent for skip layers.
|
||||
layer_cache_dtype = (
|
||||
"auto"
|
||||
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
|
||||
and not isinstance(kv_cache_spec, TQFullAttentionSpec)
|
||||
else cache_dtype
|
||||
)
|
||||
block_dim = group.backend.get_kv_cache_block_dim(
|
||||
kernel_block_sizes[group.kv_cache_group_id],
|
||||
kv_cache_spec.num_kv_heads,
|
||||
kv_cache_spec.head_size,
|
||||
cache_dtype_str=layer_cache_dtype,
|
||||
)
|
||||
# if the first dim of the kvcache's layout is already num_blocks, continue
|
||||
if block_dim == 0:
|
||||
continue
|
||||
|
||||
assert block_dim == 1, (
|
||||
"Expected the dim `num_blocks` at the second dim when updating"
|
||||
" the kvcache's layout of full attention layer"
|
||||
)
|
||||
|
||||
for layer_name in group.layer_names:
|
||||
if layer_name not in kv_caches:
|
||||
# Shared layer — will be aliased to its target after this pass.
|
||||
continue
|
||||
|
||||
kv_cache = kv_caches[layer_name]
|
||||
if kv_cache.shape[0] == 2:
|
||||
assert kv_cache.shape[1] != 2, (
|
||||
f"Cannot determine layout for tensor of shape {kv_cache.shape}"
|
||||
)
|
||||
hidden_size = kv_cache.shape[2:].numel()
|
||||
kv_cache.as_strided_(
|
||||
size=kv_cache.shape,
|
||||
stride=(
|
||||
hidden_size,
|
||||
2 * hidden_size,
|
||||
*kv_cache.stride()[2:],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def init_kv_cache(
|
||||
runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
|
||||
forward_context: dict[str, Any],
|
||||
|
||||
@@ -28,7 +28,7 @@ class EncoderCache:
|
||||
Clear the multi-modal cache that was used during profiling,
|
||||
but no longer needed during inference.
|
||||
"""
|
||||
# NOTE: v2 encoder cache profiling skips the multi-modal cache
|
||||
# TODO: Implement MM budget for encoder dummy run
|
||||
pass
|
||||
|
||||
def reset_encoder_cache(self) -> None:
|
||||
|
||||
@@ -3,9 +3,7 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
|
||||
from vllm.multimodal.encoder_budget import MultiModalBudget
|
||||
from vllm.multimodal.inputs import MultiModalKwargsItem
|
||||
from vllm.multimodal.utils import (
|
||||
get_mm_features_in_window,
|
||||
@@ -15,8 +13,6 @@ from vllm.multimodal.utils import (
|
||||
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
|
||||
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class EncoderRunner:
|
||||
def __init__(
|
||||
@@ -56,45 +52,6 @@ class EncoderRunner:
|
||||
|
||||
return mm_hashes, mm_kwargs
|
||||
|
||||
@torch.inference_mode()
|
||||
def profile_encoder_cache(
|
||||
self,
|
||||
dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
|
||||
budget: MultiModalBudget,
|
||||
) -> None:
|
||||
"""Profile multimodal encoder and temporary encoder cache memory."""
|
||||
if (encoder_budget := budget.get_encoder_budget()) <= 0:
|
||||
return
|
||||
|
||||
if not budget.mm_max_toks_per_item:
|
||||
logger.info(
|
||||
"Skipping encoder profiling for embedding-only mode "
|
||||
"(all modality limits=0 with enable_mm_embeds=True).",
|
||||
)
|
||||
return
|
||||
|
||||
assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
|
||||
dummy_modality = dummy_mm_inputs[0][0]
|
||||
max_mm_items_per_batch = len(dummy_mm_inputs)
|
||||
|
||||
logger.info_once(
|
||||
"Encoder cache will be initialized with a budget of %s tokens, "
|
||||
"and profiled with %s %s items of the maximum feature size.",
|
||||
encoder_budget,
|
||||
max_mm_items_per_batch,
|
||||
dummy_modality,
|
||||
)
|
||||
|
||||
dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
|
||||
|
||||
sanity_check_mm_encoder_outputs(
|
||||
dummy_encoder_outputs,
|
||||
expected_num_items=max_mm_items_per_batch,
|
||||
)
|
||||
self.encoder_cache.encoder_outputs.update(
|
||||
(f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def execute_mm_encoder(
|
||||
self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
|
||||
|
||||
@@ -43,10 +43,6 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
|
||||
)
|
||||
from vllm.model_executor.model_loader import get_model_loader
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.encoder_budget import (
|
||||
MultiModalBudget,
|
||||
get_dummy_encoder_profile_inputs,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.tasks import SupportedTask
|
||||
from vllm.utils.math_utils import cdiv
|
||||
@@ -675,22 +671,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
|
||||
@torch.inference_mode()
|
||||
def profile_run(self) -> None:
|
||||
if self.supports_mm_inputs and self.is_first_pp_rank:
|
||||
mm_config = self.model_config.multimodal_config
|
||||
if mm_config is not None and not mm_config.skip_mm_profiling:
|
||||
mm_budget = MultiModalBudget(
|
||||
self.vllm_config,
|
||||
self.mm_registry,
|
||||
enable_cache=False,
|
||||
)
|
||||
dummy_mm_inputs = get_dummy_encoder_profile_inputs(
|
||||
self.mm_registry,
|
||||
mm_budget,
|
||||
)
|
||||
self.model_state.encoder_runner.profile_encoder_cache(
|
||||
dummy_mm_inputs, mm_budget
|
||||
)
|
||||
|
||||
hidden_states, sample_hidden_states = self._dummy_run(
|
||||
self.max_num_tokens, skip_attn=True, is_profile=True
|
||||
)
|
||||
@@ -705,7 +685,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
del hidden_states, sample_hidden_states
|
||||
self.reset_encoder_cache()
|
||||
gc.collect()
|
||||
|
||||
def post_kv_cache_wake_up(self) -> None:
|
||||
|
||||
@@ -218,7 +218,9 @@ def _bias_kernel(
|
||||
mask=mask,
|
||||
)
|
||||
bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
|
||||
logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)
|
||||
logits = tl.load(
|
||||
logits_ptr + token_idx * logits_stride + token_ids, mask=mask
|
||||
).to(tl.float32)
|
||||
logits += bias
|
||||
tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user