Compare commits

...
Author SHA1 Message Date
Bugen ZhaoandOpenAI Codex e1a763558c [CI] Discover Rust coverage artifacts from build metadata
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 06:44:16 +00:00
Bugen ZhaoandOpenAI Codex 84aeec9f22 [CI] Simplify Rust coverage reporting
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 08:18:14 +00:00
Bugen ZhaoandOpenAI Codex 82a770ddbd [CI] Simplify Rust coverage aggregation
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:58:45 +00:00
Bugen Zhao a09a9bace1 [CI] Disable redundant Codecov file fixes 2026-07-21 13:57:38 +00:00
Bugen Zhao 6c20d467a2 [CI] Run Codecov from repository root 2026-07-21 13:34:22 +00:00
Bugen ZhaoandOpenAI Codex cb59d0a351 [CI] Collect Rust coverage in Buildkite
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 13:01:56 +00:00
Bugen ZhaoandOpenAI Codex 0ab1bded36 [CI] Instrument Rust artifacts for coverage
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 12:22:25 +00:00
Umut PolatandGitHub 040cbf95cc [Misc] Use VLLMValidationError in chat completion tool and batch validators (#49214)
Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com>
2026-07-21 11:38:20 +00:00
SyaGitHubLi, Jiang <jiang1.li@intel.com>
5b3762a7f0 [Bugfix][CPU] Fix Clang OpenMP build on macOS (#49021)
Signed-off-by: markyangcc <mmdou3@163.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-21 09:58:52 +00:00
bastefaniakandGitHub 4d30c510ce [bugfix] Fix Cosmos3 Edge checkpoint weights filtering, video loading, prompt expansion (#49190)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
2026-07-21 17:18:36 +08:00
6700813f86 [3/N][KV-Cache Layout Refactor] Standardize Mamba cache; drop get_transfer_cache_regions (#44456)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 09:16:15 +00:00
Bugen ZhaoandGitHub eb44b3aaa4 [Rust][Benchmark] Use async HTTP clients (#49295)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 16:53:57 +08:00
Nicolò LucchesiandGitHub 7a98c7a392 [Misc] Remove old now unsupported max_num_partial_prefills and max_long_partial_prefills (#49244)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-07-21 08:52:52 +00:00
Lena OnyshchenkoandGitHub 0d9e60619b [Misc][Docs] Fix XPU compute-runtime driver link version mismatch (#49299)
Signed-off-by: oonyshch <xonyshch@gmail.com>
2026-07-21 08:45:41 +00:00
1134545b6f Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) (#49033)
Co-authored-by: vllm-agent <vllm-agent@users.noreply.github.com>
2026-07-21 09:36:45 +01:00
Miłosz GrunwaldGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Harry Mellor
3e0c887511 [Bugfix] Fix Ovis2_5 special tokens for transformers v5 (#47298)
Signed-off-by: mgrunwal <milosz.grunwald@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-21 08:09:20 +00:00
Stefan KaestleandGitHub adfbbc1005 Propagate Flash Attention cache configuration to Ray workers (#49177)
Signed-off-by: Stefan Kaestle <skaestle@nvidia.com>
2026-07-21 07:47:53 +00:00
Roy WangandGitHub adc98f04d0 [Misc] Add @esmeetu to codeowners for rust/src/bench (#49298)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-07-21 07:44:24 +00:00
8def3cdde2 [Bugfix] Propagate quant_config to LFM2 ShortConv projections (#48917)
Signed-off-by: Alex Yuan <alex.yuan@liquid.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 06:51:55 +00:00
616c9bd0f4 [Frontend] Support additional sampling parameters for translation API (#45839)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-21 05:53:56 +00:00
Bugen ZhaoandGitHub 8688a06d67 [Rust][Benchmark] Use tracing for logs (#48937)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 05:18:13 +00:00
Salt SatoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f25953cc59 [Bugfix][Rust Frontend] Handle zero-column logprobs payloads without panicking (#49113)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
2026-07-21 04:30:29 +00:00
d9aa35161d Update BGE-M3 token expectations for leading spaces (#49269)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-07-21 03:49:17 +00:00
6bcda970fd [CI][NIXL] Isolate concurrent engine internal ports (#49129)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 22:28:11 -05:00
Isotr0pyandGitHub ea0e9c8f2e [MRV2] Add encoder cache profiling implementation (#47985)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-20 20:18:14 -07:00
ChaunceyandGitHub 94ed0bf4e0 [Bugfix][KV Offloading] Handle queued request aborts without allocated KV blocks (#49146)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-07-21 11:16:26 +08:00
Connor CarpenterGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
1940c8441e [Rust Frontend][gRPC] Add engine-aware health reporting (#48992)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
2026-07-21 10:54:39 +08:00
Simon MoandGitHub 72d16aee15 [CI] Exercise FA3 FP8 attention on SM90 (#49231)
Signed-off-by: Simon Mo <simon@inferact.ai>
2026-07-21 10:26:55 +08:00
Kunshang JiandGitHub e78a0c8e59 [XPU][Doc] Update XPU docker image documents (#49148)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-21 10:11:27 +08:00
Chris LeonardandGitHub 97a98006b0 Update qutlass cmake for stable abi (#47879)
Signed-off-by: Chris Leonard <chleonar@redhat.com>
2026-07-20 18:31:10 -07:00
0a684ab0c0 [Bugfix] Fix WSL circular import from pin_memory warning_once (#48444)
Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:55 -07:00
0d9210a502 Fixes non-coalesced HBM access in marlin_int4_fp8_preprocess_kernel_awq (#47268)
Signed-off-by: xjx <493337577@qq.com>
Signed-off-by: flutist-alibaba <30485581+flutist@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:38 -07:00
1d874867ea [Misc][Docs] Fix broken protocol link in speech_to_text doc (#47212)
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-21 01:05:05 +00:00
2e2e626b40 [Bugfix] Count per-group blocks in get_max_concurrency_for_kv_cache_config (#48317)
Signed-off-by: David Orman <ormandj@corenode.com>
Co-authored-by: Luke Alonso <lalonso@gmail.com>
Co-authored-by: Martin Vit <martin@voipmonitor.org>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-07-21 00:29:10 +00:00
Nick HillandGitHub af91f4b3e4 [Cleanup] Remove unused StructuredOutputRequest.status field (#49235)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 00:06:16 +00:00
2396a61108 [Attention][MLA][DCP] Query replication for MLA decode (DeepSeek-V2/R1 + Kimi-K2.5) (#45964)
Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 23:51:27 +00:00
97a668152b [RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel (#44214)
Signed-off-by: Xuanyu Zhang <xuanyu.zhang@mistral.ai>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 16:45:10 -07:00
58b2012aa2 [copy of #45208] CuMem slept-L1 fragmentation accounting (#49208)
Signed-off-by: haosdent <haosdent@gmail.com>
Signed-off-by: Justin Wood <justin.m.wood@me.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: haosdent <haosdent@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Justin Wood <jwood@me.com>
2026-07-20 23:08:11 +00:00
Ning XieandGitHub b7c20d0cfa [chore] adjust logo be more friendly to white background terminal (#48938)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
2026-07-20 15:15:28 -07:00
TJianandGitHub a2b1f9fc3b [ROCm] [Release] [Bugfix] Fix the per commit wheel release pipeline. (#49245)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-07-20 22:12:38 +00:00
642076d26c Support loading sample_from_anchor flag from speculators config (#48639)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-07-20 14:52:43 -07:00
Charlie FuandGitHub 5feb3950e5 [ROCm][CI] fix test_rocm_quick_reduce.py (#49234)
Signed-off-by: charlifu <charlifu@amd.com>
2026-07-20 16:39:33 -05:00
4ec199b66a [Bugfix][Spec-Decode] Populate draft seq_lens_cpu_upper_bound for spec-decode attention metadata (#44492)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 20:50:58 +00:00
Roberto L. CastroGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
7ca017778f [Feat][Perf] Add new warmup infrastructure for JITs (#47451)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 13:21:55 -07:00
fbfe58133d [Bugfix][KV Offload] Preserve reachable tails for hybrid SWA groups (#48911)
Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-07-20 22:12:36 +03:00
9dd62d80ab Cosmos3 FP8 ModelOpt/Diffusers remapping (#48952)
Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
Signed-off-by: wkutak <wkutak@nvidia.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-07-20 11:31:48 -07:00
f878367898 [Revert][Bugfix] Restore MiniCPM-V 4.6 ViT QKV weight loader (#49193)
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: wjinxu <1299461899@qq.com>
2026-07-20 18:17:46 +00:00
Matthew BonanniGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
bd091079cb [Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 10:53:27 -07:00
b23bd73f54 [XPU]add sycl path for Mhc (#47245)
Signed-off-by: root <xiaolong.guo@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:32:54 +00:00
Bugen ZhaoandGitHub e2d7adeb64 [Rust Frontend] Bump xgrammar-structural-tag and enable local extension (#49161)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-20 16:22:24 +01:00
173 changed files with 5876 additions and 1806 deletions
+2 -2
View File
@@ -813,8 +813,8 @@ steps:
# Download artifacts from current build
echo "Downloading artifacts from current build"
# buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
# buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# # Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh
@@ -29,6 +29,7 @@ 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 "--- $*"
@@ -106,6 +107,18 @@ 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 \
@@ -176,14 +189,41 @@ run_tests() {
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest"
cargo nextest run \
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 \
--manifest-path rust/Cargo.toml \
--workspace \
--all-features \
--locked \
--no-fail-fast
--no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
}
install_protoc
+182
View File
@@ -0,0 +1,182 @@
#!/bin/sh
RUST_CODECOV_VERSION="v11.3.1"
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
rust_coverage_repo_root() {
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
printf '%s\n' /vllm-workspace
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
else
git rev-parse --show-toplevel
fi
}
rust_coverage_start() {
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
mkdir -p "$RUST_COVERAGE_DIR"
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
export LLVM_PROFILE_FILE
trap rust_coverage_finalize 0
}
rust_coverage_objects() {
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
python3 - "$rust_cov_objects_manifest" <<'PY'
from pathlib import Path
import sys
for relative in Path(sys.argv[1]).read_text().splitlines():
for entry in sys.path:
path = Path(entry or ".").resolve() / relative
if path.is_file():
print(path)
break
else:
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
PY
}
rust_coverage_collect() {
rust_cov_collect_flag=${1:?coverage flag is required}
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
rust_cov_collect_primary=
set --
while IFS= read -r rust_cov_collect_object; do
if [ -z "$rust_cov_collect_primary" ]; then
rust_cov_collect_primary=$rust_cov_collect_object
else
set -- "$@" "--object=$rust_cov_collect_object"
fi
done <<EOF
$rust_cov_collect_objects
EOF
llvm-profdata merge \
-sparse \
"$RUST_COVERAGE_DIR"/*.profraw \
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
llvm-cov export \
"$rust_cov_collect_primary" \
"$@" \
--format=lcov \
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
> "$rust_cov_collect_lcov" || return 1
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
export RUST_COVERAGE_LCOV
}
rust_coverage_upload() {
rust_cov_upload_lcov=${1:?LCOV path is required}
rust_cov_upload_flag=${2:?coverage flag is required}
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
if [ "$(uname -m)" != "x86_64" ]; then
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
return 1
fi
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|| return 1
curl -fsSL \
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
| sha256sum -c - || return 1
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
rust_cov_upload_slug="vllm-project/vllm"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
case "$rust_cov_upload_slug" in
*/*) ;;
*) rust_cov_upload_slug="vllm-project/vllm" ;;
esac
fi
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
if [ -z "${CODECOV_TOKEN:-}" ]; then
# Codecov accepts tokenless public uploads on unprotected branch names.
# A colon-separated prefix keeps feature-branch and fork uploads from
# requiring a repository secret.
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
else
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
fi
fi
set --
set -- "$@" upload-process
set -- "$@" --file "$rust_cov_upload_lcov"
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
# source-line fix scanning, which is unrelated to path mapping.
set -- "$@" --disable-search --disable-file-fixes
set -- "$@" --fail-on-error --git-service github
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
set -- "$@" --branch "$rust_cov_upload_branch"
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
set -- "$@" --slug "$rust_cov_upload_slug"
set -- "$@" --flag "$rust_cov_upload_flag"
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
set -- "$@" --dir "$rust_cov_upload_repo_root"
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
fi
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
# E2E steps run from tests/, so execute from the repository root to resolve
# codecov.yml and repository paths consistently.
(
cd "$rust_cov_upload_repo_root" || exit 1
"$rust_cov_upload_codecov_dir/codecov" "$@"
) >"$rust_cov_upload_log" 2>&1
rust_cov_upload_rc=$?
cat "$rust_cov_upload_log"
# v11.3.1 can log API failures while returning zero even with
# --fail-on-error. Preserve the strict CI contract explicitly.
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
echo "Codecov CLI reported an upload error" >&2
rust_cov_upload_rc=1
fi
rm -rf "$rust_cov_upload_codecov_dir"
return "$rust_cov_upload_rc"
}
rust_coverage_finalize() {
rust_cov_finalize_test_rc=$?
trap - 0
set +e
rust_coverage_collect "$RUST_COVERAGE_FLAG"
rust_cov_finalize_collect_rc=$?
rust_cov_finalize_upload_rc=0
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
rust_cov_finalize_upload_rc=$?
fi
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
exit "$rust_cov_finalize_test_rc"
fi
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
exit "$rust_cov_finalize_collect_rc"
fi
exit "$rust_cov_finalize_upload_rc"
}
@@ -10,7 +10,9 @@ steps:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
@@ -34,7 +36,9 @@ steps:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
+30
View File
@@ -8,6 +8,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/benchmarks/
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
@@ -23,6 +28,7 @@ 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)"
@@ -43,6 +49,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/entrypoints/serve/
- vllm/v1/engine/
@@ -54,6 +65,7 @@ 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
@@ -72,10 +84,16 @@ 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
@@ -86,11 +104,17 @@ 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"
@@ -101,6 +125,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/distributed/
- vllm/engine/
- vllm/executor/
@@ -111,6 +140,7 @@ 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,5 +26,7 @@ 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
+1
View File
@@ -47,6 +47,7 @@
# 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
+1
View File
@@ -257,3 +257,4 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
+1 -1
View File
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
- Tool calling and reasoning parsers
- 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, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
+37
View File
@@ -8,6 +8,8 @@
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/')
@@ -30,4 +32,39 @@ 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
+5 -5
View File
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
else()
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _qutlass_fc_root)
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
QUTLASS_MINIMAL_BUILD=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
TORCH_TARGET_VERSION=0x020B000000000000ULL
USE_CUDA)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+13
View File
@@ -10,3 +10,16 @@ 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
+3 -1
View File
@@ -102,7 +102,9 @@ class TileGemm82 {
kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
auto fp32_b_regs = load_b_pair_vec(curr_b);
auto fp32_b_0_reg = fp32_b_regs.first;
auto fp32_b_1_reg = fp32_b_regs.second;
float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
// AWQ zeros: (size_k // group_size, size_n // 8)
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
int32_t group_size) {
int32_t val =
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
int32_t zero =
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
blockIdx.y];
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
// same row (stride 1) instead of striding across rows (stride size_n/8).
int col = blockIdx.y * 32 + threadIdx.x;
if (col >= size_n / 8) return;
(void)size_k;
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
int32_t new_val = 0;
#pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
zero >>= 4;
}
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
output[blockIdx.x * (size_n / 8) + col] = new_val;
}
torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
"qweight.size(0) % qzeros.size(0) != 0");
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
dim3 blocks(size_k / 32, size_n / 8);
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
+11
View File
@@ -294,6 +294,9 @@ 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 \
@@ -902,6 +905,14 @@ 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
+3
View File
@@ -46,6 +46,9 @@
"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"
},
+2 -2
View File
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
@@ -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.14.37833.4) release, to avoid potential compatibility issue.
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
--8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images]
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
vLLM offers official Docker images for deployment.
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
```bash
docker run --rm \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
vllm/vllm-openai-xpu:<tag> \
--model Qwen/Qwen3-0.6B
```
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
???+ console "Commands"
```bash
docker run --rm -it \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
--entrypoint /bin/bash \
vllm/vllm-openai-xpu:<tag>
```
--8<-- [end:pre-built-images]
--8<-- [start:build-image-from-source]
+9
View File
@@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform:
!!! tip
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
=== "Intel GPU"
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
!!! tip
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
=== "Google TPU"
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
- `response_format`: Format of the response ("json", "text") (optional)
- `temperature`: Sampling temperature between 0 and 1 (optional)
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
**Response Format:**
+29 -12
View File
@@ -3446,7 +3446,6 @@ dependencies = [
"base64 0.22.1",
"bytes",
"encoding_rs",
"futures-channel",
"futures-core",
"futures-util",
"h2",
@@ -4876,6 +4875,7 @@ dependencies = [
"futures-core",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
@@ -4937,9 +4937,9 @@ dependencies = [
[[package]]
name = "tonic"
version = "0.14.5"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
dependencies = [
"async-trait",
"axum",
@@ -4966,9 +4966,9 @@ dependencies = [
[[package]]
name = "tonic-build"
version = "0.14.5"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -4977,10 +4977,23 @@ dependencies = [
]
[[package]]
name = "tonic-prost"
version = "0.14.5"
name = "tonic-health"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
dependencies = [
"prost",
"tokio",
"tokio-stream",
"tonic",
"tonic-prost",
]
[[package]]
name = "tonic-prost"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
dependencies = [
"bytes",
"prost",
@@ -4989,9 +5002,9 @@ dependencies = [
[[package]]
name = "tonic-prost-build"
version = "0.14.5"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -5484,10 +5497,13 @@ dependencies = [
"serde",
"serde_json",
"thiserror 2.0.18",
"thiserror-ext",
"tiktoken-rs 0.9.1",
"tokenizers",
"tokio",
"tokio-stream",
"tracing",
"tracing-subscriber",
"url",
"uuid",
]
@@ -5730,6 +5746,7 @@ dependencies = [
"tokio-stream",
"tokio-util",
"tonic",
"tonic-health",
"tonic-prost",
"tonic-prost-build",
"tower",
@@ -6321,9 +6338,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
[[package]]
name = "xgrammar-structural-tag"
version = "0.1.0+xgrammar.0.2.2.4d145cc"
version = "0.2.0+xgrammar.0.2.4.dd729e7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
dependencies = [
"auto_impl",
"serde",
+6 -5
View File
@@ -118,10 +118,11 @@ 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.5"
tonic-build = "0.14.5"
tonic-prost = "0.14.5"
tonic-prost-build = "0.14.5"
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"
tool-parser = "1.2.0"
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
@@ -143,7 +144,7 @@ vllm-server = { path = "src/server" }
vllm-text = { path = "src/text" }
vllm-tokenizer = { path = "src/tokenizer" }
winnow = { version = "1.0.2", features = ["simd"] }
xgrammar-structural-tag = "0.1.0"
xgrammar-structural-tag = "0.2.0"
zeromq = { version = "0.6.0", default-features = false, features = [
"tokio-runtime",
"all-transport",
+4 -1
View File
@@ -20,16 +20,19 @@ mimalloc.workspace = true
rand.workspace = true
rand_distr.workspace = true
rayon.workspace = true
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
reqwest = { workspace = true, features = ["json", "stream", "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
+8 -7
View File
@@ -158,9 +158,10 @@ impl PoolingBackend {
// (mirrors Python async_request_vllm_rerank).
if let Some(ref list) = input.prompt_list {
if list.len() < 2 {
eprintln!(
"WARNING: vllm-rerank request has no documents \
(prompt_list needs [query, doc, ...])"
tracing::warn!(
backend = "vllm-rerank",
inputs = list.len(),
"rerank request has no documents"
);
}
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -175,10 +176,10 @@ impl PoolingBackend {
// Legacy path: text prompt as query, documents via --extra-body.
let query = input.prompt.as_ref();
if query.is_empty() && input.prompt_token_ids.is_some() {
eprintln!(
"WARNING: vllm-rerank received empty query (random dataset uses \
token IDs only). Use --dataset-name random-rerank for meaningful \
rerank benchmarks."
tracing::warn!(
backend = "vllm-rerank",
dataset = "random",
"rerank request has an empty query; use the random-rerank dataset"
);
}
serde_json::json!({
+114 -84
View File
@@ -6,6 +6,7 @@ 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};
@@ -72,12 +73,12 @@ pub fn pre_resolve_dns(
v4.extend(v6);
if !v4.is_empty() {
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
println!("Pre-resolved {host} -> {ips:?}");
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
builder = builder.resolve_to_addrs(host, &v4);
}
}
Err(e) => {
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
}
}
@@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
println!("Model not specified, fetching first model from server...");
tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) =
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}");
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name))
};
@@ -358,10 +363,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);
println!("Loading tokenizer: {tid}");
tracing::info!(tokenizer = tid, "loading tokenizer");
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)?;
println!("Tokenizer loaded successfully.");
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
Some(t)
};
let has_tokenizer = tokenizer.is_some();
@@ -421,7 +426,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.num_prompts, config.random_batch_size, config.is_reranker,
),
};
println!("Generating {dataset_label}...");
tracing::info!(
dataset = ?config.dataset_name,
prompts = config.num_prompts,
description = %dataset_label,
"generating benchmark dataset"
);
let gen_start = Instant::now();
let mut input_requests = match config.dataset_name {
@@ -472,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded.as_str()
}
};
@@ -512,7 +522,8 @@ 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()
}
};
@@ -543,7 +554,8 @@ 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,
@@ -608,18 +620,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
};
let gen_elapsed = gen_start.elapsed();
println!(
"Generated {} prompts in {:.2}s",
input_requests.len(),
gen_elapsed.as_secs_f64()
tracing::info!(
prompts = input_requests.len(),
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated benchmark dataset"
);
let filtered_count =
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
if filtered_count > 0 {
println!(
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
config.max_model_len.unwrap()
tracing::info!(
filtered_prompts = filtered_count,
max_model_len = config.max_model_len.unwrap(),
"filtered prompts above maximum model length"
);
}
if input_requests.is_empty() {
@@ -670,7 +683,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Ready check
if config.ready_check_timeout_sec > 0 {
println!("Starting initial single prompt test run...");
tracing::info!("starting initial single-prompt test run");
let test_output = wait_for_endpoint(
config.backend,
&client,
@@ -685,7 +698,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
test_output.error
)));
}
println!("Initial test run completed.");
tracing::info!("initial single-prompt test run completed");
}
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -703,12 +716,15 @@ 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() {
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
tracing::info!(
reason = "prompt_token_ids",
"skipping server tokenizer verification"
);
}
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
if is_tokenizer_verified(&cache_key) {
println!("Tokenizer verified in previous run (cached), skipping verification.");
tracing::info!(reason = "cached", "skipping server tokenizer verification");
} else {
let num_special =
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -723,14 +739,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await?
{
SampleVerifyOutcome::Passed => {
println!("Sample verification passed, skipping full verification.");
tracing::info!("tokenizer sample verification passed");
mark_tokenizer_verified(&cache_key);
}
SampleVerifyOutcome::Skipped(reason) => {
println!("Server /tokenize unavailable ({reason}), skipping verification.");
tracing::warn!(
reason = %reason,
"server tokenizer unavailable; skipping prompt verification"
);
}
SampleVerifyOutcome::Mismatch => {
println!("Sample verification found mismatch, running full verify+fix...");
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
match verify_and_fix_prompt_lengths(
&client,
&config.base_url,
@@ -742,16 +761,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await
{
Ok(()) => {
println!(
"All {} prompts verified: exact token length match.",
input_requests.len()
tracing::info!(
prompts = input_requests.len(),
"verified exact prompt token lengths"
);
mark_tokenizer_verified(&cache_key);
}
Err(BenchError::TokenizeUnavailable(reason)) => {
println!(
"Server /tokenize became unavailable during verification \
({reason}); proceeding with client-side token counts."
tracing::warn!(
reason = %reason,
"server tokenizer became unavailable; using client token counts"
);
}
Err(e) => return Err(e),
@@ -763,7 +782,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Warmup
if config.num_warmups > 0 {
println!("Warming up with {} requests...", config.num_warmups);
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
run_warmup(
config.backend,
&client,
@@ -776,7 +795,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.disable_tqdm,
)
.await;
println!("Warmup run completed.");
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -814,28 +833,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
println!("Speculative decoding detected, will collect metrics.");
tracing::info!("detected speculative decoding; collecting metrics");
}
// Main benchmark
println!("Starting main benchmark run...");
let distribution = if config.burstiness == 1.0 {
"Poisson process"
} else {
"Gamma distribution"
};
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)
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"
);
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -847,11 +860,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
);
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!(
"LoRA adapters ({}): {:?} [assignment={:?}]",
modules.len(),
names,
config.lora_assignment
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
"assigned LoRA adapters"
);
}
@@ -1125,7 +1138,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}");
tracing::error!(error = %e.as_report(), "profiler background task failed");
}
}
@@ -1289,12 +1302,14 @@ pub(crate) async fn start_profiler_immediate(
base_url: &str,
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) => println!("Profiler started"),
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
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")
}
}
}
@@ -1304,12 +1319,14 @@ 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) => println!("Profiler stopped"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
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")
}
}
}
@@ -1371,25 +1388,30 @@ pub(crate) async fn profile_on_batch_threshold(
duration_secs: f64,
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
) {
println!(
"Waiting for batch size >= {threshold} before starting profiler \
(will capture {duration_secs}s)..."
tracing::info!(
threshold,
duration_seconds = duration_secs,
"waiting for profiler batch threshold"
);
loop {
if let Some(running) = fetch_num_requests_running(client, base_url).await
&& running >= threshold
{
println!("Batch size {running} >= {threshold}, starting profiler...");
tracing::info!(
running_requests = running,
threshold,
"profiler batch threshold reached"
);
break;
}
// Wait 500ms or until the benchmark signals cancellation
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = &mut cancel_rx => {
eprintln!(
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
profiling skipped."
tracing::warn!(
threshold,
"benchmark finished before profiler batch threshold; skipping profiling"
);
return;
}
@@ -1398,13 +1420,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) => println!("Profiler started"),
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
Ok(false) => {
eprintln!("WARNING: Profiler start request returned non-success");
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
return;
}
Err(e) => {
eprintln!("WARNING: Failed to start profiler: {e}");
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
return;
}
}
@@ -1413,15 +1435,17 @@ pub(crate) async fn profile_on_batch_threshold(
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
_ = &mut cancel_rx => {
println!("Benchmark finished, stopping profiler early...");
tracing::info!("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) => println!("Profiler stopped after capturing"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
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")
}
}
}
@@ -1502,10 +1526,11 @@ 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 {
eprintln!(
"Prompt {i}: server consistently adds {excess} extra token(s) \
(likely BOS), compensating target to {}.",
expected_input_len.saturating_sub(excess),
tracing::warn!(
prompt_index = i,
extra_tokens = excess,
adjusted_target = expected_input_len.saturating_sub(excess),
"server consistently adds prompt tokens; compensating verification target"
);
}
excess
@@ -1563,7 +1588,10 @@ async fn verify_and_fix_prompt_lengths(
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
if fc > 0 {
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
tracing::info!(
fixed_prompts = fc,
"fixed prompt lengths using server tokenizer"
);
}
Ok(())
@@ -1818,7 +1846,7 @@ async fn sample_verify_prompts(
let tokenize_url = format!("{base_url}/tokenize");
let api_key = std::env::var("OPENAI_API_KEY").ok();
println!("Sampling {sample_size} prompts for verification...");
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
for (i, request) in requests.iter().enumerate().take(sample_size) {
let tokens = match server_tokenize(
@@ -1841,9 +1869,11 @@ async fn sample_verify_prompts(
let expected = request.prompt_len + num_special;
if tokens.len() != expected {
println!(
"Prompt {i}: expected {expected} tokens, server returned {}",
tokens.len()
tracing::warn!(
prompt_index = i,
expected_tokens = expected,
actual_tokens = tokens.len(),
"tokenizer verification sample mismatch"
);
return Ok(SampleVerifyOutcome::Mismatch);
}
+18 -9
View File
@@ -288,10 +288,18 @@ impl BenchConfig {
}
Some(other) => {
// extra_body was not an object — just use sampling params
eprintln!(
"Warning: --extra-body is not a JSON object, sampling params may be lost"
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"
);
let _ = other;
sampling_params
}
None => sampling_params,
@@ -489,9 +497,9 @@ impl BenchConfig {
_ => {}
}
if !args.skip_chat_template {
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
tracing::warn!(
dataset = "custom",
"client-side chat template rendering is unsupported; sending prompts raw"
);
}
}
@@ -570,9 +578,10 @@ impl BenchConfig {
}
if ignore_eos {
eprintln!(
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
respect output length limits, causing unbounded context growth."
tracing::warn!(
ignore_eos,
multi_turn = true,
"output length limits may be ignored, causing unbounded context growth"
);
}
+4 -2
View File
@@ -128,8 +128,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
}
#[test]
+74 -37
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -50,18 +51,19 @@ enum ColumnFormat {
/// Make a GET request with retry logic (3 retries with exponential backoff).
/// Returns the parsed JSON response.
fn get_with_retry(
client: &reqwest::blocking::Client,
async fn get_with_retry(
client: &reqwest::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() {
let resp = match client.get(url).send().await {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
.await;
continue;
}
return Err(BenchError::Config(format!(
@@ -80,7 +82,7 @@ fn get_with_retry(
}
if status.is_server_error() && attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
continue;
}
@@ -92,6 +94,7 @@ 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);
}
@@ -105,7 +108,7 @@ 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 fn download_hf_dataset(
pub async fn download_hf_dataset(
dataset: &str,
subset: Option<&str>,
split: Option<&str>,
@@ -115,7 +118,7 @@ pub fn download_hf_dataset(
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
let mut client_builder =
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
// Add HF_TOKEN auth header if available
if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -138,7 +141,7 @@ pub fn download_hf_dataset(
// Call /info to discover available configs and splits
let info_url =
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
let dataset_info =
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -201,7 +204,12 @@ pub fn download_hf_dataset(
(resolved_config, resolved_split)
};
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"resolved Hugging Face dataset"
);
// Check cache
let dir = cache_dir();
@@ -215,11 +223,16 @@ pub fn download_hf_dataset(
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset cached: {path_str}");
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
return Ok((path_str, resolved_config, resolved_split));
}
println!("Downloading HF dataset '{dataset}' from datasets-server...");
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"downloading Hugging Face dataset"
);
let encoded_config: String =
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -229,6 +242,7 @@ pub fn download_hf_dataset(
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -240,7 +254,7 @@ pub fn download_hf_dataset(
&length={page_size}"
);
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
let rows = data["rows"]
.as_array()
@@ -260,14 +274,14 @@ pub fn download_hf_dataset(
offset += fetched;
let total = data["num_rows_total"].as_u64().unwrap_or(0);
eprint!("\r Fetched {offset}/{total} rows...");
progress.update(offset, total);
// Stop if we have enough rows or reached end of dataset
if all_rows.len() >= num_rows_needed || fetched < page_size {
break;
}
}
eprintln!(); // newline after progress
progress.finish();
if all_rows.is_empty() {
return Err(BenchError::Config(format!(
@@ -280,7 +294,12 @@ pub fn download_hf_dataset(
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
tracing::info!(
dataset,
rows = all_rows.len(),
path = %path_str,
"saved Hugging Face dataset"
);
Ok((path_str, resolved_config, resolved_split))
}
@@ -483,21 +502,31 @@ 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) => println!("HF dataset: detected chat column '{col}'"),
ColumnFormat::Chat(col) => {
tracing::info!(
format = "chat",
column = col,
"detected Hugging Face dataset format"
);
}
ColumnFormat::Text {
prompt_col,
output_col,
} => {
let out_msg = output_col.as_deref().unwrap_or("none");
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
tracing::info!(
format = "text",
prompt_column = prompt_col,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
}
ColumnFormat::Combined { cols, output_col } => {
let out_msg = output_col.as_deref().unwrap_or("none");
println!(
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
cols
tracing::info!(
format = "combined",
prompt_columns = ?cols,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
}
}
@@ -608,9 +637,10 @@ pub fn load_hf_dataset(
if len == 0 { 128 } else { len }
} else {
if !warned_no_output {
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
);
warned_no_output = true;
}
@@ -618,9 +648,10 @@ pub fn load_hf_dataset(
}
} else {
if !warned_no_output {
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
);
warned_no_output = true;
}
@@ -640,9 +671,11 @@ pub fn load_hf_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "hf",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -652,9 +685,11 @@ pub fn load_hf_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled HF dataset from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "hf",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
@@ -1002,8 +1037,10 @@ mod tests {
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
crate::tokenizer::TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
}
/// Write JSON data to a unique temp file and return the path string.
+9 -6
View File
@@ -5,6 +5,7 @@ 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;
@@ -90,9 +91,10 @@ pub fn oversample_requests(
return;
}
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
requests.len()
tracing::info!(
samples = requests.len(),
requested = num_requests,
"skipping dataset oversampling"
);
return;
}
@@ -103,9 +105,10 @@ pub fn oversample_requests(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
requests.push(req);
}
println!(
"Oversampled requests from {original_len} to {} total samples.",
requests.len()
tracing::info!(
original_samples = original_len,
samples = requests.len(),
"oversampled dataset"
);
}
+29 -18
View File
@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
conversations.push(conv);
}
println!(
"Oversampled multi-turn conversations from {original_len} to {} total.",
conversations.len()
tracing::info!(
original_conversations = original_len,
conversations = conversations.len(),
"oversampled multi-turn conversations"
);
}
@@ -525,10 +526,12 @@ mod tests {
len
}
#[test]
#[tokio::test]
#[ignore]
fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
async fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 5,
@@ -610,10 +613,12 @@ mod tests {
println!("All prefix sharing checks passed!");
}
#[test]
#[tokio::test]
#[ignore]
fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
async fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 4,
@@ -650,10 +655,12 @@ mod tests {
println!("per_turn_input_len default-mode checks passed!");
}
#[test]
#[tokio::test]
#[ignore]
fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
async fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 50,
@@ -684,10 +691,12 @@ mod tests {
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
}
#[test]
#[tokio::test]
#[ignore]
fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
async fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 10,
@@ -709,10 +718,12 @@ mod tests {
println!("variable_turns_fixed checks passed!");
}
#[test]
#[tokio::test]
#[ignore]
fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
async fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
@@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
}
let total = prompts_per_prefix * num_prefixes;
if total != num_requests {
println!(
"prefix_repetition: generating {total} requests \
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
{} dropped to divide evenly)",
num_requests - total
tracing::info!(
requested = num_requests,
generated = total,
prefixes = num_prefixes,
prompts_per_prefix,
dropped = num_requests - total,
"adjusted prefix-repetition request count"
);
}
@@ -109,8 +111,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
}
#[test]
+77
View File
@@ -0,0 +1,77 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
/// Reports row download progress to an interactive progress bar, or through
/// periodic tracing events when the progress bar is hidden on a non-TTY.
pub(super) struct RowDownloadReporter {
progress: ProgressBar,
next_report: Instant,
}
impl RowDownloadReporter {
/// Creates a reporter that emits non-TTY updates every 10 seconds.
pub fn new() -> Self {
let progress = ProgressBar::new(0);
progress.set_style(
ProgressStyle::with_template(
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
)
.unwrap()
.progress_chars("#>-"),
);
Self {
progress,
next_report: Instant::now() + REPORT_INTERVAL,
}
}
/// Updates the current row count and reports progress when due.
pub fn update(&mut self, rows: usize, total: u64) {
let rows = rows as u64;
let total = total.max(rows);
self.progress.set_length(total);
self.progress.set_position(rows);
if self.should_report(Instant::now()) {
tracing::info!(rows, total, "fetching dataset rows");
}
}
/// Clears the interactive progress bar after the download completes.
pub fn finish(self) {
self.progress.finish_and_clear();
}
fn should_report(&mut self, now: Instant) -> bool {
if !self.progress.is_hidden() || now < self.next_report {
return false;
}
self.next_report = now + REPORT_INTERVAL;
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hidden_reporter_uses_ten_second_deadline() {
let start = Instant::now();
let mut reporter = RowDownloadReporter {
progress: ProgressBar::hidden(),
next_report: start + REPORT_INTERVAL,
};
assert!(!reporter.should_report(start + Duration::from_secs(9)));
assert!(reporter.should_report(start + Duration::from_secs(10)));
assert!(!reporter.should_report(start + Duration::from_secs(19)));
assert!(reporter.should_report(start + Duration::from_secs(20)));
}
}
+18 -12
View File
@@ -49,9 +49,12 @@ pub fn generate_random_dataset(
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
let (output_low, output_high) = range_ratio.output_bounds(output_len);
if !range_ratio.is_fixed() {
println!(
"Sampling input_len from [{input_low}, {input_high}] and \
output_len from [{output_low}, {output_high}]"
tracing::info!(
input_low,
input_high,
output_low,
output_high,
"sampling random request lengths"
);
}
@@ -305,7 +308,8 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_token_ids() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -337,7 +341,8 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_text() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -371,7 +376,8 @@ mod tests {
#[test]
#[ignore]
fn test_token_length_exact_local() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let target_len = 512;
let requests = generate_random_dataset(
&tokenizer,
@@ -405,11 +411,11 @@ mod tests {
}
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
#[test]
#[tokio::test]
#[ignore]
fn test_token_length_exact_tiktoken() {
async 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);
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
@@ -453,10 +459,10 @@ mod tests {
/// Test encode/decode roundtrip stability for tiktoken.
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
#[test]
#[tokio::test]
#[ignore]
fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
async fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
+4 -2
View File
@@ -139,8 +139,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
}
fn fixed_ratio() -> RangeRatio {
+19 -12
View File
@@ -22,18 +22,21 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
/// Download the default ShareGPT dataset from HuggingFace Hub.
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
pub fn download_sharegpt_dataset() -> Result<String> {
println!(
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
pub async fn download_sharegpt_dataset() -> Result<String> {
tracing::info!(
repository = DEFAULT_SHAREGPT_REPO,
file = DEFAULT_SHAREGPT_FILE,
"downloading ShareGPT dataset"
);
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
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| {
BenchError::Config(format!(
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
))
})?;
let path_str = path.to_string_lossy().to_string();
println!("ShareGPT dataset ready: {path_str}");
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
Ok(path_str)
}
@@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset(
// Oversample if dataset is smaller than requested
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "sharegpt",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let needed = num_requests - samples.len();
@@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled requests from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "sharegpt",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
+30 -23
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::cli::SpeedBenchConfig;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -25,7 +26,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 fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let config_name = config.as_str();
let dir = cache_dir();
@@ -35,13 +36,13 @@ pub 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();
println!("SPEED-Bench ({config_name}) cached: {path_str}");
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
return Ok(path_str);
}
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
let client = reqwest::blocking::Client::builder()
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -64,13 +66,14 @@ pub 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() {
let resp = match client.get(&url).send().await {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(
tokio::time::sleep(std::time::Duration::from_secs(
2 * (attempt as u64 + 1),
));
))
.await;
continue;
}
return Err(BenchError::Config(format!(
@@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
};
if resp.status().is_server_error() && attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
continue;
}
@@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
)));
}
data = Some(resp.json().map_err(|e| {
data = Some(resp.json().await.map_err(|e| {
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
})?);
break;
@@ -116,15 +119,14 @@ pub 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);
eprint!("\r Fetched {offset}/{total} rows...");
progress.update(offset, total);
if fetched < page_size {
break;
}
}
eprintln!(); // newline after progress
progress.finish();
if all_rows.is_empty() {
return Err(BenchError::Config(
@@ -137,9 +139,11 @@ pub 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();
println!(
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
all_rows.len()
tracing::info!(
config = config_name,
rows = all_rows.len(),
path = %path_str,
"saved SPEED-Bench dataset"
);
Ok(path_str)
}
@@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "speed-bench",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "speed-bench",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
@@ -288,7 +296,6 @@ pub fn load_speed_bench_dataset(
));
}
// Print category distribution
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
for entry in &filtered[..filtered.len().min(samples.len())] {
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -297,7 +304,7 @@ pub fn load_speed_bench_dataset(
let mut cats: Vec<_> = cat_counts.into_iter().collect();
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
println!("SPEED-Bench categories: {}", cat_str.join(", "));
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
Ok(samples)
}
+20 -35
View File
@@ -1,54 +1,39 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Sync facade over the async `hf_hub` API.
//!
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
//! backend unconditionally pulls ureq's default rustls feature. So we use the
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
//! loaders, tokenizer fallback in rayon threads) by running each download on a
//! dedicated thread with its own single-threaded runtime.
use std::path::PathBuf;
use hf_hub::Repo;
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
pub struct HubRepo {
repo: hf_hub::Repo,
repo: ApiRepo,
}
impl HubRepo {
pub fn model(model_id: String) -> Self {
Self {
repo: hf_hub::Repo::model(model_id),
}
pub fn model(model_id: String) -> Result<Self, String> {
Self::new(Repo::model(model_id))
}
pub fn dataset(repo_id: String) -> Self {
Self {
repo: hf_hub::Repo::dataset(repo_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));
}
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 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())?
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
self.repo.get(filename).await.map_err(|e| format!("{e}"))
}
}
+1 -1
View File
@@ -33,7 +33,7 @@ pub fn prepare_process() {
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
tracing::info!(soft_limit = new, "raised open-file limit");
}
}
+12
View File
@@ -19,7 +19,19 @@ struct Cli {
args: vllm_bench::BenchServeArgs,
}
// 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();
+18 -16
View File
@@ -7,6 +7,22 @@ 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.
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
let failed = outputs.len() - completed;
// 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);
}
}
log_failed_requests(outputs);
// Calculate max output tokens per second and max concurrent requests
let mut max_output_tokens_per_s = 0.0_f64;
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
let failed = outputs.len() - completed;
// 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);
}
}
log_failed_requests(outputs);
// Compute peak concurrent requests from start_time + latency windows
let successful_outputs: Vec<&RequestFuncOutput> =
+53 -39
View File
@@ -6,6 +6,7 @@ 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};
@@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
println!("Model not specified, fetching first model from server...");
tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}");
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name))
};
@@ -83,15 +88,19 @@ 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);
println!("Loading tokenizer: {tid}");
tracing::info!(tokenizer = tid, "loading tokenizer");
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)?;
println!("Tokenizer loaded successfully.");
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
Some(t)
};
// Generate/load conversations
println!("Generating multi-turn conversations...");
tracing::info!(
dataset = ?config.dataset_name,
conversations = config.num_prompts,
"generating multi-turn conversations"
);
let gen_start = Instant::now();
let mut conversations = match config.dataset_name {
@@ -131,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded.as_str()
}
};
@@ -179,8 +188,11 @@ 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 {
println!(
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
tracing::info!(
filtered_turns,
filtered_conversations,
max_model_len,
"filtered conversations above maximum model length"
);
}
if conversations.is_empty() {
@@ -192,11 +204,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let gen_elapsed = gen_start.elapsed();
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
println!(
"Generated {} conversations ({} total turns) in {:.2}s",
conversations.len(),
tracing::info!(
conversations = conversations.len(),
total_turns,
gen_elapsed.as_secs_f64()
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated multi-turn conversations"
);
// Log prefix sharing info
@@ -208,18 +220,15 @@ 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);
println!(
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
config.multi_turn_prefix_global_ratio * 100.0,
tracing::info!(
global_ratio = config.multi_turn_prefix_global_ratio,
global_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,
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
conversation_tokens = conv_tokens,
unique_tokens,
history_accumulation = false,
"configured multi-turn prefix sharing"
);
println!("No history accumulation: each turn sends fixed-length prompt only.");
}
if config.dry_run {
@@ -253,7 +262,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
..Default::default()
};
println!("Starting initial single prompt test run...");
tracing::info!("starting initial single-prompt test run");
let test_output = crate::ready_checker::wait_for_endpoint(
config.backend,
&client,
@@ -268,7 +277,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
test_output.error
)));
}
println!("Initial test run completed.");
tracing::info!("initial single-prompt test run completed");
}
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -283,9 +292,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
"min_tokens".to_string(),
serde_json::json!(config.random_output_len),
);
println!(
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
config.random_output_len
tracing::info!(
min_tokens = config.random_output_len,
dataset = "random",
"set minimum output tokens for multi-turn dataset"
);
}
Some(body)
@@ -297,7 +307,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() {
println!("Speculative decoding detected, will collect metrics.");
tracing::info!("detected speculative decoding; collecting metrics");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -330,10 +340,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
};
// Main benchmark
println!("Starting multi-turn benchmark...");
println!("Conversations: {}", conversations.len());
println!("Concurrency: {concurrency}");
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
tracing::info!(
conversations = conversations.len(),
total_turns,
concurrency,
inter_turn_delay_ms = config.multi_turn_delay_ms,
"starting multi-turn benchmark"
);
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
@@ -364,11 +377,12 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
);
if let Some(modules) = config.lora_modules.as_ref() {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!(
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
modules.len(),
names,
config.lora_assignment
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
scope = "conversation",
"assigned LoRA adapters"
);
}
@@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
match handle.await {
Ok(output) => all_outputs.push(output),
Err(e) => {
eprintln!("Conversation task panicked: {e}");
tracing::error!(error = %e.as_report(), "conversation task panicked");
}
}
}
@@ -453,7 +467,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}");
tracing::error!(error = %e.as_report(), "profiler background task failed");
}
}
+2 -2
View File
@@ -603,7 +603,7 @@ fn add_metric_stats(
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
let content = serde_json::to_string(json)?;
std::fs::write(file_path, content)?;
println!("Results saved to {file_path}");
tracing::info!(path = file_path, "saved benchmark results");
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())?;
println!("Results appended to {file_path}");
tracing::info!(path = file_path, "appended benchmark results");
Ok(())
}
+8 -2
View File
@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
let backend = get_backend(backend)?;
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
println!("Waiting for endpoint to become up in {timeout_seconds}s");
tracing::info!(
timeout_seconds,
retry_interval,
"waiting for endpoint readiness"
);
let pb = ProgressBar::new(timeout_seconds);
pb.set_style(
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
Ok(output) => {
let err = output.error.clone();
let err_last_line = err.lines().last().unwrap_or(&err);
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
pb.suspend(|| {
tracing::warn!(error = err_last_line, "endpoint is not ready");
});
last_error = err;
}
Err(e) => {
+1 -1
View File
@@ -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() {
println!("Prefix cache reset successfully.");
tracing::info!(url = %url, "reset prefix cache");
} else {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
+36 -24
View File
@@ -199,12 +199,17 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
}
};
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
tracing::info!(
encoding,
kind = "built-in-tiktoken",
vocab_size,
"loaded tokenizer"
);
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
}
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
// Phase 1: If model_id is a local directory, look for tiktoken files there
let local_dir = Path::new(model_id);
if local_dir.is_dir() {
@@ -212,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
}
// Phase 2: Fall back to HuggingFace Hub download
try_load_tiktoken_from_hf(model_id)
try_load_tiktoken_from_hf(model_id).await
}
/// Common tiktoken model filenames to search for.
@@ -247,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
}
/// Load a tiktoken tokenizer from a HuggingFace model repo.
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string());
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)?;
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 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 num_base_tokens = count_base_tokens(&model_path)?;
let config = match repo.get("tokenizer_config.json") {
let config = match repo.get("tokenizer_config.json").await {
Ok(config_path) => read_tokenizer_config(&config_path),
Err(_) => None,
};
let pattern = extract_pat_str_from_repo(&repo);
let pattern = extract_pat_str_from_repo(&repo).await;
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
}
@@ -309,15 +317,16 @@ fn build_tiktoken(
}
}
println!(
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
num_base_tokens,
all_special_tokens.len(),
if pattern.is_some() {
tracing::info!(
model = model_id,
base_tokens = num_base_tokens,
special_tokens = all_special_tokens.len(),
pattern = if pattern.is_some() {
"custom"
} else {
"default"
},
"loading tiktoken model"
);
TiktokenTokenizer::from_file(
@@ -397,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
/// Try to download the Python tokenizer source file and extract pat_str via regex.
/// Returns None if unavailable or unparsable.
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
// Try common Python tokenizer filenames
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
let py_path = match repo.get("tokenization_kimi.py").await {
Ok(path) => path,
Err(_) => repo.get("tokenizer.py").await.ok()?,
};
let source = std::fs::read_to_string(&py_path).ok()?;
@@ -438,9 +450,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
if !fragments.is_empty() {
let pattern = fragments.join("|");
println!(
"Extracted pat_str from Python source: {} fragments",
fragments.len()
tracing::debug!(
fragments = fragments.len(),
"extracted tiktoken pattern from Python source"
);
return Some(pattern);
}
+98 -21
View File
@@ -2,8 +2,10 @@
// 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};
@@ -18,7 +20,8 @@ pub enum TokenizerKind {
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
pub struct ServerTokenizer {
client: reqwest::blocking::Client,
client: reqwest::Client,
runtime: tokio::runtime::Handle,
tokenize_url: String,
detokenize_url: String,
model: String,
@@ -27,8 +30,8 @@ pub struct ServerTokenizer {
impl ServerTokenizer {
/// Create a new server tokenizer and verify connectivity.
pub fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::blocking::Client::builder()
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -38,6 +41,7 @@ impl ServerTokenizer {
let st = Self {
client,
runtime: tokio::runtime::Handle::current(),
tokenize_url,
detokenize_url,
model: model.to_string(),
@@ -45,7 +49,7 @@ impl ServerTokenizer {
};
// Probe the endpoint to verify it works and discover vocab size
let test_tokens = st.encode_inner("test")?;
let test_tokens = st.encode_async("test").await?;
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
let estimated_vocab = (max_id * 2).max(131072);
@@ -56,6 +60,10 @@ 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,
@@ -66,6 +74,7 @@ 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() {
@@ -75,7 +84,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().map_err(|e| {
let data: serde_json::Value = resp.json().await.map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
})?;
@@ -95,6 +104,10 @@ 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,
@@ -105,6 +118,7 @@ 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() {
@@ -114,7 +128,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().map_err(|e| {
let data: serde_json::Value = resp.json().await.map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
})?;
@@ -123,6 +137,26 @@ 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 ---
@@ -192,7 +226,7 @@ impl TokenizerKind {
/// 3. Server-side /tokenize + /detokenize endpoints
///
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
pub fn load_tokenizer(
pub async fn load_tokenizer(
model_id: &str,
_trust_remote_code: bool,
server_info: Option<(&str, &str)>,
@@ -212,31 +246,48 @@ pub fn load_tokenizer(
}
// 1. Try local HuggingFace tokenizer (tokenizer.json)
match try_load_local(model_id) {
match try_load_local(model_id).await {
Ok(tok) => {
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
tracing::info!(
model = model_id,
kind = "local",
vocab_size = tok.get_vocab_size(true),
"loaded tokenizer"
);
Ok(TokenizerKind::Local(Box::new(tok)))
}
Err(local_err) => {
// 2. Try tiktoken format
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
match crate::tiktoken::try_load_tiktoken(model_id) {
tracing::info!(
model = model_id,
error = %local_err.as_report(),
"local tokenizer unavailable; trying tiktoken"
);
match crate::tiktoken::try_load_tiktoken(model_id).await {
Ok(tok) => {
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
tracing::info!(
model = model_id,
kind = "tiktoken",
vocab_size = tok.vocab_size(),
"loaded tokenizer"
);
Ok(TokenizerKind::Tiktoken(tok))
}
Err(tiktoken_err) => {
// 3. Try server-side fallback
if let Some((base_url, model)) = server_info {
println!(
"Tiktoken also not available ({tiktoken_err}), \
trying server-side tokenization..."
tracing::info!(
model = model_id,
error = %tiktoken_err.as_report(),
"tiktoken unavailable; trying server-side tokenization"
);
match ServerTokenizer::new(base_url, model) {
match ServerTokenizer::new(base_url, model).await {
Ok(srv) => {
println!(
"Tokenizer: Server (vocab_size≈{})",
srv.cached_vocab_size
tracing::info!(
model = model_id,
kind = "server",
vocab_size = srv.cached_vocab_size,
"loaded tokenizer"
);
return Ok(TokenizerKind::Server(srv));
}
@@ -264,7 +315,7 @@ pub fn load_tokenizer(
}
/// Try loading tokenizer.json from local path or HuggingFace Hub.
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
// 1. Try local directory with tokenizer.json
let local_path = Path::new(model_id).join("tokenizer.json");
if local_path.exists() {
@@ -290,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result<Tokenizer> {
}
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
let repo = crate::hub::HubRepo::model(model_id.to_string());
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
let tokenizer_path = repo
.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);
}
}
+1 -1
View File
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
};
apply_structural_tag_constraint(request, parser.structural_tag_model())?;
apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
if parser.preserve_special_tokens() {
request.decode_options.skip_special_tokens = false;
@@ -7,7 +7,8 @@ use thiserror_ext::AsReport;
use vllm_engine_core_client::protocol::structured_outputs::{
StructuredOutputBackend, StructuredOutputsParams,
};
use vllm_parser::tool::StructuralTagModel;
use vllm_parser::tool::StructuralTagBuilder;
use xgrammar_structural_tag::builders::StructuralTagOptions;
use xgrammar_structural_tag::{
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
build_structural_tag,
@@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult};
/// support and the request's tool choice.
pub(super) fn apply_structural_tag_constraint(
request: &mut ChatRequest,
model: Option<StructuralTagModel>,
builder: Option<&dyn StructuralTagBuilder>,
) -> ChatResult<()> {
let Some(model) = model else {
let Some(builder) = builder else {
return Ok(());
};
let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint(
})
.collect::<Vec<_>>();
let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
let structural_tag = build_structural_tag(
builder,
&tools,
tool_choice,
StructuralTagOptions::default().with_reasoning(false),
)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
// Overwrite any existing structured output settings with the structural tag constraint.
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -141,7 +147,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -154,7 +160,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -169,7 +175,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -184,7 +190,7 @@ mod tests {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -201,7 +207,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -221,7 +227,7 @@ mod tests {
);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build");
let tag = structural_tag_value(&request).to_string();
@@ -234,7 +240,7 @@ mod tests {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -249,7 +255,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed");
let params = structured_outputs(&request);
+17 -7
View File
@@ -26,18 +26,21 @@ 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() {
pub(crate) fn init_tracing(process_label: &str) {
let filter = build_targets_filter(
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
env::var("RUST_LOG").ok().as_deref(),
);
let formatter = VllmEventFormatter::new();
let formatter = VllmEventFormatter::new(process_label);
let _ = tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
.with(
tracing_subscriber::fmt::layer()
.event_format(formatter)
.with_writer(std::io::stderr)
.with_filter(filter),
)
.try_init();
}
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
}
impl VllmEventFormatter {
fn new() -> Self {
fn new(process_label: &str) -> Self {
Self {
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
prefix: format!("({process_label} pid={})", process::id()),
timer: VllmLocalTimer::default(),
}
}
@@ -291,6 +294,13 @@ 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"));
+9 -1
View File
@@ -5,6 +5,7 @@ mod cli;
mod logging;
use std::env;
use std::ffi::OsStr;
use std::process::ExitStatus;
use anyhow::{Context, Result, anyhow, bail};
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
}
fn main() -> Result<()> {
logging::init_tracing();
let process_label =
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
"bench" => "Bench",
"serve" | "frontend" => "RustFrontend",
_ => "Rust",
};
logging::init_tracing(process_label);
let cli = Cli::parse();
let mut runtime = tokio::runtime::Builder::new_multi_thread();
@@ -460,6 +460,12 @@ 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()
+20 -1
View File
@@ -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;
use tokio::sync::{mpsc, watch};
use tracing::{debug, info, trace, warn};
use vllm_metrics::METRICS;
use zeromq::RouterSendHalf;
@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {
request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>,
health_tx: watch::Sender<bool>,
}
impl ClientInner {
@@ -57,6 +58,7 @@ 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),
}
}
@@ -169,6 +171,7 @@ 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();
@@ -191,6 +194,12 @@ 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 {
@@ -280,6 +289,11 @@ 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 {
@@ -461,13 +475,18 @@ 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,6 +269,19 @@ impl WireLogprobs {
);
}
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
if token_ids.rows == 0 {
return Ok(Logprobs {
positions: Vec::new(),
});
}
if token_ids.cols == 0 {
bail_ext_value_decode!(
"{field_prefix}: zero-column logprobs payload with {} rows",
token_ids.rows
);
}
let mut positions = Vec::with_capacity(token_ids.rows);
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
.data
@@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() {
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
);
}
#[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"
);
}
+3 -3
View File
@@ -4,7 +4,7 @@
use std::sync::Arc;
use vllm_parser::tool::{
Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
};
use vllm_parser::unified::{
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
self.inner.preserve_special_tokens()
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
self.inner.structural_tag_model()
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.inner.structural_tag_builder()
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.2 models.
///
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV32)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V4 models.
///
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV4)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -73,7 +73,7 @@ mod tests {
use super::DeepSeekV4ToolParser;
use crate::tool::test_utils::{collect_stream, test_tools};
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
use crate::tool::{ToolParser, ToolParserTestExt as _};
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
let params = params
@@ -91,13 +91,10 @@ mod tests {
}
#[test]
fn deepseek_v4_exposes_structural_tag_model() {
fn deepseek_v4_exposes_structural_tag_builder() {
let parser = DeepSeekV4ToolParser::new(&test_tools());
assert_eq!(
parser.structural_tag_model(),
Some(StructuralTagModel::DeepSeekV4)
);
assert!(parser.structural_tag_builder().is_some());
}
#[test]
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
///
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekR1)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
///
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::DeepSeekV31)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{GlmXmlToolParser, Separator};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
///
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Glm47)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Glm47.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<tool_calls>";
const TOOL_CALLS_END: &str = "</tool_calls>";
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::HyV3)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::HyV3.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Hermes",
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Hermes)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Hermes.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+5 -3
View File
@@ -12,7 +12,9 @@ use super::{
argument_delta_event, tool_call_header_event,
};
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
};
#[derive(Debug, Clone, PartialEq, Eq)]
enum LlamaJsonMode {
@@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Llama)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Llama.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Qwen XML",
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Qwen3)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Qwen3.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
true
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Kimi)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Kimi.builder())
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<minimax:tool_call>";
const TOOL_CALL_END: &str = "</minimax:tool_call>";
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Minimax)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Minimax.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
pub use seed_oss::SeedOssToolParser;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub use xgrammar_structural_tag::Model as StructuralTagModel;
pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
use crate::utils;
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
false
}
/// Return the xgrammar structural-tag model used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
/// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None
}
+7 -10
View File
@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::Tool;
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<tool_call>";
const TOOL_CALL_END: &str = "</tool_call>";
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
Ok(Box::new(Self::new(tools)))
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
Some(StructuralTagModel::Qwen3Coder)
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -294,7 +294,7 @@ mod tests {
use serde_json::{Value, json};
use thiserror_ext::AsReport;
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
use super::{Qwen3CoderToolParser, ToolParser};
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
@@ -308,13 +308,10 @@ mod tests {
}
#[test]
fn qwen_coder_exposes_structural_tag_model() {
fn qwen_coder_exposes_structural_tag_builder() {
let parser = Qwen3CoderToolParser::new(&test_tools());
assert_eq!(
parser.structural_tag_model(),
Some(StructuralTagModel::Qwen3Coder)
);
assert!(parser.structural_tag_builder().is_some());
}
#[test]
+4 -7
View File
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
use crate::reasoning::ReasoningParser;
use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Unified parser that composes existing reasoning and tool parsers.
pub struct CombinedParser {
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
|| self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
}
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -269,10 +269,7 @@ mod tests {
fn combined_parser_emits_tool_calls_from_visible_content() {
let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
let mut parser = CombinedParser::new(None, Some(tool));
assert!(matches!(
parser.structural_tag_model(),
Some(crate::tool::StructuralTagModel::Qwen3)
));
assert!(parser.structural_tag_builder().is_some());
let output = collect(
&mut parser,
+3 -3
View File
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
use crate::reasoning::ReasoningError;
use crate::tool::{
StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
};
/// Result alias for unified parser operations.
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
false
}
/// Return the xgrammar structural-tag model used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
/// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None
}
+1
View File
@@ -35,6 +35,7 @@ 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
+70
View File
@@ -0,0 +1,70 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tonic::server::NamedService;
use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::GenerateGrpcService;
pub(crate) async fn monitor_health(
mut health_reporter: HealthReporter,
mut engine_health: watch::Receiver<bool>,
shutdown: CancellationToken,
) {
let generate_service = GenerateGrpcService::NAME;
let status = ServingStatus::NotServing;
let health_event_first = tokio::select! {
result = engine_health.wait_for(|healthy| !*healthy) => {
match result {
Ok(_) => warn!(
generate_service,
overall_service = true,
status = ?status,
reason = "engine_unhealthy",
"marking gRPC health services as not serving"
),
Err(error) => warn!(
%error,
generate_service,
overall_service = true,
status = ?status,
reason = "health_channel_closed",
"engine health channel closed; marking gRPC health services as not serving"
),
}
true
}
_ = shutdown.cancelled() => {
info!(
generate_service,
overall_service = true,
status = ?status,
reason = "server_shutdown",
"server shutting down; marking gRPC health services as not serving"
);
false
}
};
health_reporter.set_not_serving::<GenerateGrpcService>().await;
// Generate is currently the only engine-backed gRPC service, so overall
// server health intentionally mirrors it.
health_reporter.set_service_status("", status).await;
if health_event_first {
shutdown.cancelled().await;
info!(
generate_service,
overall_service = true,
reason = "server_shutdown",
"server shutting down; closing gRPC health watches"
);
}
health_reporter.clear_service_status(generate_service).await;
health_reporter.clear_service_status("").await;
}
+4
View File
@@ -4,6 +4,7 @@
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
mod convert;
mod health;
use std::pin::Pin;
use std::sync::Arc;
@@ -24,8 +25,11 @@ 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;
+175 -12
View File
@@ -16,6 +16,10 @@ 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,
@@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend {
async fn setup_grpc_service(
engine_id: impl Into<EngineId>,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
) -> (
GenerateServer<GenerateServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint();
let engine_id = engine_id.into();
@@ -232,6 +240,7 @@ async fn setup_grpc_service(
)
.await
.expect("connect client");
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
test_llm(client),
@@ -240,6 +249,7 @@ 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,
)
}
@@ -254,25 +264,51 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
svc,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
(GenerateClient::new(channel), server_task, engine_task)
}
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
TonicServer::builder()
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
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");
});
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
let channel = Endpoint::from_shared(format!("http://{addr}"))
.expect("grpc endpoint")
.connect()
.await
.expect("connect grpc client");
.expect("connect grpc channel");
(grpc_client, server_task, engine_task)
(channel, server_task)
}
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -283,7 +319,7 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
let (svc, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr").to_string();
@@ -1035,3 +1072,129 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
service
};
let mut stream = health_client
.watch(HealthCheckRequest {
service: service.to_string(),
})
.await
.unwrap_or_else(|error| {
panic!("failed to start health watch for {service_label}: {error}")
})
.into_inner();
let initial = stream
.message()
.await
.unwrap_or_else(|error| {
panic!("failed to read initial health status for {service_label}: {error}")
})
.unwrap_or_else(|| {
panic!("health watch for {service_label} ended before its initial status")
});
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for {service_label}"
);
health_streams.push((service_label, stream));
}
engine_health_tx.send(false).expect("publish unhealthy engine state");
for (service_label, mut stream) in health_streams {
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))
.unwrap_or_else(|error| {
panic!("failed to read health update for {service_label}: {error}")
})
.unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected health status for {service_label}"
);
}
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
service: "vllm.Generate".to_string(),
})
.await
.expect("start health watch for vllm.Generate")
.into_inner();
let initial = stream
.message()
.await
.expect("read initial health status for vllm.Generate")
.expect("health watch ended before its initial status");
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for vllm.Generate"
);
shutdown.cancel();
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for shutdown health update for vllm.Generate")
.expect("failed to read shutdown health update for vllm.Generate")
.expect("health watch ended before its shutdown update");
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected shutdown health status for vllm.Generate"
);
let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for vllm.Generate health watch to close")
.expect("failed while closing vllm.Generate health watch");
assert!(
stream_end.is_none(),
"vllm.Generate health watch remained open"
);
tokio::time::timeout(Duration::from_secs(2), server_task)
.await
.expect("timed out waiting for gRPC server shutdown")
.expect("gRPC server task failed");
}
+28 -14
View File
@@ -39,6 +39,7 @@ use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
use tokio_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};
@@ -203,14 +204,19 @@ where
.map(tls::build_grpc_server_config)
.transpose()
.context("invalid gRPC TLS configuration")?;
let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(svc);
.add_service(health_service)
.add_service(generate_service);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls))
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
} else {
None
};
@@ -294,7 +300,8 @@ where
let server_shutdown = server_shutdown.clone();
let force_shutdown = force_shutdown.clone();
async move {
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
else {
// No gRPC configured: just wait for shutdown so we do not race the
// join! by resolving early and tripping the cancellation token.
shutdown.cancelled().await;
@@ -304,19 +311,26 @@ where
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
None => MaybeTlsListener::plain(grpc_listener),
};
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
let server =
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
let result = tokio::select! {
result = server => {
result.context("gRPC server failed")
}
_ = force_shutdown.cancelled() => {
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
Ok(())
}
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
};
server_shutdown.cancel();
let (result, ()) = tokio::join!(server, health_monitor);
result
}
};
+13 -3
View File
@@ -60,6 +60,8 @@ def _make_quick_allreduce(
qar.use_fp16_kernels = use_fp16_kernels
qar.qr_quant_level = QuickReduceRegime[quant_level]
qar.qr_max_size = qr_max_size
qar.qr_min_size = None
qar.qr_quantization_min_size = None
return qar
@@ -511,13 +513,21 @@ def test_quick_reduce_regime_values():
assert QuickReduceRegime.INT8.value == 1
assert QuickReduceRegime.INT6.value == 2
assert QuickReduceRegime.INT4.value == 3
assert QuickReduceRegime.NONE.value == 4
assert QuickReduceRegime.INT3.value == 4
assert QuickReduceRegime.NONE.value == 5
def test_quick_reduce_regime_names():
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
assert set(QuickReduceRegime.__members__) == {
"FP",
"INT8",
"INT6",
"INT4",
"INT3",
"NONE",
}
@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
@@ -693,7 +703,7 @@ def test_quick_allreduce_min_size_table():
for dtype in [torch.float16, torch.bfloat16]:
for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
assert len(min_sizes) == 4
assert len(min_sizes) == 5
assert all(size > 0 for size in min_sizes)
@@ -0,0 +1,880 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end tests for routed-expert capture on the monolithic MoE path.
These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
expert IDs picked by FlashInfer's fused router-and-experts kernels (the
"monolithic" path). When ``set_capture_fn`` is installed on
a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
should:
* allocate an int16 ``(num_tokens, top_k)`` buffer,
* pass it to FlashInfer as ``routing_replay_out``,
* invoke the callback after the kernel returns.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
FusedMoEQuantConfig,
RoutingMethodType,
fp8_w8a8_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
TrtLlmBf16ExpertsMonolithic,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
TrtLlmFp8ExpertsMonolithic,
)
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
TrtLlmNvFp4ExpertsMonolithic,
)
from vllm.platforms import current_platform
try:
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
except ImportError:
pytest.skip("flashinfer not available", allow_module_level=True)
if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
pytest.skip(
"Requires FlashInfer TRT-LLM fused MoE on CUDA",
allow_module_level=True,
)
if not current_platform.has_device_capability(100):
pytest.skip(
"TRT-LLM fused MoE kernels require SM100+",
allow_module_level=True,
)
def _shuffle_bf16_weights_block_major_k(
w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
) -> torch.Tensor:
"""Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
row shuffle.
"""
from flashinfer import shuffle_matrix_a
from flashinfer.fused_moe import convert_to_block_layout
num_experts = w.shape[0]
shuffled = []
for i in range(num_experts):
t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
shuffled.append(convert_to_block_layout(t, block_k))
return torch.stack(shuffled).view(torch.bfloat16)
def _make_bf16_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
routing_method: RoutingMethodType,
device: torch.device,
) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
"""Construct the monolithic BF16 experts plus the BlockMajorK weights
expected by ``trtllm_bf16_moe``.
"""
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=routing_method,
max_num_tokens=max(8, 1),
)
quant_config = FusedMoEQuantConfig.make(
quant_dtype=None,
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
)
experts = TrtLlmBf16ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
gemm1 = (
torch.randn(
num_experts,
2 * intermediate_size,
hidden_size,
device=device,
dtype=torch.bfloat16,
)
* 0.1
)
gemm2 = (
torch.randn(
num_experts,
hidden_size,
intermediate_size,
device=device,
dtype=torch.bfloat16,
)
* 0.1
)
w13 = _shuffle_bf16_weights_block_major_k(gemm1)
w2 = _shuffle_bf16_weights_block_major_k(gemm2)
return experts, w13, w2
def _run_bf16_monolithic(
experts: TrtLlmBf16ExpertsMonolithic,
hidden_states: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
*,
n_group: int | None = None,
topk_group: int | None = None,
routed_scaling_factor: float | None = None,
e_score_correction_bias: torch.Tensor | None = None,
) -> torch.Tensor:
return experts.apply(
hidden_states=hidden_states,
w1=w13,
w2=w2,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
a1q_scale=None,
apply_router_weight_on_input=False,
num_expert_group=n_group,
topk_group=topk_group,
e_score_correction_bias=e_score_correction_bias,
routed_scaling_factor=routed_scaling_factor,
)
_DSV3_NUM_EXPERTS = 32
_DSV3_N_GROUP = 4
_DSV3_TOPK_GROUP = 2
def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""The capture callback should receive the int16 routed-expert IDs the
kernel actually used, the values should be valid expert indices, and
each token should pick ``top_k`` distinct experts."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
captured: list[torch.Tensor] = []
def capture_fn(replay_out: torch.Tensor) -> None:
captured.append(replay_out.clone())
assert experts.supports_routing_replay_capture()
experts.set_capture_fn(capture_fn)
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
n_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
routed_scaling_factor=1.0,
e_score_correction_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize(
"routing_method",
[
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
],
)
def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
num_tokens: int,
routing_method: RoutingMethodType,
) -> None:
"""Routing replay works for non-DeepSeekV3 routing methods too.
FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = 8
top_k = 2
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=routing_method,
device=device,
)
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
"""With no callback installed the kernel should not see a
``routing_replay_out`` tensor verify the helper short-circuits."""
torch.manual_seed(0)
device = torch.device("cuda:0")
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=2,
hidden_size=1024,
intermediate_size=1024,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
# No callback installed.
buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
assert buf is None
# Dispatch is also a no-op.
experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
"""FlashInfer's ``routing_replay_out`` is supported by all routing
methods, so ``supports_routing_replay_capture`` should be True
regardless of routing method."""
device = torch.device("cuda:0")
for routing_method in (
RoutingMethodType.DeepSeekV3,
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
):
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=2,
hidden_size=1024,
intermediate_size=1024,
routing_method=routing_method,
device=device,
)
assert experts.supports_routing_replay_capture() is True, (
f"{routing_method!r} should support routing replay capture"
)
def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
"""When capture is installed, the allocated buffer is int16 and shaped
``(num_tokens, experts_per_token)``."""
device = torch.device("cuda:0")
experts, _, _ = _make_bf16_monolithic_experts(
num_experts=_DSV3_NUM_EXPERTS,
top_k=4,
hidden_size=1024,
intermediate_size=1024,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
experts.set_capture_fn(lambda r: None)
buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
assert buf is not None
assert buf.dtype == torch.int16
assert buf.shape[0] >= 11
assert buf.shape[1] == 4
assert buf.device.type == "cuda"
def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
"""End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
on the monolithic experts and verify the captured rows land in the
capturer's device buffer at the correct layer slot.
Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
for the monolithic path: a single closure is installed on the monolithic
``fused_experts`` (in addition to ``router.set_capture_fn`` on the
non-monolithic path) and the capturer routes per-layer based on the
closed-over ``layer_id``.
"""
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
RoutedExpertsCapturer,
)
torch.manual_seed(7)
device = torch.device("cuda:0")
num_tokens = 4
top_k = 2
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_bf16_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
routing_method=RoutingMethodType.DeepSeekV3,
device=device,
)
num_layers = 3
layer_id = 1
capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
capturer.dp_rank = 0
capturer.tp_size = 1
capturer.device_buffer = torch.full(
(num_tokens + 4, num_layers, top_k),
-1,
dtype=torch.int32,
device=device,
)
def capture_fn(replay_out: torch.Tensor) -> None:
capturer.capture(layer_id, replay_out)
experts.set_capture_fn(capture_fn)
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
# Patch get_forward_context to return a dp_metadata=None context so the
# capturer takes the single-DP branch.
import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
with patch.object(
rec,
"get_forward_context",
return_value=SimpleNamespace(dp_metadata=None),
):
_ = _run_bf16_monolithic(
experts,
hidden_states=hidden_states,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
n_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
routed_scaling_factor=1.0,
e_score_correction_bias=routing_bias,
)
captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
# Valid expert IDs at this layer.
assert (captured >= 0).all()
assert (captured < num_experts).all()
for t in range(num_tokens):
unique = captured[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts at layer "
f"{layer_id}, got {unique.numel()}"
)
# Other layers / trailing token rows untouched.
for other_layer in range(num_layers):
if other_layer == layer_id:
continue
assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
f"layer {other_layer} should be untouched, got writes"
)
assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
"tail rows beyond num_tokens should remain sentinel"
)
# ----------------------------------------------------------------------------
# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
# ----------------------------------------------------------------------------
def _make_fp8_block_scale_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
device: torch.device,
) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
"""Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
code path with DSV3 routing.
Weights are shuffled into the BlockMajorK layout the kernel expects
(same helper the vLLM weight loader uses for DeepSeek-FP8 models).
"""
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
_shuffle_deepseek_fp8_moe_weights,
)
block_k = 128
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=RoutingMethodType.DeepSeekV3,
max_num_tokens=max(8, 1),
)
# Random fp8 weights + ones-block scales (the kernel decoder only cares
# that the per-block scales are present and finite for routing/replay).
gemm1 = torch.randn(
num_experts, 2 * intermediate_size, hidden_size, device=device
).to(torch.float8_e4m3fn)
gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
torch.float8_e4m3fn
)
w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
w1_scale = torch.ones(
num_experts,
2 * intermediate_size // block_k,
hidden_size // block_k,
device=device,
dtype=torch.float32,
)
w2_scale = torch.ones(
num_experts,
hidden_size // block_k,
intermediate_size // block_k,
device=device,
dtype=torch.float32,
)
quant_config = fp8_w8a8_moe_quant_config(
w1_scale=w1_scale,
w2_scale=w2_scale,
block_shape=[block_k, block_k],
per_act_token_quant=False,
)
experts = TrtLlmFp8ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
return experts, w13_shuffled, w2_shuffled
def _run_fp8_block_scale_monolithic(
experts: TrtLlmFp8ExpertsMonolithic,
hidden_states_fp8: torch.Tensor,
hidden_states_scale: torch.Tensor,
w13: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
routing_bias: torch.Tensor,
) -> torch.Tensor:
return experts.apply(
hidden_states=hidden_states_fp8,
w1=w13,
w2=w2,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
# The block-scale apply path reads ``a1q_scale`` and transposes it
# to ``(hidden_size/128, num_tokens)`` for the kernel call.
a1q_scale=hidden_states_scale,
apply_router_weight_on_input=False,
num_expert_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
e_score_correction_bias=routing_bias,
routed_scaling_factor=1.0,
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
path, DSV3 routing) captures valid expert IDs."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
device=device,
)
assert experts.supports_routing_replay_capture()
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
# Per-token / per-block hidden scales (ones is fine for the routing
# path; the GEMM output isn't being asserted on).
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
).to(torch.float8_e4m3fn)
hidden_states_scale = torch.ones(
num_tokens, hidden_size // 128, device=device, dtype=torch.float32
)
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_fp8_block_scale_monolithic(
experts,
hidden_states_fp8=hidden_states,
hidden_states_scale=hidden_states_scale,
w13=w13,
w2=w2,
router_logits=router_logits,
num_experts=num_experts,
routing_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
# ----------------------------------------------------------------------------
# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
# ----------------------------------------------------------------------------
def _make_nvfp4_monolithic_experts(
num_experts: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
device: torch.device,
) -> tuple[
TrtLlmNvFp4ExpertsMonolithic,
torch.Tensor, # w13 (packed nvfp4 uint8)
torch.Tensor, # w13 block-scale (fp8)
torch.Tensor, # w2 (packed nvfp4 uint8)
torch.Tensor, # w2 block-scale (fp8)
torch.Tensor, # input global scale (per-tensor float32)
]:
"""Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
and DSV3 routing.
NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
"global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
* weights: uint8 (packed fp4) ``(E, M, K//2)``
* weight scales: fp8 ``(E, M, K//16)``
* hidden states: uint8 (packed fp4) ``(N, K//2)``
* hidden state scales: fp8 ``(N, K//16)``
"""
from flashinfer import fp4_quantize
block_size = 16
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
moe_config = FusedMoEConfig(
num_experts=num_experts,
experts_per_token=top_k,
hidden_dim=hidden_size,
intermediate_size=intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
moe_parallel_config=parallel_cfg,
in_dtype=torch.bfloat16,
activation=MoEActivation.SILU,
device=device,
routing_method=RoutingMethodType.DeepSeekV3,
max_num_tokens=max(8, 1),
)
gemm1 = torch.randn(
num_experts,
2 * intermediate_size,
hidden_size,
device=device,
dtype=torch.bfloat16,
)
gemm2 = torch.randn(
num_experts,
hidden_size,
intermediate_size,
device=device,
dtype=torch.bfloat16,
)
# Per-tensor weight scaling factor (used to build ``g1_alphas`` /
# ``g2_alphas`` below).
w_global_scale = torch.tensor(1.0, device=device)
# Per-tensor input scaling factor.
a_global_scale = torch.tensor(1.0, device=device)
w13_q, w13_scale = fp4_quantize(
gemm1,
w_global_scale,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
num_experts, 2 * intermediate_size, hidden_size // block_size
)
w2_q, w2_scale = fp4_quantize(
gemm2,
w_global_scale,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
num_experts, hidden_size, intermediate_size // block_size
)
# NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
# g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
quant_config = FusedMoEQuantConfig.make(
quant_dtype="nvfp4",
per_act_token_quant=False,
per_out_ch_quant=False,
block_shape=None,
w1_scale=w13_scale,
w2_scale=w2_scale,
g1_alphas=g_alphas,
g2_alphas=g_alphas,
a1_gscale=a_global_scale,
a2_gscale=a2_gscale,
)
experts = TrtLlmNvFp4ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
)
return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
def _run_nvfp4_monolithic(
experts: TrtLlmNvFp4ExpertsMonolithic,
hidden_states_q: torch.Tensor,
hidden_states_scale: torch.Tensor,
router_logits: torch.Tensor,
num_experts: int,
routing_bias: torch.Tensor,
) -> torch.Tensor:
"""The monolithic NVFP4 apply expects packed fp4 hidden states + the
matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
# Stash the weight tensors on the experts in the locations the apply()
# implementation reads from (it pulls them from quant_config / scales
# already; w1/w2 come in as args).
return experts.apply(
hidden_states=hidden_states_q,
w1=experts._w13_packed,
w2=experts._w2_packed,
router_logits=router_logits,
activation=MoEActivation.SILU,
global_num_experts=num_experts,
expert_map=None,
a1q_scale=hidden_states_scale,
apply_router_weight_on_input=False,
num_expert_group=_DSV3_N_GROUP,
topk_group=_DSV3_TOPK_GROUP,
e_score_correction_bias=routing_bias,
routed_scaling_factor=1.0,
)
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
@pytest.mark.parametrize("top_k", [2, 4])
def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
num_tokens: int,
top_k: int,
) -> None:
"""End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
on the DSV3 routing path."""
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
pytest.skip(
f"DSV3 requires top_k <= n_group * topk_group "
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
)
from flashinfer import fp4_quantize
torch.manual_seed(0)
device = torch.device("cuda:0")
num_experts = _DSV3_NUM_EXPERTS
hidden_size = 1024
intermediate_size = 1024
block_size = 16
experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
device=device,
)
# The apply() reads w1/w2 from its args, but we keep them on the experts
# for convenience of the helper.
experts._w13_packed = w13_q
experts._w2_packed = w2_q
assert experts.supports_routing_replay_capture()
captured: list[torch.Tensor] = []
experts.set_capture_fn(lambda r: captured.append(r.clone()))
hidden_states = (
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
)
hidden_states_q, hidden_states_scale = fp4_quantize(
hidden_states,
a_gs,
block_size,
sf_use_ue8m0=False,
is_sf_swizzled_layout=False,
)
# The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
# ``hidden_states_scale`` in its native (uint8 packed) form.
router_logits = torch.rand(
num_tokens, num_experts, device=device, dtype=torch.float32
)
routing_bias = _make_dsv3_routing_bias(num_experts, device)
_ = _run_nvfp4_monolithic(
experts,
hidden_states_q=hidden_states_q,
hidden_states_scale=hidden_states_scale,
router_logits=router_logits,
num_experts=num_experts,
routing_bias=routing_bias,
)
assert len(captured) == 1
replay = captured[0]
assert replay.dtype == torch.int16
assert replay.shape == (num_tokens, top_k)
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
for t in range(num_tokens):
unique = replay[t].unique()
assert unique.numel() == top_k, (
f"token {t}: expected {top_k} distinct experts, "
f"got {unique.numel()} ({replay[t].tolist()})"
)
-1
View File
@@ -66,7 +66,6 @@ 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(
+306
View File
@@ -0,0 +1,306 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ast
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, cast
import pytest
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
WarmupIntRange,
get_ast_full_name,
zip_inputs,
)
def _next_power_of_2(value: int) -> int:
return 1 << max(0, value - 1).bit_length()
def _round_up(value: int, *, multiple: int) -> int:
return ((value + multiple - 1) // multiple) * multiple
def _config(
*,
bias: int = 0,
disabled: bool = False,
name: str = "base",
vectorized: bool = False,
) -> SimpleNamespace:
return SimpleNamespace(
bias=bias,
disabled=disabled,
name=name,
vectorized=vectorized,
)
class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
block_size: int
work: int
vector_width: int
descriptor: tuple[object, ...]
enabled: bool
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
cfg: Any,
lanes: int = 1,
mode: str = "default",
debug: int = 0,
) -> CompileKey:
block_size = _next_power_of_2(tokens)
work: int = block_size * lanes + cfg.bias
return self.CompileKey(
block_size=block_size,
work=work,
vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
descriptor=(
cfg.name,
mode,
-block_size,
block_size % 3,
block_size**2,
),
enabled=not cfg.disabled,
)
def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(
tokens=WarmupIntRange(1, max_tokens + 1),
cfg=cfg,
# This argument is intentionally unused by dispatch expressions.
debug=WarmupIntRange(0, 100),
)
def compile(self, compile_key: CompileKey) -> None:
pass
class RecordingToyKernel(ToyKernel):
def __init__(self) -> None:
self.compiled: list[ToyKernel.CompileKey] = []
super().__init__()
def compile(self, compile_key: ToyKernel.CompileKey) -> None:
self.compiled.append(compile_key)
def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
cfg = _config()
assert ToyKernel().get_warmup_keys(5, cfg) == [
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
]
def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
assert ToyKernel().compile_key(
{
"tokens": 4,
"cfg": cfg,
"lanes": 2,
}
) == ToyKernel.CompileKey(
block_size=4,
work=11,
vector_width=4,
descriptor=("cfg", "default", -4, 1, 16),
enabled=False,
)
def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
cfg = _config(vectorized=True)
keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
zip_inputs(
dict(tokens=1, mode="small"),
dict(tokens=4, mode="wide"),
),
cfg=cfg,
lanes=(1, 2),
)
assert keys == [
ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
]
def test_zip_inputs_validates_input_rows() -> None:
with pytest.raises(ValueError, match="requires at least one"):
zip_inputs()
with pytest.raises(ValueError, match="rows must be mappings"):
zip_inputs(cast(Any, ("tokens", 1)))
with pytest.raises(ValueError, match="at least one dispatch input name"):
zip_inputs({})
with pytest.raises(ValueError, match="dispatch input names must be strings"):
zip_inputs(cast(Any, {1: 2}))
with pytest.raises(ValueError, match="same dispatch input names"):
zip_inputs({"tokens": 1}, {"mode": "small"})
def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
kernel = ToyKernel()
with pytest.raises(TypeError, match="zip_inputs"):
kernel._trace_dispatch(kernel.dispatch)(
cast(Any, {"tokens": 1}),
cfg=_config(),
)
with pytest.raises(ValueError, match="specified more than once"):
kernel._trace_dispatch(kernel.dispatch)(
zip_inputs(dict(tokens=1, mode="small")),
tokens=2,
cfg=_config(),
)
def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(value=_round_up(tokens, multiple=block_size))
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch( # type: ignore[override]
self,
*,
tokens: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
assert HelperKernel().compile_key(
{
"tokens": 5,
"block_size": 4,
}
) == HelperKernel.CompileKey(value=8)
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
if value > 0:
value = 1
return self.CompileKey(value=value)
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
return self.CompileKey(**{"value": value})
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
with pytest.raises(ValueError, match="local assignments"):
BranchKernel()
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
KwargsReturnKernel()
def test_dispatch_reports_unsupported_expression_with_context() -> None:
class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: object
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
return self.CompileKey(value={value})
def get_warmup_keys(self) -> list[CompileKey]:
return []
def compile(self, compile_key: CompileKey) -> None:
pass
with pytest.raises(ValueError) as exc_info:
UnsupportedKernel().compile_key({"value": 1})
message = str(exc_info.value)
assert "Unsupported dispatch expression" in message
assert "{value}" in message
assert "Supported dispatch expressions" in message
def test_warmup_compiles_all_returned_keys_in_order() -> None:
kernel = RecordingToyKernel()
cfg = _config()
kernel.warmup(3, cfg)
assert kernel.compiled == [
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
]
def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
dotted_expr = ast.parse("foo.bar.baz").body[0]
call_expr = ast.parse("foo()").body[0]
assert isinstance(dotted_expr, ast.Expr)
assert isinstance(call_expr, ast.Expr)
assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
assert get_ast_full_name(call_expr.value) is None
@@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
)
def _make_modular_routed_experts():
return types.SimpleNamespace(
quant_method=types.SimpleNamespace(is_monolithic=False),
)
def test_base_router_capture_pre_eplb_mapping():
router = _make_router()
captured = []
@@ -122,6 +128,8 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
def __init__(self):
self.layer_id = 7
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
class DummyCapturer:
def __init__(self):
@@ -160,6 +168,8 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
def __init__(self):
self.layer_id = 11
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
class DummyCapturer:
def __init__(self):
@@ -197,6 +207,8 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
def __init__(self, layer_id):
self.layer_id = layer_id
self.router = _make_router()
self.routed_experts = _make_modular_routed_experts()
self._quant_method = self.routed_experts.quant_method
target_module = DummyFusedMoE(layer_id=7)
draft_module = DummyFusedMoE(layer_id=0)
@@ -222,6 +234,49 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
assert draft_module.router.capture_fn is None
def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
from vllm.v1.worker import gpu_model_runner as gmr
class DummyFusedMoE:
def __init__(self):
self.layer_id = 3
self.router = _make_router()
# Use a concrete monolithic expert and override its capability
# instead of instantiating the abstract base class directly.
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
CPUExpertsFp8,
)
fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
self.routed_experts = types.SimpleNamespace(
quant_method=types.SimpleNamespace(
is_monolithic=True,
moe_kernel=types.SimpleNamespace(
impl=types.SimpleNamespace(fused_experts=fused_experts)
),
)
)
self._quant_method = self.routed_experts.quant_method
self._quant_method.moe_kernel.impl.fused_experts = fused_experts
fused_experts.supports_routing_replay_capture = lambda: False
class DummyCapturer:
def capture(self, layer_id, topk_ids):
pass
dummy_module = DummyFusedMoE()
import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
dummy_self = types.SimpleNamespace(
model=types.SimpleNamespace(modules=lambda: [dummy_module])
)
with pytest.raises(ValueError, match="monolithic MoE kernel"):
gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
def test_routed_experts_capturer_single_dp_no_metadata():
"""dp_metadata is None: capture writes the full topk_ids rows."""
capturer = _capturer_with_buffer(dp_rank=0)
+13 -2
View File
@@ -9,6 +9,7 @@ import torch.nn.functional as F
from transformers import AutoModel
from vllm.platforms import current_platform
from vllm.utils.mem_constants import MiB_bytes
from ....conftest import HfRunner
from ....utils import VLLM_PATH
@@ -106,7 +107,12 @@ def test_prm_models(
if current_platform.is_cpu():
pytest.skip("CPU only supports V1")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
if not FIXTURE_REWARD_RESULT.get(model):
pytest.skip(f"No available golden outputs for {model}.")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
@@ -70,7 +70,19 @@ 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
assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
prompt_token_ids = processed["prompt_token_ids"]
assert prompt_token_ids.count(video_token_id) == expected_tokens
hf_processor = processor.info.get_hf_processor()
expected_frame_wrappers = int(grid_thw[:, 0].sum())
assert (
prompt_token_ids.count(hf_processor.vision_start_token_id)
== expected_frame_wrappers
)
assert (
prompt_token_ids.count(hf_processor.vision_end_token_id)
== expected_frame_wrappers
)
@pytest.mark.parametrize("num_images", [1, 2])
+39
View File
@@ -74,6 +74,44 @@ def test_cosmos3_new_checkpoint_weights_mapper():
)
def test_cosmos3_modelopt_quantizer_weights_mapper():
"""ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
(``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
(it has no parameter for them) while keeping the latter."""
from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
# Native ModelOpt quantizer buffers are dropped.
assert (
mapper.apply_list(
[
"layers.0.self_attn.to_q.input_quantizer._amax",
"layers.0.self_attn.to_q.weight_quantizer._amax",
"layers.0.self_attn.to_q.weight_quantizer._scale",
"layers.0.mlp.down_proj.output_quantizer._amax",
]
)
== []
)
# The FP8 scale sidecars vLLM actually consumes are kept and remapped.
assert mapper.apply_list(
[
"layers.0.self_attn.to_q.weight",
"layers.0.self_attn.to_q.weight_scale",
"layers.0.self_attn.to_q.input_scale",
"layers.0.mlp.down_proj.input_scale",
]
) == [
"language_model.model.layers.0.self_attn.q_proj.weight",
"language_model.model.layers.0.self_attn.q_proj.weight_scale",
"language_model.model.layers.0.self_attn.q_proj.input_scale",
"language_model.model.layers.0.mlp.down_proj.input_scale",
]
def test_cosmos3_edge_checkpoint_weights_mapper():
from vllm.model_executor.models.cosmos3_edge import (
Cosmos3EdgeForConditionalGeneration,
@@ -132,6 +170,7 @@ 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",
+2 -8
View File
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
import pytest
from tests.quantization.utils import is_quant_method_supported
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
from vllm.platforms import current_platform
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
from ..utils import check_logprobs_close
@@ -70,13 +70,7 @@ def test_models(
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
if not (
current_platform.is_xpu()
or (
get_flash_attn_version() == 3
and current_platform.is_device_capability_family(90)
)
):
if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
pytest.skip(
f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
)
+7
View File
@@ -304,6 +304,13 @@ def test_pynvvideocodec_decoder_slot_retains_simple_decoder():
# ============================================================================
def test_cosmos3_edge_uses_qwen3_vl_video_backend():
backend = get_video_loader_backend_for_processor("Cosmos3EdgeVideoProcessor")
assert backend == "qwen3_vl"
assert isinstance(VIDEO_LOADER_REGISTRY.load(backend), Qwen3VLVideoBackend)
@pytest.mark.parametrize(
"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}
+6 -3
View File
@@ -93,9 +93,6 @@ 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")
@@ -105,9 +102,15 @@ 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:
+17
View File
@@ -67,6 +67,23 @@ def test_memory_profiling():
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
assert abs(non_torch_ratio - 1) <= 0.05
assert result.torch_peak_increase == 1024 * 1024 * 1024
expected_total_consumed = (256 + 512) * 1024 * 1024
total_consumed_ratio = result.total_consumed / expected_total_consumed
assert abs(total_consumed_ratio - 1) <= 0.05, (
f"total_consumed={result.total_consumed}, "
f"expected={expected_total_consumed}, "
f"ratio={total_consumed_ratio}"
)
expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
assert abs(non_kv_ratio - 1) <= 0.05, (
f"non_kv_cache_memory={result.non_kv_cache_memory}, "
f"expected={expected_non_kv}, "
f"ratio={non_kv_ratio}"
)
del weights
lib.cudaFree(handle1)
lib.cudaFree(handle2)
+64 -7
View File
@@ -21,6 +21,7 @@ from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import (
STR_DTYPE_TO_TORCH_DTYPE,
is_quantized_kv_cache,
is_torch_equal_or_newer,
set_random_seed,
)
@@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [
DEVICE_TYPE = current_platform.device_type
FP8_KV_CACHE_DTYPES = {
"fp8": torch.float8_e4m3fn,
"fp8_e4m3": torch.float8_e4m3fn,
}
# Remove flashinfer from the list if it's not available
try:
import flashinfer # noqa: F401
@@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache(
num_blocks: int,
common_attn_metadata: CommonAttentionMetadata,
randomize_blocks: bool = True,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
"""Create and prepopulate a KV cache with context data.
@@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache(
block_table = common_attn_metadata.block_table_tensor
slot_mapping = common_attn_metadata.slot_mapping
# For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
# the higher-precision context tensors quantizes them, mirroring runtime.
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
kv_cache = torch.zeros(
num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
num_blocks,
block_size,
num_kv_heads,
2 * head_size,
dtype=storage_dtype,
device=device,
)
kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
@@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache(
] * block_size + token_inter_block_offsets.to(device)
# Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
return kv_cache.transpose(1, 2).contiguous()
kv_cache = kv_cache.transpose(1, 2).contiguous()
if fp8_kv_cache:
kv_cache = kv_cache.view(torch.uint8)
return kv_cache
class MockAttentionLayer:
@@ -224,6 +246,7 @@ def run_attention_backend(
kv_cache: torch.Tensor,
attn_type: AttentionType = AttentionType.DECODER,
sliding_window: int | None = None,
kv_cache_dtype: str = "auto",
) -> torch.Tensor:
"""Run attention computation using the specified backend's AttentionImpl."""
@@ -291,13 +314,16 @@ def run_attention_backend(
alibi_slopes=None,
sliding_window=sliding_window,
attn_type=attn_type,
kv_cache_dtype="auto",
kv_cache_dtype=kv_cache_dtype,
)
# Create mock layer and output buffer
mock_layer = MockAttentionLayer(device)
output = torch.empty_like(query)
if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
query = query.to(current_platform.fp8_dtype())
# Run forward pass
# NOTE: The query, key, and value are already shaped correctly
# in the calling test function.
@@ -324,6 +350,7 @@ def _test_backend_correctness(
atol: float = 1e-2,
rtol: float = 1e-2,
tensor_parallel_size: int = 1,
kv_cache_dtype: str = "auto",
):
"""
Test that all backends produce similar outputs to a reference implementation
@@ -372,6 +399,7 @@ def _test_backend_correctness(
num_gpu_blocks=8192,
hf_config_override=hf_config_override,
)
vllm_config.cache_config.cache_dtype = kv_cache_dtype
device = torch.device(f"{DEVICE_TYPE}:0")
kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
@@ -392,6 +420,13 @@ def _test_backend_correctness(
block_size = vllm_config.cache_config.block_size
scale = 1.0 / (head_size**0.5)
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
if fp8_kv_cache:
query_fp8_dtype = current_platform.fp8_dtype()
kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
atol = max(atol, 6e-2)
rtol = max(rtol, 1e-1)
# 2. Generate data and compute SDPA reference output
all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
all_sdpa_outputs = []
@@ -407,10 +442,17 @@ def _test_backend_correctness(
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
if fp8_kv_cache:
q_ref = q.to(query_fp8_dtype).to(dtype)
k_ref = k_full.to(kv_fp8_dtype).to(dtype)
v_ref = v_full.to(kv_fp8_dtype).to(dtype)
else:
q_ref, k_ref, v_ref = q, k_full, v_full
# SDPA expects (N, H, L, D), so unsqueeze batch and permute
q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)
if num_q_heads != num_kv_heads:
assert num_q_heads % num_kv_heads == 0, (
@@ -471,6 +513,7 @@ def _test_backend_correctness(
num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
common_attn_metadata=common_attn_metadata,
randomize_blocks=True,
kv_cache_dtype=kv_cache_dtype,
)
# 4. Run vLLM backends and compare
@@ -488,6 +531,12 @@ def _test_backend_correctness(
else:
backend_cls = None
if is_quantized_kv_cache(kv_cache_dtype) and (
backend_cls is None
or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
):
continue
if backend_name == AttentionBackendEnum.FLASHINFER:
set_kv_cache_layout("HND")
reset_kv_cache_layout = True
@@ -521,6 +570,7 @@ def _test_backend_correctness(
kv_cache_for_backend,
sliding_window=sliding_window,
attn_type=attn_type,
kv_cache_dtype=kv_cache_dtype,
)
finally:
if reset_kv_cache_layout:
@@ -570,8 +620,13 @@ def _test_backend_correctness(
)
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
def test_causal_backend_correctness(
default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
default_vllm_config,
batch_spec_name: str,
model: str,
tensor_parallel_size: int,
kv_cache_dtype: str,
):
"""Test backend's correctness with causal attention."""
@@ -612,6 +667,7 @@ def test_causal_backend_correctness(
SMALL_BLOCK_BACKENDS,
causal_mask_mod,
tensor_parallel_size=tensor_parallel_size,
kv_cache_dtype=kv_cache_dtype,
)
# Fast FlexAttention needs to run with block_size=128
@@ -623,6 +679,7 @@ def test_causal_backend_correctness(
causal_mask_mod,
block_size=128,
tensor_parallel_size=tensor_parallel_size,
kv_cache_dtype=kv_cache_dtype,
)
+1
View File
@@ -119,6 +119,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
layer.kv_b_proj.quant_method = None
layer.is_aiter_triton_fp4_bmm_enabled = False
layer.is_aiter_triton_fp8_bmm_enabled = False
layer.dcp_q_replicate = False
layer.quant_config = None
layer.layer_name = "test"
+2
View File
@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
EncoderOnlyAttentionSpec,
FullAttentionSpec,
MambaSpec,
get_kv_quant_mode,
)
@@ -178,6 +179,7 @@ def create_standard_kv_cache_spec(
head_size=vllm_config.model_config.get_head_size(),
dtype=vllm_config.model_config.dtype,
sliding_window=vllm_config.model_config.get_sliding_window(),
kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
)
+123
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
import hashlib
import importlib
from collections.abc import Callable
@@ -1477,6 +1478,128 @@ def test_get_max_concurrency_for_kv_cache_config():
assert num_tokens == max_concurrency_hybrid_model * max_model_len
assert max_concurrency == max_concurrency_hybrid_model
# Unequal group sizes in the standard layout: each group's pages cost
# whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
# same as the equal-hybrid case above, regardless of the second group
# holding only 2 layers.
kv_cache_config_unequal_groups = KVCacheConfig(
num_blocks=1153 * 3,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
],
)
assert (
get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_unequal_groups
)
== 3
)
# UniformTypeKVCacheSpecs group (worker config shape): the aggregated
# spec's memory/page ratio equals a single layer's page count, so the
# group needs 1024 blocks and the request 1153 in total. The previous
# formula normalized both groups' memory by the first group's page size,
# reporting 3459/1057 = 3.27 here instead of 3 — and a different value
# again for the scheduler-config shape below.
uniform_full_spec = UniformTypeKVCacheSpecs(
block_size=full_attention_spec.block_size,
kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
)
kv_cache_config_uniform_group = KVCacheConfig(
num_blocks=1153 * 3,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
],
)
assert (
get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_uniform_group
)
== 3
)
# Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
# uniform-type group's spec with a representative per-layer spec.
# Capacity must not change between the two shapes (the engine computes
# on the scheduler config, the worker loop on the worker config).
kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
[copy.deepcopy(kv_cache_config_uniform_group)]
)
assert get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_scheduler_shape
) == get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_uniform_group
)
def test_get_max_concurrency_packed_kv_cache_config():
from vllm.v1.core.kv_cache_utils import (
_get_kv_cache_config_packed,
_use_packed_kv_cache_config,
)
model_config = ModelConfig(
"Qwen/Qwen1.5-7B",
runner="generate",
dtype="float16",
max_model_len=16384,
)
scheduler_config = SchedulerConfig(
max_num_batched_tokens=1024,
enable_chunked_prefill=True,
max_model_len=model_config.max_model_len,
is_encoder_decoder=model_config.is_encoder_decoder,
async_scheduling=False,
)
vllm_config = VllmConfig(
model_config=model_config,
scheduler_config=scheduler_config,
)
# All-UniformTypeKVCacheSpecs groups select the packed layout.
mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
swa_specs = {
f"layer_{i}": SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
sliding_window=128,
)
for i in range(4, 6)
}
kv_cache_groups = [
KVCacheGroupSpec(
list(mla_specs),
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
),
KVCacheGroupSpec(
list(swa_specs),
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
),
]
assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
vllm_config, kv_cache_groups, 2 * GiB_bytes
)
assert num_blocks > 0
kv_cache_config_packed = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=kv_cache_groups,
)
# Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
# the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
# previous formula normalized by the first group's page size and gave
# 1061 blocks per request instead of 1097.
assert get_max_concurrency_for_kv_cache_config(
vllm_config, kv_cache_config_packed
) == num_blocks / (1024 + 73)
def test_allocate_with_lookahead():
"""Verify that lookahead tokens correctly affect block allocation"""
@@ -83,6 +83,11 @@ 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
@@ -154,12 +159,14 @@ 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 \
@@ -208,12 +215,18 @@ 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,6 +12,9 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
)
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
from 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,
@@ -19,6 +22,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
OffloadingConnectorScheduler,
RequestOffloadState,
is_store_reachable_swa_chunk,
)
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_cache_interface import (
@@ -108,6 +112,40 @@ def test_last_block_offloaded_at_request_finish(
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_abort_queued_request_does_not_build_store_job(
request_runner, async_scheduling: bool
):
"""Aborting a never-scheduled request must not store unallocated KV."""
block_size = 4
runner = request_runner(
block_size=block_size,
num_gpu_blocks=8,
async_scheduling=async_scheduling,
)
runner.new_request(token_ids=[0] * (block_size * 4))
runner.scheduler.schedule()
runner.new_request(token_ids=[1] * (block_size * 4))
queued_req_id = str(runner.req_id)
assert any(
request.request_id == queued_req_id for request in runner.scheduler.waiting
)
runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
req_status = runner.connector_scheduler._req_status[queued_req_id]
assert all(group_state.offload_keys for group_state in req_status.group_states)
assert all(not group_state.block_ids for group_state in req_status.group_states)
scheduler_output = runner.scheduler.schedule()
metadata = scheduler_output.kv_connector_metadata
assert isinstance(metadata, OffloadingConnectorMetadata)
assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
assert queued_req_id not in runner.connector_scheduler._req_status
def test_scheduler_reports_lookup_sync_delay(request_runner):
runner = request_runner(
block_size=4,
@@ -126,6 +164,55 @@ def test_scheduler_reports_lookup_sync_delay(request_runner):
assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0
@pytest.mark.parametrize(
(
"absolute_chunk_index",
"storable_chunk_count",
"alignment_chunk_count",
"sliding_window_chunks",
"is_eagle_group",
"expected",
),
[
# Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61.
(61, 64, 64, 2, False, False),
(62, 64, 64, 2, False, True),
(60, 64, 64, 2, True, False),
(61, 64, 64, 2, True, True),
# Partial 48-of-64 segment: the reachable tail ends at chunk 47.
(45, 48, 64, 2, False, False),
(46, 48, 64, 2, False, True),
(44, 48, 64, 2, True, False),
(45, 48, 64, 2, True, True),
# A later partial segment uses its own actual end (chunks 64-79).
(76, 80, 64, 3, False, False),
(77, 80, 64, 3, False, True),
# No alignment means no store-pruning optimization.
(0, 1, None, None, False, True),
# A tail at least as large as the segment keeps every chunk.
(0, 2, 64, 2, False, True),
],
)
def test_is_store_reachable_swa_chunk(
absolute_chunk_index: int,
storable_chunk_count: int,
alignment_chunk_count: int | None,
sliding_window_chunks: int | None,
is_eagle_group: bool,
expected: bool,
):
assert (
is_store_reachable_swa_chunk(
absolute_chunk_index,
storable_chunk_count,
alignment_chunk_count,
sliding_window_chunks,
is_eagle_group,
)
is expected
)
def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
"""A deferred lookup reports its async delay once it resolves."""
runner = request_runner(
@@ -125,8 +125,7 @@ def test_register_kv_caches(backend):
own dedicated tensors.
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
kv_caches, which automatically applies
_update_hybrid_attention_mamba_layout for hybrid models.
the raw per-layer kv_caches registered by the connector.
Verifies that the canonicalized CanonicalKVCaches has the correct
block tensors, tensor_idx references, and page sizes across all groups.
-38
View File
@@ -5,13 +5,11 @@ import torch
from torch import Generator
from 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
@@ -1045,39 +1043,3 @@ 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))
@@ -0,0 +1,128 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the EAGLE speculator's draft attention metadata builder.
These tests guard the regression where ``_build_draft_attn_metadata`` did
not populate ``seq_lens_cpu_upper_bound`` on the per-step
``CommonAttentionMetadata``. Several downstream attention backends and
helpers (``split_decodes_prefills_and_extends``, the MLA indexer,
flex-attention, cross-attention) assert this field is non-None, so
omitting it caused crashes at the start of draft decode for certain
backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode):
AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
"""
from types import SimpleNamespace
from unittest.mock import patch
import torch
from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator
from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
def _make_fake_speculator(
*,
max_num_reqs: int = 8,
max_num_tokens: int = 16,
max_model_len: int = 1024,
draft_max_seq_len: int = 1024,
) -> SimpleNamespace:
"""Build a fake EagleSpeculator with just the attributes used by
``_build_draft_attn_metadata``. We deliberately avoid constructing a
real ``EagleSpeculator`` because that requires a full ``VllmConfig``
and a draft model.
"""
fake_input_buffers = SimpleNamespace(
query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
)
fake_block_tables = SimpleNamespace(
input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
)
return SimpleNamespace(
arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
block_tables=fake_block_tables,
input_buffers=fake_input_buffers,
attn_groups=[],
kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
max_model_len=max_model_len,
draft_max_seq_len=draft_max_seq_len,
)
def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step):
captured: dict[str, object] = {}
def fake_build_attn_metadata(**kwargs):
captured.update(kwargs)
return {}
with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):
EagleSpeculator._build_draft_attn_metadata(
fake, # type: ignore[arg-type]
num_reqs=num_reqs,
num_reqs_padded=num_reqs_padded,
num_tokens_padded=num_tokens_padded,
seq_lens_cpu_upper_bound=base,
step=step,
)
return captured
def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():
"""The fix: every per-step ``CommonAttentionMetadata`` carries a non-None
``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus
the current draft-step offset. Padded entries are zeroed (matching the
main model runner's convention)."""
fake = _make_fake_speculator()
base = torch.tensor([100, 200, 300, 0], dtype=torch.int32)
captured = _run_build(
fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2
)
bound = captured["seq_lens_cpu_upper_bound"]
assert isinstance(bound, torch.Tensor), (
"seq_lens_cpu_upper_bound must be a tensor, not None"
)
assert bound.shape == (4,), (
f"expected shape (num_reqs_padded=4,), got {bound.shape}"
)
assert bound.device.type == "cpu"
assert bound.dtype == torch.int32
# base[:num_reqs] + step, padded tail zeroed.
assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32))
def test_build_draft_attn_metadata_handles_zero_unpadded_reqs():
"""Edge case: when ``num_reqs == 0`` the upper-bound tensor must
still be a valid all-zero tensor of length ``num_reqs_padded``."""
fake = _make_fake_speculator()
base = torch.zeros(2, dtype=torch.int32)
captured = _run_build(
fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1
)
bound = captured["seq_lens_cpu_upper_bound"]
assert isinstance(bound, torch.Tensor)
assert bound.shape == (2,)
assert torch.equal(bound, torch.zeros(2, dtype=torch.int32))
def test_build_draft_attn_metadata_clamps_to_max_model_len():
"""The per-request upper bound (target bound + step) is clamped to the
model length so it never exceeds the allocated KV range."""
fake = _make_fake_speculator(max_model_len=1024)
base = torch.tensor([1023, 500], dtype=torch.int32)
captured = _run_build(
fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3
)
bound = captured["seq_lens_cpu_upper_bound"]
# 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
+20 -1
View File
@@ -36,6 +36,22 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
]
def write_coverage_objects(extensions: list[RustExtension], output: Path) -> None:
artifacts = []
for extension in extensions:
for target in sorted(set(extension.target.values())):
target_path = ROOT_DIR.joinpath(*target.split("."))
if extension.binding == Binding.Exec:
matches = [target_path]
else:
matches = sorted(target_path.parent.glob(f"{target_path.name}*.so"))
if len(matches) != 1 or not matches[0].is_file():
raise RuntimeError(f"unable to locate Rust artifact for {target}")
artifacts.append(matches[0].relative_to(ROOT_DIR).as_posix())
output.write_text("\n".join(artifacts) + "\n")
def rust_py_extension_module_names() -> list[str]:
module_names = []
for extension in rust_extensions():
@@ -52,12 +68,15 @@ def rust_py_extension_module_names() -> list[str]:
def build_binary(build_rust_args: list[str]) -> None:
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=rust_extensions(optional=False),
rust_extensions=extensions,
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:
@@ -1083,7 +1083,8 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
return {}
# Analyze the functions to determine FA3/FA4-specific features
fa3_supports_fp8 = True
fa3_supports_fp8 = False
fa4_supports_fp8 = False
fa3_supports_sinks = False
fa4_supports_sinks = False
fa3_compute_cap: str | None = None
@@ -1093,6 +1094,44 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
if not isinstance(node, ast.FunctionDef):
continue
# Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version.
# Accept both equality checks and membership checks such as `in (3, 4)`.
if node.name == "flash_attn_supports_kv_cache_dtype":
for n in ast.walk(node):
if not (
isinstance(n, ast.Compare)
and len(n.ops) == 1
and len(n.comparators) == 1
):
continue
is_version_compare = (
isinstance(n.left, ast.Name) and n.left.id == "fa_version"
) or (
isinstance(n.left, ast.Call)
and isinstance(n.left.func, ast.Name)
and n.left.func.id == "get_flash_attn_version"
)
if not is_version_compare:
continue
versions: list[Any] = []
comparator = n.comparators[0]
if isinstance(n.ops[0], ast.Eq) and isinstance(
comparator, ast.Constant
):
versions = [comparator.value]
elif isinstance(n.ops[0], ast.In) and isinstance(
comparator, (ast.Tuple, ast.List, ast.Set)
):
versions = [
elt.value
for elt in comparator.elts
if isinstance(elt, ast.Constant)
]
fa3_supports_fp8 |= 3 in versions
fa4_supports_fp8 |= 4 in versions
# Check flash_attn_supports_sinks - looks for `fa_version == 3/4`
# or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)
if node.name == "flash_attn_supports_sinks":
@@ -1199,7 +1238,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
},
"fa4": {
"compute_capability": fa4_compute_cap,
"supports_fp8": False,
"supports_fp8": fa4_supports_fp8,
"supports_sink": fa4_supports_sinks,
},
}
@@ -1255,6 +1294,22 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
# ---------------------------------------------------------------------------
def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str:
"""Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string.
The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``,
but actual support varies by FA version (e.g. FA2 has no fp8 support), so each
variant's row is normalized to match its own capability.
"""
fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"}
dtypes = kv_cache_dtypes.split(", ")
base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes]
if supports_fp8:
advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes]
return ", ".join(base_dtypes + advertised_fp8_dtypes)
return ", ".join(base_dtypes)
def _expand_flash_attn_variants(
all_backends: list[dict[str, Any]],
fa_features: dict[str, dict[str, Any]],
@@ -1275,6 +1330,9 @@ def _expand_flash_attn_variants(
fa2["_sort_key"] = "FLASH_ATTN"
fa2["_sort_order"] = 0
fa2["supports_sink"] = fa_features["fa2"]["supports_sink"]
fa2["kv_cache_dtypes"] = _apply_fp8_support(
backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"]
)
# Create FA3 entry (uses parsed compute_capability from fa_utils)
fa3 = backend.copy()
@@ -1284,11 +1342,9 @@ def _expand_flash_attn_variants(
if fa_features["fa3"]["compute_capability"]:
fa3["compute_capability"] = fa_features["fa3"]["compute_capability"]
fa3["supports_sink"] = fa_features["fa3"]["supports_sink"]
if fa_features["fa3"]["supports_fp8"]:
base_dtypes = backend["kv_cache_dtypes"].split(", ")
fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"]
new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes]
fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes)
fa3["kv_cache_dtypes"] = _apply_fp8_support(
backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"]
)
expanded.append(fa2)
expanded.append(fa3)
@@ -1302,6 +1358,9 @@ def _expand_flash_attn_variants(
if fa_features["fa4"].get("compute_capability"):
fa4["compute_capability"] = fa_features["fa4"]["compute_capability"]
fa4["supports_sink"] = fa_features["fa4"]["supports_sink"]
fa4["kv_cache_dtypes"] = _apply_fp8_support(
backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"]
)
expanded.append(fa4)
return expanded

Some files were not shown because too many files have changed in this diff Show More