Compare commits

..
Author SHA1 Message Date
Shengqi ChenandGitHub a02155c787 Merge branch 'main' into cuda-arch-fixup 2026-07-21 09:33:48 +08:00
Shengqi ChenandGitHub 319db65b68 Merge branch 'main' into cuda-arch-fixup 2026-07-09 01:55:33 +08:00
Shengqi ChenandGitHub 83a7669827 Merge branch 'main' into cuda-arch-fixup
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-07 00:13:53 +08:00
Shengqi ChenandGitHub 401bed48ad Merge branch 'main' into cuda-arch-fixup
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-03 15:54:56 +08:00
Shengqi ChenandGitHub d492d1e697 Merge branch 'main' into cuda-arch-fixup 2026-06-30 21:30:29 +08:00
Shengqi Chen f37e113590 [Build] Apply ruff format to CUDA arch regex
Keep the compiled CUDA arch regex on one line to match ruff-format output.

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 20:37:59 +08:00
Shengqi Chen a10e369f06 [Build] Address CUDA arch review comments
Fix CUDA arch warning tests so they do not depend on importing the stable libtorch extension, which keeps the warning coverage active in lightweight CI environments and avoids mypy treating a fixture value as a base class.

Use regex for the compiled-arch parser, apply formatter output, and make the CUTLASS grouped GEMM Python support query fall back to false when the op is unavailable or unimplemented in the current build.

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 20:03:28 +08:00
Shengqi Chen aa3f2efe42 [Temp] Cherry-pick #47139 to fix build
Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:38:28 +08:00
Shengqi ChenandCodex a32d1bff95 [Doc] Document CUDA wheel architecture coverage
Explain that pre-built CUDA wheels use the architecture lists selected by the release and build pipelines, which may be narrower than the full set vLLM can build from source.

Call out CUDA 12.9 architecture-specific wheel coverage, CUDA 13 family-specific targets, and the no-kernel-image error users may see when a wheel does not cover their GPU.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex ee47a21fcd [Build] Warn on uncovered CUDA device architectures
Expose the compiled CUDA arch list from the stable extension and check visible CUDA devices against it during CUDA platform startup. The runtime check distinguishes exact architecture targets from CUDA 13 family targets so users get an early warning before hitting missing kernel images.

The startup warning is limited to the NVML-backed CUDA platform path to preserve the existing no-CUDA-init import behavior for non-NVML environments.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex 0f471a3088 [Build] Scope stable CUDA kernel feature macros
Keep optional stable CUDA kernel feature macros on the source files that consume them instead of adding them to VLLM_GPU_FLAGS. This avoids perturbing unrelated compile commands and invalidating more cache entries when optional kernel families change.

Also align CUTLASS grouped MoE support with the SM10x/SM11x family so Thor works under both CUDA 12 SM101 and CUDA 13 SM110 reporting, and remove the stale ENABLE_CUTLASS_MLA definition left after the old CUTLASS MLA path was deleted.

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
103 changed files with 1255 additions and 1983 deletions
@@ -29,7 +29,6 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
log_section() {
echo "--- $*"
@@ -107,18 +106,6 @@ install_cargo_nextest() {
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
}
install_cargo_llvm_cov() {
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
local toolchain
toolchain="$(rust_toolchain)"
rustup component add --toolchain "$toolchain" llvm-tools-preview
cargo binstall \
--no-confirm \
--force \
--secure \
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
}
install_uv() {
log_section "Installing uv ${UV_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf \
@@ -189,41 +176,14 @@ run_tests() {
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest with Rust coverage"
mkdir -p artifacts
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
cargo llvm-cov clean \
--manifest-path rust/Cargo.toml \
--profraw-only
set +e
cargo llvm-cov nextest \
log_section "Running cargo nextest"
cargo nextest run \
--manifest-path rust/Cargo.toml \
--workspace \
--all-features \
--locked \
--no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
--no-fail-fast
}
install_protoc
-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"
}
-30
View File
@@ -8,11 +8,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/benchmarks/
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
@@ -28,7 +23,6 @@ steps:
- tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
@@ -49,11 +43,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
- vllm/v1/engine/
@@ -65,7 +54,6 @@ steps:
# - tests/entrypoints/serve/dev/test_sleep.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
@@ -84,16 +72,10 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
@@ -104,17 +86,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
@@ -125,11 +101,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/distributed/
- vllm/engine/
- vllm/executor/
@@ -140,7 +111,6 @@ steps:
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -26,7 +26,5 @@ steps:
- rust-toolchain.toml
- .buildkite/test_areas/rust_frontend_cargo.yaml
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
commands:
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
-1
View File
@@ -47,7 +47,6 @@
# Rust Frontend
/rust/ @BugenZhao @njhill
/rust/src/bench @esmeetu
/build_rust.sh @BugenZhao @njhill
/rust-toolchain.toml @BugenZhao @njhill
/.buildkite/test_areas/rust* @BugenZhao @njhill
-1
View File
@@ -257,4 +257,3 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
+86 -42
View File
@@ -227,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
cuda_archs_loose_intersection(CUDA_ARCHS
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
set(VLLM_COMPILED_CUDA_ARCHS)
foreach(_ARCH ${CUDA_ARCHS})
set(_COMPILED_ARCH "${_ARCH}")
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
if(_ARCH MATCHES "^(10|11|12)\\.0$")
set(_COMPILED_ARCH "${_ARCH}f")
endif()
elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
set(_COMPILED_ARCH "${_ARCH}a")
endif()
endif()
list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
endforeach()
list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
else()
#
# For other GPU targets override the GPU architectures detected by cmake/torch
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
#
# Shared entry sources are part of the base extension source list, but some
# optional kernel families below append source-local feature macros to them.
set(VLLM_STABLE_TORCH_BINDINGS_SRC
"csrc/libtorch_stable/torch_bindings.cpp")
set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
set(SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
set(NVFP4_QUANT_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
set(NVFP4_SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
"csrc/libtorch_stable/activation_kernels.cu"
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/sampler.cu"
"csrc/libtorch_stable/topk.cu"
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"${CACHE_KERNELS_SRC}"
"csrc/libtorch_stable/cache_kernels_fused.cu"
"csrc/libtorch_stable/custom_all_reduce.cu"
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(COOPERATIVE_TOPK_ARCHS)
list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
endif()
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
"${SCALED_MM_ENTRY_SRC}"
"${NVFP4_QUANT_ENTRY_SRC}"
"${NVFP4_SCALED_MM_ENTRY_SRC}"
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")
#
# Machete kernels
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
CUDA_ARCHS "${CUDA_ARCHS}")
if(COOPERATIVE_TOPK_ARCHS)
set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cooperative_topk.cu")
"${COOPERATIVE_TOPK_SRC}")
set_gencode_flags_for_srcs(
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
SRCS "${COOPERATIVE_TOPK_SRC}"
CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Only build Marlin kernels if we are building for at least some compatible archs.
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM120_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_C2X_SRCS}"
CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_C2X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
else()
if (SCALED_MM_3X_ARCHS)
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
# The implementation is named sm100 historically, but it is built for the
# SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
# SM110. Keep the compile-time macro aligned with runtime dispatch.
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
#
# SM12x FP4 kernels. These share some generic NVFP4 quantization entry
# sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
# per-source flags, so shared files accumulate both SM12x and SM10x/11x
# gencodes when both families are requested.
# Shared FP4 implementation sources live next to the FP4 arch logic because
# both SM12x and SM10x/11x append family-specific gencodes and feature macros
# to them.
set(FP4_SHARED_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
set(FP4_SM120_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
)
set_gencode_flags_for_srcs(
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM120=1)
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
else()
message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
set(FP4_SM100_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
message(STATUS
"Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
else()
message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${CUTLASS_MLA_SRCS}"
CUDA_ARCHS "${MLA_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
# Add MLA-specific include directories only to MLA source files
set_source_files_properties(${CUTLASS_MLA_SRCS}
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# Needed to use cuda/hip APIs from C-shim
if(VLLM_GPU_LANG STREQUAL "CUDA")
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
if(COOPERATIVE_TOPK_ARCHS)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Needed by CUTLASS kernels
target_compile_definitions(_C_stable_libtorch PRIVATE
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+1 -1
View File
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
- Tool calling and reasoning parsers
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
- Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, 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:
-37
View File
@@ -8,8 +8,6 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
CARGO_LLVM_COV_VERSION="0.8.7"
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -32,39 +30,4 @@ else
PROFILE_ARG="--release"
fi
rm -rf "$COVERAGE_TOOLS_DIR"
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
# Keep instrumented intermediates isolated when local builds switch modes.
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
cargo +"$TOOLCHAIN" install \
--locked \
--version "$CARGO_LLVM_COV_VERSION" \
cargo-llvm-cov
eval "$(
cargo +"$TOOLCHAIN" llvm-cov show-env \
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
--sh
)"
# Build scripts and proc macros can run during compilation. Their profiles
# are unrelated to runtime coverage and would otherwise pollute the tree.
export LLVM_PROFILE_FILE=/dev/null
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
--print target-libdir)")/bin"
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
fi
+22
View File
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
endif()
endmacro()
#
# For a list of source files append preprocessor definitions to file-specific
# compile options. Use this for optional kernel feature macros so toggling one
# kernel family does not perturb the compile command for every source in the
# extension target.
#
macro(set_compile_definitions_for_srcs)
set(options)
set(oneValueArgs)
set(multiValueArgs SRCS DEFINITIONS)
cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
"${multiValueArgs}" ${ARGN})
foreach(_DEF ${arg_DEFINITIONS})
set_property(
SOURCE ${arg_SRCS}
APPEND PROPERTY
COMPILE_DEFINITIONS "${_DEF}"
)
endforeach()
endmacro()
#
# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
# `<major>.<minor>[letter]` compute the "loose intersection" with the
-13
View File
@@ -10,16 +10,3 @@ fixes:
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
# Map Rust sources built in the E2E image and on Buildkite agents.
- "/workspace/rust/::rust/"
- "/var/lib/buildkite-agent/.*/rust/::rust/"
flags:
rust-unit:
paths:
- rust/
carryforward: false
rust-e2e:
paths:
- rust/
carryforward: false
+1 -3
View File
@@ -102,9 +102,7 @@ class TileGemm82 {
kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto fp32_b_regs = load_b_pair_vec(curr_b);
auto fp32_b_0_reg = fp32_b_regs.first;
auto fp32_b_1_reg = fp32_b_regs.second;
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
+2
View File
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
torch::stable::Tensor const& perm);
#ifndef USE_ROCM
std::string get_compiled_cuda_archs();
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
#endif
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& a_tensors,
torch::stable::Tensor const& b_tensors,
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
std::optional<torch::stable::Tensor> const& bias);
#endif
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
void get_cutlass_moe_mm_data_caller(
const torch::stable::Tensor& topk_ids,
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
// CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
// actual cutlass_moe_mm dispatch compiled into this file.
// or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
// SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
// implementation. Only report archs that have an actual cutlass_moe_mm
// dispatch compiled into this file.
#if defined CUDA_VERSION
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
return CUDA_VERSION >= 12080;
}
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& c_strides, bool per_act_token,
bool per_out_ch) {
int32_t version_num = get_sm_version_num();
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
if (version_num >= 100 && version_num < 110) {
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
// Keep runtime dispatch aligned with the CMake arch list and support query:
// CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
if (version_num >= 100 && version_num < 120) {
cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
expert_offsets, problem_sizes, a_strides, b_strides,
c_strides, per_act_token, per_out_ch);
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
problem_sizes2, input_permutation,
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
const bool swap_ab) {
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
problem_sizes2, expert_num_tokens,
+6
View File
@@ -4,6 +4,10 @@
#include <torch/csrc/stable/library.h>
#ifndef USE_ROCM
std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
#endif
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
// Note: We register under namespace "_C" so ops are accessible as
// torch.ops._C.<op_name> for compatibility with existing code.
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.def("get_compiled_cuda_archs() -> str");
#ifndef USE_ROCM
@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
#ifndef USE_ROCM
ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
ops.impl("cutlass_scaled_mm_supports_fp8",
TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
ops.impl("cutlass_group_gemm_supported",
-11
View File
@@ -294,9 +294,6 @@ FROM base AS rust-build
ARG BUILD_OS
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
# Temporary default for the initial CI validation. Set this back to 0 when
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
ARG VLLM_RUST_COVERAGE=1
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -905,14 +902,6 @@ COPY ./vllm/collect_env.py .
# note that this uses vllm installed by `pip`
FROM vllm-base AS test
COPY --from=rust-build \
/workspace/rust-coverage-tools/ \
/opt/vllm-rust-coverage/
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
ENV LLVM_PROFILE_FILE=/dev/null
ADD . /vllm-workspace/
ARG PYTHON_VERSION
-3
View File
@@ -46,9 +46,6 @@
"TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
},
"VLLM_RUST_COVERAGE": {
"default": "1"
},
"MAX_JOBS": {
"default": "2"
},
@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
```
!!! warning "CUDA architecture coverage"
Pre-built CUDA wheels are compiled for the CUDA architectures selected by
vLLM's release and build pipelines. This list is intentionally smaller than
every architecture CMake can build, because each additional architecture
increases wheel size.
In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
architecture-specific Blackwell/Thor targets, such as `sm_103` or
`sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
cover the corresponding major-version GPU family.
If vLLM logs a warning that your visible CUDA device is not covered by the
wheel's compiled CUDA architectures, or if you see a CUDA error such as
`no kernel image is available for execution on the device`, install a CUDA
13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
that includes your GPU.
#### Install the latest code
LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
@@ -58,40 +58,7 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
--8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images]
vLLM offers official Docker images for deployment.
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
```bash
docker run --rm \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
vllm/vllm-openai-xpu:<tag> \
--model Qwen/Qwen3-0.6B
```
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
???+ console "Commands"
```bash
docker run --rm -it \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
--entrypoint /bin/bash \
vllm/vllm-openai-xpu:<tag>
```
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
--8<-- [end:pre-built-images]
--8<-- [start:build-image-from-source]
-9
View File
@@ -65,15 +65,6 @@ This guide will help you quickly get started with vLLM to perform:
!!! tip
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
=== "Intel GPU"
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
!!! tip
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
=== "Google TPU"
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
+27
View File
@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc
If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.
## CUDA architecture not covered by the wheel
If vLLM logs a warning that the current wheel was built for a set of CUDA
architectures but one of your visible CUDA devices is not covered, the installed
wheel may not contain native CUDA kernels for that GPU. The same issue may also
surface later as a CUDA runtime error such as `no kernel image is available for
execution on the device`.
The architecture list in a pre-built wheel is determined by the vLLM release and
build pipelines, then filtered by CMake before individual kernels choose their
own per-kernel architectures. It is not the same as the full set of
architectures that vLLM can build from source.
CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
them from source. CUDA 13 wheels can use family-specific targets such as
`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
major-version GPU family.
If you see this warning or error, install a CUDA 13 wheel when possible.
Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
your GPU's compute capability. See the CUDA installation guide's
[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
sections for details.
## Python multiprocessing
### `RuntimeError` Exception
+9 -26
View File
@@ -3446,6 +3446,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"encoding_rs",
"futures-channel",
"futures-core",
"futures-util",
"h2",
@@ -4875,7 +4876,6 @@ dependencies = [
"futures-core",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
@@ -4937,9 +4937,9 @@ dependencies = [
[[package]]
name = "tonic"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
dependencies = [
"async-trait",
"axum",
@@ -4966,9 +4966,9 @@ dependencies = [
[[package]]
name = "tonic-build"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -4976,24 +4976,11 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "tonic-health"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
dependencies = [
"prost",
"tokio",
"tokio-stream",
"tonic",
"tonic-prost",
]
[[package]]
name = "tonic-prost"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
dependencies = [
"bytes",
"prost",
@@ -5002,9 +4989,9 @@ dependencies = [
[[package]]
name = "tonic-prost-build"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -5497,13 +5484,10 @@ dependencies = [
"serde",
"serde_json",
"thiserror 2.0.18",
"thiserror-ext",
"tiktoken-rs 0.9.1",
"tokenizers",
"tokio",
"tokio-stream",
"tracing",
"tracing-subscriber",
"url",
"uuid",
]
@@ -5746,7 +5730,6 @@ dependencies = [
"tokio-stream",
"tokio-util",
"tonic",
"tonic-health",
"tonic-prost",
"tonic-prost-build",
"tower",
+4 -5
View File
@@ -118,11 +118,10 @@ tokio = { version = "1.47.1", features = [
tokio-openssl = "0.6"
tokio-stream = "0.1"
tokio-util = { version = "0.7.18", features = ["rt"] }
tonic = "0.14.6"
tonic-build = "0.14.6"
tonic-health = "0.14.6"
tonic-prost = "0.14.6"
tonic-prost-build = "0.14.6"
tonic = "0.14.5"
tonic-build = "0.14.5"
tonic-prost = "0.14.5"
tonic-prost-build = "0.14.5"
tool-parser = "1.2.0"
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
+1 -4
View File
@@ -20,19 +20,16 @@ mimalloc.workspace = true
rand.workspace = true
rand_distr.workspace = true
rayon.workspace = true
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
rlimit.workspace = true
rustc-hash.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json = { workspace = true, features = ["raw_value"] }
thiserror.workspace = true
thiserror-ext.workspace = true
tiktoken-rs.workspace = true
tokenizers.workspace = true
tokio.workspace = true
tokio-stream.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true
uuid.workspace = true
+7 -8
View File
@@ -158,10 +158,9 @@ impl PoolingBackend {
// (mirrors Python async_request_vllm_rerank).
if let Some(ref list) = input.prompt_list {
if list.len() < 2 {
tracing::warn!(
backend = "vllm-rerank",
inputs = list.len(),
"rerank request has no documents"
eprintln!(
"WARNING: vllm-rerank request has no documents \
(prompt_list needs [query, doc, ...])"
);
}
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -176,10 +175,10 @@ impl PoolingBackend {
// Legacy path: text prompt as query, documents via --extra-body.
let query = input.prompt.as_ref();
if query.is_empty() && input.prompt_token_ids.is_some() {
tracing::warn!(
backend = "vllm-rerank",
dataset = "random",
"rerank request has an empty query; use the random-rerank dataset"
eprintln!(
"WARNING: vllm-rerank received empty query (random dataset uses \
token IDs only). Use --dataset-name random-rerank for meaningful \
rerank benchmarks."
);
}
serde_json::json!({
+84 -114
View File
@@ -6,7 +6,6 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,12 +72,12 @@ pub fn pre_resolve_dns(
v4.extend(v6);
if !v4.is_empty() {
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
println!("Pre-resolved {host} -> {ips:?}");
builder = builder.resolve_to_addrs(host, &v4);
}
}
Err(e) => {
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
}
}
@@ -347,14 +346,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
tracing::info!(base_url = %config.base_url, "fetching first model from server");
println!("Model not specified, fetching first model from server...");
let (name, id) =
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
println!("First model name: {name}, first model id: {id}");
(id, Some(name))
};
@@ -363,10 +358,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
tracing::info!(tokenizer = tid, "loading tokenizer");
println!("Loading tokenizer: {tid}");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
let has_tokenizer = tokenizer.is_some();
@@ -426,12 +421,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.num_prompts, config.random_batch_size, config.is_reranker,
),
};
tracing::info!(
dataset = ?config.dataset_name,
prompts = config.num_prompts,
description = %dataset_label,
"generating benchmark dataset"
);
println!("Generating {dataset_label}...");
let gen_start = Instant::now();
let mut input_requests = match config.dataset_name {
@@ -482,7 +472,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded.as_str()
}
};
@@ -522,8 +512,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None => {
downloaded = crate::datasets::speed_bench::download_speed_bench(
config.speed_bench_config,
)
.await?;
)?;
downloaded.as_str()
}
};
@@ -554,8 +543,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.hf_subset.as_deref(),
config.hf_split.as_deref(),
config.num_prompts,
)
.await?;
)?;
crate::datasets::hf_dataset::load_hf_dataset(
tok,
&downloaded_path,
@@ -620,19 +608,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
};
let gen_elapsed = gen_start.elapsed();
tracing::info!(
prompts = input_requests.len(),
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated benchmark dataset"
println!(
"Generated {} prompts in {:.2}s",
input_requests.len(),
gen_elapsed.as_secs_f64()
);
let filtered_count =
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
if filtered_count > 0 {
tracing::info!(
filtered_prompts = filtered_count,
max_model_len = config.max_model_len.unwrap(),
"filtered prompts above maximum model length"
println!(
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
config.max_model_len.unwrap()
);
}
if input_requests.is_empty() {
@@ -683,7 +670,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Ready check
if config.ready_check_timeout_sec > 0 {
tracing::info!("starting initial single-prompt test run");
println!("Starting initial single prompt test run...");
let test_output = wait_for_endpoint(
config.backend,
&client,
@@ -698,7 +685,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
test_output.error
)));
}
tracing::info!("initial single-prompt test run completed");
println!("Initial test run completed.");
}
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -716,15 +703,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
DatasetName::Random | DatasetName::PrefixRepetition
);
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
tracing::info!(
reason = "prompt_token_ids",
"skipping server tokenizer verification"
);
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
}
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
if is_tokenizer_verified(&cache_key) {
tracing::info!(reason = "cached", "skipping server tokenizer verification");
println!("Tokenizer verified in previous run (cached), skipping verification.");
} else {
let num_special =
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -739,17 +723,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await?
{
SampleVerifyOutcome::Passed => {
tracing::info!("tokenizer sample verification passed");
println!("Sample verification passed, skipping full verification.");
mark_tokenizer_verified(&cache_key);
}
SampleVerifyOutcome::Skipped(reason) => {
tracing::warn!(
reason = %reason,
"server tokenizer unavailable; skipping prompt verification"
);
println!("Server /tokenize unavailable ({reason}), skipping verification.");
}
SampleVerifyOutcome::Mismatch => {
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
println!("Sample verification found mismatch, running full verify+fix...");
match verify_and_fix_prompt_lengths(
&client,
&config.base_url,
@@ -761,16 +742,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await
{
Ok(()) => {
tracing::info!(
prompts = input_requests.len(),
"verified exact prompt token lengths"
println!(
"All {} prompts verified: exact token length match.",
input_requests.len()
);
mark_tokenizer_verified(&cache_key);
}
Err(BenchError::TokenizeUnavailable(reason)) => {
tracing::warn!(
reason = %reason,
"server tokenizer became unavailable; using client token counts"
println!(
"Server /tokenize became unavailable during verification \
({reason}); proceeding with client-side token counts."
);
}
Err(e) => return Err(e),
@@ -782,7 +763,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Warmup
if config.num_warmups > 0 {
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
println!("Warming up with {} requests...", config.num_warmups);
run_warmup(
config.backend,
&client,
@@ -795,7 +776,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.disable_tqdm,
)
.await;
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
println!("Warmup run completed.");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -833,22 +814,28 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
tracing::info!("detected speculative decoding; collecting metrics");
println!("Speculative decoding detected, will collect metrics.");
}
// Main benchmark
println!("Starting main benchmark run...");
let distribution = if config.burstiness == 1.0 {
"Poisson process"
} else {
"Gamma distribution"
};
tracing::info!(
request_rate = config.request_rate,
burstiness = config.burstiness,
distribution,
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
prompts = config.num_prompts,
"starting main benchmark run"
println!(
"Traffic request rate: {}",
if config.request_rate.is_infinite() {
"inf".to_string()
} else {
format!("{}", config.request_rate)
}
);
println!("Burstiness factor: {} ({distribution})", config.burstiness);
println!(
"Maximum request concurrency: {}",
config.max_concurrency.unwrap_or(config.num_prompts)
);
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -860,11 +847,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
);
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
"assigned LoRA adapters"
println!(
"LoRA adapters ({}): {:?} [assignment={:?}]",
modules.len(),
names,
config.lora_assignment
);
}
@@ -1138,7 +1125,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
tracing::error!(error = %e.as_report(), "profiler background task failed");
eprintln!("WARNING: Profile background task failed: {e}");
}
}
@@ -1302,14 +1289,12 @@ pub(crate) async fn start_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Starting profiler...");
let profile_url = format!("{base_url}/start_profile");
tracing::info!(url = %profile_url, "starting profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
}
Ok(true) => println!("Profiler started"),
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
}
}
@@ -1319,14 +1304,12 @@ pub(crate) async fn stop_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Stopping profiler...");
let profile_url = format!("{base_url}/stop_profile");
tracing::info!(url = %profile_url, "stopping profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
}
Ok(true) => println!("Profiler stopped"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
}
}
@@ -1388,30 +1371,25 @@ pub(crate) async fn profile_on_batch_threshold(
duration_secs: f64,
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
) {
tracing::info!(
threshold,
duration_seconds = duration_secs,
"waiting for profiler batch threshold"
println!(
"Waiting for batch size >= {threshold} before starting profiler \
(will capture {duration_secs}s)..."
);
loop {
if let Some(running) = fetch_num_requests_running(client, base_url).await
&& running >= threshold
{
tracing::info!(
running_requests = running,
threshold,
"profiler batch threshold reached"
);
println!("Batch size {running} >= {threshold}, starting profiler...");
break;
}
// Wait 500ms or until the benchmark signals cancellation
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = &mut cancel_rx => {
tracing::warn!(
threshold,
"benchmark finished before profiler batch threshold; skipping profiling"
eprintln!(
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
profiling skipped."
);
return;
}
@@ -1420,13 +1398,13 @@ pub(crate) async fn profile_on_batch_threshold(
let start_url = format!("{base_url}/start_profile");
match send_profile_request(client, &start_url, extra_headers).await {
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
Ok(true) => println!("Profiler started"),
Ok(false) => {
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
eprintln!("WARNING: Profiler start request returned non-success");
return;
}
Err(e) => {
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
eprintln!("WARNING: Failed to start profiler: {e}");
return;
}
}
@@ -1435,17 +1413,15 @@ pub(crate) async fn profile_on_batch_threshold(
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
_ = &mut cancel_rx => {
tracing::info!("benchmark finished; stopping profiler early");
println!("Benchmark finished, stopping profiler early...");
}
}
let stop_url = format!("{base_url}/stop_profile");
match send_profile_request(client, &stop_url, extra_headers).await {
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
}
Ok(true) => println!("Profiler stopped after capturing"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
}
}
@@ -1526,11 +1502,10 @@ async fn verify_and_fix_prompt_lengths(
let excess = tokens.len().saturating_sub(expected_input_len);
let compensate = if excess > 0 && last_excess == Some(excess) {
if _iter == 1 {
tracing::warn!(
prompt_index = i,
extra_tokens = excess,
adjusted_target = expected_input_len.saturating_sub(excess),
"server consistently adds prompt tokens; compensating verification target"
eprintln!(
"Prompt {i}: server consistently adds {excess} extra token(s) \
(likely BOS), compensating target to {}.",
expected_input_len.saturating_sub(excess),
);
}
excess
@@ -1588,10 +1563,7 @@ async fn verify_and_fix_prompt_lengths(
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
if fc > 0 {
tracing::info!(
fixed_prompts = fc,
"fixed prompt lengths using server tokenizer"
);
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
}
Ok(())
@@ -1846,7 +1818,7 @@ async fn sample_verify_prompts(
let tokenize_url = format!("{base_url}/tokenize");
let api_key = std::env::var("OPENAI_API_KEY").ok();
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
println!("Sampling {sample_size} prompts for verification...");
for (i, request) in requests.iter().enumerate().take(sample_size) {
let tokens = match server_tokenize(
@@ -1869,11 +1841,9 @@ async fn sample_verify_prompts(
let expected = request.prompt_len + num_special;
if tokens.len() != expected {
tracing::warn!(
prompt_index = i,
expected_tokens = expected,
actual_tokens = tokens.len(),
"tokenizer verification sample mismatch"
println!(
"Prompt {i}: expected {expected} tokens, server returned {}",
tokens.len()
);
return Ok(SampleVerifyOutcome::Mismatch);
}
+9 -18
View File
@@ -288,18 +288,10 @@ impl BenchConfig {
}
Some(other) => {
// extra_body was not an object — just use sampling params
let value_type = match &other {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => unreachable!(),
};
tracing::warn!(
value_type,
"sampling parameters may be lost because --extra-body is not a JSON object"
eprintln!(
"Warning: --extra-body is not a JSON object, sampling params may be lost"
);
let _ = other;
sampling_params
}
None => sampling_params,
@@ -497,9 +489,9 @@ impl BenchConfig {
_ => {}
}
if !args.skip_chat_template {
tracing::warn!(
dataset = "custom",
"client-side chat template rendering is unsupported; sending prompts raw"
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
);
}
}
@@ -578,10 +570,9 @@ impl BenchConfig {
}
if ignore_eos {
tracing::warn!(
ignore_eos,
multi_turn = true,
"output length limits may be ignored, causing unbounded context growth"
eprintln!(
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
respect output length limits, causing unbounded context growth."
);
}
+2 -4
View File
@@ -128,10 +128,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
#[test]
+37 -74
View File
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -51,19 +50,18 @@ enum ColumnFormat {
/// Make a GET request with retry logic (3 retries with exponential backoff).
/// Returns the parsed JSON response.
async fn get_with_retry(
client: &reqwest::Client,
fn get_with_retry(
client: &reqwest::blocking::Client,
url: &str,
label: &str,
) -> Result<serde_json::Value> {
let max_retries = 3;
for attempt in 0..=max_retries {
let resp = match client.get(url).send().await {
let resp = match client.get(url).send() {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
.await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
return Err(BenchError::Config(format!(
@@ -82,7 +80,7 @@ async fn get_with_retry(
}
if status.is_server_error() && attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
@@ -94,7 +92,6 @@ async fn get_with_retry(
let data: serde_json::Value = resp
.json()
.await
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
return Ok(data);
}
@@ -108,7 +105,7 @@ async fn get_with_retry(
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
/// or the dataset is exhausted.
pub async fn download_hf_dataset(
pub fn download_hf_dataset(
dataset: &str,
subset: Option<&str>,
split: Option<&str>,
@@ -118,7 +115,7 @@ pub async fn download_hf_dataset(
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
let mut client_builder =
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
// Add HF_TOKEN auth header if available
if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -141,7 +138,7 @@ pub async fn download_hf_dataset(
// Call /info to discover available configs and splits
let info_url =
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
let dataset_info =
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -204,12 +201,7 @@ pub async fn download_hf_dataset(
(resolved_config, resolved_split)
};
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"resolved Hugging Face dataset"
);
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
// Check cache
let dir = cache_dir();
@@ -223,16 +215,11 @@ pub async fn download_hf_dataset(
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
println!("HF dataset cached: {path_str}");
return Ok((path_str, resolved_config, resolved_split));
}
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"downloading Hugging Face dataset"
);
println!("Downloading HF dataset '{dataset}' from datasets-server...");
let encoded_config: String =
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -242,7 +229,6 @@ pub async fn download_hf_dataset(
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -254,7 +240,7 @@ pub async fn download_hf_dataset(
&length={page_size}"
);
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
let rows = data["rows"]
.as_array()
@@ -274,14 +260,14 @@ pub async fn download_hf_dataset(
offset += fetched;
let total = data["num_rows_total"].as_u64().unwrap_or(0);
progress.update(offset, total);
eprint!("\r Fetched {offset}/{total} rows...");
// Stop if we have enough rows or reached end of dataset
if all_rows.len() >= num_rows_needed || fetched < page_size {
break;
}
}
progress.finish();
eprintln!(); // newline after progress
if all_rows.is_empty() {
return Err(BenchError::Config(format!(
@@ -294,12 +280,7 @@ pub async fn download_hf_dataset(
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(
dataset,
rows = all_rows.len(),
path = %path_str,
"saved Hugging Face dataset"
);
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
Ok((path_str, resolved_config, resolved_split))
}
@@ -502,31 +483,21 @@ pub fn load_hf_dataset(
// Detect column format from first row
let format = detect_column_format(&entries[0], text_column_override)?;
// Print detected format
match &format {
ColumnFormat::Chat(col) => {
tracing::info!(
format = "chat",
column = col,
"detected Hugging Face dataset format"
);
}
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
ColumnFormat::Text {
prompt_col,
output_col,
} => {
tracing::info!(
format = "text",
prompt_column = prompt_col,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
let out_msg = output_col.as_deref().unwrap_or("none");
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
}
ColumnFormat::Combined { cols, output_col } => {
tracing::info!(
format = "combined",
prompt_columns = ?cols,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
let out_msg = output_col.as_deref().unwrap_or("none");
println!(
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
cols
);
}
}
@@ -637,10 +608,9 @@ pub fn load_hf_dataset(
if len == 0 { 128 } else { len }
} else {
if !warned_no_output {
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
);
warned_no_output = true;
}
@@ -648,10 +618,9 @@ pub fn load_hf_dataset(
}
} else {
if !warned_no_output {
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
);
warned_no_output = true;
}
@@ -671,11 +640,9 @@ pub fn load_hf_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "hf",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -685,11 +652,9 @@ pub fn load_hf_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "hf",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled HF dataset from {original_len} to {} total samples.",
samples.len()
);
}
}
@@ -1037,10 +1002,8 @@ mod tests {
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
crate::tokenizer::TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
/// Write JSON data to a unique temp file and return the path string.
+6 -9
View File
@@ -5,7 +5,6 @@ pub mod custom;
pub mod hf_dataset;
pub mod multi_turn;
pub mod prefix_repetition;
mod progress;
pub mod random;
pub mod random_mm;
pub mod random_rerank;
@@ -91,10 +90,9 @@ pub fn oversample_requests(
return;
}
if no_oversample {
tracing::info!(
samples = requests.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
requests.len()
);
return;
}
@@ -105,10 +103,9 @@ pub fn oversample_requests(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
requests.push(req);
}
tracing::info!(
original_samples = original_len,
samples = requests.len(),
"oversampled dataset"
println!(
"Oversampled requests from {original_len} to {} total samples.",
requests.len()
);
}
+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);
conversations.push(conv);
}
tracing::info!(
original_conversations = original_len,
conversations = conversations.len(),
"oversampled multi-turn conversations"
println!(
"Oversampled multi-turn conversations from {original_len} to {} total.",
conversations.len()
);
}
@@ -526,12 +525,10 @@ mod tests {
len
}
#[tokio::test]
#[test]
#[ignore]
async fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 5,
@@ -613,12 +610,10 @@ mod tests {
println!("All prefix sharing checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 4,
@@ -655,12 +650,10 @@ mod tests {
println!("per_turn_input_len default-mode checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 50,
@@ -691,12 +684,10 @@ mod tests {
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 10,
@@ -718,12 +709,10 @@ mod tests {
println!("variable_turns_fixed checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
@@ -41,13 +41,11 @@ pub fn generate_prefix_repetition_dataset(
}
let total = prompts_per_prefix * num_prefixes;
if total != num_requests {
tracing::info!(
requested = num_requests,
generated = total,
prefixes = num_prefixes,
prompts_per_prefix,
dropped = num_requests - total,
"adjusted prefix-repetition request count"
println!(
"prefix_repetition: generating {total} requests \
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
{} dropped to divide evenly)",
num_requests - total
);
}
@@ -111,10 +109,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
#[test]
-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 (output_low, output_high) = range_ratio.output_bounds(output_len);
if !range_ratio.is_fixed() {
tracing::info!(
input_low,
input_high,
output_low,
output_high,
"sampling random request lengths"
println!(
"Sampling input_len from [{input_low}, {input_high}] and \
output_len from [{output_low}, {output_high}]"
);
}
@@ -308,8 +305,7 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_token_ids() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -341,8 +337,7 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_text() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -376,8 +371,7 @@ mod tests {
#[test]
#[ignore]
fn test_token_length_exact_local() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let target_len = 512;
let requests = generate_random_dataset(
&tokenizer,
@@ -411,11 +405,11 @@ mod tests {
}
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
#[tokio::test]
#[test]
#[ignore]
async fn test_token_length_exact_tiktoken() {
fn test_token_length_exact_tiktoken() {
// Use Qwen2.5 which has a tiktoken-format tokenizer
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
@@ -459,10 +453,10 @@ mod tests {
/// Test encode/decode roundtrip stability for tiktoken.
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
#[tokio::test]
#[test]
#[ignore]
async fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
+2 -4
View File
@@ -139,10 +139,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
fn fixed_ratio() -> RangeRatio {
+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.
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
pub async fn download_sharegpt_dataset() -> Result<String> {
tracing::info!(
repository = DEFAULT_SHAREGPT_REPO,
file = DEFAULT_SHAREGPT_FILE,
"downloading ShareGPT dataset"
pub fn download_sharegpt_dataset() -> Result<String> {
println!(
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
);
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
.map_err(BenchError::Config)?;
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
BenchError::Config(format!(
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
))
})?;
let path_str = path.to_string_lossy().to_string();
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
println!("ShareGPT dataset ready: {path_str}");
Ok(path_str)
}
@@ -138,11 +135,9 @@ pub fn load_sharegpt_dataset(
// Oversample if dataset is smaller than requested
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "sharegpt",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let needed = num_requests - samples.len();
@@ -152,11 +147,9 @@ pub fn load_sharegpt_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "sharegpt",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled requests from {original_len} to {} total samples.",
samples.len()
);
}
}
+23 -30
View File
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::cli::SpeedBenchConfig;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -26,7 +25,7 @@ fn cache_dir() -> std::path::PathBuf {
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
/// Results are cached as JSON locally for subsequent runs.
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let config_name = config.as_str();
let dir = cache_dir();
@@ -36,13 +35,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
// Return cached file if it exists
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
println!("SPEED-Bench ({config_name}) cached: {path_str}");
return Ok(path_str);
}
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
let client = reqwest::Client::builder()
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -50,7 +49,6 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -66,14 +64,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let max_retries = 3;
let mut data: Option<serde_json::Value> = None;
for attempt in 0..=max_retries {
let resp = match client.get(&url).send().await {
let resp = match client.get(&url).send() {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(
std::thread::sleep(std::time::Duration::from_secs(
2 * (attempt as u64 + 1),
))
.await;
));
continue;
}
return Err(BenchError::Config(format!(
@@ -83,7 +80,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
};
if resp.status().is_server_error() && attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
@@ -94,7 +91,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
)));
}
data = Some(resp.json().await.map_err(|e| {
data = Some(resp.json().map_err(|e| {
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
})?);
break;
@@ -119,14 +116,15 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let fetched = rows.len();
offset += fetched;
// Print progress
let total = data["num_rows_total"].as_u64().unwrap_or(0);
progress.update(offset, total);
eprint!("\r Fetched {offset}/{total} rows...");
if fetched < page_size {
break;
}
}
progress.finish();
eprintln!(); // newline after progress
if all_rows.is_empty() {
return Err(BenchError::Config(
@@ -139,11 +137,9 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(
config = config_name,
rows = all_rows.len(),
path = %path_str,
"saved SPEED-Bench dataset"
println!(
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
all_rows.len()
);
Ok(path_str)
}
@@ -267,11 +263,9 @@ pub fn load_speed_bench_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "speed-bench",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -281,11 +275,9 @@ pub fn load_speed_bench_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "speed-bench",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
samples.len()
);
}
}
@@ -296,6 +288,7 @@ pub fn load_speed_bench_dataset(
));
}
// Print category distribution
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
for entry in &filtered[..filtered.len().min(samples.len())] {
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -304,7 +297,7 @@ pub fn load_speed_bench_dataset(
let mut cats: Vec<_> = cat_counts.into_iter().collect();
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
println!("SPEED-Bench categories: {}", cat_str.join(", "));
Ok(samples)
}
+35 -20
View File
@@ -1,39 +1,54 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::path::PathBuf;
//! Sync facade over the async `hf_hub` API.
//!
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
//! backend unconditionally pulls ureq's default rustls feature. So we use the
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
//! loaders, tokenizer fallback in rayon threads) by running each download on a
//! dedicated thread with its own single-threaded runtime.
use hf_hub::Repo;
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
use std::path::PathBuf;
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
pub struct HubRepo {
repo: ApiRepo,
repo: hf_hub::Repo,
}
impl HubRepo {
pub fn model(model_id: String) -> Result<Self, String> {
Self::new(Repo::model(model_id))
pub fn model(model_id: String) -> Self {
Self {
repo: hf_hub::Repo::model(model_id),
}
}
pub fn dataset(repo_id: String) -> Result<Self, String> {
Self::new(Repo::dataset(repo_id))
}
fn new(repo: Repo) -> Result<Self, String> {
let mut builder = ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
pub fn dataset(repo_id: String) -> Self {
Self {
repo: hf_hub::Repo::dataset(repo_id),
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
Ok(Self {
repo: api.repo(repo),
})
}
/// Download (or fetch from cache) a single file from the repo.
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
self.repo.get(filename).await.map_err(|e| format!("{e}"))
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
let repo = self.repo.clone();
let filename = filename.to_string();
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
.join()
.map_err(|_| "HF Hub download thread panicked".to_string())?
}
}
+1 -1
View File
@@ -33,7 +33,7 @@ pub fn prepare_process() {
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
tracing::info!(soft_limit = new, "raised open-file limit");
eprintln!("Open-file limit: {new}");
}
}
-12
View File
@@ -19,19 +19,7 @@ struct Cli {
args: vllm_bench::BenchServeArgs,
}
// TODO: unify the tracing subscriber used by different binaries.
fn init_tracing() {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
let _ = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_writer(std::io::stderr)
.try_init();
}
fn main() -> anyhow::Result<()> {
init_tracing();
let cli = Cli::parse();
vllm_bench::prepare_process();
+16 -18
View File
@@ -7,22 +7,6 @@ use crate::datasets::SampleRequest;
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
use crate::multi_turn::ConversationOutput;
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
if failed_outputs.is_empty() {
return;
}
tracing::warn!(
failed_requests = failed_outputs.len(),
displayed_errors = failed_outputs.len().min(10),
"benchmark requests failed"
);
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
tracing::warn!(index, error = %output.error, "benchmark request failed");
}
}
/// Calculate benchmark metrics from request outputs.
///
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
@@ -79,7 +63,14 @@ pub fn calculate_metrics(
let failed = outputs.len() - completed;
log_failed_requests(outputs);
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Calculate max output tokens per second and max concurrent requests
let mut max_output_tokens_per_s = 0.0_f64;
@@ -304,7 +295,14 @@ pub fn calculate_embedding_metrics(
let failed = outputs.len() - completed;
log_failed_requests(outputs);
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Compute peak concurrent requests from start_time + latency windows
let successful_outputs: Vec<&RequestFuncOutput> =
+39 -53
View File
@@ -6,7 +6,6 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,13 +72,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
tracing::info!(base_url = %config.base_url, "fetching first model from server");
println!("Model not specified, fetching first model from server...");
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
println!("First model name: {name}, first model id: {id}");
(id, Some(name))
};
@@ -88,19 +83,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
tracing::info!(tokenizer = tid, "loading tokenizer");
println!("Loading tokenizer: {tid}");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
// Generate/load conversations
tracing::info!(
dataset = ?config.dataset_name,
conversations = config.num_prompts,
"generating multi-turn conversations"
);
println!("Generating multi-turn conversations...");
let gen_start = Instant::now();
let mut conversations = match config.dataset_name {
@@ -140,7 +131,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded.as_str()
}
};
@@ -188,11 +179,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (filtered_conversations, filtered_turns) =
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
if filtered_turns > 0 || filtered_conversations > 0 {
tracing::info!(
filtered_turns,
filtered_conversations,
max_model_len,
"filtered conversations above maximum model length"
println!(
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
);
}
if conversations.is_empty() {
@@ -204,11 +192,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let gen_elapsed = gen_start.elapsed();
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
tracing::info!(
conversations = conversations.len(),
println!(
"Generated {} conversations ({} total turns) in {:.2}s",
conversations.len(),
total_turns,
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated multi-turn conversations"
gen_elapsed.as_secs_f64()
);
// Log prefix sharing info
@@ -220,15 +208,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let conv_tokens =
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
tracing::info!(
global_ratio = config.multi_turn_prefix_global_ratio,
println!(
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
config.multi_turn_prefix_global_ratio * 100.0,
global_tokens,
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
conversation_tokens = conv_tokens,
config.multi_turn_prefix_conversation_ratio * 100.0,
conv_tokens,
(1.0 - config.multi_turn_prefix_global_ratio
- config.multi_turn_prefix_conversation_ratio)
* 100.0,
unique_tokens,
history_accumulation = false,
"configured multi-turn prefix sharing"
);
println!("No history accumulation: each turn sends fixed-length prompt only.");
}
if config.dry_run {
@@ -262,7 +253,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
..Default::default()
};
tracing::info!("starting initial single-prompt test run");
println!("Starting initial single prompt test run...");
let test_output = crate::ready_checker::wait_for_endpoint(
config.backend,
&client,
@@ -277,7 +268,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
test_output.error
)));
}
tracing::info!("initial single-prompt test run completed");
println!("Initial test run completed.");
}
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -292,10 +283,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
"min_tokens".to_string(),
serde_json::json!(config.random_output_len),
);
tracing::info!(
min_tokens = config.random_output_len,
dataset = "random",
"set minimum output tokens for multi-turn dataset"
println!(
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
config.random_output_len
);
}
Some(body)
@@ -307,7 +297,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
tracing::info!("detected speculative decoding; collecting metrics");
println!("Speculative decoding detected, will collect metrics.");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -340,13 +330,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
};
// Main benchmark
tracing::info!(
conversations = conversations.len(),
total_turns,
concurrency,
inter_turn_delay_ms = config.multi_turn_delay_ms,
"starting multi-turn benchmark"
);
println!("Starting multi-turn benchmark...");
println!("Conversations: {}", conversations.len());
println!("Concurrency: {concurrency}");
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
@@ -377,12 +364,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
);
if let Some(modules) = config.lora_modules.as_ref() {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
scope = "conversation",
"assigned LoRA adapters"
println!(
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
modules.len(),
names,
config.lora_assignment
);
}
@@ -447,7 +433,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
match handle.await {
Ok(output) => all_outputs.push(output),
Err(e) => {
tracing::error!(error = %e.as_report(), "conversation task panicked");
eprintln!("Conversation task panicked: {e}");
}
}
}
@@ -467,7 +453,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
tracing::error!(error = %e.as_report(), "profiler background task failed");
eprintln!("WARNING: Profile background task failed: {e}");
}
}
+2 -2
View File
@@ -603,7 +603,7 @@ fn add_metric_stats(
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
let content = serde_json::to_string(json)?;
std::fs::write(file_path, content)?;
tracing::info!(path = file_path, "saved benchmark results");
println!("Results saved to {file_path}");
Ok(())
}
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
file.write_all(b"\n")?;
}
file.write_all(content.as_bytes())?;
tracing::info!(path = file_path, "appended benchmark results");
println!("Results appended to {file_path}");
Ok(())
}
+2 -8
View File
@@ -23,11 +23,7 @@ pub async fn wait_for_endpoint(
let backend = get_backend(backend)?;
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
tracing::info!(
timeout_seconds,
retry_interval,
"waiting for endpoint readiness"
);
println!("Waiting for endpoint to become up in {timeout_seconds}s");
let pb = ProgressBar::new(timeout_seconds);
pb.set_style(
@@ -57,9 +53,7 @@ pub async fn wait_for_endpoint(
Ok(output) => {
let err = output.error.clone();
let err_last_line = err.lines().last().unwrap_or(&err);
pb.suspend(|| {
tracing::warn!(error = err_last_line, "endpoint is not ready");
});
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
last_error = err;
}
Err(e) => {
+1 -1
View File
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
.await
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
if resp.status().is_success() {
tracing::info!(url = %url, "reset prefix cache");
println!("Prefix cache reset successfully.");
} else {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
+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}")))?;
tracing::info!(
encoding,
kind = "built-in-tiktoken",
vocab_size,
"loaded tokenizer"
);
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
}
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
// Phase 1: If model_id is a local directory, look for tiktoken files there
let local_dir = Path::new(model_id);
if local_dir.is_dir() {
@@ -217,7 +212,7 @@ pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
}
// Phase 2: Fall back to HuggingFace Hub download
try_load_tiktoken_from_hf(model_id).await
try_load_tiktoken_from_hf(model_id)
}
/// Common tiktoken model filenames to search for.
@@ -252,28 +247,25 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
}
/// Load a tiktoken tokenizer from a HuggingFace model repo.
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string());
let mut model_path = None;
for filename in TIKTOKEN_MODEL_FILENAMES {
if let Ok(path) = repo.get(filename).await {
model_path = Some(path);
break;
}
}
let model_path = model_path.ok_or_else(|| {
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
})?;
let model_path = repo
.get("tiktoken.model")
.or_else(|_| repo.get("qwen.tiktoken"))
.or_else(|_| repo.get("vocab.tiktoken"))
.map_err(|_| {
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
})?;
let num_base_tokens = count_base_tokens(&model_path)?;
let config = match repo.get("tokenizer_config.json").await {
let config = match repo.get("tokenizer_config.json") {
Ok(config_path) => read_tokenizer_config(&config_path),
Err(_) => None,
};
let pattern = extract_pat_str_from_repo(&repo).await;
let pattern = extract_pat_str_from_repo(&repo);
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
}
@@ -317,16 +309,15 @@ fn build_tiktoken(
}
}
tracing::info!(
model = model_id,
base_tokens = num_base_tokens,
special_tokens = all_special_tokens.len(),
pattern = if pattern.is_some() {
println!(
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
num_base_tokens,
all_special_tokens.len(),
if pattern.is_some() {
"custom"
} else {
"default"
},
"loading tiktoken model"
);
TiktokenTokenizer::from_file(
@@ -406,12 +397,9 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
/// Try to download the Python tokenizer source file and extract pat_str via regex.
/// Returns None if unavailable or unparsable.
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
// Try common Python tokenizer filenames
let py_path = match repo.get("tokenization_kimi.py").await {
Ok(path) => path,
Err(_) => repo.get("tokenizer.py").await.ok()?,
};
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
let source = std::fs::read_to_string(&py_path).ok()?;
@@ -450,9 +438,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
if !fragments.is_empty() {
let pattern = fragments.join("|");
tracing::debug!(
fragments = fragments.len(),
"extracted tiktoken pattern from Python source"
println!(
"Extracted pat_str from Python source: {} fragments",
fragments.len()
);
return Some(pattern);
}
+21 -98
View File
@@ -2,10 +2,8 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::HashSet;
use std::future::Future;
use std::path::Path;
use thiserror_ext::AsReport as _;
use tokenizers::Tokenizer;
use crate::error::{BenchError, Result};
@@ -20,8 +18,7 @@ pub enum TokenizerKind {
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
pub struct ServerTokenizer {
client: reqwest::Client,
runtime: tokio::runtime::Handle,
client: reqwest::blocking::Client,
tokenize_url: String,
detokenize_url: String,
model: String,
@@ -30,8 +27,8 @@ pub struct ServerTokenizer {
impl ServerTokenizer {
/// Create a new server tokenizer and verify connectivity.
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::Client::builder()
pub fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -41,7 +38,6 @@ impl ServerTokenizer {
let st = Self {
client,
runtime: tokio::runtime::Handle::current(),
tokenize_url,
detokenize_url,
model: model.to_string(),
@@ -49,7 +45,7 @@ impl ServerTokenizer {
};
// Probe the endpoint to verify it works and discover vocab size
let test_tokens = st.encode_async("test").await?;
let test_tokens = st.encode_inner("test")?;
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
let estimated_vocab = (max_id * 2).max(131072);
@@ -60,10 +56,6 @@ impl ServerTokenizer {
}
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
self.block_on(self.encode_async(text))
}
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
let payload = serde_json::json!({
"model": self.model,
"prompt": text,
@@ -74,7 +66,6 @@ impl ServerTokenizer {
.post(&self.tokenize_url)
.json(&payload)
.send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
if !resp.status().is_success() {
@@ -84,7 +75,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().await.map_err(|e| {
let data: serde_json::Value = resp.json().map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
})?;
@@ -104,10 +95,6 @@ impl ServerTokenizer {
}
fn decode_inner(&self, ids: &[u32]) -> Result<String> {
self.block_on(self.decode_async(ids))
}
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
let payload = serde_json::json!({
"model": self.model,
"tokens": ids,
@@ -118,7 +105,6 @@ impl ServerTokenizer {
.post(&self.detokenize_url)
.json(&payload)
.send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
if !resp.status().is_success() {
@@ -128,7 +114,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().await.map_err(|e| {
let data: serde_json::Value = resp.json().map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
})?;
@@ -137,26 +123,6 @@ impl ServerTokenizer {
.map(|s| s.to_string())
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
}
fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
if matches!(
self.runtime.runtime_flavor(),
tokio::runtime::RuntimeFlavor::CurrentThread
) {
return Err(BenchError::Tokenizer(
"Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
));
}
// Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
// Tokio workers must enter a blocking region before re-entering the runtime;
// Rayon workers can drive the future directly with the saved runtime handle.
if tokio::runtime::Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
}
// --- TokenizerKind methods ---
@@ -226,7 +192,7 @@ impl TokenizerKind {
/// 3. Server-side /tokenize + /detokenize endpoints
///
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
pub async fn load_tokenizer(
pub fn load_tokenizer(
model_id: &str,
_trust_remote_code: bool,
server_info: Option<(&str, &str)>,
@@ -246,48 +212,31 @@ pub async fn load_tokenizer(
}
// 1. Try local HuggingFace tokenizer (tokenizer.json)
match try_load_local(model_id).await {
match try_load_local(model_id) {
Ok(tok) => {
tracing::info!(
model = model_id,
kind = "local",
vocab_size = tok.get_vocab_size(true),
"loaded tokenizer"
);
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
Ok(TokenizerKind::Local(Box::new(tok)))
}
Err(local_err) => {
// 2. Try tiktoken format
tracing::info!(
model = model_id,
error = %local_err.as_report(),
"local tokenizer unavailable; trying tiktoken"
);
match crate::tiktoken::try_load_tiktoken(model_id).await {
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
match crate::tiktoken::try_load_tiktoken(model_id) {
Ok(tok) => {
tracing::info!(
model = model_id,
kind = "tiktoken",
vocab_size = tok.vocab_size(),
"loaded tokenizer"
);
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
Ok(TokenizerKind::Tiktoken(tok))
}
Err(tiktoken_err) => {
// 3. Try server-side fallback
if let Some((base_url, model)) = server_info {
tracing::info!(
model = model_id,
error = %tiktoken_err.as_report(),
"tiktoken unavailable; trying server-side tokenization"
println!(
"Tiktoken also not available ({tiktoken_err}), \
trying server-side tokenization..."
);
match ServerTokenizer::new(base_url, model).await {
match ServerTokenizer::new(base_url, model) {
Ok(srv) => {
tracing::info!(
model = model_id,
kind = "server",
vocab_size = srv.cached_vocab_size,
"loaded tokenizer"
println!(
"Tokenizer: Server (vocab_size≈{})",
srv.cached_vocab_size
);
return Ok(TokenizerKind::Server(srv));
}
@@ -315,7 +264,7 @@ pub async fn load_tokenizer(
}
/// Try loading tokenizer.json from local path or HuggingFace Hub.
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
// 1. Try local directory with tokenizer.json
let local_path = Path::new(model_id).join("tokenizer.json");
if local_path.exists() {
@@ -341,37 +290,11 @@ async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
}
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
let repo = crate::hub::HubRepo::model(model_id.to_string());
let tokenizer_path = repo
.get("tokenizer.json")
.await
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
Tokenizer::from_file(&tokenizer_path)
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_tokenizer_sync_bridge() {
let tokenizer = std::sync::Arc::new(ServerTokenizer {
client: reqwest::Client::new(),
runtime: tokio::runtime::Handle::current(),
tokenize_url: String::new(),
detokenize_url: String::new(),
model: String::new(),
cached_vocab_size: 0,
});
assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
let (tx, rx) = tokio::sync::oneshot::channel();
rayon::spawn(move || {
let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
});
assert_eq!(rx.await.unwrap().unwrap(), 2);
}
}
+7 -17
View File
@@ -26,21 +26,18 @@ const RESET: &str = "\x1b[0m";
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
format_description!("[month]-[day] [hour]:[minute]:[second]");
const PROCESS_LABEL: &str = "RustFrontend";
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
pub(crate) fn init_tracing(process_label: &str) {
pub(crate) fn init_tracing() {
let filter = build_targets_filter(
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
env::var("RUST_LOG").ok().as_deref(),
);
let formatter = VllmEventFormatter::new(process_label);
let formatter = VllmEventFormatter::new();
let _ = tracing_subscriber::registry()
.with(
tracing_subscriber::fmt::layer()
.event_format(formatter)
.with_writer(std::io::stderr)
.with_filter(filter),
)
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
.try_init();
}
@@ -97,9 +94,9 @@ struct VllmEventFormatter {
}
impl VllmEventFormatter {
fn new(process_label: &str) -> Self {
fn new() -> Self {
Self {
prefix: format!("({process_label} pid={})", process::id()),
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
timer: VllmLocalTimer::default(),
}
}
@@ -294,13 +291,6 @@ fn map_python_log_level(level: &str) -> LevelFilter {
mod tests {
use super::*;
#[test]
fn formatter_prefix_uses_process_label() {
let formatter = VllmEventFormatter::new("Bench");
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
}
#[test]
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
+1 -9
View File
@@ -5,7 +5,6 @@ mod cli;
mod logging;
use std::env;
use std::ffi::OsStr;
use std::process::ExitStatus;
use anyhow::{Context, Result, anyhow, bail};
@@ -83,14 +82,7 @@ fn shutdown_signal() -> CancellationToken {
}
fn main() -> Result<()> {
let process_label =
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
"bench" => "Bench",
"serve" | "frontend" => "RustFrontend",
_ => "Rust",
};
logging::init_tracing(process_label);
logging::init_tracing();
let cli = Cli::parse();
let mut runtime = tokio::runtime::Builder::new_multi_thread();
@@ -460,12 +460,6 @@ impl EngineCoreClient {
self.inner.is_healthy()
}
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
self.inner.subscribe_health()
}
/// Return the first persistent health error observed by the client, if any.
pub fn health_error(&self) -> Option<Arc<Error>> {
self.inner.health_error()
+1 -20
View File
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
use parking_lot::Mutex;
use thiserror_ext::AsReport as _;
use tokio::runtime::Handle;
use tokio::sync::{mpsc, watch};
use tokio::sync::mpsc;
use tracing::{debug, info, trace, warn};
use vllm_metrics::METRICS;
use zeromq::RouterSendHalf;
@@ -36,7 +36,6 @@ pub(crate) struct ClientInner {
request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>,
health_tx: watch::Sender<bool>,
}
impl ClientInner {
@@ -58,7 +57,6 @@ impl ClientInner {
request_reg: Mutex::new(RequestRegistry::new(engines)),
utility_reg: Mutex::new(UtilityRegistry::default()),
health_error: ArcSwapOption::empty(),
health_tx: watch::Sender::new(true),
}
}
@@ -171,7 +169,6 @@ impl ClientInner {
/// persistent health error.
pub fn close_registries(&self, error: Arc<Error>) {
let persistent_error = self.record_health_error(error);
self.publish_unhealthy();
let request_senders = self.request_reg.lock().close();
let utility_senders = self.utility_reg.lock().close();
@@ -194,12 +191,6 @@ impl ClientInner {
self.health_error.load().is_none()
}
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> watch::Receiver<bool> {
self.health_tx.subscribe()
}
/// Resolve one utility output to the waiting caller. Returns `true` if a
/// waiting caller existed.
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
@@ -289,11 +280,6 @@ impl ClientInner {
.expect("health error must be recorded before registries close")
}
/// Publish the sticky healthy-to-unhealthy transition.
fn publish_unhealthy(&self) {
self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
}
/// Assert there is a recorded health error and return a `Shared` variant
/// wrapping it for error returns when the client is already closed.
fn closed_error(&self) -> Error {
@@ -475,18 +461,13 @@ mod tests {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn close_registries_records_first_health_error_only() {
let inner = test_inner().await;
let mut health = inner.subscribe_health();
assert!(*health.borrow());
inner.close_registries(Arc::new(Error::EngineCoreDead));
health.changed().await.expect("health sender remains open");
assert!(!inner.is_healthy());
assert!(!*health.borrow());
assert!(matches!(
inner.health_error().as_deref(),
Some(Error::EngineCoreDead)
));
assert!(!*inner.subscribe_health().borrow());
inner.close_registries(Arc::new(client_closed!("shutdown")));
assert!(matches!(
@@ -269,19 +269,6 @@ impl WireLogprobs {
);
}
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
if token_ids.rows == 0 {
return Ok(Logprobs {
positions: Vec::new(),
});
}
if token_ids.cols == 0 {
bail_ext_value_decode!(
"{field_prefix}: zero-column logprobs payload with {} rows",
token_ids.rows
);
}
let mut positions = Vec::with_capacity(token_ids.rows);
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
.data
@@ -303,49 +303,3 @@ fn rejects_non_none_cu_num_generated_tokens() {
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
);
}
#[test]
fn decodes_zero_row_logprobs_as_empty() {
for shape in [[0usize, 0], [0, 3]] {
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
None,
Some(Value::Array(vec![
ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
Value::Nil,
])),
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0]
.new_prompt_logprobs_tensors
.clone()
.unwrap()
.into_direct()
.unwrap();
assert!(logprobs.is_empty());
}
}
#[test]
fn rejects_zero_column_logprobs_with_rows() {
let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
Some(Value::Array(vec![
ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[2], ranks),
Value::Nil,
])),
None,
)))];
let error = decode_engine_core_outputs(&frames).unwrap_err();
let crate::error::Error::ExtValueDecode { message } = &error else {
panic!("expected ExtValueDecode");
};
assert_eq!(
message,
"new_logprobs: zero-column logprobs payload with 2 rows"
);
}
-1
View File
@@ -35,7 +35,6 @@ tokio-openssl.workspace = true
tokio-stream.workspace = true
tokio-util.workspace = true
tonic.workspace = true
tonic-health.workspace = true
tonic-prost.workspace = true
tower.workspace = true
tower-http.workspace = true
-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.
mod convert;
mod health;
use std::pin::Pin;
use std::sync::Arc;
@@ -25,11 +24,8 @@ pub mod pb {
tonic::include_proto!("vllm");
}
pub(crate) use health::monitor_health;
pub use pb::generate_server::GenerateServer;
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
#[cfg(test)]
mod tests;
+12 -175
View File
@@ -16,10 +16,6 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use tokio::net::TcpStream;
use tokio_openssl::SslStream;
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
use tonic_health::pb::HealthCheckRequest;
use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
use tonic_health::pb::health_client::HealthClient;
use tonic_health::server::health_reporter;
use tower::service_fn;
use vllm_chat::{
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
@@ -204,11 +200,7 @@ impl ChatRenderer for FakeTextBackend {
async fn setup_grpc_service(
engine_id: impl Into<EngineId>,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (
GenerateServer<GenerateServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint();
let engine_id = engine_id.into();
@@ -240,7 +232,6 @@ async fn setup_grpc_service(
)
.await
.expect("connect client");
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
test_llm(client),
@@ -249,7 +240,6 @@ async fn setup_grpc_service(
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
(
GenerateServer::new(GenerateServiceImpl::new(state)),
engine_health,
engine_task,
)
}
@@ -264,51 +254,25 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
svc,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
(GenerateClient::new(channel), server_task, engine_task)
}
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
let server = TonicServer::builder()
.add_service(health_service)
.add_service(generate_service)
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor =
super::monitor_health(health_reporter, engine_health, shutdown.clone());
let server = async move {
let result = server.await;
shutdown.cancel();
result
};
let (server_result, ()) = tokio::join!(server, health_monitor);
server_result.expect("grpc server");
TonicServer::builder()
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
});
let channel = Endpoint::from_shared(format!("http://{addr}"))
.expect("grpc endpoint")
.connect()
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
.await
.expect("connect grpc channel");
.expect("connect grpc client");
(channel, server_task)
(grpc_client, server_task, engine_task)
}
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -319,7 +283,7 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -409,8 +373,7 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr").to_string();
@@ -1072,129 +1035,3 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
service
};
let mut stream = health_client
.watch(HealthCheckRequest {
service: service.to_string(),
})
.await
.unwrap_or_else(|error| {
panic!("failed to start health watch for {service_label}: {error}")
})
.into_inner();
let initial = stream
.message()
.await
.unwrap_or_else(|error| {
panic!("failed to read initial health status for {service_label}: {error}")
})
.unwrap_or_else(|| {
panic!("health watch for {service_label} ended before its initial status")
});
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for {service_label}"
);
health_streams.push((service_label, stream));
}
engine_health_tx.send(false).expect("publish unhealthy engine state");
for (service_label, mut stream) in health_streams {
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))
.unwrap_or_else(|error| {
panic!("failed to read health update for {service_label}: {error}")
})
.unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected health status for {service_label}"
);
}
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
service: "vllm.Generate".to_string(),
})
.await
.expect("start health watch for vllm.Generate")
.into_inner();
let initial = stream
.message()
.await
.expect("read initial health status for vllm.Generate")
.expect("health watch ended before its initial status");
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for vllm.Generate"
);
shutdown.cancel();
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for shutdown health update for vllm.Generate")
.expect("failed to read shutdown health update for vllm.Generate")
.expect("health watch ended before its shutdown update");
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected shutdown health status for vllm.Generate"
);
let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for vllm.Generate health watch to close")
.expect("failed while closing vllm.Generate health watch");
assert!(
stream_end.is_none(),
"vllm.Generate health watch remained open"
);
tokio::time::timeout(Duration::from_secs(2), server_task)
.await
.expect("timed out waiting for gRPC server shutdown")
.expect("gRPC server task failed");
}
+14 -28
View File
@@ -39,7 +39,6 @@ use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
use tokio_util::sync::CancellationToken;
use tonic::transport::Server as TonicServer;
use tonic_health::server::health_reporter;
use tower::ServiceExt as _;
use tracing::{info, trace, warn};
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
@@ -204,19 +203,14 @@ where
.map(tls::build_grpc_server_config)
.transpose()
.context("invalid gRPC TLS configuration")?;
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(health_service)
.add_service(generate_service);
.add_service(svc);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
Some((grpc_listener, svc, grpc_tls))
} else {
None
};
@@ -300,8 +294,7 @@ where
let server_shutdown = server_shutdown.clone();
let force_shutdown = force_shutdown.clone();
async move {
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
else {
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
// No gRPC configured: just wait for shutdown so we do not race the
// join! by resolving early and tripping the cancellation token.
shutdown.cancelled().await;
@@ -311,26 +304,19 @@ where
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
None => MaybeTlsListener::plain(grpc_listener),
};
let server =
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
let server = async move {
let result = tokio::select! {
result = server => {
result.context("gRPC server failed")
}
_ = force_shutdown.cancelled() => {
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
Ok(())
}
};
server_shutdown.cancel();
result
let result = tokio::select! {
result = server => {
result.context("gRPC server failed")
}
_ = force_shutdown.cancelled() => {
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
Ok(())
}
};
let (result, ()) = tokio::join!(server, health_monitor);
server_shutdown.cancel();
result
}
};
+131
View File
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from types import ModuleType, SimpleNamespace
from typing import Any
import pytest
from vllm.platforms.interface import DeviceCapability
@pytest.fixture
def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
monkeypatch.setitem(
sys.modules,
"vllm._C_stable_libtorch",
stable_libtorch_module,
)
from vllm.platforms.cuda import CudaPlatformBase
return CudaPlatformBase
def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
assert cuda_platform_base._compiled_arch_covers_device(
"12.1a", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"12.0a", DeviceCapability(12, 1)
)
assert cuda_platform_base._compiled_arch_covers_device(
"12.0f", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"10.0f", DeviceCapability(12, 1)
)
def test_warn_if_device_arch_not_compiled(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 2
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
capabilities = {
0: DeviceCapability(12, 1),
1: DeviceCapability(10, 3),
}
return capabilities[device_id]
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
return f"GPU {device_id}"
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(
_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert len(warnings) == 1
assert warnings[0][1] == "12.0f, 10.0a"
assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
def test_warn_if_device_arch_not_compiled_no_warning(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 1
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
return DeviceCapability(10, 3)
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
raise AssertionError("get_device_name should not be called for covered devices")
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert warnings == []
@@ -0,0 +1,31 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import torch
from vllm import _custom_ops as ops
def test_cutlass_group_gemm_python_guard_allows_thor(monkeypatch):
seen_capabilities: list[int] = []
def fake_cutlass_group_gemm_supported(capability: int) -> bool:
seen_capabilities.append(capability)
return True
monkeypatch.setattr(
torch.ops,
"_C",
SimpleNamespace(cutlass_group_gemm_supported=fake_cutlass_group_gemm_supported),
)
# CUDA 12 reports Thor as SM101 and CUDA 13 reports it as SM110. Both
# should reach the C++ query for the SM10x/SM11x CUTLASS MoE kernel.
assert ops.cutlass_group_gemm_supported(101)
assert ops.cutlass_group_gemm_supported(110)
# SM120 uses separate kernels and must not be advertised by this path.
assert not ops.cutlass_group_gemm_supported(120)
assert seen_capabilities == [101, 110]
+1
View File
@@ -66,6 +66,7 @@ def test_worker_apply_lora(qwen3_lora_files):
runner_type="generate",
max_num_batched_tokens=32,
max_num_seqs=32,
max_num_partial_prefills=32,
),
device_config=DeviceConfig(DEVICE_TYPE),
cache_config=CacheConfig(
@@ -70,19 +70,7 @@ def _assert_video_outputs(processor, processed) -> None:
merge_size = processor.info.get_hf_config().vision_config.spatial_merge_size
expected_tokens = int(grid_thw.prod()) // merge_size**2
video_token_id = processor.info.get_hf_config().video_token_id
prompt_token_ids = processed["prompt_token_ids"]
assert prompt_token_ids.count(video_token_id) == expected_tokens
hf_processor = processor.info.get_hf_processor()
expected_frame_wrappers = int(grid_thw[:, 0].sum())
assert (
prompt_token_ids.count(hf_processor.vision_start_token_id)
== expected_frame_wrappers
)
assert (
prompt_token_ids.count(hf_processor.vision_end_token_id)
== expected_frame_wrappers
)
assert processed["prompt_token_ids"].count(video_token_id) == expected_tokens
@pytest.mark.parametrize("num_images", [1, 2])
-1
View File
@@ -170,7 +170,6 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
"layers.0.self_attn.to_add_out.weight",
"layers.0.self_attn.norm_added_q.weight",
"layers.0.self_attn.norm_added_k.weight",
"layers.0.self_attn.k_norm_und_for_gen.weight",
"layers.0.self_attn.q_proj_moe_gen.weight",
"layers.0.mlp_moe_gen.up_proj.weight",
"norm_moe_gen.weight",
-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(
"model_repo, expected_loader_cls, hf_sample_kwargs",
[
@@ -48,12 +48,12 @@ def _check_dense_embedding(data, index=0):
def _check_sparse_embedding(data, check_tokens=False):
expected_weights = [
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
{"token_id": 70, "weight": 0.09808349609375, "token": " the"},
{"token_id": 83, "weight": 0.08154296875, "token": " is"},
{"token_id": 111, "weight": 0.11810302734375, "token": " of"},
{"token_id": 4865, "weight": 0.1171875, "token": " What"},
{"token_id": 9942, "weight": 0.292236328125, "token": " France"},
{"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
{"token_id": 70, "weight": 0.09808349609375, "token": "the"},
{"token_id": 83, "weight": 0.08154296875, "token": "is"},
{"token_id": 111, "weight": 0.11810302734375, "token": "of"},
{"token_id": 4865, "weight": 0.1171875, "token": "What"},
{"token_id": 9942, "weight": 0.292236328125, "token": "France"},
{"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
]
expected_embed = {x["token_id"]: x for x in expected_weights}
+3 -6
View File
@@ -93,6 +93,9 @@ def test_online_quantization(
use_rocm_aiter: bool,
monkeypatch,
) -> None:
if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
pytest.skip("FA3 currently rejects FP8 KV cache output dtype on SM90")
if use_rocm_aiter:
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
@@ -102,15 +105,9 @@ def test_online_quantization(
if force_marlin:
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")
model_dtype = "auto"
if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
# FA3 requires BF16 output when the query input is FP8.
model_dtype = "bfloat16"
with vllm_runner(
"facebook/opt-125m",
quantization="fp8",
dtype=model_dtype,
enforce_eager=True,
kv_cache_dtype=kv_cache_dtype,
) as llm:
@@ -83,11 +83,6 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
ENFORCE_EAGER=${ENFORCE_EAGER:-1}
# Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
# Pin concurrent prefiller and non-DP decoder engines to separate internal
# port windows. DP decoder ranks retain their existing internal port selection.
PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000}
DECODER_INTERNAL_PORT_BASE=${DECODER_INTERNAL_PORT_BASE:-30000}
INTERNAL_PORT_STRIDE=${INTERNAL_PORT_STRIDE:-100}
# Resolve the repository root from the script location instead of `.git`.
# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
@@ -159,14 +154,12 @@ run_tests_for_model() {
PORT=$((8100 + i))
# Calculate side channel port. Avoid clash with with TP workers.
SIDE_CHANNEL_PORT=$((5559 + i))
INTERNAL_PORT=$((PREFILLER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
# Build the command with or without model-specific args
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
VLLM_KV_CACHE_LAYOUT='HND' \
VLLM_PORT=$INTERNAL_PORT \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
vllm serve $model_name \
@@ -215,18 +208,12 @@ run_tests_for_model() {
PORT=$((8200 + i))
# Calculate side channel port
SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
INTERNAL_PORT=$((DECODER_INTERNAL_PORT_BASE + i * INTERNAL_PORT_STRIDE))
DECODER_INTERNAL_PORT_ENV=
if [[ -z "${DP_EP:-}" ]]; then
DECODER_INTERNAL_PORT_ENV="VLLM_PORT=$INTERNAL_PORT"
fi
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
# Build the command with or without model-specific args
BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
$DECODER_INTERNAL_PORT_ENV \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
vllm serve $model_name \
@@ -12,9 +12,6 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
)
from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
_ConnectorMetricName,
@@ -112,40 +109,6 @@ def test_last_block_offloaded_at_request_finish(
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_abort_queued_request_does_not_build_store_job(
request_runner, async_scheduling: bool
):
"""Aborting a never-scheduled request must not store unallocated KV."""
block_size = 4
runner = request_runner(
block_size=block_size,
num_gpu_blocks=8,
async_scheduling=async_scheduling,
)
runner.new_request(token_ids=[0] * (block_size * 4))
runner.scheduler.schedule()
runner.new_request(token_ids=[1] * (block_size * 4))
queued_req_id = str(runner.req_id)
assert any(
request.request_id == queued_req_id for request in runner.scheduler.waiting
)
runner.scheduler.finish_requests(queued_req_id, RequestStatus.FINISHED_ABORTED)
req_status = runner.connector_scheduler._req_status[queued_req_id]
assert all(group_state.offload_keys for group_state in req_status.group_states)
assert all(not group_state.block_ids for group_state in req_status.group_states)
scheduler_output = runner.scheduler.schedule()
metadata = scheduler_output.kv_connector_metadata
assert isinstance(metadata, OffloadingConnectorMetadata)
assert all(job.req_id != queued_req_id for job in metadata.store_jobs.values())
assert queued_req_id not in runner.connector_scheduler._req_status
def test_scheduler_reports_lookup_sync_delay(request_runner):
runner = request_runner(
block_size=4,
@@ -125,7 +125,8 @@ def test_register_kv_caches(backend):
own dedicated tensors.
Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
the raw per-layer kv_caches registered by the connector.
kv_caches, which automatically applies
_update_hybrid_attention_mamba_layout for hybrid models.
Verifies that the canonicalized CanonicalKVCaches has the correct
block tensors, tensor_idx references, and page sizes across all groups.
+38
View File
@@ -5,11 +5,13 @@ import torch
from torch import Generator
from tests.utils import large_gpu_mark
from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.sample.ops.topk_topp_sampler import (
apply_top_k_top_p_pytorch,
flashinfer_sample,
random_sample,
)
from vllm.v1.sample.sampler import Sampler
@@ -1043,3 +1045,39 @@ class TestFlashInferDistributionMatch:
f"{label}: distribution differs from theoretical: "
f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
)
@pytest.mark.skipif(
not FLASHINFER_TOPK_TOPP_SUPPORTED,
reason="FlashInfer top-k/top-p sampler is not available on this platform.",
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)])
def test_flashinfer_sample_padded_vocab(
dtype: torch.dtype, k: int | None, p: float | None
):
"""flashinfer_sample must accept the logits the sampler actually hands it.
compute_logits slices the padding off the vocab, so for a vocab that isn't a
multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model
dtype, while FlashInfer requires contiguous fp32.
"""
torch.set_default_device(DEVICE_TYPE)
batch_size = 8
org_vocab_size = 50272
padded_vocab_size = pad_vocab_size(org_vocab_size)
assert padded_vocab_size != org_vocab_size
logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[
..., :org_vocab_size
]
# A single row stays contiguous despite the padded stride, hence batch_size > 1.
assert not logits.is_contiguous()
token_ids = flashinfer_sample(
logits,
torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None,
torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None,
)
assert token_ids.shape == (batch_size,)
assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size))
+1 -20
View File
@@ -36,22 +36,6 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]:
]
def write_coverage_objects(extensions: list[RustExtension], output: Path) -> None:
artifacts = []
for extension in extensions:
for target in sorted(set(extension.target.values())):
target_path = ROOT_DIR.joinpath(*target.split("."))
if extension.binding == Binding.Exec:
matches = [target_path]
else:
matches = sorted(target_path.parent.glob(f"{target_path.name}*.so"))
if len(matches) != 1 or not matches[0].is_file():
raise RuntimeError(f"unable to locate Rust artifact for {target}")
artifacts.append(matches[0].relative_to(ROOT_DIR).as_posix())
output.write_text("\n".join(artifacts) + "\n")
def rust_py_extension_module_names() -> list[str]:
module_names = []
for extension in rust_extensions():
@@ -68,15 +52,12 @@ def rust_py_extension_module_names() -> list[str]:
def build_binary(build_rust_args: list[str]) -> None:
os.chdir(ROOT_DIR)
(ROOT_DIR / "vllm").mkdir(exist_ok=True)
extensions = rust_extensions(optional=False)
setup(
name="vllm-rust-frontend-build",
packages=[],
rust_extensions=extensions,
rust_extensions=rust_extensions(optional=False),
script_args=["build_rust", "--quiet", "--inplace", *build_rust_args],
)
if output := os.getenv("VLLM_RUST_COVERAGE_OBJECTS"):
write_coverage_objects(extensions, Path(output))
def main() -> None:
+7 -3
View File
@@ -835,12 +835,16 @@ def cutlass_scaled_mm_azp(
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
if cuda_device_capability < 90 or cuda_device_capability >= 110:
# CUTLASS grouped FP8 MoE uses the same sm100-named implementation for
# SM10x and SM11x. Thor reports SM101 with CUDA 12 and SM110 with CUDA 13,
# so keep this Python guard aligned with the C++ support query/dispatch.
if cuda_device_capability < 90 or cuda_device_capability >= 120:
return False
try:
return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
except AttributeError:
# Return False on non-CUDA platforms where it is not available
except (AttributeError, RuntimeError):
# Return False on non-CUDA platforms where it is not available or not
# implemented for the current build.
return False
+40 -4
View File
@@ -67,6 +67,16 @@ class SchedulerConfig:
In real usage, this should be set in `EngineArgs.create_engine_config`.
"""
max_num_partial_prefills: int = Field(default=1, ge=1)
"""For chunked prefill, the maximum number of sequences that can be
partially prefilled concurrently."""
max_long_partial_prefills: int = Field(default=1, ge=1)
"""For chunked prefill, the maximum number of prompts longer than
long_prefill_token_threshold that will be prefilled concurrently. Setting
this less than max_num_partial_prefills will allow shorter prompts to jump
the queue in front of longer prompts in some cases, improving latency."""
long_prefill_token_threshold: int = Field(default=0, ge=0)
"""For chunked prefill, a request is considered long if the prompt is
longer than this number of tokens. 0 disables the cap (default)."""
@@ -244,6 +254,19 @@ class SchedulerConfig:
self.max_num_batched_tokens,
)
if self.max_num_partial_prefills > 1:
if self.long_prefill_token_threshold == 0:
self.long_prefill_token_threshold = int(max_model_len * 0.04)
logger.info(
"Concurrent partial prefills enabled with "
"max_num_partial_prefills=%d, max_long_partial_prefills=%d, "
"long_prefill_token_threshold=%d",
self.max_num_partial_prefills,
self.max_long_partial_prefills,
self.long_prefill_token_threshold,
)
self.verify_max_model_len(max_model_len)
def verify_max_model_len(self, max_model_len: int) -> Self:
@@ -275,11 +298,24 @@ class SchedulerConfig:
self.max_num_seqs * max_model_len,
)
if self.long_prefill_token_threshold > max_model_len:
if self.max_num_partial_prefills > 1:
if not self.enable_chunked_prefill:
raise ValueError(
"Chunked prefill must be enabled to set "
"max_num_partial_prefills > 1."
)
if self.long_prefill_token_threshold > max_model_len:
raise ValueError(
"long_prefill_token_threshold "
f"({self.long_prefill_token_threshold}) cannot be greater "
f"than the max_model_len ({max_model_len})."
)
if self.max_long_partial_prefills > self.max_num_partial_prefills:
raise ValueError(
"long_prefill_token_threshold "
f"({self.long_prefill_token_threshold}) cannot be greater "
f"than the max_model_len ({max_model_len})."
f"{self.max_long_partial_prefills=} must be less than or equal to "
f"{self.max_num_partial_prefills=}."
)
return self
@@ -21,10 +21,12 @@ from vllm.logger import init_logger
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.platforms import current_platform
from vllm.v1.attention.backend import AttentionBackend
from vllm.v1.kv_cache_interface import MambaSpec
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
if TYPE_CHECKING:
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
from vllm.v1.kv_cache_interface import KVCacheSpec
logger = init_logger(__name__)
@@ -592,6 +594,32 @@ class TransferTopology:
abs_ratio = -tp_ratio
return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
def get_transfer_cache_regions(
self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
) -> list[torch.Tensor] | torch.Tensor:
"""Return the cache tensor(s) to register as NIXL memory regions,
also accounting for hybrid SSM models specificities.
"""
if isinstance(layer_spec, MambaSpec):
# Register the whole kv cache shared tensor, including
# SSM/Conv.
conv, ssm = cache
return [conv]
# Check may be hacky but it's matching
# `_update_hybrid_attention_mamba_layout`.
if self.is_mamba and cache.shape[0] == 2:
# When MAMBA is present, all backends are blocks first, so
# that blocks can be shared between attention layers and mamba
# layers. Runner already adjusted strides for FlashAttn-like
# backends so its num_blocks first.
# Swap [2<>num_blocks] dims for hybrid SSM layout.
cache = cache.transpose(0, 1)
# K and V are packed into one tensor (content dim), so each layer
# registers as a single region.
return [cache]
def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
"""One-line summary of transfer config for logging."""
info = self._engines[(remote_engine_id, remote_pp_rank)]
@@ -1678,9 +1678,9 @@ class MooncakeConnectorWorker:
conv, _ = cache_or_caches
cache_list = [conv]
else:
# K and V are packed into one blocks-first tensor per layer,
# so each layer registers as a single region.
cache_list = [cache_or_caches]
cache_list = self.transfer_topo.get_transfer_cache_regions(
cache_or_caches, layer_spec
)
logger.debug(
"registering layer %s with %d cache tensor(s)",
@@ -1092,7 +1092,7 @@ class NixlBaseConnectorWorker:
# to better exploit the memory layout (ie num_blocks is the first dim).
tensor_size_bytes = None
for layer_name, cache in xfer_buffers.items():
for layer_name, cache_or_caches in xfer_buffers.items():
# NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
# that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
# However, physical page_size may differ when kernel requires a specific
@@ -1109,6 +1109,9 @@ class NixlBaseConnectorWorker:
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
# MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
layer_spec = layer_spec.kv_cache_specs[layer_name]
cache_list = self.transfer_topo.get_transfer_cache_regions(
cache_or_caches, layer_spec
)
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
# the page_size assuming constant `self._logical_num_blocks`.
physical_page_size = (
@@ -1117,6 +1120,8 @@ class NixlBaseConnectorWorker:
else layer_spec.page_size_bytes
// self._physical_blocks_per_logical_kv_block
)
# For when registering multiple tensors eg K/V in separate regions.
physical_page_size = physical_page_size // len(cache_list)
if self.transfer_topo._cross_layers_blocks:
# When cross-layers blocks are used, multiply by number of layers
physical_page_size = physical_page_size * len(
@@ -1131,61 +1136,66 @@ class NixlBaseConnectorWorker:
# [`num_blocks` * `page_size`]
curr_tensor_size_bytes = num_blocks * physical_page_size
base_addr = cache.data_ptr()
if base_addr in seen_base_addresses:
# NOTE (NickLucche) HMA employs memory pooling to share tensors
# across groups. This results in skipping all tensors but the ones
# pointed to by group0. Also, generally we will have more blocks
# per tensor but fewer regions.
logger.debug("Skipping %s because it's already seen", layer_name)
continue
logger.debug(
"Registering layer %s with cache shape: %s", layer_name, cache.shape
)
seen_base_addresses.append(base_addr)
# Only record non-Mamba page sizes.
if isinstance(layer_spec, MambaSpec):
self.block_len_per_layer.append(
physical_page_size // self._physical_blocks_per_logical_kv_block
# TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
# registering a single tensor for both K/V and splitting logically like FI.
for cache in cache_list:
base_addr = cache.data_ptr()
if base_addr in seen_base_addresses:
# NOTE (NickLucche) HMA employs memory pooling to share tensors
# across groups. This results in skipping all tensors but the ones
# pointed to by group0. Also, generally we will have more blocks
# per tensor but fewer regions.
logger.debug("Skipping %s because it's already seen", layer_name)
continue
logger.debug(
"Registering layer %s with cache shape: %s", layer_name, cache.shape
)
else:
self.block_len_per_layer.append(physical_page_size)
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
)
self._region_is_mla.append(is_mla_region)
if not is_mla_region:
if tensor_size_bytes is None:
tensor_size_bytes = curr_tensor_size_bytes
assert tensor_size_bytes == curr_tensor_size_bytes, (
"All non-MLA kv cache tensors must have the same size"
seen_base_addresses.append(base_addr)
# Only record non-Mamba page sizes.
if isinstance(layer_spec, MambaSpec):
self.block_len_per_layer.append(
physical_page_size // self._physical_blocks_per_logical_kv_block
)
else:
self.block_len_per_layer.append(physical_page_size)
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
)
self._region_is_mla.append(is_mla_region)
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
if (
self._physical_blocks_per_logical_kv_block == 1
and cache.shape[0] != num_blocks
):
raise AssertionError(
"All kv cache tensors must have the same number of "
f"blocks; layer={layer_name}, "
f"expected_num_blocks={num_blocks}, "
f"cache_shape={tuple(cache.shape)}, "
f"cache_stride={tuple(cache.stride())}, "
f"layer_spec={type(layer_spec).__name__}, "
f"backend={self.backend_name}, "
"all_backends="
f"{[backend.get_name() for backend in self.attn_backends]}, "
f"kv_cache_layout={self.kv_cache_layout}"
if not is_mla_region:
if tensor_size_bytes is None:
tensor_size_bytes = curr_tensor_size_bytes
assert tensor_size_bytes == curr_tensor_size_bytes, (
"All non-MLA kv cache tensors must have the same size"
)
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
if (
self._physical_blocks_per_logical_kv_block == 1
and cache.shape[0] != num_blocks
):
raise AssertionError(
"All kv cache tensors must have the same number of "
f"blocks; layer={layer_name}, "
f"expected_num_blocks={num_blocks}, "
f"cache_shape={tuple(cache.shape)}, "
f"cache_stride={tuple(cache.stride())}, "
f"layer_spec={type(layer_spec).__name__}, "
f"backend={self.backend_name}, "
"all_backends="
f"{[backend.get_name() for backend in self.attn_backends]}, "
f"kv_cache_layout={self.kv_cache_layout}"
)
# Need to make sure the device ID is non-negative for NIXL,
# Torch uses -1 to indicate CPU tensors.
self.device_id = max(cache.get_device(), 0)
caches_data.append(
(base_addr, curr_tensor_size_bytes, self.device_id, "")
)
# Need to make sure the device ID is non-negative for NIXL,
# Torch uses -1 to indicate CPU tensors.
self.device_id = max(cache.get_device(), 0)
caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))
logger.debug(
"Different block lengths collected: %s", set(self.block_len_per_layer)
)
@@ -326,12 +326,9 @@ class RequestOffloadState:
group_state.block_ids.extend(new_blocks)
def storable_chunks(
self,
group_config: "GroupOffloadConfig",
group_state: RequestGroupState,
num_offloadable_tokens: int,
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
) -> int:
"""Number of allocated leading offloaded chunks eligible for store.
"""Number of leading offloaded chunks eligible for store.
For eagle/MTP groups the volatile trailing chunk of the offloadable
range is excluded while decoding: the draft-layer KV of the last
@@ -348,10 +345,7 @@ class RequestOffloadState:
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
if group_config.is_eagle_group and is_decoding:
num_chunks = max(0, num_chunks - 1)
num_allocated_chunks = (
len(group_state.block_ids) // self.config.blocks_per_chunk
)
return min(num_chunks, num_allocated_chunks)
return num_chunks
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
# max(): at the prefill->decode transition of a chunk-aligned prompt,
@@ -362,7 +356,7 @@ class RequestOffloadState:
):
group_state.next_stored_chunk_idx = max(
group_state.next_stored_chunk_idx,
self.storable_chunks(group_config, group_state, num_offloadable_tokens),
self.storable_chunks(group_config, num_offloadable_tokens),
)
def update_num_hit_chunks(self, num_cached_tokens: int) -> None:
@@ -997,7 +991,7 @@ class OffloadingConnectorScheduler:
self.config.kv_group_configs, req_status.group_states
):
num_chunks = req_status.storable_chunks(
group_config, group_state, num_offloadable_tokens
group_config, num_offloadable_tokens
)
start_chunk_idx = group_state.next_stored_chunk_idx
@@ -1074,7 +1068,7 @@ class OffloadingConnectorScheduler:
group_config.sliding_window_size_in_chunks is not None
)
num_chunks = req_status.storable_chunks(
group_config, group_state, num_offloadable_tokens
group_config, num_offloadable_tokens
)
start_chunk_idx = group_state.next_stored_chunk_idx
block_ids = group_state.block_ids
@@ -56,7 +56,9 @@ class OffloadingConnectorWorker:
def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
self.worker = self.spec.get_worker(kv_caches)
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
def register_kv_caches(
self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
):
kv_cache_config = self.kv_cache_config
num_blocks = kv_cache_config.num_blocks
@@ -118,13 +120,24 @@ class OffloadingConnectorWorker:
)
elif isinstance(layer_kv_cache_spec, MambaSpec):
layer_kv_cache = kv_caches[layer_name]
assert layer_kv_cache.dtype == torch.int8
tensors_per_block[layer_name] = (
layer_kv_cache.view(
num_blocks, layer_kv_cache_spec.page_size_bytes
),
state_tensors = kv_caches[layer_name]
assert isinstance(state_tensors, list)
# re-construct the raw (num_blocks, page_size) tensor
# from the first state tensor
assert len(state_tensors) > 0
first_state_tensor = state_tensors[0]
assert first_state_tensor.storage_offset() == 0
tensor = (
torch.tensor(
[],
dtype=torch.int8,
device=first_state_tensor.device,
)
.set_(first_state_tensor.untyped_storage())
.view((num_blocks, layer_kv_cache_spec.page_size_bytes))
)
tensors_per_block[layer_name] = (tensor,)
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
unpadded_page_size_bytes[layer_name] = replace(
+19
View File
@@ -527,6 +527,8 @@ class EngineArgs:
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
max_num_batched_tokens: int | None = None
max_num_scheduled_tokens: int | None = None
max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills
max_long_partial_prefills: int = SchedulerConfig.max_long_partial_prefills
long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold
max_num_seqs: int | None = None
max_logprobs: int = ModelConfig.max_logprobs
@@ -1438,6 +1440,13 @@ class EngineArgs:
"default": None,
},
)
scheduler_group.add_argument(
"--max-num-partial-prefills", **scheduler_kwargs["max_num_partial_prefills"]
)
scheduler_group.add_argument(
"--max-long-partial-prefills",
**scheduler_kwargs["max_long_partial_prefills"],
)
scheduler_group.add_argument(
"--long-prefill-token-threshold",
**scheduler_kwargs["long_prefill_token_threshold"],
@@ -2181,6 +2190,8 @@ class EngineArgs:
is_encoder_decoder=model_config.is_encoder_decoder,
policy=self.scheduling_policy,
scheduler_cls=self.scheduler_cls,
max_num_partial_prefills=self.max_num_partial_prefills,
max_long_partial_prefills=self.max_long_partial_prefills,
long_prefill_token_threshold=self.long_prefill_token_threshold,
scheduler_reserve_full_isl=self.scheduler_reserve_full_isl,
watermark=self.watermark,
@@ -2391,6 +2402,14 @@ class EngineArgs:
def _check_feature_supported(self):
"""Raise an error if the feature is not supported."""
# No Concurrent Partial Prefills so far.
if (
self.max_num_partial_prefills != SchedulerConfig.max_num_partial_prefills
or self.max_long_partial_prefills
!= SchedulerConfig.max_long_partial_prefills
):
_raise_unsupported_error(feature_name="Concurrent Partial Prefill")
if self.pipeline_parallel_size > 1:
supports_pp = getattr(
self.distributed_executor_backend, "supports_pp", False
@@ -687,9 +687,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
skip_special_tokens=self.skip_special_tokens,
spaces_between_special_tokens=self.spaces_between_special_tokens,
include_stop_str_in_output=self.include_stop_str_in_output,
output_kind=(
RequestOutputKind.DELTA if self.stream else RequestOutputKind.FINAL_ONLY
),
output_kind=RequestOutputKind.DELTA
if self.stream
else RequestOutputKind.FINAL_ONLY,
structured_outputs=self.extract_structured_outputs(),
logit_bias=self.logit_bias,
bad_words=self.bad_words,
@@ -849,10 +849,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
# Reject empty tools array, matching OpenAI API behavior
if data.get("tools") == []:
raise VLLMValidationError(
raise ValueError(
"`tools` must not be an empty array. "
"Either provide at least one tool or omit the field entirely.",
parameter="tools",
"Either provide at least one tool or omit the field entirely."
)
# if "tool_choice" is not specified but tools are provided,
@@ -1076,15 +1075,13 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
if isinstance(data, BatchChatCompletionRequest):
data = data.model_dump(exclude_unset=True)
if data.get("use_beam_search"):
raise VLLMValidationError(
raise ValueError(
"Batch chat completions do not support beam search. "
"Please set `use_beam_search` to False.",
parameter="use_beam_search",
"Please set `use_beam_search` to False."
)
if data.get("logprob_token_ids") and not data.get("logprobs"):
raise VLLMValidationError(
"when using `logprob_token_ids`, `logprobs` must be set to true.",
parameter="logprob_token_ids",
raise ValueError(
"when using `logprob_token_ids`, `logprobs` must be set to true."
)
response_format = data.get("response_format")
rf_type = (
@@ -1098,10 +1095,8 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
validate_structured_outputs_structural_tag(structured_outputs)
n = data.get("n", 1)
if n is not None and n != 1:
raise VLLMValidationError(
"Batch chat completions do not support `n > 1`. Please set `n` to 1.",
parameter="n",
value=n,
raise ValueError(
"Batch chat completions do not support `n > 1`. Please set `n` to 1."
)
return data
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import time
from typing import TYPE_CHECKING, Literal, TypeAlias
@@ -79,6 +78,7 @@ class TranslationRequest(OpenAIBaseModel):
`verbose_json`, or `vtt`.
"""
# TODO support additional sampling parameters
# --8<-- [start:translation-sampling-params]
use_beam_search: bool = False
"""Whether or not beam search should be used."""
@@ -103,28 +103,6 @@ class TranslationRequest(OpenAIBaseModel):
will use [log probability](https://en.wikipedia.org/wiki/Log_probability)
to automatically increase the temperature until certain thresholds are hit.
"""
top_p: float | None = None
"""Enables nucleus (top-p) sampling, where tokens are selected from the
smallest possible set whose cumulative probability exceeds `p`.
"""
top_k: int | None = None
"""Limits sampling to the `k` most probable tokens at each step."""
min_p: float | None = None
"""Filters out tokens with a probability lower than `min_p`, ensuring a
minimum likelihood threshold during sampling.
"""
frequency_penalty: float | None = 0.0
"""The frequency penalty to use for sampling."""
repetition_penalty: float | None = None
"""The repetition penalty to use for sampling."""
presence_penalty: float | None = 0.0
"""The presence penalty to use for sampling."""
# --8<-- [end:translation-sampling-params]
# --8<-- [start:translation-extra-params]
@@ -161,23 +139,11 @@ class TranslationRequest(OpenAIBaseModel):
max_completion_tokens: int | None = None
"""The maximum number of tokens to generate."""
vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(
default=None,
description=(
"Additional request parameters with (list of) string or "
"numeric values, used by custom extensions."
),
)
# --8<-- [end:translation-extra-params]
# Default sampling parameters for translation requests.
_DEFAULT_SAMPLING_PARAMS: dict = {
"repetition_penalty": 1.0,
"temperature": 0,
"top_p": 1.0,
"top_k": 0,
"min_p": 0.0,
}
def build_stt_params(
@@ -233,38 +199,14 @@ class TranslationRequest(OpenAIBaseModel):
temperature = default_sampling_params.get(
"temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]
)
if (top_p := self.top_p) is None:
top_p = default_sampling_params.get(
"top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]
)
if (top_k := self.top_k) is None:
top_k = default_sampling_params.get(
"top_k", self._DEFAULT_SAMPLING_PARAMS["top_k"]
)
if (min_p := self.min_p) is None:
min_p = default_sampling_params.get(
"min_p", self._DEFAULT_SAMPLING_PARAMS["min_p"]
)
if (repetition_penalty := self.repetition_penalty) is None:
repetition_penalty = default_sampling_params.get(
"repetition_penalty",
self._DEFAULT_SAMPLING_PARAMS["repetition_penalty"],
)
return SamplingParams.from_optional(
temperature=temperature,
max_tokens=max_tokens,
seed=self.seed,
top_p=top_p,
top_k=top_k,
min_p=min_p,
frequency_penalty=self.frequency_penalty,
repetition_penalty=repetition_penalty,
presence_penalty=self.presence_penalty,
output_kind=RequestOutputKind.DELTA
if self.stream
else RequestOutputKind.FINAL_ONLY,
extra_args=self.vllm_xargs,
skip_clone=True, # Created fresh per request, safe to skip clone
)
@@ -284,16 +226,6 @@ class TranslationRequest(OpenAIBaseModel):
parameter=invalid_param,
)
xargs = data.get("vllm_xargs")
if isinstance(xargs, str):
try:
data["vllm_xargs"] = json.loads(xargs)
except json.JSONDecodeError as e:
raise VLLMValidationError(
f"Failed to parse vllm_xargs. Must be valid JSON: {e}",
parameter="vllm_xargs",
) from e
return data
@@ -4,8 +4,6 @@
from abc import ABC, abstractmethod
import torch
from vllm.config import VllmConfig
from vllm.v1.attention.backend import AttentionBackend, AttentionImpl
from vllm.v1.kv_cache_interface import KVCacheSpec
@@ -23,14 +21,6 @@ class AttentionLayerBase(ABC):
impl: "AttentionImpl"
supports_dcp: bool = True
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
"""Bind the allocated KV cache tensor to this layer.
The default stores the cache view as-is; subclasses (e.g. Mamba)
override this to unpack the raw buffer into per-state views.
"""
self.kv_cache = kv_cache
@abstractmethod
def get_attn_backend(self) -> type[AttentionBackend]:
"""Get the attention backend class for this layer."""
@@ -2,13 +2,11 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from collections.abc import Iterable
from math import prod
import torch
from vllm.config import VllmConfig
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.utils.torch_utils import get_dtype_size
from vllm.v1.attention.backend import AttentionBackend
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
from vllm.v1.attention.selector import get_mamba_attn_backend
@@ -26,22 +24,6 @@ class MambaBase(AttentionLayerBase):
kv_cache: tuple[torch.Tensor, ...]
supports_dcp: bool = False
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
"""Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.
Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
packed contiguously; slice them out and reinterpret per dtype/shape.
"""
pages = kv_cache.squeeze(dim=(1, 2))
states: list[torch.Tensor] = []
offset = 0
for shape, dtype in zip(self.get_state_shape(), self.get_state_dtype()):
nbytes = prod(shape) * get_dtype_size(dtype)
state = pages[:, offset : offset + nbytes].view(dtype)
states.append(state.view(-1, *shape))
offset += nbytes
self.kv_cache = tuple(states)
@abstractmethod
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
"""
@@ -23,7 +23,6 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
causal_conv1d_fn,
causal_conv1d_update,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.platforms import current_platform
from vllm.utils.torch_utils import direct_register_custom_op
from vllm.v1.attention.backend import AttentionMetadata
@@ -43,7 +42,6 @@ class ShortConv(MambaBase, CustomOp):
layer_idx: int,
model_config: ModelConfig | None = None,
cache_config: CacheConfig | None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
@@ -69,14 +67,12 @@ class ShortConv(MambaBase, CustomOp):
input_size=dim,
output_sizes=[dim] * 3,
bias=self.bias,
quant_config=quant_config,
prefix=f"{prefix}.in_proj",
)
self.out_proj = RowParallelLinear(
input_size=dim,
output_size=dim,
bias=self.bias,
quant_config=quant_config,
prefix=f"{prefix}.out_proj",
)
+1 -8
View File
@@ -276,13 +276,7 @@ class Cosmos3EdgeProcessingInfo(Qwen3VLProcessingInfo):
class Cosmos3EdgeMultiModalProcessor(Qwen3VLMultiModalProcessor):
@staticmethod
def _expands_only_video_token(_hf_processor: ProcessorMixin) -> bool:
# Cosmos renders each video as a full
# <|vision_start|><|video_pad|><|vision_end|> placeholder, and its
# reference processor replaces that entire triplet with timestamped,
# per-frame vision sequences.
return False
pass
class Cosmos3EdgeDummyInputsBuilder(Qwen3VLDummyInputsBuilder):
@@ -523,7 +517,6 @@ class Cosmos3EdgeForConditionalGeneration(
},
orig_to_new_substr={
"_moe_gen": None,
"k_norm_und_for_gen": None,
".add_q_proj.": None,
".add_k_proj.": None,
".add_v_proj.": None,
-1
View File
@@ -257,7 +257,6 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
layer_idx=layer_idx,
model_config=model_config,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.conv",
)
-1
View File
@@ -351,7 +351,6 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):
layer_idx=layer_idx,
model_config=model_config,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.conv",
)
+4 -3
View File
@@ -402,11 +402,12 @@ class Ovis2_5MultiModalProcessor(BaseMultiModalProcessor[Ovis2_5ProcessingInfo])
hf_processor_mm_kwargs: Mapping[str, object],
out_mm_kwargs: MultiModalKwargsItems,
) -> list[PromptReplacement]:
hf_processor = self.info.get_hf_processor()
tokenizer = self.info.get_tokenizer()
vocab = tokenizer.get_vocab()
placeholder = {
"image": hf_processor.get_token_value("image_token"),
"video": hf_processor.get_token_value("video_token"),
"image": vocab[IMAGE_TOKEN],
"video": vocab[VIDEO_TOKEN],
}
def get_replacement_ovis(item_idx, modality: str):
+1 -28
View File
@@ -4,7 +4,6 @@ from collections.abc import Mapping
from vllm.config import ModelConfig, VllmConfig
from vllm.logger import init_logger
from vllm.multimodal.inputs import MultiModalKwargsItem
from vllm.multimodal.processing import BaseMultiModalProcessor
from vllm.multimodal.registry import MultiModalRegistry
from vllm.utils.torch_utils import set_default_torch_num_threads
@@ -49,8 +48,6 @@ class MultiModalBudget:
self,
vllm_config: VllmConfig,
mm_registry: MultiModalRegistry,
*,
enable_cache: bool = True,
) -> None:
super().__init__()
@@ -61,11 +58,7 @@ class MultiModalBudget:
self.max_num_reqs = scheduler_config.max_num_seqs
with set_default_torch_num_threads(): # Avoid hang during startup
cache = (
mm_registry.processor_only_cache_from_config(vllm_config)
if enable_cache
else None
)
cache = mm_registry.processor_only_cache_from_config(vllm_config)
processor = mm_registry.create_processor(model_config, cache=cache)
self.cache = cache
@@ -198,23 +191,3 @@ class MultiModalBudget:
def reset_cache(self) -> None:
if self.cache is not None:
self.cache.clear_cache()
def get_dummy_encoder_profile_inputs(
mm_registry: MultiModalRegistry,
budget: MultiModalBudget,
) -> list[tuple[str, MultiModalKwargsItem]]:
if budget.get_encoder_budget() <= 0 or not budget.mm_max_toks_per_item:
return []
modality = budget.get_modality_with_max_tokens()
max_items_per_batch = budget.mm_max_items_per_batch[modality]
dummy_mm_inputs = mm_registry.get_dummy_mm_inputs(
budget.model_config,
mm_counts={modality: 1},
processor=budget.processor,
)
dummy_mm_item = dummy_mm_inputs["mm_kwargs"][modality][0]
assert dummy_mm_item is not None, "Dummy item should be generated"
return [(modality, dummy_mm_item)] * max_items_per_batch
+1 -1
View File
@@ -1177,7 +1177,7 @@ class PyNvVideoCodecVideoBackend(VideoBackend):
@VIDEO_LOADER_REGISTRY.register(
"qwen3_vl",
video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"),
video_processor="Qwen3VLVideoProcessor",
)
class Qwen3VLVideoBackend(VideoBackend):
@classmethod
+60
View File
@@ -14,6 +14,7 @@ from datetime import timedelta
from functools import cache, lru_cache, wraps
from typing import TYPE_CHECKING, NamedTuple, TypeVar
import regex as re
import torch
from torch.distributed import PrefixStore, ProcessGroup
from torch.distributed.distributed_c10d import is_nccl_available
@@ -48,6 +49,8 @@ _R = TypeVar("_R")
pynvml = import_pynvml()
_COMPILED_CUDA_ARCH_RE = re.compile(r"^(?P<major>\d+)\.(?P<minor>\d+)(?P<kind>[af])?$")
# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
# see https://github.com/huggingface/diffusers/issues/9704 for details
torch.backends.cuda.enable_cudnn_sdp(False)
@@ -282,6 +285,62 @@ class CudaPlatformBase(Platform):
def log_warnings(cls):
pass
@staticmethod
def _compiled_arch_covers_device(arch: str, capability: DeviceCapability) -> bool:
match = _COMPILED_CUDA_ARCH_RE.match(arch)
if match is None:
return False
major = int(match.group("major"))
minor = int(match.group("minor"))
kind = match.group("kind")
if kind == "f":
return major == capability.major
return major == capability.major and minor == capability.minor
@classmethod
def _warn_if_device_arch_not_compiled(cls) -> None:
try:
compiled_archs_str = torch.ops._C.get_compiled_cuda_archs()
except (AttributeError, RuntimeError):
return
compiled_archs = [
arch.strip() for arch in compiled_archs_str.split(",") if arch.strip()
]
if not compiled_archs:
return
unsupported_devices: list[str] = []
for device_id in range(cls.device_count()):
capability = cls.get_device_capability(device_id)
if capability is None:
continue
if any(
cls._compiled_arch_covers_device(arch, capability)
for arch in compiled_archs
):
continue
device_name = cls.get_device_name(device_id)
unsupported_devices.append(
f"{device_id}: {device_name} (compute capability "
f"{capability.as_version_str()})"
)
if unsupported_devices:
logger.warning_once(
"The current vLLM wheel was built for CUDA architectures %s, "
"but the following visible CUDA device(s) are not covered: %s. "
"This may cause CUDA kernel launch failures or force slower "
"fallback paths. Note that CUDA 13 family-specific targets "
"such as 10.0f and 12.0f cover the corresponding major-version "
"GPU family, while architecture-specific targets such as 12.1a "
"cover only that exact compute capability.",
", ".join(compiled_archs),
"; ".join(unsupported_devices),
)
@classmethod
def is_pin_memory_available(cls) -> bool:
if in_wsl():
@@ -949,6 +1008,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
@classmethod
@with_nvml_context
def log_warnings(cls):
cls._warn_if_device_arch_not_compiled()
device_ids: int = pynvml.nvmlDeviceGetCount()
if device_ids > 1:
device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
-1
View File
@@ -35,7 +35,6 @@ except json.JSONDecodeError:
# ---------------------------------------------------------------------------
DEFAULT_ENV_VAR_PREFIXES: set[str] = {
"VLLM_",
"FLASH_ATTENTION_",
"LMCACHE_",
"NCCL_",
"UCX_",
+15 -9
View File
@@ -78,6 +78,7 @@ class Ovis2_5Processor(ProcessorMixin):
@cached_property
def extra_special_tokens(self):
vocab = self.tokenizer.get_vocab()
required_tokens = {
"image_token": "<image>",
"video_token": "<video>",
@@ -89,16 +90,21 @@ class Ovis2_5Processor(ProcessorMixin):
"image_pad": "<|image_pad|>",
}
# The checkpoint defines both `additional_special_tokens` and
# `extra_special_tokens`, with the latter empty. Transformers ignores
# the former because the latter is explicitly empty, so the tokens are
# missing from the vocab. Re-add them to restore the expected ids.
self.tokenizer.add_tokens(list(required_tokens.values()), special_tokens=True)
extra_special_tokens = {}
suggestion = (
"please add '<image>', '<video>', '<ovis_visual_atom>', "
"'<ovis_image_start>', '<ovis_image_end>', '<ovis_video_start>', "
"'<ovis_video_end>' in 'additional_special_tokens' of "
"tokenizer_config.json, You can refer to "
"https://huggingface.co/AIDC-AI/Ovis2.6-30B-A3B/blob/main/tokenizer_config.json"
)
return {
key: self.tokenizer.convert_tokens_to_ids(token_name)
for key, token_name in required_tokens.items()
}
for key, token_name in required_tokens.items():
if token_name not in vocab:
raise ValueError(f"Can not find {token_name}, {suggestion}")
extra_special_tokens[key] = vocab[token_name]
return extra_special_tokens
def __call__(
self,
+6 -4
View File
@@ -392,8 +392,9 @@ def apply_top_k_top_p_pytorch(
logits_sort.masked_fill_(top_k_mask, -float("inf"))
if p is not None:
# Apply top-p.
probs_sort = logits_sort.softmax(dim=-1)
# Apply top-p. The cumsum below runs over the whole vocab, so accumulating
# in a low-precision dtype makes the nucleus undershoot p.
probs_sort = logits_sort.softmax(dim=-1, dtype=torch.float32)
probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
# at least one
@@ -500,9 +501,10 @@ def flashinfer_sample(
probs, k, deterministic=True
)
else:
# Both top-k and top-p.
# Both top-k and top-p. FlashInfer requires contiguous fp32 logits; the
# branches above get that from softmax().
next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
logits, k, p, deterministic=True
logits.float().contiguous(), k, p, deterministic=True
)
return next_token_ids.view(-1)
+19 -13
View File
@@ -131,7 +131,7 @@ def _topk_topp_kernel(
mask_n = offs < VOCAB_SIZE
logits_blk0 = tl.load(
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
# Exclude -inf values (e.g. from grammar bitmasks) from
# statistics to avoid NaN in pivot computation.
finite_mask = (logits_blk0 > -float("inf")) & mask_n
@@ -164,7 +164,7 @@ def _topk_topp_kernel(
mask_n = offs_n < VOCAB_SIZE
logits_blk = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
# Exclude -inf from min to keep binary search bounds
@@ -305,7 +305,7 @@ def _topk_topp_kernel(
mask_n = offs_n < VOCAB_SIZE
logits_blk2 = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
above_0 = logits_blk2 > k_pivot_0
above_1 = logits_blk2 > k_pivot_1
@@ -457,7 +457,7 @@ def _topk_topp_kernel(
LOGITS_ROW + offs_n,
mask=mask_n,
other=-float("inf"),
)
).to(tl.float32)
outlier_mask = (probs_blk > min_logit) & mask_n
@@ -600,7 +600,7 @@ def _topk_topp_kernel(
mask_n = offs < VOCAB_SIZE
logits_blk0 = tl.load(
LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
# Exclude -inf values (e.g. from grammar bitmasks) from
# statistics to avoid NaN in pivot computation.
finite_mask = (logits_blk0 > -float("inf")) & mask_n
@@ -626,7 +626,7 @@ def _topk_topp_kernel(
mask_n = offs_n < VOCAB_SIZE
logits_blk = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
max_logit = tl.maximum(max_logit, tl.max(logits_blk))
# Exclude -inf from min to keep binary search bounds
# finite (avoids NaN pivots).
@@ -660,7 +660,7 @@ def _topk_topp_kernel(
probs_blk = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
probs_blk = tl.exp(probs_blk - max_sample)
probs_blk = probs_blk / sum_exp_logits
@@ -754,7 +754,7 @@ def _topk_topp_kernel(
probs_blk = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
probs_blk = tl.exp(probs_blk - max_sample)
probs_blk = probs_blk / sum_exp_logits
tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)
@@ -835,7 +835,7 @@ def _topk_topp_kernel(
mask_n = offs_n < VOCAB_SIZE
logits_blk = tl.load(
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
)
).to(tl.float32)
keep_mask = (logits_blk > final_pivot) & mask_n
# Duplicate logit handling
@@ -878,7 +878,7 @@ def apply_top_k_top_p_triton(
The masked logits tensor. It may or may not be modified in-place.
"""
assert logits.ndim == 2
assert logits.dtype == torch.float32
assert logits.dtype in (torch.float32, torch.bfloat16, torch.float16)
batch_size, vocab_size = logits.shape
topk_enabled = k is not None
topp_enabled = p is not None
@@ -911,7 +911,9 @@ def apply_top_k_top_p_triton(
buffer = _TRITON_BUFFER_CACHE.get(buf_key)
if buffer is None or buffer.shape[0] < NUM_PROGRAMS:
size = min(next_power_of_2(NUM_PROGRAMS), num_sm)
buffer = logits.new_empty((size, vocab_size))
buffer = torch.empty(
(size, vocab_size), dtype=torch.float32, device=logits.device
)
_TRITON_BUFFER_CACHE[buf_key] = buffer
if buffer.shape[0] > NUM_PROGRAMS:
buffer = buffer[:NUM_PROGRAMS]
@@ -919,8 +921,12 @@ def apply_top_k_top_p_triton(
# Cache lookup table entries on each device.
tables = _TRITON_TABLE_CACHE.get(logits.device)
if tables is None:
normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)
percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)
normal_cdf_to_sigma_table = torch.tensor(
_NORMAL_CDF_TO_SIGMA_TABLE, dtype=torch.float32, device=logits.device
)
percentile_to_std_table = torch.tensor(
_PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=logits.device
)
_TRITON_TABLE_CACHE[logits.device] = (
normal_cdf_to_sigma_table,
percentile_to_std_table,
+87 -11
View File
@@ -262,7 +262,7 @@ def _reshape_kv_cache(
kv_cache_config: "KVCacheConfig | None" = None,
) -> dict[str, Any]:
kv_caches: dict[str, Any] = {}
has_attn = False
has_attn, has_mamba = False, False
layer_packing: dict[str, tuple[int, int]] = {}
if kv_cache_config is not None:
@@ -340,21 +340,38 @@ def _reshape_kv_cache(
)
elif isinstance(kv_cache_spec, MambaSpec):
page_size_bytes = kv_cache_spec.page_size_bytes
# Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
# int8 page view per layer; the layer's bind_kv_cache unpacks
# each block's bytes into its conv/ssm state views. Keeping
# one tensor per layer lets the KV connector register it
# without special-casing Mamba.
kv_caches[layer_name] = kv_raw_tensor[
: num_blocks * page_size_bytes
].view(num_blocks, 1, 1, page_size_bytes)
has_mamba = True
state_tensors = []
storage_offset_bytes = 0
for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
dtype_size = get_dtype_size(dtype)
num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size
target_shape = (num_blocks, *shape)
stride = torch.empty(target_shape).stride()
target_stride = (num_element_per_page, *stride[1:])
assert storage_offset_bytes % dtype_size == 0
tensor = torch.as_strided(
kv_raw_tensor.view(dtype),
size=target_shape,
stride=target_stride,
storage_offset=storage_offset_bytes // dtype_size,
)
state_tensors.append(tensor)
storage_offset_bytes += stride[0] * dtype_size
kv_caches[layer_name] = state_tensors
else:
raise NotImplementedError(
f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
)
if has_attn and kv_cache_config is not None:
if has_attn and has_mamba:
_update_hybrid_attention_layout(
attn_groups=attn_groups,
kv_caches=kv_caches,
kernel_block_sizes=kernel_block_sizes,
cache_dtype=cache_dtype,
)
elif has_attn and kv_cache_config is not None:
_align_mixed_attention_kv_cache_views(
attn_groups=attn_groups,
kv_caches=kv_caches,
@@ -441,6 +458,65 @@ def _restride_blocks_first_kv_cache_to_kv_first_storage(
)
def _update_hybrid_attention_layout(
attn_groups: Iterable[AttentionGroup],
kv_caches: dict[str, Any],
kernel_block_sizes: list[int],
cache_dtype: str,
) -> None:
for group in attn_groups:
if group.kv_cache_group_id >= len(kernel_block_sizes):
continue
kv_cache_spec = group.kv_cache_spec
if not isinstance(kv_cache_spec, AttentionSpec):
continue
# Mirror the per-layer dtype selection used when building the shape
# above. The block-dim index is dtype-independent for current backends
# (quantization only changes the last dim), so this is a no-op today,
# but it keeps both call sites consistent for skip layers.
layer_cache_dtype = (
"auto"
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
and not isinstance(kv_cache_spec, TQFullAttentionSpec)
else cache_dtype
)
block_dim = group.backend.get_kv_cache_block_dim(
kernel_block_sizes[group.kv_cache_group_id],
kv_cache_spec.num_kv_heads,
kv_cache_spec.head_size,
cache_dtype_str=layer_cache_dtype,
)
# if the first dim of the kvcache's layout is already num_blocks, continue
if block_dim == 0:
continue
assert block_dim == 1, (
"Expected the dim `num_blocks` at the second dim when updating"
" the kvcache's layout of full attention layer"
)
for layer_name in group.layer_names:
if layer_name not in kv_caches:
# Shared layer — will be aliased to its target after this pass.
continue
kv_cache = kv_caches[layer_name]
if kv_cache.shape[0] == 2:
assert kv_cache.shape[1] != 2, (
f"Cannot determine layout for tensor of shape {kv_cache.shape}"
)
hidden_size = kv_cache.shape[2:].numel()
kv_cache.as_strided_(
size=kv_cache.shape,
stride=(
hidden_size,
2 * hidden_size,
*kv_cache.stride()[2:],
),
)
def init_kv_cache(
runner_kv_caches: list[torch.Tensor | list[torch.Tensor]],
forward_context: dict[str, Any],
+1 -1
View File
@@ -28,7 +28,7 @@ class EncoderCache:
Clear the multi-modal cache that was used during profiling,
but no longer needed during inference.
"""
# NOTE: v2 encoder cache profiling skips the multi-modal cache
# TODO: Implement MM budget for encoder dummy run
pass
def reset_encoder_cache(self) -> None:
-43
View File
@@ -3,9 +3,7 @@
import numpy as np
import torch
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
from vllm.multimodal.encoder_budget import MultiModalBudget
from vllm.multimodal.inputs import MultiModalKwargsItem
from vllm.multimodal.utils import (
get_mm_features_in_window,
@@ -15,8 +13,6 @@ from vllm.multimodal.utils import (
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
logger = init_logger(__name__)
class EncoderRunner:
def __init__(
@@ -56,45 +52,6 @@ class EncoderRunner:
return mm_hashes, mm_kwargs
@torch.inference_mode()
def profile_encoder_cache(
self,
dummy_mm_inputs: list[tuple[str, MultiModalKwargsItem]],
budget: MultiModalBudget,
) -> None:
"""Profile multimodal encoder and temporary encoder cache memory."""
if (encoder_budget := budget.get_encoder_budget()) <= 0:
return
if not budget.mm_max_toks_per_item:
logger.info(
"Skipping encoder profiling for embedding-only mode "
"(all modality limits=0 with enable_mm_embeds=True).",
)
return
assert dummy_mm_inputs, "Dummy inputs should be generated for encoder profiling"
dummy_modality = dummy_mm_inputs[0][0]
max_mm_items_per_batch = len(dummy_mm_inputs)
logger.info_once(
"Encoder cache will be initialized with a budget of %s tokens, "
"and profiled with %s %s items of the maximum feature size.",
encoder_budget,
max_mm_items_per_batch,
dummy_modality,
)
dummy_encoder_outputs = self.execute_mm_encoder(dummy_mm_inputs)
sanity_check_mm_encoder_outputs(
dummy_encoder_outputs,
expected_num_items=max_mm_items_per_batch,
)
self.encoder_cache.encoder_outputs.update(
(f"tmp_{i}", output) for i, output in enumerate(dummy_encoder_outputs)
)
@torch.inference_mode()
def execute_mm_encoder(
self, mm_kwargs: list[tuple[str, MultiModalKwargsItem]]
-21
View File
@@ -43,10 +43,6 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
)
from vllm.model_executor.model_loader import get_model_loader
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.encoder_budget import (
MultiModalBudget,
get_dummy_encoder_profile_inputs,
)
from vllm.sequence import IntermediateTensors
from vllm.tasks import SupportedTask
from vllm.utils.math_utils import cdiv
@@ -675,22 +671,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
@torch.inference_mode()
def profile_run(self) -> None:
if self.supports_mm_inputs and self.is_first_pp_rank:
mm_config = self.model_config.multimodal_config
if mm_config is not None and not mm_config.skip_mm_profiling:
mm_budget = MultiModalBudget(
self.vllm_config,
self.mm_registry,
enable_cache=False,
)
dummy_mm_inputs = get_dummy_encoder_profile_inputs(
self.mm_registry,
mm_budget,
)
self.model_state.encoder_runner.profile_encoder_cache(
dummy_mm_inputs, mm_budget
)
hidden_states, sample_hidden_states = self._dummy_run(
self.max_num_tokens, skip_attn=True, is_profile=True
)
@@ -705,7 +685,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
torch.accelerator.synchronize()
del hidden_states, sample_hidden_states
self.reset_encoder_cache()
gc.collect()
def post_kv_cache_wake_up(self) -> None:
+3 -1
View File
@@ -218,7 +218,9 @@ def _bias_kernel(
mask=mask,
)
bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)
logits = tl.load(
logits_ptr + token_idx * logits_stride + token_ids, mask=mask
).to(tl.float32)
logits += bias
tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)

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