Compare commits

..
Author SHA1 Message Date
Nick Hill 5ea7cac55b revert inadvertent change to .pre-commit-config.yaml
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-22 16:12:01 +01:00
Nick HillandClaude Opus 4.8 be3476447f [Doc] register_kv_caches: views are authoritative, not storage nbytes
Two connectors (NIXL packed registration, SimpleCPUOffload) derived KV
geometry from untyped_storage().nbytes() and broke under the extensible
KV cache, where storages span reserved capacity. Document the contract
so out-of-tree connectors avoid the same pattern.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-22 15:59:16 +01:00
Nick HillandClaude Opus 4.8 f1473092c4 [Core] Make SimpleCPUOffloadConnector geometry extensible-KV-cache aware
The worker derived per-block sizes from storage.nbytes() // num_blocks
and viewed whole storages as (num_blocks, block_bytes). With the
extensible KV cache, registration-view storages span the reserved
capacity while num_blocks is the committed count, so block strides were
wrong and tail rows pointed into unmapped virtual memory. Derive the
per-block size from the registration views' committed extent instead
(summing a layer's state tensors for Mamba), keep the bounded-storage
size for packed layouts, and slice each segment to its committed block
prefix. Byte-identical behavior when committed == capacity.

No sleep/wake override is needed for this connector: it holds VA-stable
views plus its own pinned CPU pool (default no-op hooks are correct,
like OffloadingConnector).

Validated on GPU: cold-vs-CPU-reload greedy outputs match 5/5 with the
extensible cache (and 5/5 baseline), incl. kv_cache_memory_bytes mode.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-22 15:59:16 +01:00
Nick HillandClaude Opus 4.8 4dcbae8670 [Core] Tighten packed extensible KV cache invariants
- Assert one packed row per logical block at allocation: the reshape
  view construction, NIXL's packed registration math, and the packed
  storage bounding all rely on bytes_per_block == block_stride, so make
  the constraint explicit at the source instead of implicit in three
  places.
- Make kv_cache_config a required argument of
  narrow_kv_caches_to_num_blocks so future callers cannot silently skip
  the packed storage bounding.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-22 14:40:35 +01:00
zjy0516 65dac3a770 Fix extensible KV cache lint errors
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-07-22 08:30:35 +00:00
zjy0516 0ba2500ef0 Fix extensible KV cache connector registrations
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-07-22 08:22:15 +00:00
Nick HillandClaude Opus 4.8 ef576befd2 [Core] Defragment extensible KV cache before KV-transfer registration
Root-caused via a minimal 2-process NIXL repro on GB200: UCX transfers
succeed for VMM-backed regions mapped as a single physical allocation
but fail (remote-endpoint invalidation, NIXL_ERR_REMOTE_DISCONNECT) for
regions spanning multiple incrementally-committed cuMemCreate handles -
exactly what the 1-block -> warmup-prefix -> final commit sequence
produces. Before deferred connector registration, extend_kv_cache now
releases the warmup-time chunks and re-commits each segment prefix as
one physical allocation (contents at that point are only warmup garbage;
no requests have been served).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:57 +01:00
Nick HillandClaude Opus 4.8 35e4a36107 [Core] Allocate shareable (IPC-exportable) VMM memory for KV connectors
NIXL 1P1D validation on GB200 showed the decode side invalidating the
prefill agent on the very first KV pull: intra-node UCX uses CUDA IPC,
and cuMemCreate allocations are only exportable to other processes when
created with requestedHandleTypes=POSIX_FILE_DESCRIPTOR, which the alloc
props did not set. When a KV connector is configured, request the POSIX
FD handle type alongside the GPU-direct-RDMA-capable flag (renamed
rdma_capable -> shareable), keeping the fallback-with-warning where such
allocations are unavailable.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 da5803d46e [ROCm] Add HIP VMM driver backend for the extensible KV cache
HIP mirrors the CUDA driver's VMM API (hipMemAddressReserve /
hipMemCreate / hipMemMap / hipMemSetAccess / ...) with identical call
signatures, struct layouts, and constants, so the backend only supplies
the library, symbol names, error-string convention, and implicit-context
handling; DLPack views use kDLROCM. The worker-side probe gates actual
use, so unsupported ROCm stacks still fall back gracefully.

Untested on AMD hardware so far.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 75ddfaf909 [ModelRunner V2] Support KV connectors with extensible KV cache
Connectors must not register KV cache memory (e.g. RDMA memory regions)
before the final size is physically committed. With the V2 runner:

- Defer ensure_kv_transfer_initialized + connector creation/registration
  from initialize_from_config to extend_kv_cache, which now receives the
  final (pristine, post-warmup-sizing) per-rank kv_cache_config through
  the executor RPC instead of a bare block count.
- Register views narrowed along each layer's block dim to the committed
  block count (narrow_kv_caches_to_num_blocks), so connectors only see
  physically backed memory. Committed blocks form a prefix of each layout
  segment, so a narrow covers exactly the committed bytes (e.g. NIXL's
  separate K/V regions land on the two committed prefixes).
- Allocate physical chunks with the gpuDirectRDMACapable flag when a KV
  connector is configured, falling back with a warning where GDR-capable
  VMM allocations are unavailable.
- Warmup runs against the no-op connector (it is disabled during warmup
  anyway); V1 runner + connectors + extensible remains rejected.

Validated e2e on GPU with ExampleConnector (shared-storage): deferred
registration, then a real external-cache save + hit through the narrowed
registered views with identical output.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 f36fe52add [Core] Support sleep mode with extensible KV cache
The VMM-backed KV cache lives outside the torch/CuMem allocators, so
sleep now discards its physical pages directly (release_physical: unmap
and release handles, keeping the VA reservation so tensor views and
captured graphs stay pointer-valid) and wake_up recommits the same block
count with freshly zeroed pages, matching the CuMem discard semantics.

Validated e2e on GPU: sleep(level=1) frees weights + KV physical memory
(0.17 GiB residual), wake_up restores and generation output is identical.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 391d918d4d [ModelRunner V2] Support packed KV cache layouts with extensible KV cache
The packed (block_stride) backing is block-major by construction: block b
occupies the b-th block_stride-byte row, holding every layer's page. Back
it with one shared single-segment ExtensibleTensor so a prefix of blocks
commits naturally, instead of rejecting the layout.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 fca040885b [Core] Extensible KV cache: VMM driver probe/fallback, manual size support
- Extract the driver ctypes bindings into vllm/utils/vmm_driver.py behind
  a small VmmDriver interface (CUDA implementation; struct layouts and
  call signatures are shared with HIP for a future ROCm backend).
- Probe VMM support on the workers (driver loads, VA reservation works)
  and fall back to standard KV cache allocation with a warning instead of
  failing on platforms without VMM (e.g. WSL2, non-GPU workers).
- Support kv_cache_memory_bytes: the requested size is committed as-is
  after warmup (single sizing pass), still avoiding warmup-time OOM.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:56 +01:00
Nick HillandClaude Opus 4.8 a7fd4c7482 [Core] Unify extensible KV cache state on ExtensibleKVCacheBuffers
Move the grow-only buffer collection from the V2 attn_utils module to
vllm/utils/extensible_tensor.py and use it from the V1 runner as well
(replacing the _extensible_kv_cache_* attribute trio). Both runners now
expose the same `extensible_kv_buffers` attribute, so worker-level
features (memory measurement, sleep, connector deferral) can treat the
runners uniformly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:55 +01:00
5131691063 [ModelRunner V2] Support extensible KV cache; size KV from measured warmup memory
Adapt #47363's extensible KV cache to the V2 model runner:

- V2 allocation (gpu/attn_utils.py): reserve each KV cache tensor's full
  virtual range with ExtensibleTensor, committing a per-segment block
  prefix. Segment counts are derived from each backend's physical layout
  (block dim / stride order), with hybrid attention+Mamba forced
  block-major to match the re-strided layout.
- V2 warmup writes to real block IDs (a contiguous prefix starting at 1),
  unlike V1's all-zero dummy block tables, so warmup_kernels and
  run_mixed_prefill_decode_warmup now commit exactly the block prefix
  they touch via a new ensure_kv_cache_blocks() hook.
- Post-warmup measurement: instead of only the CUDA graph pool bytes,
  the worker measures actual non-KV memory in use after ALL warmup
  (retained worst-case activation segments, NCCL buffers, CUDA graphs)
  and reports the excess over the profiling estimate
  (CompilationTimes.cuda_graph renamed to warmup_memory). The engine's
  second sizing pass then commits a KV cache that leaves room for the
  real runtime working set - including the worst-case spec-decode
  logits all-gather that memory profiling misses today.
- Gate extensible mode against KV connectors and sleep mode; drop it
  from the V2-unsupported feature list.
- Extend tests/v1/worker/test_extensible_kv_cache.py with V2 coverage
  (segment inference, staged prefix commits, hybrid re-stride layout).

Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
2026-07-20 14:42:55 +01:00
Nick HillandZhuohan Li 80e00e5ac6 [Core] Pick extensible KV cache memory from #47363
Reserve the KV cache address range with CUDA virtual memory, commit a
minimal prefix before CUDA graph capture, measure real post-capture
memory usage, then commit the final KV cache size with stable tensor
addresses. Opt-in via --enable-extensible-kv-cache.

Squashed pick of vllm-project/vllm#47363.

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

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