forked from Karylab-cklius/vllm
Compare commits
107
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d98b72bd3 | ||
|
|
75ccdf3145 | ||
|
|
c6fe94b4d5 | ||
|
|
46f01a50ac | ||
|
|
f00efc5265 | ||
|
|
b0cb1da1bd | ||
|
|
0e36e3bbd1 | ||
|
|
494845e79f | ||
|
|
0416dab275 | ||
|
|
c8db00b16c | ||
|
|
80c7683923 | ||
|
|
638d6e9757 | ||
|
|
1ad84fea86 | ||
|
|
12213c6795 | ||
|
|
10c75477b0 | ||
|
|
ac36a7a1e7 | ||
|
|
521aa80f71 | ||
|
|
a76df87db8 | ||
|
|
a4904ba903 | ||
|
|
f83de6d44c | ||
|
|
239fc73553 | ||
|
|
76bf55240c | ||
|
|
9a698f3255 | ||
|
|
4080263bb2 | ||
|
|
fc5fda105f | ||
|
|
b07ec92faa | ||
|
|
27ffbfde8d | ||
|
|
229e01e9e1 | ||
|
|
191146dba5 | ||
|
|
f3a920a076 | ||
|
|
149daf0d72 | ||
|
|
917fdb5bf7 | ||
|
|
4b594b4aa1 | ||
|
|
7d10a4cfce | ||
|
|
910cc8543a | ||
|
|
431934522b | ||
|
|
61a09532f2 | ||
|
|
3de4b2bf3c | ||
|
|
b44311b6ef | ||
|
|
b0d7875180 | ||
|
|
53c2f20dd9 | ||
|
|
37e370fe93 | ||
|
|
2dc5a72e7e | ||
|
|
c79ff5f918 | ||
|
|
1a659a0c37 | ||
|
|
0f6cf7f628 | ||
|
|
c79ad3ae21 | ||
|
|
61c9ef986a | ||
|
|
d6dbdb9b0d | ||
|
|
06da482fb4 | ||
|
|
2f75e7f712 | ||
|
|
7c21548ce3 | ||
|
|
9df2f91232 | ||
|
|
387189c429 | ||
|
|
75576c63be | ||
|
|
16aca639b7 | ||
|
|
6049424b7e | ||
|
|
060b5f61dc | ||
|
|
ec59c1579f | ||
|
|
1750e443f2 | ||
|
|
ba18929079 | ||
|
|
0500ca6a58 | ||
|
|
a1c15bcb0f | ||
|
|
4809de7317 | ||
|
|
05781e21dd | ||
|
|
85f638a2b8 | ||
|
|
08e5067561 | ||
|
|
a7d00ec051 | ||
|
|
b8fb56d970 | ||
|
|
96a739289e | ||
|
|
60d443f738 | ||
|
|
1dca300653 | ||
|
|
fca252d59e | ||
|
|
33178f9006 | ||
|
|
b2b8f679d0 | ||
|
|
de6ec294ef | ||
|
|
61e10f0116 | ||
|
|
6e96891ba0 | ||
|
|
47f1b47a73 | ||
|
|
5aab491bc9 | ||
|
|
5812e1a66b | ||
|
|
8950394e0a | ||
|
|
7bb49be4d1 | ||
|
|
c67650f04b | ||
|
|
f890e1dbe2 | ||
|
|
040cbf95cc | ||
|
|
5b3762a7f0 | ||
|
|
4d30c510ce | ||
|
|
6700813f86 | ||
|
|
eb44b3aaa4 | ||
|
|
7a98c7a392 | ||
|
|
0d9e60619b | ||
|
|
1134545b6f | ||
|
|
3e0c887511 | ||
|
|
adfbbc1005 | ||
|
|
adc98f04d0 | ||
|
|
8def3cdde2 | ||
|
|
616c9bd0f4 | ||
|
|
8688a06d67 | ||
|
|
f25953cc59 | ||
|
|
d9aa35161d | ||
|
|
6bcda970fd | ||
|
|
ea0e9c8f2e | ||
|
|
94ed0bf4e0 | ||
|
|
1940c8441e | ||
|
|
72d16aee15 | ||
|
|
e78a0c8e59 |
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
|
||||
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
|
||||
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
|
||||
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_METADATA_VERSION="1"
|
||||
IMAGE_EXISTED_BEFORE_BUILD=0
|
||||
|
||||
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
|
||||
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
|
||||
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
|
||||
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
|
||||
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
|
||||
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
|
||||
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
|
||||
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
|
||||
|
||||
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
|
||||
return 0
|
||||
fi
|
||||
|
||||
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
if [[ -n "${!var:-}" ]]; then
|
||||
echo "Using provided ${var}: ${!var}"
|
||||
continue
|
||||
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
|
||||
compute_dependency_cache_keys() {
|
||||
local bake_dir=""
|
||||
local dockerfile_rocm=""
|
||||
local rixl_branch=""
|
||||
local nixl_branch=""
|
||||
local ucx_branch=""
|
||||
local rocshmem_branch=""
|
||||
local deepep_branch=""
|
||||
local rixl_material=""
|
||||
local nixl_material=""
|
||||
local rocshmem_material=""
|
||||
local deepep_material=""
|
||||
|
||||
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
|
||||
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
|
||||
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
|
||||
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
|
||||
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
|
||||
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
|
||||
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
|
||||
|
||||
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
|
||||
RIXL_CACHE_KEY=$(
|
||||
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
|
||||
NIXL_CACHE_KEY=$(
|
||||
compose_dependency_cache_key \
|
||||
"${rixl_branch}-ucx-${ucx_branch}" \
|
||||
"${rixl_material}"
|
||||
"${nixl_branch}-ucx-${ucx_branch}" \
|
||||
"${nixl_material}"
|
||||
)
|
||||
export RIXL_CACHE_KEY
|
||||
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
|
||||
export NIXL_CACHE_KEY
|
||||
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
|
||||
fi
|
||||
|
||||
if [[ -n "${rocshmem_branch}" ]]; then
|
||||
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
|
||||
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
|
||||
|
||||
case "${target}" in
|
||||
rixl-rocm-ci)
|
||||
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
|
||||
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
nixl-rocm-ci)
|
||||
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
|
||||
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
fi
|
||||
;;
|
||||
rocshmem-rocm-ci)
|
||||
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
|
||||
|
||||
resolve_ci_base_dependency_targets() {
|
||||
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
|
||||
local rixl_ref=""
|
||||
local nixl_ref=""
|
||||
local rocshmem_ref=""
|
||||
local deepep_ref=""
|
||||
|
||||
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
|
||||
case "${mode}" in
|
||||
always)
|
||||
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
|
||||
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
|
||||
add_dependency_cache_target "${target}"
|
||||
fi
|
||||
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
|
||||
;;
|
||||
esac
|
||||
|
||||
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${rixl_ref}"; then
|
||||
echo "RIXL dependency cache exists: ${rixl_ref}"
|
||||
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${nixl_ref}"; then
|
||||
echo "NIXL dependency cache exists: ${nixl_ref}"
|
||||
else
|
||||
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
|
||||
add_dependency_cache_target "rixl-rocm-ci"
|
||||
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
|
||||
add_dependency_cache_target "nixl-rocm-ci"
|
||||
fi
|
||||
fi
|
||||
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -euox pipefail
|
||||
|
||||
export VLLM_CPU_KVCACHE_SPACE=1
|
||||
export VLLM_CPU_KVCACHE_SPACE=1
|
||||
export VLLM_CPU_CI_ENV=1
|
||||
# Reduce sub-processes for acceleration
|
||||
export TORCH_COMPILE_DISABLE=1
|
||||
# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
|
||||
# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
|
||||
# when callers specify fullgraph=True.
|
||||
export VLLM_ENABLE_V1_MULTIPROCESSING=0
|
||||
|
||||
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
|
||||
@@ -49,15 +50,15 @@ wait_for_pid_and_check_log() {
|
||||
}
|
||||
|
||||
# Test Sky Lake (AVX512F)
|
||||
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
|
||||
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
|
||||
PID_TEST_0=$!
|
||||
|
||||
# Test Cascade Lake (AVX512F + VNNI)
|
||||
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
|
||||
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
|
||||
PID_TEST_1=$!
|
||||
|
||||
# Test Cooper Lake (AVX512F + VNNI + BF16)
|
||||
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
|
||||
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
|
||||
PID_TEST_2=$!
|
||||
|
||||
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -26,7 +26,10 @@ steps:
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- vllm/config/compilation.py
|
||||
- vllm/compilation
|
||||
- vllm/v1/worker/encoder_cudagraph.py
|
||||
- vllm/v1/worker/encoder_cudagraph_defs.py
|
||||
commands:
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 35
|
||||
timeout_in_minutes: 53
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
|
||||
@@ -39,7 +39,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -59,7 +59,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -78,7 +78,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -156,7 +156,7 @@ steps:
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
@@ -90,6 +90,7 @@ steps:
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/simple_kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/streaming_input
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/ec_connector/unit
|
||||
- tests/v1/metrics
|
||||
@@ -103,6 +104,7 @@ steps:
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s v1/streaming_input
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -143,6 +145,7 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
|
||||
@@ -46,6 +46,19 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Inkling Unit Tests (B200)
|
||||
key: inkling-unit-tests-b200
|
||||
timeout_in_minutes: 40
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- vllm/models/inkling/
|
||||
- vllm/cute_utils/
|
||||
- cmake/external_projects/tml_fa4.cmake
|
||||
- tests/models/inkling/
|
||||
commands:
|
||||
# FA4 kernel tests require SM100; the suite skips them elsewhere.
|
||||
- pytest -v -s models/inkling
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
depends_on:
|
||||
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 98
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
device: h200_35gb
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 150
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
|
||||
@@ -170,3 +170,19 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
|
||||
|
||||
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
|
||||
key: spec-decode-deepseek-mtp-parallel-load-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/llm_base_proposer.py
|
||||
- vllm/v1/spec_decode/eagle.py
|
||||
- vllm/v1/worker/gpu/spec_decode/eagle/
|
||||
- vllm/model_executor/models/deepseek_mtp.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
|
||||
@@ -47,6 +47,7 @@
|
||||
|
||||
# Rust Frontend
|
||||
/rust/ @BugenZhao @njhill
|
||||
/rust/src/bench @esmeetu
|
||||
/build_rust.sh @BugenZhao @njhill
|
||||
/rust-toolchain.toml @BugenZhao @njhill
|
||||
/.buildkite/test_areas/rust* @BugenZhao @njhill
|
||||
|
||||
@@ -181,6 +181,18 @@ pull_request_rules:
|
||||
add:
|
||||
- performance
|
||||
|
||||
- name: label-quantization
|
||||
description: Automatically apply quantization label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm/model_executor/layers/quantization/
|
||||
- title~=(?i)quant
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- quantization
|
||||
|
||||
- name: label-qwen
|
||||
description: Automatically apply qwen label
|
||||
conditions:
|
||||
|
||||
@@ -130,6 +130,47 @@ jobs:
|
||||
},
|
||||
],
|
||||
},
|
||||
quantization: {
|
||||
keywords: [
|
||||
{
|
||||
term: "quantization",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "quantized",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
"intel-gpu": {
|
||||
// Keyword search - matches whole words only (with word boundaries)
|
||||
keywords: [
|
||||
{
|
||||
term: "B50",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B60",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B70",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "intel gpu",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "Arc GPU",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "BMG",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
// Add more label configurations here as needed
|
||||
// example: {
|
||||
// keywords: [...],
|
||||
@@ -491,4 +532,4 @@ jobs:
|
||||
issue_number: context.issue.number,
|
||||
body: message,
|
||||
});
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
|
||||
+44
-88
@@ -68,8 +68,8 @@ endif()
|
||||
# requirements.txt files and should be kept consistent. The ROCm torch
|
||||
# versions are derived from docker/Dockerfile.rocm
|
||||
#
|
||||
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
|
||||
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
|
||||
set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
|
||||
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
|
||||
# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
|
||||
# version check would always warn. Only treat it as a nightly build when the
|
||||
# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
|
||||
@@ -227,22 +227,6 @@ 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
|
||||
@@ -401,20 +385,8 @@ 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
|
||||
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
|
||||
"csrc/libtorch_stable/torch_bindings.cpp"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/activation_kernels.cu"
|
||||
@@ -437,7 +409,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"
|
||||
"${CACHE_KERNELS_SRC}"
|
||||
"csrc/libtorch_stable/cache_kernels.cu"
|
||||
"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")
|
||||
@@ -453,6 +425,11 @@ 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")
|
||||
@@ -490,14 +467,11 @@ 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"
|
||||
"${SCALED_MM_ENTRY_SRC}"
|
||||
"${NVFP4_QUANT_ENTRY_SRC}"
|
||||
"${NVFP4_SCALED_MM_ENTRY_SRC}"
|
||||
"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"
|
||||
"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
|
||||
@@ -573,15 +547,11 @@ 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
|
||||
"${COOPERATIVE_TOPK_SRC}")
|
||||
"csrc/libtorch_stable/cooperative_topk.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${COOPERATIVE_TOPK_SRC}"
|
||||
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
|
||||
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.
|
||||
@@ -790,10 +760,8 @@ 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}")
|
||||
@@ -826,10 +794,8 @@ 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}")
|
||||
@@ -862,10 +828,8 @@ 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}")
|
||||
@@ -894,10 +858,8 @@ 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)
|
||||
@@ -922,10 +884,8 @@ 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)
|
||||
@@ -948,13 +908,8 @@ 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}")
|
||||
# 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.
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
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)
|
||||
@@ -995,16 +950,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
|
||||
#
|
||||
|
||||
# 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")
|
||||
|
||||
# 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.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
@@ -1012,19 +961,18 @@ 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
|
||||
${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_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.")
|
||||
@@ -1039,10 +987,14 @@ 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
|
||||
${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_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/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.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}).")
|
||||
@@ -1050,10 +1002,9 @@ 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.")
|
||||
@@ -1107,6 +1058,7 @@ 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")
|
||||
@@ -1154,6 +1106,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
VLLM_ENABLE_COOPERATIVE_TOPK=1)
|
||||
endif()
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
|
||||
- Tool calling and reasoning parsers
|
||||
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
|
||||
- Efficient multi-LoRA support for dense and MoE layers
|
||||
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||
|
||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
fmha_sm100
|
||||
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
|
||||
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
|
||||
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
|
||||
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -344,28 +344,6 @@ 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
|
||||
|
||||
@@ -102,7 +102,9 @@ class TileGemm82 {
|
||||
kv_cache_t* __restrict__ curr_b = b_tile;
|
||||
|
||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
|
||||
auto fp32_b_regs = load_b_pair_vec(curr_b);
|
||||
auto fp32_b_0_reg = fp32_b_regs.first;
|
||||
auto fp32_b_1_reg = fp32_b_regs.second;
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
|
||||
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
|
||||
batch_offset += offset;
|
||||
int32_t block_table_id = batch_offset / block_size;
|
||||
int32_t slot_id = batch_offset % block_size;
|
||||
// seq_starts may push the block index past the end of the batch's block
|
||||
// table row.
|
||||
if (block_table_id >= block_table_stride) continue;
|
||||
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
|
||||
int32_t block_id = block_table[block_table_offset];
|
||||
int64_t cache_offset =
|
||||
|
||||
@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias);
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_sigmoid(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor);
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
|
||||
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
|
||||
double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid);
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding);
|
||||
|
||||
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> topk_ids,
|
||||
|
||||
@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int end_expert,
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor)
|
||||
const double routed_scaling_factor,
|
||||
const bool* is_padding)
|
||||
{
|
||||
|
||||
using cub_kvp = cub::KeyValuePair<int, float>;
|
||||
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int expert = result_kvp.key;
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
|
||||
|
||||
const int idx = k * block_row + k_idx;
|
||||
// Return the unbiased scores for output weights
|
||||
output[idx] = inputs_after_softmax[thread_read_offset + expert];
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
|
||||
assert(indices[idx] >= 0);
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : num_experts);
|
||||
assert(is_pad_row || indices[idx] >= 0);
|
||||
source_rows[idx] = k_idx * num_rows + block_row;
|
||||
if (renormalize) {
|
||||
selected_sum += inputs_after_softmax[thread_read_offset + expert];
|
||||
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
|
||||
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor)
|
||||
const float* bias, const double routed_scaling_factor, const bool* is_padding)
|
||||
{
|
||||
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
|
||||
// single) thread per row of the input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = max_val;
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -605,7 +610,7 @@ struct TopkConstants
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
|
||||
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
|
||||
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
|
||||
{
|
||||
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
|
||||
}
|
||||
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
|
||||
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor,
|
||||
cudaStream_t stream)
|
||||
cudaStream_t stream,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const float* bias_ptr = nullptr;
|
||||
if (bias.has_value()) {
|
||||
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
|
||||
bias_ptr = bias_tensor.const_data_ptr<float>();
|
||||
}
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
|
||||
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
|
||||
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
|
||||
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
|
||||
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -806,7 +825,8 @@ void topk_softmax(
|
||||
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias)
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -825,15 +845,15 @@ void topk_softmax(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
@@ -846,7 +866,8 @@ void topk_sigmoid(
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor)
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -865,15 +886,15 @@ void topk_sigmoid(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
|
||||
@@ -80,22 +80,27 @@ __launch_bounds__(128) __global__
|
||||
OutIndType* indices, int num_rows,
|
||||
int num_experts, float routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid) {
|
||||
const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
if (warp >= num_rows) return;
|
||||
const int64_t token_id = load_index_as_int64(input_ids, warp);
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
int expert = 0;
|
||||
float weight = 0.f;
|
||||
if (lane < 6) {
|
||||
if (lane < 6 && !is_pad_row) {
|
||||
// only load and calculate for 6 experts
|
||||
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
|
||||
const float x = input[warp * num_experts + expert];
|
||||
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
|
||||
if (isnan(weight)) {
|
||||
weight = 0.f;
|
||||
}
|
||||
}
|
||||
float weight_sum = weight;
|
||||
#pragma unroll
|
||||
@@ -111,7 +116,8 @@ __launch_bounds__(128) __global__
|
||||
const int offset = warp * 6 + lane;
|
||||
output[offset] =
|
||||
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
|
||||
indices[offset] = static_cast<OutIndType>(expert);
|
||||
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
|
||||
: static_cast<OutIndType>(-1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
int num_rows, int num_experts,
|
||||
double routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
const HashIndType* tid2eid, cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
if (num_rows == 0) return;
|
||||
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
|
||||
cudaLaunchConfig_t config = {};
|
||||
@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
config.numAttrs = 1;
|
||||
const float scale = static_cast<float>(routed_scaling_factor);
|
||||
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
|
||||
num_experts, scale, input_ids, tid2eid);
|
||||
num_experts, scale, input_ids, tid2eid, is_padding);
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid) {
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
return;
|
||||
}
|
||||
const bool row_is_active = finished ? !finished[thread_row] : true;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// We finally start setting up the read pointers for each thread. First, each
|
||||
// thread jumps to the start of the row it will read.
|
||||
@@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
// to float
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = 0.f;
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, float>) {
|
||||
using VecType = AlignedArray<float, ELTS_PER_LDG>;
|
||||
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
@@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
if constexpr (USE_HASH) {
|
||||
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
|
||||
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
row_chunk[ii] = sqrtf(val);
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
}
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
@@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = static_cast<IndType>(expert);
|
||||
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
|
||||
: static_cast<IndType>(-1);
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
@@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
#endif
|
||||
return;
|
||||
} else {
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
|
||||
// Original TopK path: find top-k experts by score
|
||||
@@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert =
|
||||
expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool should_process_row =
|
||||
row_is_active && node_uses_expert && !is_pad_row;
|
||||
|
||||
// The lead thread from each sub-group will write out the final results
|
||||
// to global memory. (This will be a single) thread per row of the
|
||||
// input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
if (correction_bias != nullptr) {
|
||||
if (correction_bias != nullptr && should_process_row) {
|
||||
max_val -= correction_bias[expert];
|
||||
}
|
||||
output[idx] = max_val;
|
||||
indices[idx] =
|
||||
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
@@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
|
||||
indices, source_row, k, start_expert, end_expert,
|
||||
renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#else
|
||||
kernel<<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert,
|
||||
end_expert, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#endif
|
||||
})
|
||||
}
|
||||
@@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
|
||||
stream);
|
||||
stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream, const bool* is_padding) {
|
||||
#ifndef USE_ROCM
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (use_hash && topk == 6 && renormalize &&
|
||||
(num_experts == 256 || num_experts == 384)) {
|
||||
launchDsv4HashTopk<IndType, HashIndType>(
|
||||
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
|
||||
routed_scaling_factor, input_ids, tid2eid, stream);
|
||||
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().const_data_ptr<float>();
|
||||
@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
auto launch = [&](auto* topk_indices_ptr) {
|
||||
using OutIndType =
|
||||
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
|
||||
"is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
|
||||
"is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (tid2eid.has_value()) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(),
|
||||
"input_ids is required for hash MoE");
|
||||
@@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream);
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
|
||||
torch::headeronly::ScalarType::Int);
|
||||
@@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
|
||||
tid2eid.value().const_data_ptr<int>(), stream);
|
||||
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
|
||||
}
|
||||
} else {
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
|
||||
@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
|
||||
static_cast<const OutIndType*>(nullptr), stream);
|
||||
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -793,7 +842,8 @@ void topk_softplus_sqrt(
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
@@ -806,21 +856,22 @@ void topk_softplus_sqrt(
|
||||
dispatch_topk_softplus_sqrt_launch<float>(
|
||||
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
|
||||
token_expert_indices, num_tokens, num_experts, topk, renormalize,
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
|
||||
is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_softplus_sqrt_launch<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
|
||||
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
m.def(
|
||||
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
|
||||
"bias) -> ()");
|
||||
"bias, Tensor? is_padding) -> ()");
|
||||
|
||||
// Apply topk sigmoid to the gating outputs.
|
||||
m.def(
|
||||
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, "
|
||||
"Tensor? bias, float routed_scaling_factor) -> ()");
|
||||
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
|
||||
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
|
||||
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts. topk_ids/expert_map are optional and, when
|
||||
|
||||
@@ -47,8 +47,6 @@ 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,8 +51,7 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
|
||||
|
||||
#endif
|
||||
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
|
||||
torch::stable::Tensor const& a_tensors,
|
||||
torch::stable::Tensor const& b_tensors,
|
||||
@@ -84,9 +83,8 @@ 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_SM10X_OR_SM11X) && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
|
||||
void get_cutlass_moe_mm_data_caller(
|
||||
const torch::stable::Tensor& topk_ids,
|
||||
@@ -177,14 +175,11 @@ 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 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.
|
||||
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
|
||||
// actual cutlass_moe_mm dispatch compiled into this file.
|
||||
|
||||
#if defined CUDA_VERSION
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
|
||||
return CUDA_VERSION >= 12080;
|
||||
}
|
||||
@@ -286,11 +281,8 @@ 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_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) {
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
if (version_num >= 100 && version_num < 110) {
|
||||
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);
|
||||
@@ -324,9 +316,8 @@ 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_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
(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,
|
||||
@@ -347,9 +338,8 @@ 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_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
(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);
|
||||
@@ -372,9 +362,8 @@ 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_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
|
||||
problem_sizes2, expert_num_tokens,
|
||||
|
||||
@@ -4,10 +4,6 @@
|
||||
|
||||
#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.
|
||||
@@ -34,7 +30,6 @@ 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
|
||||
|
||||
@@ -768,7 +763,6 @@ 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",
|
||||
|
||||
+2
-2
@@ -22,7 +22,7 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=13.0.2
|
||||
ARG CUDA_VERSION=13.0.3
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
@@ -793,7 +793,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
# https://docs.flashinfer.ai/installation.html
|
||||
# From versions.json: .flashinfer.version
|
||||
ARG FLASHINFER_VERSION=0.6.14
|
||||
ARG FLASHINFER_VERSION=0.6.15.post1
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
+53
-31
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="39be1de8"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="bfb51733"
|
||||
# NIXL/UCX build stages
|
||||
FROM base AS build_nixl
|
||||
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
|
||||
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
|
||||
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
|
||||
ENV NIXL_HOME=/usr/local/nixl
|
||||
|
||||
# RIXL build system dependences and RDMA support
|
||||
# NIXL build system dependencies and RDMA support
|
||||
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
libgrpc-dev \
|
||||
libgrpc++-dev \
|
||||
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system meson auditwheel patchelf tomlkit
|
||||
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
|
||||
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
cd /usr/local/src && \
|
||||
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
|
||||
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
git clone ${RIXL_REPO} /opt/rixl && \
|
||||
cd /opt/rixl && \
|
||||
git checkout ${RIXL_BRANCH} && \
|
||||
git clone ${NIXL_REPO} /opt/nixl && \
|
||||
cd /opt/nixl && \
|
||||
git checkout ${NIXL_BRANCH} && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
meson setup build --prefix=${RIXL_HOME} \
|
||||
meson setup build --prefix=${NIXL_HOME} \
|
||||
-Ducx_path=${UCX_HOME} \
|
||||
-Drocm_path=${ROCM_PATH} && \
|
||||
-Dwheel_variant=rocm \
|
||||
-Dbuild_tests=false \
|
||||
-Dbuild_examples=false && \
|
||||
cd build && \
|
||||
ninja -j$(nproc) && \
|
||||
ninja install
|
||||
ninja install && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
|
||||
> /etc/ld.so.conf.d/nixl.conf && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
|
||||
>> /etc/ld.so.conf.d/nixl.conf && \
|
||||
ldconfig
|
||||
|
||||
# Generate RIXL wheel
|
||||
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
|
||||
# so configure the ROCm wheel variant directly through Meson.
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
_ucx_install_dir=${UCX_HOME} \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
RUN cd /opt/nixl && \
|
||||
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
|
||||
--python ${PYTHON_VERSION} \
|
||||
-Csetup-args=-Ducx_path=${UCX_HOME} \
|
||||
-Csetup-args=-Dwheel_variant=rocm \
|
||||
-Csetup-args=-Dbuild_tests=false \
|
||||
-Csetup-args=-Dbuild_examples=false && \
|
||||
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
|
||||
auditwheel repair \
|
||||
--exclude 'libamdhip64*' \
|
||||
--exclude 'libcore*' \
|
||||
--exclude 'libpull*' \
|
||||
/tmp/nixl_wheels/nixl_rocm*.whl \
|
||||
--plat manylinux_2_34_$(uname -m) \
|
||||
--wheel-dir /tmp/nixl_wheels/repaired && \
|
||||
./contrib/wheel_add_ucx_plugins.py \
|
||||
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
|
||||
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
|
||||
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
|
||||
/tmp/nixl_wheels/repaired/*.whl && \
|
||||
cp /tmp/nixl_wheels/repaired/*.whl /app/install
|
||||
|
||||
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
|
||||
# invalidate the slow ROCShmem build.
|
||||
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
|
||||
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
|
||||
fi
|
||||
|
||||
# Install RIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
# Install NIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
|
||||
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
|
||||
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
|
||||
|
||||
# Copy ROCShmem runtime libraries.
|
||||
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
|
||||
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
# ROCm profiler limits workaround.
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
&& pip uninstall -y vllm \
|
||||
&& uv pip install --system *.whl
|
||||
|
||||
# Install RIXL wheel
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
# Install NIXL ROCm wheel
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
uv pip install --system /nixl_install/*.whl
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
ARG BASE_IMAGE
|
||||
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.16.post3"
|
||||
ARG AITER_BRANCH="v0.1.16.post5"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
+25
-1
@@ -86,6 +86,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p /tmp/hf-xet/dist && \
|
||||
cp dist/*.whl /tmp/hf-xet/dist/
|
||||
|
||||
# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
|
||||
# llvmlite v0.47 does not support; only SystemZ target is needed).
|
||||
FROM base AS llvm20-build
|
||||
ARG LLVM_VERSION=20.1.8
|
||||
WORKDIR /tmp
|
||||
RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
|
||||
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
|
||||
tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
|
||||
cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
|
||||
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
|
||||
-DLLVM_ENABLE_RTTI=ON \
|
||||
-DLLVM_BUILD_TOOLS=OFF \
|
||||
-DLLVM_BUILD_UTILS=ON \
|
||||
-DLLVM_BUILD_EXAMPLES=OFF \
|
||||
-DLLVM_BUILD_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_EXAMPLES=OFF \
|
||||
-DLLVM_INCLUDE_BENCHMARKS=OFF && \
|
||||
ninja -C build install && \
|
||||
rm -rf build llvm-project-${LLVM_VERSION}.src*
|
||||
|
||||
# Build numba
|
||||
FROM python-install AS numba-builder
|
||||
|
||||
@@ -96,11 +119,13 @@ WORKDIR /tmp
|
||||
|
||||
# Clone all required dependencies
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
|
||||
microdnf install ninja-build gcc gcc-c++ -y && \
|
||||
git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
|
||||
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
|
||||
cd llvmlite && \
|
||||
uv pip install 'cmake<4' 'setuptools<70' numpy && \
|
||||
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
|
||||
python setup.py bdist_wheel && \
|
||||
cd ../numba && \
|
||||
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
|
||||
@@ -158,7 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
|
||||
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
|
||||
uv pip install -v \
|
||||
$ARROW_WHL_FILE \
|
||||
$VISION_WHL_FILE \
|
||||
$HF_XET_WHL_FILE \
|
||||
$LLVM_WHL_FILE \
|
||||
|
||||
+13
-13
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
|
||||
}
|
||||
|
||||
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
|
||||
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
|
||||
# stage and is irrelevant when building --target ci_base itself.
|
||||
variable "CI_BASE_IMAGE" {
|
||||
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
|
||||
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
|
||||
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
|
||||
# functions produce no entries when the variable is empty.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "RIXL_CACHE_KEY" {
|
||||
variable "NIXL_CACHE_KEY" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
|
||||
])
|
||||
}
|
||||
|
||||
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
|
||||
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
|
||||
# These stages are pinned to specific upstream commit hashes, so cache keys use
|
||||
# those hashes rather than the Buildkite commit. This means the cache persists
|
||||
# across all vLLM commits as long as the upstream dependency pins don't change.
|
||||
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
|
||||
function "get_cache_from_rocm_deps" {
|
||||
params = []
|
||||
result = compact([
|
||||
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
|
||||
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
|
||||
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
function "get_cache_to_rocm_rixl" {
|
||||
function "get_cache_to_rocm_nixl" {
|
||||
params = []
|
||||
result = compact([
|
||||
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
|
||||
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
|
||||
# (e.g., requirements change), these stages are cache hits if their upstream
|
||||
# pins haven't changed -- saving ~35min of compilation.
|
||||
target "rixl-rocm-ci" {
|
||||
target "nixl-rocm-ci" {
|
||||
inherits = ["_common-rocm", "_ci-rocm"]
|
||||
target = "build_rixl"
|
||||
target = "build_nixl"
|
||||
cache-from = get_cache_from_rocm_deps()
|
||||
cache-to = get_cache_to_rocm_rixl()
|
||||
cache-to = get_cache_to_rocm_nixl()
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
|
||||
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
|
||||
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
|
||||
# files drifts from the remote image label. Per-PR builds then pull the result
|
||||
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
|
||||
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
|
||||
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
|
||||
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
|
||||
]),
|
||||
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
|
||||
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
|
||||
# are cache hits even when ci_base itself needs rebuilding.
|
||||
get_cache_from_rocm_deps(),
|
||||
)
|
||||
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
|
||||
# Group for ci_base builds -- exports dependency stage caches alongside the
|
||||
# ci_base image so future rebuilds can reuse them independently.
|
||||
group "ci-base-rocm-ci-with-deps" {
|
||||
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
|
||||
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
|
||||
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
|
||||
# overrides) and writes a small HCL override before invoking CI targets.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ target "test-rocm" {
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
|
||||
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
|
||||
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
|
||||
# the auto-rebuild trigger when requirements change in a PR.
|
||||
target "ci-base-rocm" {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "13.0.2"
|
||||
"default": "13.0.3"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
@@ -11,10 +11,10 @@
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
|
||||
},
|
||||
"BUILD_OS": {
|
||||
"default": "ubuntu"
|
||||
@@ -68,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.14"
|
||||
"default": "0.6.15.post1"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
|
||||
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
|
||||
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
|
||||
|
||||
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
|
||||
|
||||
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
|
||||
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
|
||||
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
|
||||
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
|
||||
|
||||
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
|
||||
|
||||
|
||||
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
|
||||
- vLLM CPU supports quantizations:
|
||||
- AWQ (x86 only)
|
||||
- GPTQ (x86 only)
|
||||
- compressed-tensor INT8 W8A8 (x86, s390x)
|
||||
- compressed-tensor INT8 W8A8 (x86 only)
|
||||
|
||||
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
|
||||
- OS: `Linux`
|
||||
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
|
||||
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
|
||||
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
|
||||
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
|
||||
|
||||
--8<-- [end:requirements]
|
||||
--8<-- [start:set-up-using-python]
|
||||
@@ -28,13 +28,24 @@ Install the following packages from the package manager before building the vLLM
|
||||
|
||||
```bash
|
||||
dnf install -y \
|
||||
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
|
||||
which procps findutils tar vim git patch xz ninja-build \
|
||||
gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
|
||||
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
|
||||
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
|
||||
clang llvm-devel llvm-static clang-devel
|
||||
```
|
||||
|
||||
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
|
||||
Build and install `numactl` from source:
|
||||
|
||||
```bash
|
||||
curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
|
||||
tar -xvzf v2.0.19.tar.gz
|
||||
cd numactl-2.0.19
|
||||
./autogen.sh && ./configure && make && make install
|
||||
cd ..
|
||||
```
|
||||
|
||||
Install rust>=1.80 which is needed for `outlines-core`, `uvloop`, and `hf-xet` python packages installation.
|
||||
|
||||
```bash
|
||||
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
@@ -44,26 +55,79 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
Execute the following commands to build and install vLLM from source.
|
||||
|
||||
!!! tip
|
||||
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
|
||||
Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`.
|
||||
See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build.
|
||||
|
||||
!!! note "LLVM 20 required for llvmlite"
|
||||
`llvmlite v0.47` requires LLVM 20, but UBI 9.6 repos ship LLVM 21 which is
|
||||
not compatible. You must build LLVM 20 from source before building `llvmlite`:
|
||||
|
||||
```bash
|
||||
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-20.1.8/llvm-project-20.1.8.src.tar.xz
|
||||
tar -xf llvm-project-20.1.8.src.tar.xz
|
||||
cmake -G Ninja -S llvm-project-20.1.8.src/llvm -B llvm-build \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
|
||||
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
|
||||
-DLLVM_ENABLE_RTTI=ON \
|
||||
-DLLVM_BUILD_TOOLS=OFF \
|
||||
-DLLVM_BUILD_UTILS=ON \
|
||||
-DLLVM_BUILD_EXAMPLES=OFF \
|
||||
-DLLVM_BUILD_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_EXAMPLES=OFF \
|
||||
-DLLVM_INCLUDE_BENCHMARKS=OFF
|
||||
ninja -C llvm-build install
|
||||
```
|
||||
|
||||
Then build `llvmlite` pointing to LLVM 20:
|
||||
|
||||
```bash
|
||||
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
|
||||
python setup.py bdist_wheel
|
||||
```
|
||||
|
||||
```bash
|
||||
uv pip install -v \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match && \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
uv pip install -v \
|
||||
/path/to/torchvision.whl \
|
||||
/path/to/llvmlite.whl \
|
||||
/path/to/numba.whl \
|
||||
/path/to/opencv_python_headless.whl \
|
||||
/path/to/hf_xet.whl \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match && \
|
||||
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
```
|
||||
|
||||
??? console "pip"
|
||||
```bash
|
||||
pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
pip install dist/*.whl
|
||||
pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
/path/to/torchvision.whl \
|
||||
/path/to/llvmlite.whl \
|
||||
/path/to/numba.whl \
|
||||
/path/to/opencv_python_headless.whl \
|
||||
/path/to/hf_xet.whl \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt && \
|
||||
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
|
||||
pip install dist/*.whl
|
||||
```
|
||||
|
||||
!!! warning "Protobuf workaround for s390x"
|
||||
The C++ protobuf extension crashes on s390x. After installation, set the
|
||||
following environment variable and remove the C++ extensions:
|
||||
|
||||
```bash
|
||||
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
|
||||
|
||||
# Remove C++ protobuf extensions that crash on s390x
|
||||
SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
rm -rf "$SITE_PKGS/google/_upb/"*.so \
|
||||
"$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true
|
||||
```
|
||||
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
@@ -80,19 +144,20 @@ docker build -f docker/Dockerfile.s390x \
|
||||
|
||||
# Launch OpenAI server
|
||||
docker run --rm \
|
||||
--privileged true \
|
||||
--security-opt seccomp=unconfined \
|
||||
--cap-add SYS_NICE \
|
||||
--shm-size 4g \
|
||||
-p 8000:8000 \
|
||||
-e VLLM_CPU_KVCACHE_SPACE=<KV cache space> \
|
||||
-e VLLM_CPU_OMP_THREADS_BIND=<CPU cores for inference> \
|
||||
vllm-cpu-env \
|
||||
--model meta-llama/Llama-3.2-1B-Instruct \
|
||||
--dtype float \
|
||||
--dtype bfloat16 \
|
||||
other vLLM OpenAI server arguments
|
||||
```
|
||||
|
||||
!!! tip
|
||||
An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
|
||||
Alternatively, `--privileged=true` also works but is broader and not generally recommended.
|
||||
|
||||
--8<-- [end:build-image-from-source]
|
||||
--8<-- [start:extra-information]
|
||||
|
||||
@@ -46,26 +46,6 @@ 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.14.37833.4) release, to avoid potential compatibility issue.
|
||||
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
--8<-- [start:pre-built-images]
|
||||
|
||||
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
|
||||
vLLM offers official Docker images for deployment.
|
||||
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||
|
||||
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
|
||||
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--network=host \
|
||||
--device /dev/dri:/dev/dri \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--ipc=host \
|
||||
--privileged \
|
||||
vllm/vllm-openai-xpu:<tag> \
|
||||
--model Qwen/Qwen3-0.6B
|
||||
```
|
||||
|
||||
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
|
||||
|
||||
???+ console "Commands"
|
||||
```bash
|
||||
docker run --rm -it \
|
||||
--network=host \
|
||||
--device /dev/dri:/dev/dri \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--ipc=host \
|
||||
--privileged \
|
||||
--entrypoint /bin/bash \
|
||||
vllm/vllm-openai-xpu:<tag>
|
||||
```
|
||||
|
||||
--8<-- [end:pre-built-images]
|
||||
--8<-- [start:build-image-from-source]
|
||||
|
||||
@@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
!!! tip
|
||||
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
|
||||
|
||||
=== "Intel GPU"
|
||||
|
||||
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
|
||||
|
||||
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||
|
||||
!!! tip
|
||||
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
|
||||
|
||||
=== "Google TPU"
|
||||
|
||||
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
// Reo.Dev documentation tracking
|
||||
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
|
||||
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
|
||||
@@ -174,33 +174,6 @@ 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
|
||||
|
||||
@@ -160,3 +160,4 @@ extra_javascript:
|
||||
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
|
||||
- mkdocs/javascript/edit_and_feedback.js
|
||||
- mkdocs/javascript/slack_and_forum.js
|
||||
- mkdocs/javascript/reo.js
|
||||
|
||||
+1
-1
@@ -7,7 +7,7 @@ requires = [
|
||||
"setuptools>=77.0.3,<81.0.0",
|
||||
"setuptools-scm>=8.0",
|
||||
"setuptools-rust>=1.9.0",
|
||||
"torch == 2.11.0",
|
||||
"torch == 2.13.0",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
]
|
||||
|
||||
@@ -4,8 +4,8 @@ packaging>=24.2
|
||||
setuptools==77.0.3 # this version can reuse CMake build dir
|
||||
setuptools-scm>=8
|
||||
setuptools-rust>=1.9.0
|
||||
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
|
||||
@@ -5,7 +5,7 @@ packaging>=24.2
|
||||
setuptools>=77.0.3,<81.0.0
|
||||
setuptools-scm>=8
|
||||
setuptools-rust>=1.9.0
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
|
||||
@@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir
|
||||
numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding
|
||||
|
||||
# Dependencies for CPUs
|
||||
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
|
||||
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
|
||||
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
|
||||
|
||||
@@ -4,18 +4,18 @@
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
|
||||
# Dependencies for NVIDIA GPUs
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
torchaudio==2.11.0
|
||||
# These must be updated alongside torch
|
||||
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
torchcodec >= 0.14
|
||||
PyNvVideoCodec==2.0.4
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
|
||||
# install_requires so the published wheel does not carry an unresolvable pin
|
||||
--extra-index-url https://flashinfer.ai/whl/
|
||||
flashinfer-python==0.6.14
|
||||
flashinfer-cubin==0.6.14
|
||||
flashinfer-python==0.6.15.post1
|
||||
flashinfer-cubin==0.6.15.post1
|
||||
apache-tvm-ffi==0.1.10
|
||||
tilelang==0.1.9
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.6.0
|
||||
quack-kernels>=0.4.0 # Required for tml-fa4
|
||||
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
@@ -1107,7 +1107,7 @@ tokenizers==0.22.2
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
# transformers
|
||||
torch==2.11.0+cpu
|
||||
torch==2.13.0+cpu
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
# accelerate
|
||||
@@ -1134,7 +1134,7 @@ torchaudio==2.11.0+cpu
|
||||
# vocos
|
||||
torchcodec==0.14.0+cpu
|
||||
# via -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cpu
|
||||
torchvision==0.28.0+cpu
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
# open-clip-torch
|
||||
|
||||
@@ -28,9 +28,9 @@ soundfile # required for audio tests
|
||||
jiwer # required for audio tests
|
||||
tblib # for pickling test exceptions
|
||||
timm >=1.0.17 # required for internvl and gemma3n-mm test
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
torchaudio==2.11.0
|
||||
torchvision==0.26.0
|
||||
torchvision==0.28.0
|
||||
transformers_stream_generator # required for qwen-vl test
|
||||
matplotlib # required for qwen-vl test
|
||||
mistral_common[image,audio] >= 1.11.5 # required for voxtral test
|
||||
|
||||
@@ -159,7 +159,7 @@ cuda-bindings==13.0.3
|
||||
# via torch
|
||||
cuda-pathfinder==1.3.3
|
||||
# via cuda-bindings
|
||||
cuda-toolkit==13.0.2
|
||||
cuda-toolkit==13.0.3.0
|
||||
# via torch
|
||||
cupy-cuda12x==13.6.0
|
||||
# via ray
|
||||
@@ -599,7 +599,7 @@ numpy==2.2.6
|
||||
# tritonclient
|
||||
# vocos
|
||||
# xgrammar
|
||||
nvidia-cublas==13.1.0.3
|
||||
nvidia-cublas==13.1.1.3
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cudnn-cu13
|
||||
@@ -607,10 +607,12 @@ nvidia-cublas==13.1.0.3
|
||||
nvidia-cuda-cupti==13.0.85
|
||||
# via cuda-toolkit
|
||||
nvidia-cuda-nvrtc==13.0.88
|
||||
# via cuda-toolkit
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cublas
|
||||
nvidia-cuda-runtime==13.0.96
|
||||
# via cuda-toolkit
|
||||
nvidia-cudnn-cu13==9.19.0.56
|
||||
nvidia-cudnn-cu13==9.20.0.48
|
||||
# via torch
|
||||
nvidia-cufft==12.0.0.61
|
||||
# via cuda-toolkit
|
||||
@@ -624,9 +626,9 @@ nvidia-cusparse==12.6.3.3
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cusolver
|
||||
nvidia-cusparselt-cu13==0.8.0
|
||||
nvidia-cusparselt-cu13==0.8.1
|
||||
# via torch
|
||||
nvidia-nccl-cu13==2.28.9
|
||||
nvidia-nccl-cu13==2.29.7
|
||||
# via torch
|
||||
nvidia-nvjitlink==13.0.88
|
||||
# via
|
||||
@@ -1202,7 +1204,7 @@ tokenizers==0.22.2
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
# transformers
|
||||
torch==2.11.0+cu130
|
||||
torch==2.13.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1233,7 +1235,7 @@ torchcodec==0.14.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cu130
|
||||
torchvision==0.28.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1270,7 +1272,7 @@ transformers==5.13.1
|
||||
# xgrammar
|
||||
transformers-stream-generator==0.0.5
|
||||
# via -r requirements/test/cuda.in
|
||||
triton==3.6.0
|
||||
triton==3.7.1
|
||||
# via
|
||||
# torch
|
||||
# xgrammar
|
||||
|
||||
@@ -12,4 +12,4 @@ ray[data]
|
||||
setuptools==78.1.0
|
||||
setuptools-rust>=1.9.0
|
||||
nixl==0.3.0
|
||||
tpu-inference==0.24.0
|
||||
tpu-inference==0.25.0
|
||||
|
||||
Generated
+27
-10
@@ -3446,7 +3446,6 @@ dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
@@ -4876,6 +4875,7 @@ dependencies = [
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4937,9 +4937,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
@@ -4966,9 +4966,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-build"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
|
||||
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -4977,10 +4977,23 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.5"
|
||||
name = "tonic-health"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
|
||||
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
|
||||
dependencies = [
|
||||
"prost",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"prost",
|
||||
@@ -4989,9 +5002,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost-build"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
|
||||
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -5484,10 +5497,13 @@ dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tiktoken-rs 0.9.1",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
@@ -5730,6 +5746,7 @@ dependencies = [
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tonic",
|
||||
"tonic-health",
|
||||
"tonic-prost",
|
||||
"tonic-prost-build",
|
||||
"tower",
|
||||
|
||||
+5
-4
@@ -118,10 +118,11 @@ tokio = { version = "1.47.1", features = [
|
||||
tokio-openssl = "0.6"
|
||||
tokio-stream = "0.1"
|
||||
tokio-util = { version = "0.7.18", features = ["rt"] }
|
||||
tonic = "0.14.5"
|
||||
tonic-build = "0.14.5"
|
||||
tonic-prost = "0.14.5"
|
||||
tonic-prost-build = "0.14.5"
|
||||
tonic = "0.14.6"
|
||||
tonic-build = "0.14.6"
|
||||
tonic-health = "0.14.6"
|
||||
tonic-prost = "0.14.6"
|
||||
tonic-prost-build = "0.14.6"
|
||||
tool-parser = "1.2.0"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
|
||||
|
||||
@@ -14,6 +14,10 @@ service Generate {
|
||||
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
|
||||
}
|
||||
|
||||
service Control {
|
||||
rpc Abort (AbortRequest) returns (AbortResponse) {}
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Generate Request
|
||||
// ======================================================================================
|
||||
@@ -201,3 +205,12 @@ message TokenIds {
|
||||
repeated uint32 ids = 1;
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Control
|
||||
// ======================================================================================
|
||||
|
||||
message AbortRequest {
|
||||
repeated string request_ids = 1;
|
||||
}
|
||||
|
||||
message AbortResponse {}
|
||||
|
||||
@@ -20,16 +20,19 @@ mimalloc.workspace = true
|
||||
rand.workspace = true
|
||||
rand_distr.workspace = true
|
||||
rayon.workspace = true
|
||||
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
|
||||
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
|
||||
rlimit.workspace = true
|
||||
rustc-hash.workspace = true
|
||||
serde = { workspace = true, features = ["rc"] }
|
||||
serde_json = { workspace = true, features = ["raw_value"] }
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
tiktoken-rs.workspace = true
|
||||
tokenizers.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
|
||||
@@ -158,9 +158,10 @@ impl PoolingBackend {
|
||||
// (mirrors Python async_request_vllm_rerank).
|
||||
if let Some(ref list) = input.prompt_list {
|
||||
if list.len() < 2 {
|
||||
eprintln!(
|
||||
"WARNING: vllm-rerank request has no documents \
|
||||
(prompt_list needs [query, doc, ...])"
|
||||
tracing::warn!(
|
||||
backend = "vllm-rerank",
|
||||
inputs = list.len(),
|
||||
"rerank request has no documents"
|
||||
);
|
||||
}
|
||||
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
|
||||
@@ -175,10 +176,10 @@ impl PoolingBackend {
|
||||
// Legacy path: text prompt as query, documents via --extra-body.
|
||||
let query = input.prompt.as_ref();
|
||||
if query.is_empty() && input.prompt_token_ids.is_some() {
|
||||
eprintln!(
|
||||
"WARNING: vllm-rerank received empty query (random dataset uses \
|
||||
token IDs only). Use --dataset-name random-rerank for meaningful \
|
||||
rerank benchmarks."
|
||||
tracing::warn!(
|
||||
backend = "vllm-rerank",
|
||||
dataset = "random",
|
||||
"rerank request has an empty query; use the random-rerank dataset"
|
||||
);
|
||||
}
|
||||
serde_json::json!({
|
||||
|
||||
+114
-84
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
|
||||
@@ -72,12 +73,12 @@ pub fn pre_resolve_dns(
|
||||
v4.extend(v6);
|
||||
if !v4.is_empty() {
|
||||
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
|
||||
println!("Pre-resolved {host} -> {ips:?}");
|
||||
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
|
||||
builder = builder.resolve_to_addrs(host, &v4);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
|
||||
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||
(m.clone(), config.model_name.clone())
|
||||
} else {
|
||||
println!("Model not specified, fetching first model from server...");
|
||||
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||
let (name, id) =
|
||||
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
|
||||
println!("First model name: {name}, first model id: {id}");
|
||||
tracing::info!(
|
||||
model_name = name,
|
||||
model_id = id,
|
||||
"selected first model from server"
|
||||
);
|
||||
(id, Some(name))
|
||||
};
|
||||
|
||||
@@ -358,10 +363,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
None
|
||||
} else {
|
||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||
println!("Loading tokenizer: {tid}");
|
||||
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
||||
println!("Tokenizer loaded successfully.");
|
||||
let t =
|
||||
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
|
||||
Some(t)
|
||||
};
|
||||
let has_tokenizer = tokenizer.is_some();
|
||||
@@ -421,7 +426,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.num_prompts, config.random_batch_size, config.is_reranker,
|
||||
),
|
||||
};
|
||||
println!("Generating {dataset_label}...");
|
||||
tracing::info!(
|
||||
dataset = ?config.dataset_name,
|
||||
prompts = config.num_prompts,
|
||||
description = %dataset_label,
|
||||
"generating benchmark dataset"
|
||||
);
|
||||
let gen_start = Instant::now();
|
||||
|
||||
let mut input_requests = match config.dataset_name {
|
||||
@@ -472,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let path = match config.dataset_path.as_deref() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -512,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
None => {
|
||||
downloaded = crate::datasets::speed_bench::download_speed_bench(
|
||||
config.speed_bench_config,
|
||||
)?;
|
||||
)
|
||||
.await?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -543,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.hf_subset.as_deref(),
|
||||
config.hf_split.as_deref(),
|
||||
config.num_prompts,
|
||||
)?;
|
||||
)
|
||||
.await?;
|
||||
crate::datasets::hf_dataset::load_hf_dataset(
|
||||
tok,
|
||||
&downloaded_path,
|
||||
@@ -608,18 +620,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
};
|
||||
|
||||
let gen_elapsed = gen_start.elapsed();
|
||||
println!(
|
||||
"Generated {} prompts in {:.2}s",
|
||||
input_requests.len(),
|
||||
gen_elapsed.as_secs_f64()
|
||||
tracing::info!(
|
||||
prompts = input_requests.len(),
|
||||
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||
"generated benchmark dataset"
|
||||
);
|
||||
|
||||
let filtered_count =
|
||||
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
|
||||
if filtered_count > 0 {
|
||||
println!(
|
||||
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
|
||||
config.max_model_len.unwrap()
|
||||
tracing::info!(
|
||||
filtered_prompts = filtered_count,
|
||||
max_model_len = config.max_model_len.unwrap(),
|
||||
"filtered prompts above maximum model length"
|
||||
);
|
||||
}
|
||||
if input_requests.is_empty() {
|
||||
@@ -670,7 +683,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
|
||||
// Ready check
|
||||
if config.ready_check_timeout_sec > 0 {
|
||||
println!("Starting initial single prompt test run...");
|
||||
tracing::info!("starting initial single-prompt test run");
|
||||
let test_output = wait_for_endpoint(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -685,7 +698,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
test_output.error
|
||||
)));
|
||||
}
|
||||
println!("Initial test run completed.");
|
||||
tracing::info!("initial single-prompt test run completed");
|
||||
}
|
||||
|
||||
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
|
||||
@@ -703,12 +716,15 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
DatasetName::Random | DatasetName::PrefixRepetition
|
||||
);
|
||||
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
|
||||
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
|
||||
tracing::info!(
|
||||
reason = "prompt_token_ids",
|
||||
"skipping server tokenizer verification"
|
||||
);
|
||||
}
|
||||
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
|
||||
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
|
||||
if is_tokenizer_verified(&cache_key) {
|
||||
println!("Tokenizer verified in previous run (cached), skipping verification.");
|
||||
tracing::info!(reason = "cached", "skipping server tokenizer verification");
|
||||
} else {
|
||||
let num_special =
|
||||
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
|
||||
@@ -723,14 +739,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
.await?
|
||||
{
|
||||
SampleVerifyOutcome::Passed => {
|
||||
println!("Sample verification passed, skipping full verification.");
|
||||
tracing::info!("tokenizer sample verification passed");
|
||||
mark_tokenizer_verified(&cache_key);
|
||||
}
|
||||
SampleVerifyOutcome::Skipped(reason) => {
|
||||
println!("Server /tokenize unavailable ({reason}), skipping verification.");
|
||||
tracing::warn!(
|
||||
reason = %reason,
|
||||
"server tokenizer unavailable; skipping prompt verification"
|
||||
);
|
||||
}
|
||||
SampleVerifyOutcome::Mismatch => {
|
||||
println!("Sample verification found mismatch, running full verify+fix...");
|
||||
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
|
||||
match verify_and_fix_prompt_lengths(
|
||||
&client,
|
||||
&config.base_url,
|
||||
@@ -742,16 +761,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
println!(
|
||||
"All {} prompts verified: exact token length match.",
|
||||
input_requests.len()
|
||||
tracing::info!(
|
||||
prompts = input_requests.len(),
|
||||
"verified exact prompt token lengths"
|
||||
);
|
||||
mark_tokenizer_verified(&cache_key);
|
||||
}
|
||||
Err(BenchError::TokenizeUnavailable(reason)) => {
|
||||
println!(
|
||||
"Server /tokenize became unavailable during verification \
|
||||
({reason}); proceeding with client-side token counts."
|
||||
tracing::warn!(
|
||||
reason = %reason,
|
||||
"server tokenizer became unavailable; using client token counts"
|
||||
);
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
@@ -763,7 +782,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
|
||||
// Warmup
|
||||
if config.num_warmups > 0 {
|
||||
println!("Warming up with {} requests...", config.num_warmups);
|
||||
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
|
||||
run_warmup(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -776,7 +795,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
config.disable_tqdm,
|
||||
)
|
||||
.await;
|
||||
println!("Warmup run completed.");
|
||||
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
|
||||
}
|
||||
|
||||
// Start profiler if requested (immediate mode — no batch threshold)
|
||||
@@ -814,28 +833,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
let spec_decode_before =
|
||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||
if spec_decode_before.is_some() {
|
||||
println!("Speculative decoding detected, will collect metrics.");
|
||||
tracing::info!("detected speculative decoding; collecting metrics");
|
||||
}
|
||||
|
||||
// Main benchmark
|
||||
println!("Starting main benchmark run...");
|
||||
let distribution = if config.burstiness == 1.0 {
|
||||
"Poisson process"
|
||||
} else {
|
||||
"Gamma distribution"
|
||||
};
|
||||
println!(
|
||||
"Traffic request rate: {}",
|
||||
if config.request_rate.is_infinite() {
|
||||
"inf".to_string()
|
||||
} else {
|
||||
format!("{}", config.request_rate)
|
||||
}
|
||||
);
|
||||
println!("Burstiness factor: {} ({distribution})", config.burstiness);
|
||||
println!(
|
||||
"Maximum request concurrency: {}",
|
||||
config.max_concurrency.unwrap_or(config.num_prompts)
|
||||
tracing::info!(
|
||||
request_rate = config.request_rate,
|
||||
burstiness = config.burstiness,
|
||||
distribution,
|
||||
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
|
||||
prompts = config.num_prompts,
|
||||
"starting main benchmark run"
|
||||
);
|
||||
|
||||
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
|
||||
@@ -847,11 +860,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
);
|
||||
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
|
||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||
println!(
|
||||
"LoRA adapters ({}): {:?} [assignment={:?}]",
|
||||
modules.len(),
|
||||
names,
|
||||
config.lora_assignment
|
||||
tracing::info!(
|
||||
adapters = modules.len(),
|
||||
names = ?names,
|
||||
assignment = ?config.lora_assignment,
|
||||
"assigned LoRA adapters"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1125,7 +1138,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
|
||||
if let Some((cancel_tx, task)) = profile_task {
|
||||
let _ = cancel_tx.send(());
|
||||
if let Err(e) = task.await {
|
||||
eprintln!("WARNING: Profile background task failed: {e}");
|
||||
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1289,12 +1302,14 @@ pub(crate) async fn start_profiler_immediate(
|
||||
base_url: &str,
|
||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||
) {
|
||||
println!("Starting profiler...");
|
||||
let profile_url = format!("{base_url}/start_profile");
|
||||
tracing::info!(url = %profile_url, "starting profiler");
|
||||
match send_profile_request(client, &profile_url, extra_headers).await {
|
||||
Ok(true) => println!("Profiler started"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
|
||||
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
|
||||
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1304,12 +1319,14 @@ pub(crate) async fn stop_profiler_immediate(
|
||||
base_url: &str,
|
||||
extra_headers: &Option<std::collections::HashMap<String, String>>,
|
||||
) {
|
||||
println!("Stopping profiler...");
|
||||
let profile_url = format!("{base_url}/stop_profile");
|
||||
tracing::info!(url = %profile_url, "stopping profiler");
|
||||
match send_profile_request(client, &profile_url, extra_headers).await {
|
||||
Ok(true) => println!("Profiler stopped"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
||||
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
|
||||
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1371,25 +1388,30 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
duration_secs: f64,
|
||||
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
|
||||
) {
|
||||
println!(
|
||||
"Waiting for batch size >= {threshold} before starting profiler \
|
||||
(will capture {duration_secs}s)..."
|
||||
tracing::info!(
|
||||
threshold,
|
||||
duration_seconds = duration_secs,
|
||||
"waiting for profiler batch threshold"
|
||||
);
|
||||
|
||||
loop {
|
||||
if let Some(running) = fetch_num_requests_running(client, base_url).await
|
||||
&& running >= threshold
|
||||
{
|
||||
println!("Batch size {running} >= {threshold}, starting profiler...");
|
||||
tracing::info!(
|
||||
running_requests = running,
|
||||
threshold,
|
||||
"profiler batch threshold reached"
|
||||
);
|
||||
break;
|
||||
}
|
||||
// Wait 500ms or until the benchmark signals cancellation
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
|
||||
_ = &mut cancel_rx => {
|
||||
eprintln!(
|
||||
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
|
||||
profiling skipped."
|
||||
tracing::warn!(
|
||||
threshold,
|
||||
"benchmark finished before profiler batch threshold; skipping profiling"
|
||||
);
|
||||
return;
|
||||
}
|
||||
@@ -1398,13 +1420,13 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
|
||||
let start_url = format!("{base_url}/start_profile");
|
||||
match send_profile_request(client, &start_url, extra_headers).await {
|
||||
Ok(true) => println!("Profiler started"),
|
||||
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
|
||||
Ok(false) => {
|
||||
eprintln!("WARNING: Profiler start request returned non-success");
|
||||
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
|
||||
return;
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("WARNING: Failed to start profiler: {e}");
|
||||
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
|
||||
return;
|
||||
}
|
||||
}
|
||||
@@ -1413,15 +1435,17 @@ pub(crate) async fn profile_on_batch_threshold(
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
|
||||
_ = &mut cancel_rx => {
|
||||
println!("Benchmark finished, stopping profiler early...");
|
||||
tracing::info!("benchmark finished; stopping profiler early");
|
||||
}
|
||||
}
|
||||
|
||||
let stop_url = format!("{base_url}/stop_profile");
|
||||
match send_profile_request(client, &stop_url, extra_headers).await {
|
||||
Ok(true) => println!("Profiler stopped after capturing"),
|
||||
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
|
||||
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
|
||||
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
|
||||
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
|
||||
Err(e) => {
|
||||
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1502,10 +1526,11 @@ async fn verify_and_fix_prompt_lengths(
|
||||
let excess = tokens.len().saturating_sub(expected_input_len);
|
||||
let compensate = if excess > 0 && last_excess == Some(excess) {
|
||||
if _iter == 1 {
|
||||
eprintln!(
|
||||
"Prompt {i}: server consistently adds {excess} extra token(s) \
|
||||
(likely BOS), compensating target to {}.",
|
||||
expected_input_len.saturating_sub(excess),
|
||||
tracing::warn!(
|
||||
prompt_index = i,
|
||||
extra_tokens = excess,
|
||||
adjusted_target = expected_input_len.saturating_sub(excess),
|
||||
"server consistently adds prompt tokens; compensating verification target"
|
||||
);
|
||||
}
|
||||
excess
|
||||
@@ -1563,7 +1588,10 @@ async fn verify_and_fix_prompt_lengths(
|
||||
|
||||
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
|
||||
if fc > 0 {
|
||||
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
|
||||
tracing::info!(
|
||||
fixed_prompts = fc,
|
||||
"fixed prompt lengths using server tokenizer"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -1818,7 +1846,7 @@ async fn sample_verify_prompts(
|
||||
let tokenize_url = format!("{base_url}/tokenize");
|
||||
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
|
||||
println!("Sampling {sample_size} prompts for verification...");
|
||||
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
|
||||
|
||||
for (i, request) in requests.iter().enumerate().take(sample_size) {
|
||||
let tokens = match server_tokenize(
|
||||
@@ -1841,9 +1869,11 @@ async fn sample_verify_prompts(
|
||||
|
||||
let expected = request.prompt_len + num_special;
|
||||
if tokens.len() != expected {
|
||||
println!(
|
||||
"Prompt {i}: expected {expected} tokens, server returned {}",
|
||||
tokens.len()
|
||||
tracing::warn!(
|
||||
prompt_index = i,
|
||||
expected_tokens = expected,
|
||||
actual_tokens = tokens.len(),
|
||||
"tokenizer verification sample mismatch"
|
||||
);
|
||||
return Ok(SampleVerifyOutcome::Mismatch);
|
||||
}
|
||||
|
||||
@@ -288,10 +288,18 @@ impl BenchConfig {
|
||||
}
|
||||
Some(other) => {
|
||||
// extra_body was not an object — just use sampling params
|
||||
eprintln!(
|
||||
"Warning: --extra-body is not a JSON object, sampling params may be lost"
|
||||
let value_type = match &other {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "boolean",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => unreachable!(),
|
||||
};
|
||||
tracing::warn!(
|
||||
value_type,
|
||||
"sampling parameters may be lost because --extra-body is not a JSON object"
|
||||
);
|
||||
let _ = other;
|
||||
sampling_params
|
||||
}
|
||||
None => sampling_params,
|
||||
@@ -489,9 +497,9 @@ impl BenchConfig {
|
||||
_ => {}
|
||||
}
|
||||
if !args.skip_chat_template {
|
||||
eprintln!(
|
||||
"NOTE: client-side chat template rendering is not supported; custom \
|
||||
dataset prompts are sent raw (equivalent to --skip-chat-template)."
|
||||
tracing::warn!(
|
||||
dataset = "custom",
|
||||
"client-side chat template rendering is unsupported; sending prompts raw"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -570,9 +578,10 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
if ignore_eos {
|
||||
eprintln!(
|
||||
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
|
||||
respect output length limits, causing unbounded context growth."
|
||||
tracing::warn!(
|
||||
ignore_eos,
|
||||
multi_turn = true,
|
||||
"output length limits may be ignored, causing unbounded context growth"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -128,8 +128,10 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::SampleRequest;
|
||||
use super::progress::RowDownloadReporter;
|
||||
use crate::error::{BenchError, Result};
|
||||
use crate::tokenizer::TokenizerKind;
|
||||
|
||||
@@ -50,18 +51,19 @@ enum ColumnFormat {
|
||||
|
||||
/// Make a GET request with retry logic (3 retries with exponential backoff).
|
||||
/// Returns the parsed JSON response.
|
||||
fn get_with_retry(
|
||||
client: &reqwest::blocking::Client,
|
||||
async fn get_with_retry(
|
||||
client: &reqwest::Client,
|
||||
url: &str,
|
||||
label: &str,
|
||||
) -> Result<serde_json::Value> {
|
||||
let max_retries = 3;
|
||||
for attempt in 0..=max_retries {
|
||||
let resp = match client.get(url).send() {
|
||||
let resp = match client.get(url).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if attempt < max_retries {
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -80,7 +82,7 @@ fn get_with_retry(
|
||||
}
|
||||
|
||||
if status.is_server_error() && attempt < max_retries {
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -92,6 +94,7 @@ fn get_with_retry(
|
||||
|
||||
let data: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
|
||||
return Ok(data);
|
||||
}
|
||||
@@ -105,7 +108,7 @@ fn get_with_retry(
|
||||
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
|
||||
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
|
||||
/// or the dataset is exhausted.
|
||||
pub fn download_hf_dataset(
|
||||
pub async fn download_hf_dataset(
|
||||
dataset: &str,
|
||||
subset: Option<&str>,
|
||||
split: Option<&str>,
|
||||
@@ -115,7 +118,7 @@ pub fn download_hf_dataset(
|
||||
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
|
||||
|
||||
let mut client_builder =
|
||||
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
|
||||
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
|
||||
|
||||
// Add HF_TOKEN auth header if available
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
@@ -138,7 +141,7 @@ pub fn download_hf_dataset(
|
||||
// Call /info to discover available configs and splits
|
||||
let info_url =
|
||||
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
|
||||
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
|
||||
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
|
||||
|
||||
let dataset_info =
|
||||
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
|
||||
@@ -201,7 +204,12 @@ pub fn download_hf_dataset(
|
||||
(resolved_config, resolved_split)
|
||||
};
|
||||
|
||||
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
|
||||
tracing::info!(
|
||||
dataset,
|
||||
config = resolved_config,
|
||||
split = resolved_split,
|
||||
"resolved Hugging Face dataset"
|
||||
);
|
||||
|
||||
// Check cache
|
||||
let dir = cache_dir();
|
||||
@@ -215,11 +223,16 @@ pub fn download_hf_dataset(
|
||||
|
||||
if cache_path.exists() {
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
println!("HF dataset cached: {path_str}");
|
||||
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
|
||||
return Ok((path_str, resolved_config, resolved_split));
|
||||
}
|
||||
|
||||
println!("Downloading HF dataset '{dataset}' from datasets-server...");
|
||||
tracing::info!(
|
||||
dataset,
|
||||
config = resolved_config,
|
||||
split = resolved_split,
|
||||
"downloading Hugging Face dataset"
|
||||
);
|
||||
|
||||
let encoded_config: String =
|
||||
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
|
||||
@@ -229,6 +242,7 @@ pub fn download_hf_dataset(
|
||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||
let mut offset = 0usize;
|
||||
let page_size = 100usize;
|
||||
let mut progress = RowDownloadReporter::new();
|
||||
|
||||
loop {
|
||||
let url = format!(
|
||||
@@ -240,7 +254,7 @@ pub fn download_hf_dataset(
|
||||
&length={page_size}"
|
||||
);
|
||||
|
||||
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
|
||||
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
|
||||
|
||||
let rows = data["rows"]
|
||||
.as_array()
|
||||
@@ -260,14 +274,14 @@ pub fn download_hf_dataset(
|
||||
offset += fetched;
|
||||
|
||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||
eprint!("\r Fetched {offset}/{total} rows...");
|
||||
progress.update(offset, total);
|
||||
|
||||
// Stop if we have enough rows or reached end of dataset
|
||||
if all_rows.len() >= num_rows_needed || fetched < page_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
eprintln!(); // newline after progress
|
||||
progress.finish();
|
||||
|
||||
if all_rows.is_empty() {
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -280,7 +294,12 @@ pub fn download_hf_dataset(
|
||||
std::fs::write(&cache_path, &json_str)?;
|
||||
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
|
||||
tracing::info!(
|
||||
dataset,
|
||||
rows = all_rows.len(),
|
||||
path = %path_str,
|
||||
"saved Hugging Face dataset"
|
||||
);
|
||||
Ok((path_str, resolved_config, resolved_split))
|
||||
}
|
||||
|
||||
@@ -483,21 +502,31 @@ pub fn load_hf_dataset(
|
||||
// Detect column format from first row
|
||||
let format = detect_column_format(&entries[0], text_column_override)?;
|
||||
|
||||
// Print detected format
|
||||
match &format {
|
||||
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
|
||||
ColumnFormat::Chat(col) => {
|
||||
tracing::info!(
|
||||
format = "chat",
|
||||
column = col,
|
||||
"detected Hugging Face dataset format"
|
||||
);
|
||||
}
|
||||
ColumnFormat::Text {
|
||||
prompt_col,
|
||||
output_col,
|
||||
} => {
|
||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
||||
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
|
||||
tracing::info!(
|
||||
format = "text",
|
||||
prompt_column = prompt_col,
|
||||
output_column = output_col.as_deref().unwrap_or("none"),
|
||||
"detected Hugging Face dataset format"
|
||||
);
|
||||
}
|
||||
ColumnFormat::Combined { cols, output_col } => {
|
||||
let out_msg = output_col.as_deref().unwrap_or("none");
|
||||
println!(
|
||||
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
|
||||
cols
|
||||
tracing::info!(
|
||||
format = "combined",
|
||||
prompt_columns = ?cols,
|
||||
output_column = output_col.as_deref().unwrap_or("none"),
|
||||
"detected Hugging Face dataset format"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -608,9 +637,10 @@ pub fn load_hf_dataset(
|
||||
if len == 0 { 128 } else { len }
|
||||
} else {
|
||||
if !warned_no_output {
|
||||
eprintln!(
|
||||
"WARNING: No output column detected and --hf-output-len not set. \
|
||||
Using default output length of 128 tokens."
|
||||
tracing::warn!(
|
||||
path = dataset_path,
|
||||
default_output_tokens = 128,
|
||||
"no dataset output column or --hf-output-len; using default output length"
|
||||
);
|
||||
warned_no_output = true;
|
||||
}
|
||||
@@ -618,9 +648,10 @@ pub fn load_hf_dataset(
|
||||
}
|
||||
} else {
|
||||
if !warned_no_output {
|
||||
eprintln!(
|
||||
"WARNING: No output column detected and --hf-output-len not set. \
|
||||
Using default output length of 128 tokens."
|
||||
tracing::warn!(
|
||||
path = dataset_path,
|
||||
default_output_tokens = 128,
|
||||
"no dataset output column or --hf-output-len; using default output length"
|
||||
);
|
||||
warned_no_output = true;
|
||||
}
|
||||
@@ -640,9 +671,11 @@ pub fn load_hf_dataset(
|
||||
// Oversample if needed
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "hf",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let original_len = samples.len();
|
||||
@@ -652,9 +685,11 @@ pub fn load_hf_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
println!(
|
||||
"Oversampled HF dataset from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "hf",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1002,8 +1037,10 @@ mod tests {
|
||||
|
||||
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
|
||||
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
crate::tokenizer::TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
}
|
||||
|
||||
/// Write JSON data to a unique temp file and return the path string.
|
||||
|
||||
@@ -5,6 +5,7 @@ pub mod custom;
|
||||
pub mod hf_dataset;
|
||||
pub mod multi_turn;
|
||||
pub mod prefix_repetition;
|
||||
mod progress;
|
||||
pub mod random;
|
||||
pub mod random_mm;
|
||||
pub mod random_rerank;
|
||||
@@ -90,9 +91,10 @@ pub fn oversample_requests(
|
||||
return;
|
||||
}
|
||||
if no_oversample {
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
requests.len()
|
||||
tracing::info!(
|
||||
samples = requests.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
);
|
||||
return;
|
||||
}
|
||||
@@ -103,9 +105,10 @@ pub fn oversample_requests(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
requests.push(req);
|
||||
}
|
||||
println!(
|
||||
"Oversampled requests from {original_len} to {} total samples.",
|
||||
requests.len()
|
||||
tracing::info!(
|
||||
original_samples = original_len,
|
||||
samples = requests.len(),
|
||||
"oversampled dataset"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
|
||||
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
|
||||
conversations.push(conv);
|
||||
}
|
||||
println!(
|
||||
"Oversampled multi-turn conversations from {original_len} to {} total.",
|
||||
conversations.len()
|
||||
tracing::info!(
|
||||
original_conversations = original_len,
|
||||
conversations = conversations.len(),
|
||||
"oversampled multi-turn conversations"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -525,10 +526,12 @@ mod tests {
|
||||
len
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_prefix_sharing_structure() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
async fn test_prefix_sharing_structure() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 5,
|
||||
@@ -610,10 +613,12 @@ mod tests {
|
||||
println!("All prefix sharing checks passed!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_per_turn_input_len_default_mode() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
async fn test_per_turn_input_len_default_mode() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 4,
|
||||
@@ -650,10 +655,12 @@ mod tests {
|
||||
println!("per_turn_input_len default-mode checks passed!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_variable_turns_range() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
async fn test_variable_turns_range() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 50,
|
||||
@@ -684,10 +691,12 @@ mod tests {
|
||||
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_variable_turns_fixed() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
async fn test_variable_turns_fixed() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cfg = MultiTurnRandomConfig {
|
||||
num_conversations: 10,
|
||||
@@ -709,10 +718,12 @@ mod tests {
|
||||
println!("variable_turns_fixed checks passed!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_per_turn_input_len_prefix_sharing() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
|
||||
async fn test_per_turn_input_len_prefix_sharing() {
|
||||
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
|
||||
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
|
||||
|
||||
@@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
|
||||
}
|
||||
let total = prompts_per_prefix * num_prefixes;
|
||||
if total != num_requests {
|
||||
println!(
|
||||
"prefix_repetition: generating {total} requests \
|
||||
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
|
||||
{} dropped to divide evenly)",
|
||||
num_requests - total
|
||||
tracing::info!(
|
||||
requested = num_requests,
|
||||
generated = total,
|
||||
prefixes = num_prefixes,
|
||||
prompts_per_prefix,
|
||||
dropped = num_requests - total,
|
||||
"adjusted prefix-repetition request count"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -109,8 +111,10 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
|
||||
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Reports row download progress to an interactive progress bar, or through
|
||||
/// periodic tracing events when the progress bar is hidden on a non-TTY.
|
||||
pub(super) struct RowDownloadReporter {
|
||||
progress: ProgressBar,
|
||||
next_report: Instant,
|
||||
}
|
||||
|
||||
impl RowDownloadReporter {
|
||||
/// Creates a reporter that emits non-TTY updates every 10 seconds.
|
||||
pub fn new() -> Self {
|
||||
let progress = ProgressBar::new(0);
|
||||
progress.set_style(
|
||||
ProgressStyle::with_template(
|
||||
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
|
||||
)
|
||||
.unwrap()
|
||||
.progress_chars("#>-"),
|
||||
);
|
||||
Self {
|
||||
progress,
|
||||
next_report: Instant::now() + REPORT_INTERVAL,
|
||||
}
|
||||
}
|
||||
|
||||
/// Updates the current row count and reports progress when due.
|
||||
pub fn update(&mut self, rows: usize, total: u64) {
|
||||
let rows = rows as u64;
|
||||
let total = total.max(rows);
|
||||
self.progress.set_length(total);
|
||||
self.progress.set_position(rows);
|
||||
|
||||
if self.should_report(Instant::now()) {
|
||||
tracing::info!(rows, total, "fetching dataset rows");
|
||||
}
|
||||
}
|
||||
|
||||
/// Clears the interactive progress bar after the download completes.
|
||||
pub fn finish(self) {
|
||||
self.progress.finish_and_clear();
|
||||
}
|
||||
|
||||
fn should_report(&mut self, now: Instant) -> bool {
|
||||
if !self.progress.is_hidden() || now < self.next_report {
|
||||
return false;
|
||||
}
|
||||
self.next_report = now + REPORT_INTERVAL;
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn hidden_reporter_uses_ten_second_deadline() {
|
||||
let start = Instant::now();
|
||||
let mut reporter = RowDownloadReporter {
|
||||
progress: ProgressBar::hidden(),
|
||||
next_report: start + REPORT_INTERVAL,
|
||||
};
|
||||
|
||||
assert!(!reporter.should_report(start + Duration::from_secs(9)));
|
||||
assert!(reporter.should_report(start + Duration::from_secs(10)));
|
||||
assert!(!reporter.should_report(start + Duration::from_secs(19)));
|
||||
assert!(reporter.should_report(start + Duration::from_secs(20)));
|
||||
}
|
||||
}
|
||||
@@ -49,9 +49,12 @@ pub fn generate_random_dataset(
|
||||
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
|
||||
let (output_low, output_high) = range_ratio.output_bounds(output_len);
|
||||
if !range_ratio.is_fixed() {
|
||||
println!(
|
||||
"Sampling input_len from [{input_low}, {input_high}] and \
|
||||
output_len from [{output_low}, {output_high}]"
|
||||
tracing::info!(
|
||||
input_low,
|
||||
input_high,
|
||||
output_low,
|
||||
output_high,
|
||||
"sampling random request lengths"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -305,7 +308,8 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_generate_random_dataset_token_ids() {
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
10, // num_requests
|
||||
@@ -337,7 +341,8 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_generate_random_dataset_text() {
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
10, // num_requests
|
||||
@@ -371,7 +376,8 @@ mod tests {
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn test_token_length_exact_local() {
|
||||
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
|
||||
let tokenizer =
|
||||
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
|
||||
let target_len = 512;
|
||||
let requests = generate_random_dataset(
|
||||
&tokenizer,
|
||||
@@ -405,11 +411,11 @@ mod tests {
|
||||
}
|
||||
|
||||
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_token_length_exact_tiktoken() {
|
||||
async fn test_token_length_exact_tiktoken() {
|
||||
// Use Qwen2.5 which has a tiktoken-format tokenizer
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||
let tokenizer = match tokenizer {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
@@ -453,10 +459,10 @@ mod tests {
|
||||
|
||||
/// Test encode/decode roundtrip stability for tiktoken.
|
||||
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
|
||||
#[test]
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
fn test_tiktoken_roundtrip_stability() {
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
|
||||
async fn test_tiktoken_roundtrip_stability() {
|
||||
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
|
||||
let tokenizer = match tokenizer {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
|
||||
@@ -139,8 +139,10 @@ mod tests {
|
||||
|
||||
/// gpt2 via built-in tiktoken encoding — loads without network access.
|
||||
fn test_tokenizer() -> TokenizerKind {
|
||||
crate::tokenizer::load_tokenizer("gpt2", false, None)
|
||||
.expect("gpt2 built-in tiktoken should always load without network")
|
||||
TokenizerKind::Tiktoken(
|
||||
crate::tiktoken::load_builtin_tiktoken("gpt2")
|
||||
.expect("gpt2 built-in tiktoken should always load without network"),
|
||||
)
|
||||
}
|
||||
|
||||
fn fixed_ratio() -> RangeRatio {
|
||||
|
||||
@@ -22,18 +22,21 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
|
||||
|
||||
/// Download the default ShareGPT dataset from HuggingFace Hub.
|
||||
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
|
||||
pub fn download_sharegpt_dataset() -> Result<String> {
|
||||
println!(
|
||||
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
|
||||
pub async fn download_sharegpt_dataset() -> Result<String> {
|
||||
tracing::info!(
|
||||
repository = DEFAULT_SHAREGPT_REPO,
|
||||
file = DEFAULT_SHAREGPT_FILE,
|
||||
"downloading ShareGPT dataset"
|
||||
);
|
||||
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
|
||||
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
|
||||
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
|
||||
.map_err(BenchError::Config)?;
|
||||
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
|
||||
BenchError::Config(format!(
|
||||
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
|
||||
))
|
||||
})?;
|
||||
let path_str = path.to_string_lossy().to_string();
|
||||
println!("ShareGPT dataset ready: {path_str}");
|
||||
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
|
||||
Ok(path_str)
|
||||
}
|
||||
|
||||
@@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset(
|
||||
// Oversample if dataset is smaller than requested
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "sharegpt",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let needed = num_requests - samples.len();
|
||||
@@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
println!(
|
||||
"Oversampled requests from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "sharegpt",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::SampleRequest;
|
||||
use super::progress::RowDownloadReporter;
|
||||
use crate::cli::SpeedBenchConfig;
|
||||
use crate::error::{BenchError, Result};
|
||||
use crate::tokenizer::TokenizerKind;
|
||||
@@ -25,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf {
|
||||
|
||||
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
|
||||
/// Results are cached as JSON locally for subsequent runs.
|
||||
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let config_name = config.as_str();
|
||||
|
||||
let dir = cache_dir();
|
||||
@@ -35,13 +36,13 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
// Return cached file if it exists
|
||||
if cache_path.exists() {
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
println!("SPEED-Bench ({config_name}) cached: {path_str}");
|
||||
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
|
||||
return Ok(path_str);
|
||||
}
|
||||
|
||||
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
|
||||
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
|
||||
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
|
||||
@@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||
let mut offset = 0usize;
|
||||
let page_size = 100usize;
|
||||
let mut progress = RowDownloadReporter::new();
|
||||
|
||||
loop {
|
||||
let url = format!(
|
||||
@@ -64,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let max_retries = 3;
|
||||
let mut data: Option<serde_json::Value> = None;
|
||||
for attempt in 0..=max_retries {
|
||||
let resp = match client.get(&url).send() {
|
||||
let resp = match client.get(&url).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if attempt < max_retries {
|
||||
std::thread::sleep(std::time::Duration::from_secs(
|
||||
tokio::time::sleep(std::time::Duration::from_secs(
|
||||
2 * (attempt as u64 + 1),
|
||||
));
|
||||
))
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
};
|
||||
|
||||
if resp.status().is_server_error() && attempt < max_retries {
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
)));
|
||||
}
|
||||
|
||||
data = Some(resp.json().map_err(|e| {
|
||||
data = Some(resp.json().await.map_err(|e| {
|
||||
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
|
||||
})?);
|
||||
break;
|
||||
@@ -116,15 +119,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let fetched = rows.len();
|
||||
offset += fetched;
|
||||
|
||||
// Print progress
|
||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||
eprint!("\r Fetched {offset}/{total} rows...");
|
||||
progress.update(offset, total);
|
||||
|
||||
if fetched < page_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
eprintln!(); // newline after progress
|
||||
progress.finish();
|
||||
|
||||
if all_rows.is_empty() {
|
||||
return Err(BenchError::Config(
|
||||
@@ -137,9 +139,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
std::fs::write(&cache_path, &json_str)?;
|
||||
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
println!(
|
||||
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
|
||||
all_rows.len()
|
||||
tracing::info!(
|
||||
config = config_name,
|
||||
rows = all_rows.len(),
|
||||
path = %path_str,
|
||||
"saved SPEED-Bench dataset"
|
||||
);
|
||||
Ok(path_str)
|
||||
}
|
||||
@@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset(
|
||||
// Oversample if needed
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let original_len = samples.len();
|
||||
@@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
println!(
|
||||
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -288,7 +296,6 @@ pub fn load_speed_bench_dataset(
|
||||
));
|
||||
}
|
||||
|
||||
// Print category distribution
|
||||
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
|
||||
for entry in &filtered[..filtered.len().min(samples.len())] {
|
||||
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
|
||||
@@ -297,7 +304,7 @@ pub fn load_speed_bench_dataset(
|
||||
let mut cats: Vec<_> = cat_counts.into_iter().collect();
|
||||
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
|
||||
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
|
||||
println!("SPEED-Bench categories: {}", cat_str.join(", "));
|
||||
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
|
||||
|
||||
Ok(samples)
|
||||
}
|
||||
|
||||
+20
-35
@@ -1,54 +1,39 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Sync facade over the async `hf_hub` API.
|
||||
//!
|
||||
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
|
||||
//! backend unconditionally pulls ureq's default rustls feature. So we use the
|
||||
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
|
||||
//! loaders, tokenizer fallback in rayon threads) by running each download on a
|
||||
//! dedicated thread with its own single-threaded runtime.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use hf_hub::Repo;
|
||||
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
|
||||
|
||||
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
|
||||
pub struct HubRepo {
|
||||
repo: hf_hub::Repo,
|
||||
repo: ApiRepo,
|
||||
}
|
||||
|
||||
impl HubRepo {
|
||||
pub fn model(model_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::model(model_id),
|
||||
}
|
||||
pub fn model(model_id: String) -> Result<Self, String> {
|
||||
Self::new(Repo::model(model_id))
|
||||
}
|
||||
|
||||
pub fn dataset(repo_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::dataset(repo_id),
|
||||
pub fn dataset(repo_id: String) -> Result<Self, String> {
|
||||
Self::new(Repo::dataset(repo_id))
|
||||
}
|
||||
|
||||
fn new(repo: Repo) -> Result<Self, String> {
|
||||
let mut builder = ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
Ok(Self {
|
||||
repo: api.repo(repo),
|
||||
})
|
||||
}
|
||||
|
||||
/// Download (or fetch from cache) a single file from the repo.
|
||||
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
|
||||
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
let repo = self.repo.clone();
|
||||
let filename = filename.to_string();
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
|
||||
rt.block_on(async move {
|
||||
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
|
||||
})
|
||||
})
|
||||
.join()
|
||||
.map_err(|_| "HF Hub download thread panicked".to_string())?
|
||||
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
self.repo.get(filename).await.map_err(|e| format!("{e}"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ pub fn prepare_process() {
|
||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||
&& new > 1024
|
||||
{
|
||||
eprintln!("Open-file limit: {new}");
|
||||
tracing::info!(soft_limit = new, "raised open-file limit");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,7 +19,19 @@ struct Cli {
|
||||
args: vllm_bench::BenchServeArgs,
|
||||
}
|
||||
|
||||
// TODO: unify the tracing subscriber used by different binaries.
|
||||
fn init_tracing() {
|
||||
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(filter)
|
||||
.with_writer(std::io::stderr)
|
||||
.try_init();
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
init_tracing();
|
||||
|
||||
let cli = Cli::parse();
|
||||
vllm_bench::prepare_process();
|
||||
|
||||
|
||||
@@ -7,6 +7,22 @@ use crate::datasets::SampleRequest;
|
||||
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
|
||||
use crate::multi_turn::ConversationOutput;
|
||||
|
||||
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
|
||||
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
|
||||
if failed_outputs.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
tracing::warn!(
|
||||
failed_requests = failed_outputs.len(),
|
||||
displayed_errors = failed_outputs.len().min(10),
|
||||
"benchmark requests failed"
|
||||
);
|
||||
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
|
||||
tracing::warn!(index, error = %output.error, "benchmark request failed");
|
||||
}
|
||||
}
|
||||
|
||||
/// Calculate benchmark metrics from request outputs.
|
||||
///
|
||||
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
|
||||
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
|
||||
|
||||
let failed = outputs.len() - completed;
|
||||
|
||||
// Print failed request errors (capped to 10)
|
||||
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
|
||||
if !failed_outputs.is_empty() {
|
||||
eprintln!("Failed requests during benchmark run detected (capping to 10):");
|
||||
for (i, err) in failed_outputs.iter().take(10).enumerate() {
|
||||
eprintln!("Error {i}: {}", err.error);
|
||||
}
|
||||
}
|
||||
log_failed_requests(outputs);
|
||||
|
||||
// Calculate max output tokens per second and max concurrent requests
|
||||
let mut max_output_tokens_per_s = 0.0_f64;
|
||||
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
|
||||
|
||||
let failed = outputs.len() - completed;
|
||||
|
||||
// Print failed request errors (capped to 10)
|
||||
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
|
||||
if !failed_outputs.is_empty() {
|
||||
eprintln!("Failed requests during benchmark run detected (capping to 10):");
|
||||
for (i, err) in failed_outputs.iter().take(10).enumerate() {
|
||||
eprintln!("Error {i}: {}", err.error);
|
||||
}
|
||||
}
|
||||
log_failed_requests(outputs);
|
||||
|
||||
// Compute peak concurrent requests from start_time + latency windows
|
||||
let successful_outputs: Vec<&RequestFuncOutput> =
|
||||
|
||||
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use indicatif::{ProgressBar, ProgressStyle};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
|
||||
@@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let (model_id, model_name) = if let Some(ref m) = config.model {
|
||||
(m.clone(), config.model_name.clone())
|
||||
} else {
|
||||
println!("Model not specified, fetching first model from server...");
|
||||
tracing::info!(base_url = %config.base_url, "fetching first model from server");
|
||||
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
|
||||
println!("First model name: {name}, first model id: {id}");
|
||||
tracing::info!(
|
||||
model_name = name,
|
||||
model_id = id,
|
||||
"selected first model from server"
|
||||
);
|
||||
(id, Some(name))
|
||||
};
|
||||
|
||||
@@ -83,15 +88,19 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
None
|
||||
} else {
|
||||
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
|
||||
println!("Loading tokenizer: {tid}");
|
||||
tracing::info!(tokenizer = tid, "loading tokenizer");
|
||||
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
|
||||
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
|
||||
println!("Tokenizer loaded successfully.");
|
||||
let t =
|
||||
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
|
||||
Some(t)
|
||||
};
|
||||
|
||||
// Generate/load conversations
|
||||
println!("Generating multi-turn conversations...");
|
||||
tracing::info!(
|
||||
dataset = ?config.dataset_name,
|
||||
conversations = config.num_prompts,
|
||||
"generating multi-turn conversations"
|
||||
);
|
||||
let gen_start = Instant::now();
|
||||
|
||||
let mut conversations = match config.dataset_name {
|
||||
@@ -131,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let path = match config.dataset_path.as_deref() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
|
||||
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
|
||||
downloaded.as_str()
|
||||
}
|
||||
};
|
||||
@@ -179,8 +188,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let (filtered_conversations, filtered_turns) =
|
||||
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
|
||||
if filtered_turns > 0 || filtered_conversations > 0 {
|
||||
println!(
|
||||
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
|
||||
tracing::info!(
|
||||
filtered_turns,
|
||||
filtered_conversations,
|
||||
max_model_len,
|
||||
"filtered conversations above maximum model length"
|
||||
);
|
||||
}
|
||||
if conversations.is_empty() {
|
||||
@@ -192,11 +204,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
|
||||
let gen_elapsed = gen_start.elapsed();
|
||||
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
|
||||
println!(
|
||||
"Generated {} conversations ({} total turns) in {:.2}s",
|
||||
conversations.len(),
|
||||
tracing::info!(
|
||||
conversations = conversations.len(),
|
||||
total_turns,
|
||||
gen_elapsed.as_secs_f64()
|
||||
elapsed_seconds = gen_elapsed.as_secs_f64(),
|
||||
"generated multi-turn conversations"
|
||||
);
|
||||
|
||||
// Log prefix sharing info
|
||||
@@ -208,18 +220,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let conv_tokens =
|
||||
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
|
||||
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
|
||||
println!(
|
||||
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
|
||||
config.multi_turn_prefix_global_ratio * 100.0,
|
||||
tracing::info!(
|
||||
global_ratio = config.multi_turn_prefix_global_ratio,
|
||||
global_tokens,
|
||||
config.multi_turn_prefix_conversation_ratio * 100.0,
|
||||
conv_tokens,
|
||||
(1.0 - config.multi_turn_prefix_global_ratio
|
||||
- config.multi_turn_prefix_conversation_ratio)
|
||||
* 100.0,
|
||||
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
|
||||
conversation_tokens = conv_tokens,
|
||||
unique_tokens,
|
||||
history_accumulation = false,
|
||||
"configured multi-turn prefix sharing"
|
||||
);
|
||||
println!("No history accumulation: each turn sends fixed-length prompt only.");
|
||||
}
|
||||
|
||||
if config.dry_run {
|
||||
@@ -253,7 +262,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
println!("Starting initial single prompt test run...");
|
||||
tracing::info!("starting initial single-prompt test run");
|
||||
let test_output = crate::ready_checker::wait_for_endpoint(
|
||||
config.backend,
|
||||
&client,
|
||||
@@ -268,7 +277,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
test_output.error
|
||||
)));
|
||||
}
|
||||
println!("Initial test run completed.");
|
||||
tracing::info!("initial single-prompt test run completed");
|
||||
}
|
||||
|
||||
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
|
||||
@@ -283,9 +292,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
"min_tokens".to_string(),
|
||||
serde_json::json!(config.random_output_len),
|
||||
);
|
||||
println!(
|
||||
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
|
||||
config.random_output_len
|
||||
tracing::info!(
|
||||
min_tokens = config.random_output_len,
|
||||
dataset = "random",
|
||||
"set minimum output tokens for multi-turn dataset"
|
||||
);
|
||||
}
|
||||
Some(body)
|
||||
@@ -297,7 +307,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
let spec_decode_before =
|
||||
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
|
||||
if spec_decode_before.is_some() {
|
||||
println!("Speculative decoding detected, will collect metrics.");
|
||||
tracing::info!("detected speculative decoding; collecting metrics");
|
||||
}
|
||||
|
||||
// Start profiler if requested (immediate mode — no batch threshold)
|
||||
@@ -330,10 +340,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
};
|
||||
|
||||
// Main benchmark
|
||||
println!("Starting multi-turn benchmark...");
|
||||
println!("Conversations: {}", conversations.len());
|
||||
println!("Concurrency: {concurrency}");
|
||||
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
|
||||
tracing::info!(
|
||||
conversations = conversations.len(),
|
||||
total_turns,
|
||||
concurrency,
|
||||
inter_turn_delay_ms = config.multi_turn_delay_ms,
|
||||
"starting multi-turn benchmark"
|
||||
);
|
||||
|
||||
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
|
||||
|
||||
@@ -364,11 +377,12 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
);
|
||||
if let Some(modules) = config.lora_modules.as_ref() {
|
||||
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
|
||||
println!(
|
||||
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
|
||||
modules.len(),
|
||||
names,
|
||||
config.lora_assignment
|
||||
tracing::info!(
|
||||
adapters = modules.len(),
|
||||
names = ?names,
|
||||
assignment = ?config.lora_assignment,
|
||||
scope = "conversation",
|
||||
"assigned LoRA adapters"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
match handle.await {
|
||||
Ok(output) => all_outputs.push(output),
|
||||
Err(e) => {
|
||||
eprintln!("Conversation task panicked: {e}");
|
||||
tracing::error!(error = %e.as_report(), "conversation task panicked");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -453,7 +467,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
if let Some((cancel_tx, task)) = profile_task {
|
||||
let _ = cancel_tx.send(());
|
||||
if let Err(e) = task.await {
|
||||
eprintln!("WARNING: Profile background task failed: {e}");
|
||||
tracing::error!(error = %e.as_report(), "profiler background task failed");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -603,7 +603,7 @@ fn add_metric_stats(
|
||||
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
|
||||
let content = serde_json::to_string(json)?;
|
||||
std::fs::write(file_path, content)?;
|
||||
println!("Results saved to {file_path}");
|
||||
tracing::info!(path = file_path, "saved benchmark results");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
|
||||
file.write_all(b"\n")?;
|
||||
}
|
||||
file.write_all(content.as_bytes())?;
|
||||
println!("Results appended to {file_path}");
|
||||
tracing::info!(path = file_path, "appended benchmark results");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
|
||||
let backend = get_backend(backend)?;
|
||||
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
|
||||
|
||||
println!("Waiting for endpoint to become up in {timeout_seconds}s");
|
||||
tracing::info!(
|
||||
timeout_seconds,
|
||||
retry_interval,
|
||||
"waiting for endpoint readiness"
|
||||
);
|
||||
|
||||
let pb = ProgressBar::new(timeout_seconds);
|
||||
pb.set_style(
|
||||
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
|
||||
Ok(output) => {
|
||||
let err = output.error.clone();
|
||||
let err_last_line = err.lines().last().unwrap_or(&err);
|
||||
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
|
||||
pb.suspend(|| {
|
||||
tracing::warn!(error = err_last_line, "endpoint is not ready");
|
||||
});
|
||||
last_error = err;
|
||||
}
|
||||
Err(e) => {
|
||||
|
||||
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
|
||||
.await
|
||||
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
|
||||
if resp.status().is_success() {
|
||||
println!("Prefix cache reset successfully.");
|
||||
tracing::info!(url = %url, "reset prefix cache");
|
||||
} else {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
|
||||
@@ -199,12 +199,17 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
|
||||
}
|
||||
};
|
||||
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
|
||||
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
|
||||
tracing::info!(
|
||||
encoding,
|
||||
kind = "built-in-tiktoken",
|
||||
vocab_size,
|
||||
"loaded tokenizer"
|
||||
);
|
||||
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
|
||||
}
|
||||
|
||||
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
|
||||
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
// Phase 1: If model_id is a local directory, look for tiktoken files there
|
||||
let local_dir = Path::new(model_id);
|
||||
if local_dir.is_dir() {
|
||||
@@ -212,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
}
|
||||
|
||||
// Phase 2: Fall back to HuggingFace Hub download
|
||||
try_load_tiktoken_from_hf(model_id)
|
||||
try_load_tiktoken_from_hf(model_id).await
|
||||
}
|
||||
|
||||
/// Common tiktoken model filenames to search for.
|
||||
@@ -247,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
|
||||
}
|
||||
|
||||
/// Load a tiktoken tokenizer from a HuggingFace model repo.
|
||||
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
||||
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||
|
||||
let model_path = repo
|
||||
.get("tiktoken.model")
|
||||
.or_else(|_| repo.get("qwen.tiktoken"))
|
||||
.or_else(|_| repo.get("vocab.tiktoken"))
|
||||
.map_err(|_| {
|
||||
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
||||
})?;
|
||||
let mut model_path = None;
|
||||
for filename in TIKTOKEN_MODEL_FILENAMES {
|
||||
if let Ok(path) = repo.get(filename).await {
|
||||
model_path = Some(path);
|
||||
break;
|
||||
}
|
||||
}
|
||||
let model_path = model_path.ok_or_else(|| {
|
||||
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
|
||||
})?;
|
||||
|
||||
let num_base_tokens = count_base_tokens(&model_path)?;
|
||||
|
||||
let config = match repo.get("tokenizer_config.json") {
|
||||
let config = match repo.get("tokenizer_config.json").await {
|
||||
Ok(config_path) => read_tokenizer_config(&config_path),
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
let pattern = extract_pat_str_from_repo(&repo);
|
||||
let pattern = extract_pat_str_from_repo(&repo).await;
|
||||
|
||||
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
|
||||
}
|
||||
@@ -309,15 +317,16 @@ fn build_tiktoken(
|
||||
}
|
||||
}
|
||||
|
||||
println!(
|
||||
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
|
||||
num_base_tokens,
|
||||
all_special_tokens.len(),
|
||||
if pattern.is_some() {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
base_tokens = num_base_tokens,
|
||||
special_tokens = all_special_tokens.len(),
|
||||
pattern = if pattern.is_some() {
|
||||
"custom"
|
||||
} else {
|
||||
"default"
|
||||
},
|
||||
"loading tiktoken model"
|
||||
);
|
||||
|
||||
TiktokenTokenizer::from_file(
|
||||
@@ -397,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
|
||||
|
||||
/// Try to download the Python tokenizer source file and extract pat_str via regex.
|
||||
/// Returns None if unavailable or unparsable.
|
||||
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
||||
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
|
||||
// Try common Python tokenizer filenames
|
||||
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
|
||||
let py_path = match repo.get("tokenization_kimi.py").await {
|
||||
Ok(path) => path,
|
||||
Err(_) => repo.get("tokenizer.py").await.ok()?,
|
||||
};
|
||||
|
||||
let source = std::fs::read_to_string(&py_path).ok()?;
|
||||
|
||||
@@ -438,9 +450,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
|
||||
|
||||
if !fragments.is_empty() {
|
||||
let pattern = fragments.join("|");
|
||||
println!(
|
||||
"Extracted pat_str from Python source: {} fragments",
|
||||
fragments.len()
|
||||
tracing::debug!(
|
||||
fragments = fragments.len(),
|
||||
"extracted tiktoken pattern from Python source"
|
||||
);
|
||||
return Some(pattern);
|
||||
}
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::future::Future;
|
||||
use std::path::Path;
|
||||
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
use crate::error::{BenchError, Result};
|
||||
@@ -18,7 +20,8 @@ pub enum TokenizerKind {
|
||||
|
||||
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
|
||||
pub struct ServerTokenizer {
|
||||
client: reqwest::blocking::Client,
|
||||
client: reqwest::Client,
|
||||
runtime: tokio::runtime::Handle,
|
||||
tokenize_url: String,
|
||||
detokenize_url: String,
|
||||
model: String,
|
||||
@@ -27,8 +30,8 @@ pub struct ServerTokenizer {
|
||||
|
||||
impl ServerTokenizer {
|
||||
/// Create a new server tokenizer and verify connectivity.
|
||||
pub fn new(base_url: &str, model: &str) -> Result<Self> {
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
|
||||
@@ -38,6 +41,7 @@ impl ServerTokenizer {
|
||||
|
||||
let st = Self {
|
||||
client,
|
||||
runtime: tokio::runtime::Handle::current(),
|
||||
tokenize_url,
|
||||
detokenize_url,
|
||||
model: model.to_string(),
|
||||
@@ -45,7 +49,7 @@ impl ServerTokenizer {
|
||||
};
|
||||
|
||||
// Probe the endpoint to verify it works and discover vocab size
|
||||
let test_tokens = st.encode_inner("test")?;
|
||||
let test_tokens = st.encode_async("test").await?;
|
||||
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
|
||||
let estimated_vocab = (max_id * 2).max(131072);
|
||||
|
||||
@@ -56,6 +60,10 @@ impl ServerTokenizer {
|
||||
}
|
||||
|
||||
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
|
||||
self.block_on(self.encode_async(text))
|
||||
}
|
||||
|
||||
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model,
|
||||
"prompt": text,
|
||||
@@ -66,6 +74,7 @@ impl ServerTokenizer {
|
||||
.post(&self.tokenize_url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
@@ -75,7 +84,7 @@ impl ServerTokenizer {
|
||||
)));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
||||
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
|
||||
})?;
|
||||
|
||||
@@ -95,6 +104,10 @@ impl ServerTokenizer {
|
||||
}
|
||||
|
||||
fn decode_inner(&self, ids: &[u32]) -> Result<String> {
|
||||
self.block_on(self.decode_async(ids))
|
||||
}
|
||||
|
||||
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
|
||||
let payload = serde_json::json!({
|
||||
"model": self.model,
|
||||
"tokens": ids,
|
||||
@@ -105,6 +118,7 @@ impl ServerTokenizer {
|
||||
.post(&self.detokenize_url)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
@@ -114,7 +128,7 @@ impl ServerTokenizer {
|
||||
)));
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().map_err(|e| {
|
||||
let data: serde_json::Value = resp.json().await.map_err(|e| {
|
||||
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
|
||||
})?;
|
||||
|
||||
@@ -123,6 +137,26 @@ impl ServerTokenizer {
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
|
||||
}
|
||||
|
||||
fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
|
||||
if matches!(
|
||||
self.runtime.runtime_flavor(),
|
||||
tokio::runtime::RuntimeFlavor::CurrentThread
|
||||
) {
|
||||
return Err(BenchError::Tokenizer(
|
||||
"Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
|
||||
// Tokio workers must enter a blocking region before re-entering the runtime;
|
||||
// Rayon workers can drive the future directly with the saved runtime handle.
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
tokio::task::block_in_place(|| self.runtime.block_on(future))
|
||||
} else {
|
||||
self.runtime.block_on(future)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- TokenizerKind methods ---
|
||||
@@ -192,7 +226,7 @@ impl TokenizerKind {
|
||||
/// 3. Server-side /tokenize + /detokenize endpoints
|
||||
///
|
||||
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
|
||||
pub fn load_tokenizer(
|
||||
pub async fn load_tokenizer(
|
||||
model_id: &str,
|
||||
_trust_remote_code: bool,
|
||||
server_info: Option<(&str, &str)>,
|
||||
@@ -212,31 +246,48 @@ pub fn load_tokenizer(
|
||||
}
|
||||
|
||||
// 1. Try local HuggingFace tokenizer (tokenizer.json)
|
||||
match try_load_local(model_id) {
|
||||
match try_load_local(model_id).await {
|
||||
Ok(tok) => {
|
||||
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "local",
|
||||
vocab_size = tok.get_vocab_size(true),
|
||||
"loaded tokenizer"
|
||||
);
|
||||
Ok(TokenizerKind::Local(Box::new(tok)))
|
||||
}
|
||||
Err(local_err) => {
|
||||
// 2. Try tiktoken format
|
||||
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
|
||||
match crate::tiktoken::try_load_tiktoken(model_id) {
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
error = %local_err.as_report(),
|
||||
"local tokenizer unavailable; trying tiktoken"
|
||||
);
|
||||
match crate::tiktoken::try_load_tiktoken(model_id).await {
|
||||
Ok(tok) => {
|
||||
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "tiktoken",
|
||||
vocab_size = tok.vocab_size(),
|
||||
"loaded tokenizer"
|
||||
);
|
||||
Ok(TokenizerKind::Tiktoken(tok))
|
||||
}
|
||||
Err(tiktoken_err) => {
|
||||
// 3. Try server-side fallback
|
||||
if let Some((base_url, model)) = server_info {
|
||||
println!(
|
||||
"Tiktoken also not available ({tiktoken_err}), \
|
||||
trying server-side tokenization..."
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
error = %tiktoken_err.as_report(),
|
||||
"tiktoken unavailable; trying server-side tokenization"
|
||||
);
|
||||
match ServerTokenizer::new(base_url, model) {
|
||||
match ServerTokenizer::new(base_url, model).await {
|
||||
Ok(srv) => {
|
||||
println!(
|
||||
"Tokenizer: Server (vocab_size≈{})",
|
||||
srv.cached_vocab_size
|
||||
tracing::info!(
|
||||
model = model_id,
|
||||
kind = "server",
|
||||
vocab_size = srv.cached_vocab_size,
|
||||
"loaded tokenizer"
|
||||
);
|
||||
return Ok(TokenizerKind::Server(srv));
|
||||
}
|
||||
@@ -264,7 +315,7 @@ pub fn load_tokenizer(
|
||||
}
|
||||
|
||||
/// Try loading tokenizer.json from local path or HuggingFace Hub.
|
||||
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
// 1. Try local directory with tokenizer.json
|
||||
let local_path = Path::new(model_id).join("tokenizer.json");
|
||||
if local_path.exists() {
|
||||
@@ -290,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result<Tokenizer> {
|
||||
}
|
||||
|
||||
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string());
|
||||
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
|
||||
let tokenizer_path = repo
|
||||
.get("tokenizer.json")
|
||||
.await
|
||||
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
|
||||
|
||||
Tokenizer::from_file(&tokenizer_path)
|
||||
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn test_server_tokenizer_sync_bridge() {
|
||||
let tokenizer = std::sync::Arc::new(ServerTokenizer {
|
||||
client: reqwest::Client::new(),
|
||||
runtime: tokio::runtime::Handle::current(),
|
||||
tokenize_url: String::new(),
|
||||
detokenize_url: String::new(),
|
||||
model: String::new(),
|
||||
cached_vocab_size: 0,
|
||||
});
|
||||
|
||||
assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
rayon::spawn(move || {
|
||||
let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
|
||||
});
|
||||
assert_eq!(rx.await.unwrap().unwrap(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
+106
-61
@@ -54,6 +54,7 @@ mod stream;
|
||||
|
||||
use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
@@ -88,6 +89,92 @@ pub fn validate_parser_overrides(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Chat request preparation shared by inference and render-only frontends.
|
||||
pub struct ChatRequestProcessor {
|
||||
backend: DynChatBackend,
|
||||
/// Effective model dtype reported by the engine.
|
||||
/// Absent for text-only frontends without an engine handshake.
|
||||
model_dtype: Option<ModelDtype>,
|
||||
}
|
||||
|
||||
impl ChatRequestProcessor {
|
||||
/// Create a processor with multimodal support using the effective model
|
||||
/// dtype reported by the engine.
|
||||
fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
model_dtype: Some(model_dtype),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a render-only processor that rejects multimodal requests.
|
||||
pub fn render_only(backend: DynChatBackend) -> Self {
|
||||
Self {
|
||||
backend,
|
||||
model_dtype: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn finalize_rendered_prompt(
|
||||
&self,
|
||||
request: &ChatRequest,
|
||||
rendered: RenderedPrompt,
|
||||
) -> Result<(Prompt, Option<MmFeatures>)> {
|
||||
match self.model_dtype {
|
||||
Some(model_dtype) => {
|
||||
multimodal::finalize_rendered_prompt(
|
||||
request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
model_dtype,
|
||||
)
|
||||
.await
|
||||
}
|
||||
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
|
||||
None => Err(Error::UnsupportedMultimodalRenderer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Prepare one chat request without submitting it to an engine.
|
||||
pub async fn prepare(
|
||||
&self,
|
||||
mut request: ChatRequest,
|
||||
options: NewChatOutputProcessorOptions<'_>,
|
||||
) -> Result<(TextRequest, DynChatOutputProcessor)> {
|
||||
request.validate()?;
|
||||
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
let output_processor = self.backend.new_chat_output_processor(&mut request, options)?;
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let reasoning_parser_kwargs =
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.is_some()
|
||||
.then(|| ReasoningParserKwargs {
|
||||
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
|
||||
});
|
||||
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
|
||||
let text_request = TextRequest {
|
||||
request_id: request.request_id,
|
||||
prompt,
|
||||
mm_features,
|
||||
sampling_params: request.sampling_params,
|
||||
decode_options: request.decode_options,
|
||||
intermediate: request.intermediate,
|
||||
priority: request.priority,
|
||||
cache_salt: request.cache_salt,
|
||||
add_special_tokens: request.add_special_tokens,
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
reasoning_parser_kwargs,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
};
|
||||
Ok((text_request, output_processor))
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured chat facade above [`TextLlm`].
|
||||
///
|
||||
/// This layer stays above raw text semantics: it takes care of chat-template
|
||||
@@ -95,9 +182,7 @@ pub fn validate_parser_overrides(
|
||||
/// request semantics such as tool calls.
|
||||
pub struct ChatLlm {
|
||||
text: TextLlm,
|
||||
backend: DynChatBackend,
|
||||
/// Effective model dtype reported by the engine.
|
||||
model_dtype: ModelDtype,
|
||||
processor: ChatRequestProcessor,
|
||||
/// Tool-call parser selection.
|
||||
tool_call_parser: ParserSelection,
|
||||
/// Reasoning parser selection.
|
||||
@@ -112,8 +197,7 @@ impl ChatLlm {
|
||||
|
||||
Self {
|
||||
text,
|
||||
backend,
|
||||
model_dtype,
|
||||
processor: ChatRequestProcessor::new(backend, model_dtype),
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
}
|
||||
@@ -140,7 +224,7 @@ impl ChatLlm {
|
||||
|
||||
/// Override the effective model dtype used for multimodal tensor encoding.
|
||||
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
|
||||
self.model_dtype = model_dtype;
|
||||
self.processor.model_dtype = Some(model_dtype);
|
||||
self
|
||||
}
|
||||
|
||||
@@ -172,57 +256,23 @@ impl ChatLlm {
|
||||
}
|
||||
|
||||
/// Render, tokenize, and submit one chat request.
|
||||
pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
|
||||
request.validate()?;
|
||||
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
|
||||
let output_processor = self.backend.new_chat_output_processor(
|
||||
&mut request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &self.tool_call_parser,
|
||||
reasoning_parser: &self.reasoning_parser,
|
||||
},
|
||||
)?;
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let reasoning_parser_kwargs =
|
||||
request
|
||||
.sampling_params
|
||||
.structured_outputs
|
||||
.is_some()
|
||||
.then(|| ReasoningParserKwargs {
|
||||
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
|
||||
});
|
||||
|
||||
let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
|
||||
&request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
self.model_dtype,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let text_request = TextRequest {
|
||||
request_id: request.request_id.clone(),
|
||||
prompt,
|
||||
mm_features,
|
||||
sampling_params: request.sampling_params,
|
||||
decode_options: request.decode_options,
|
||||
intermediate: request.intermediate,
|
||||
priority: request.priority,
|
||||
cache_salt: request.cache_salt,
|
||||
add_special_tokens: request.add_special_tokens,
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
reasoning_parser_kwargs,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
};
|
||||
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
|
||||
let (text_request, output_processor) = self
|
||||
.processor
|
||||
.prepare(
|
||||
request,
|
||||
NewChatOutputProcessorOptions {
|
||||
tool_call_parser: &self.tool_call_parser,
|
||||
reasoning_parser: &self.reasoning_parser,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let request_id = text_request.request_id.clone();
|
||||
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
|
||||
|
||||
let structured_stream = output_processor.process(decoded_stream)?;
|
||||
|
||||
Ok(ChatEventStream::new(request.request_id, structured_stream))
|
||||
Ok(ChatEventStream::new(request_id, structured_stream))
|
||||
}
|
||||
|
||||
/// Render through the chat template and tokenize, without submitting to the engine.
|
||||
@@ -233,14 +283,9 @@ impl ChatLlm {
|
||||
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
|
||||
request.validate()?;
|
||||
|
||||
let rendered = self.backend.chat_renderer().render(&request)?;
|
||||
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
|
||||
&request,
|
||||
rendered,
|
||||
self.backend.multimodal_model_info(),
|
||||
self.model_dtype,
|
||||
)
|
||||
.await?;
|
||||
let rendered = self.processor.backend.chat_renderer().render(&request)?;
|
||||
let (prompt, _mm_features) =
|
||||
self.processor.finalize_rendered_prompt(&request, rendered).await?;
|
||||
|
||||
let tokenizer = self.text.tokenizer();
|
||||
let token_ids = match prompt {
|
||||
|
||||
@@ -26,18 +26,21 @@ const RESET: &str = "\x1b[0m";
|
||||
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
|
||||
format_description!("[month]-[day] [hour]:[minute]:[second]");
|
||||
|
||||
const PROCESS_LABEL: &str = "RustFrontend";
|
||||
|
||||
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
|
||||
pub(crate) fn init_tracing() {
|
||||
pub(crate) fn init_tracing(process_label: &str) {
|
||||
let filter = build_targets_filter(
|
||||
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
|
||||
env::var("RUST_LOG").ok().as_deref(),
|
||||
);
|
||||
let formatter = VllmEventFormatter::new();
|
||||
let formatter = VllmEventFormatter::new(process_label);
|
||||
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.event_format(formatter)
|
||||
.with_writer(std::io::stderr)
|
||||
.with_filter(filter),
|
||||
)
|
||||
.try_init();
|
||||
}
|
||||
|
||||
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
|
||||
}
|
||||
|
||||
impl VllmEventFormatter {
|
||||
fn new() -> Self {
|
||||
fn new(process_label: &str) -> Self {
|
||||
Self {
|
||||
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
|
||||
prefix: format!("({process_label} pid={})", process::id()),
|
||||
timer: VllmLocalTimer::default(),
|
||||
}
|
||||
}
|
||||
@@ -291,6 +294,13 @@ fn map_python_log_level(level: &str) -> LevelFilter {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn formatter_prefix_uses_process_label() {
|
||||
let formatter = VllmEventFormatter::new("Bench");
|
||||
|
||||
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
|
||||
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
|
||||
|
||||
@@ -5,6 +5,7 @@ mod cli;
|
||||
mod logging;
|
||||
|
||||
use std::env;
|
||||
use std::ffi::OsStr;
|
||||
use std::process::ExitStatus;
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
logging::init_tracing();
|
||||
let process_label =
|
||||
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
|
||||
"bench" => "Bench",
|
||||
"serve" | "frontend" => "RustFrontend",
|
||||
_ => "Rust",
|
||||
};
|
||||
logging::init_tracing(process_label);
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
let mut runtime = tokio::runtime::Builder::new_multi_thread();
|
||||
|
||||
@@ -460,6 +460,12 @@ impl EngineCoreClient {
|
||||
self.inner.is_healthy()
|
||||
}
|
||||
|
||||
/// Subscribe to engine health changes. The current value is `true` while
|
||||
/// the client is healthy and changes permanently to `false` on failure.
|
||||
pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
|
||||
self.inner.subscribe_health()
|
||||
}
|
||||
|
||||
/// Return the first persistent health error observed by the client, if any.
|
||||
pub fn health_error(&self) -> Option<Arc<Error>> {
|
||||
self.inner.health_error()
|
||||
|
||||
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
|
||||
use parking_lot::Mutex;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::runtime::Handle;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tracing::{debug, info, trace, warn};
|
||||
use vllm_metrics::METRICS;
|
||||
use zeromq::RouterSendHalf;
|
||||
@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {
|
||||
request_reg: Mutex<RequestRegistry>,
|
||||
utility_reg: Mutex<UtilityRegistry>,
|
||||
health_error: ArcSwapOption<Error>,
|
||||
health_tx: watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl ClientInner {
|
||||
@@ -57,6 +58,7 @@ impl ClientInner {
|
||||
request_reg: Mutex::new(RequestRegistry::new(engines)),
|
||||
utility_reg: Mutex::new(UtilityRegistry::default()),
|
||||
health_error: ArcSwapOption::empty(),
|
||||
health_tx: watch::Sender::new(true),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,6 +171,7 @@ impl ClientInner {
|
||||
/// persistent health error.
|
||||
pub fn close_registries(&self, error: Arc<Error>) {
|
||||
let persistent_error = self.record_health_error(error);
|
||||
self.publish_unhealthy();
|
||||
let request_senders = self.request_reg.lock().close();
|
||||
let utility_senders = self.utility_reg.lock().close();
|
||||
|
||||
@@ -191,6 +194,12 @@ impl ClientInner {
|
||||
self.health_error.load().is_none()
|
||||
}
|
||||
|
||||
/// Subscribe to engine health changes. The current value is `true` while
|
||||
/// the client is healthy and changes permanently to `false` on failure.
|
||||
pub fn subscribe_health(&self) -> watch::Receiver<bool> {
|
||||
self.health_tx.subscribe()
|
||||
}
|
||||
|
||||
/// Resolve one utility output to the waiting caller. Returns `true` if a
|
||||
/// waiting caller existed.
|
||||
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
|
||||
@@ -280,6 +289,11 @@ impl ClientInner {
|
||||
.expect("health error must be recorded before registries close")
|
||||
}
|
||||
|
||||
/// Publish the sticky healthy-to-unhealthy transition.
|
||||
fn publish_unhealthy(&self) {
|
||||
self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
|
||||
}
|
||||
|
||||
/// Assert there is a recorded health error and return a `Shared` variant
|
||||
/// wrapping it for error returns when the client is already closed.
|
||||
fn closed_error(&self) -> Error {
|
||||
@@ -461,13 +475,18 @@ mod tests {
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn close_registries_records_first_health_error_only() {
|
||||
let inner = test_inner().await;
|
||||
let mut health = inner.subscribe_health();
|
||||
assert!(*health.borrow());
|
||||
|
||||
inner.close_registries(Arc::new(Error::EngineCoreDead));
|
||||
health.changed().await.expect("health sender remains open");
|
||||
assert!(!inner.is_healthy());
|
||||
assert!(!*health.borrow());
|
||||
assert!(matches!(
|
||||
inner.health_error().as_deref(),
|
||||
Some(Error::EngineCoreDead)
|
||||
));
|
||||
assert!(!*inner.subscribe_health().borrow());
|
||||
|
||||
inner.close_registries(Arc::new(client_closed!("shutdown")));
|
||||
assert!(matches!(
|
||||
|
||||
@@ -269,6 +269,19 @@ impl WireLogprobs {
|
||||
);
|
||||
}
|
||||
|
||||
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
|
||||
if token_ids.rows == 0 {
|
||||
return Ok(Logprobs {
|
||||
positions: Vec::new(),
|
||||
});
|
||||
}
|
||||
if token_ids.cols == 0 {
|
||||
bail_ext_value_decode!(
|
||||
"{field_prefix}: zero-column logprobs payload with {} rows",
|
||||
token_ids.rows
|
||||
);
|
||||
}
|
||||
|
||||
let mut positions = Vec::with_capacity(token_ids.rows);
|
||||
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
|
||||
.data
|
||||
|
||||
@@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() {
|
||||
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decodes_zero_row_logprobs_as_empty() {
|
||||
for shape in [[0usize, 0], [0, 3]] {
|
||||
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
|
||||
None,
|
||||
Some(Value::Array(vec![
|
||||
ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
|
||||
ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
|
||||
ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
|
||||
Value::Nil,
|
||||
])),
|
||||
)))];
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||
let logprobs = decoded.outputs[0]
|
||||
.new_prompt_logprobs_tensors
|
||||
.clone()
|
||||
.unwrap()
|
||||
.into_direct()
|
||||
.unwrap();
|
||||
assert!(logprobs.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_column_logprobs_with_rows() {
|
||||
let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
|
||||
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
|
||||
Some(Value::Array(vec![
|
||||
ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
|
||||
ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
|
||||
ndarray_value("<i8", &[2], ranks),
|
||||
Value::Nil,
|
||||
])),
|
||||
None,
|
||||
)))];
|
||||
|
||||
let error = decode_engine_core_outputs(&frames).unwrap_err();
|
||||
let crate::error::Error::ExtValueDecode { message } = &error else {
|
||||
panic!("expected ExtValueDecode");
|
||||
};
|
||||
assert_eq!(
|
||||
message,
|
||||
"new_logprobs: zero-column logprobs payload with 2 rows"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -35,6 +35,7 @@ tokio-openssl.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
tokio-util.workspace = true
|
||||
tonic.workspace = true
|
||||
tonic-health.workspace = true
|
||||
tonic-prost.workspace = true
|
||||
tower.workspace = true
|
||||
tower-http.workspace = true
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
// 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::{ControlGrpcService, 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 control_service = ControlGrpcService::NAME;
|
||||
let status = ServingStatus::NotServing;
|
||||
let health_event_first = tokio::select! {
|
||||
result = engine_health.wait_for(|healthy| !*healthy) => {
|
||||
match result {
|
||||
Ok(_) => warn!(
|
||||
status = ?status,
|
||||
reason = "engine_unhealthy",
|
||||
"marking gRPC health services as not serving"
|
||||
),
|
||||
Err(error) => warn!(
|
||||
%error,
|
||||
status = ?status,
|
||||
reason = "health_channel_closed",
|
||||
"engine health channel closed; marking gRPC health services as not serving"
|
||||
),
|
||||
}
|
||||
true
|
||||
}
|
||||
_ = shutdown.cancelled() => {
|
||||
info!(
|
||||
status = ?status,
|
||||
reason = "server_shutdown",
|
||||
"server shutting down; marking gRPC health services as not serving"
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
health_reporter.set_not_serving::<GenerateGrpcService>().await;
|
||||
health_reporter.set_not_serving::<ControlGrpcService>().await;
|
||||
// Both gRPC services use the same engine client, so overall server health
|
||||
// mirrors their shared engine health.
|
||||
health_reporter.set_service_status("", status).await;
|
||||
|
||||
if health_event_first {
|
||||
shutdown.cancelled().await;
|
||||
info!(
|
||||
reason = "server_shutdown",
|
||||
"server shutting down; closing gRPC health watches"
|
||||
);
|
||||
}
|
||||
|
||||
health_reporter.clear_service_status(generate_service).await;
|
||||
health_reporter.clear_service_status(control_service).await;
|
||||
health_reporter.clear_service_status("").await;
|
||||
}
|
||||
@@ -4,6 +4,7 @@
|
||||
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
|
||||
|
||||
mod convert;
|
||||
mod health;
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
@@ -24,8 +25,13 @@ pub mod pb {
|
||||
tonic::include_proto!("vllm");
|
||||
}
|
||||
|
||||
pub(crate) use health::monitor_health;
|
||||
pub use pb::control_server::ControlServer;
|
||||
pub use pb::generate_server::GenerateServer;
|
||||
|
||||
pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
|
||||
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
@@ -40,6 +46,36 @@ impl GenerateServiceImpl {
|
||||
}
|
||||
}
|
||||
|
||||
/// gRPC control service backed by the shared application state.
|
||||
pub struct ControlServiceImpl {
|
||||
state: Arc<AppState>,
|
||||
}
|
||||
|
||||
impl ControlServiceImpl {
|
||||
pub fn new(state: Arc<AppState>) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl pb::control_server::Control for ControlServiceImpl {
|
||||
async fn abort(
|
||||
&self,
|
||||
request: Request<pb::AbortRequest>,
|
||||
) -> Result<Response<pb::AbortResponse>, Status> {
|
||||
let request_ids = request.into_inner().request_ids;
|
||||
if request_ids.is_empty() {
|
||||
return Ok(Response::new(pb::AbortResponse {}));
|
||||
}
|
||||
self.state
|
||||
.chat
|
||||
.abort(&request_ids)
|
||||
.await
|
||||
.map_err(|error| Status::internal(error.to_report_string()))?;
|
||||
Ok(Response::new(pb::AbortResponse {}))
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl pb::generate_server::Generate for GenerateServiceImpl {
|
||||
type GenerateStreamStream =
|
||||
|
||||
@@ -16,6 +16,10 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_openssl::SslStream;
|
||||
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
|
||||
use tonic_health::pb::HealthCheckRequest;
|
||||
use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
|
||||
use tonic_health::pb::health_client::HealthClient;
|
||||
use tonic_health::server::health_reporter;
|
||||
use tower::service_fn;
|
||||
use vllm_chat::{
|
||||
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
|
||||
@@ -34,8 +38,9 @@ use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
use super::pb::control_client::ControlClient;
|
||||
use super::pb::generate_client::GenerateClient;
|
||||
use super::{GenerateServer, GenerateServiceImpl, pb};
|
||||
use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
|
||||
use crate::listener::{Listener, MaybeTlsListener};
|
||||
use crate::state::AppState;
|
||||
use crate::tls;
|
||||
@@ -149,10 +154,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
|
||||
dealer.recv().await.expect("recv engine message").into_vec()
|
||||
}
|
||||
|
||||
fn test_llm(client: EngineCoreClient) -> Llm {
|
||||
Llm::new(client).with_request_id_randomization(false)
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct FakeTextBackend;
|
||||
|
||||
@@ -200,7 +201,12 @@ impl ChatRenderer for FakeTextBackend {
|
||||
async fn setup_grpc_service(
|
||||
engine_id: impl Into<EngineId>,
|
||||
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
||||
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
|
||||
) -> (
|
||||
GenerateServer<GenerateServiceImpl>,
|
||||
ControlServer<ControlServiceImpl>,
|
||||
tokio::sync::watch::Receiver<bool>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = engine_id.into();
|
||||
@@ -232,14 +238,17 @@ async fn setup_grpc_service(
|
||||
)
|
||||
.await
|
||||
.expect("connect client");
|
||||
let engine_health = client.subscribe_health();
|
||||
|
||||
let chat = ChatLlm::from_shared_backend(
|
||||
test_llm(client),
|
||||
Llm::new(client),
|
||||
Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
|
||||
);
|
||||
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
|
||||
(
|
||||
GenerateServer::new(GenerateServiceImpl::new(state)),
|
||||
GenerateServer::new(GenerateServiceImpl::new(state.clone())),
|
||||
ControlServer::new(ControlServiceImpl::new(state)),
|
||||
engine_health,
|
||||
engine_task,
|
||||
)
|
||||
}
|
||||
@@ -254,25 +263,56 @@ async fn grpc_test_server(
|
||||
tokio::task::JoinHandle<()>,
|
||||
MockEngineTask,
|
||||
) {
|
||||
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let (generate_service, control_service, engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, output_specs).await;
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
generate_service,
|
||||
control_service,
|
||||
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>,
|
||||
control_service: ControlServer<ControlServiceImpl>,
|
||||
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;
|
||||
health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc server");
|
||||
let server = TonicServer::builder()
|
||||
.add_service(health_service)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||
let health_monitor =
|
||||
super::monitor_health(health_reporter, engine_health, shutdown.clone());
|
||||
let server = async move {
|
||||
let result = server.await;
|
||||
shutdown.cancel();
|
||||
result
|
||||
};
|
||||
let (server_result, ()) = tokio::join!(server, health_monitor);
|
||||
server_result.expect("grpc server");
|
||||
});
|
||||
|
||||
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
|
||||
let channel = Endpoint::from_shared(format!("http://{addr}"))
|
||||
.expect("grpc endpoint")
|
||||
.connect()
|
||||
.await
|
||||
.expect("connect grpc client");
|
||||
.expect("connect grpc channel");
|
||||
|
||||
(grpc_client, server_task, engine_task)
|
||||
(channel, server_task)
|
||||
}
|
||||
|
||||
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
|
||||
@@ -283,7 +323,8 @@ async fn grpc_tls_test_server(
|
||||
certs: &TestCerts,
|
||||
cert_reqs: i32,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
|
||||
let (generate_service, control_service, _engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, output_specs).await;
|
||||
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
|
||||
.expect("build grpc tls config");
|
||||
|
||||
@@ -293,7 +334,8 @@ async fn grpc_tls_test_server(
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc tls server");
|
||||
@@ -373,7 +415,8 @@ async fn grpc_server_with_keepalive(
|
||||
engine_id: impl Into<EngineId>,
|
||||
keepalive: Option<Duration>,
|
||||
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
|
||||
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
|
||||
let (generate_service, control_service, _engine_health, engine_task) =
|
||||
setup_grpc_service(engine_id, default_stream_output_specs()).await;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
|
||||
let addr = listener.local_addr().expect("local addr").to_string();
|
||||
@@ -388,7 +431,8 @@ async fn grpc_server_with_keepalive(
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
builder
|
||||
.add_service(svc)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service)
|
||||
.serve_with_incoming(incoming)
|
||||
.await
|
||||
.expect("grpc server");
|
||||
@@ -1035,3 +1079,226 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
|
||||
|
||||
server_task.abort();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn control_abort_resolves_external_id_and_empty_is_noop() {
|
||||
let (generate_service, control_service, engine_health, engine_task) =
|
||||
setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
|
||||
let (channel, server_task) = start_grpc_test_server(
|
||||
generate_service,
|
||||
control_service,
|
||||
engine_health,
|
||||
tokio_util::sync::CancellationToken::new(),
|
||||
)
|
||||
.await;
|
||||
let mut generate_client = GenerateClient::new(channel.clone());
|
||||
let mut control_client = ControlClient::new(channel);
|
||||
let request_id = "test-abort-active";
|
||||
|
||||
let mut stream = generate_client
|
||||
.generate_stream(pb::GenerateRequest {
|
||||
request_id: request_id.to_string(),
|
||||
model: "test-model".to_string(),
|
||||
prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
|
||||
stopping: Some(pb::StoppingCriteria {
|
||||
max_new_tokens: 10,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect("start generation")
|
||||
.into_inner();
|
||||
|
||||
loop {
|
||||
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||
.await
|
||||
.expect("timed out waiting for active generation output")
|
||||
.expect("read active generation output")
|
||||
.expect("generation ended before producing output");
|
||||
if let Some(output) = response.outputs {
|
||||
assert!(
|
||||
output.finish_info.is_none(),
|
||||
"generation finished before abort behavior was exercised"
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest::default())
|
||||
.await
|
||||
.expect("empty abort should be a no-op");
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(100), stream.message())
|
||||
.await
|
||||
.is_err(),
|
||||
"empty abort unexpectedly ended the active generation"
|
||||
);
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest {
|
||||
request_ids: vec![
|
||||
request_id.to_string(),
|
||||
request_id.to_string(),
|
||||
"unknown".to_string(),
|
||||
],
|
||||
})
|
||||
.await
|
||||
.expect("abort active generation");
|
||||
|
||||
let finish_reason = loop {
|
||||
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
|
||||
.await
|
||||
.expect("timed out waiting for aborted generation")
|
||||
.expect("read aborted generation")
|
||||
.expect("generation ended without an aborted response");
|
||||
if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
|
||||
break finish_info.finish_reason;
|
||||
}
|
||||
};
|
||||
assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
|
||||
|
||||
control_client
|
||||
.abort(pb::AbortRequest {
|
||||
request_ids: vec![request_id.to_string()],
|
||||
})
|
||||
.await
|
||||
.expect("repeated abort should be idempotent");
|
||||
|
||||
engine_task.await.expect("mock engine task");
|
||||
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, control_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,
|
||||
control_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", "vllm.Control", ""] {
|
||||
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, control_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,
|
||||
control_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");
|
||||
}
|
||||
|
||||
+32
-14
@@ -39,6 +39,7 @@ use tokio::net::TcpListener;
|
||||
use tokio::time::{Instant, sleep_until};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tonic::transport::Server as TonicServer;
|
||||
use tonic_health::server::health_reporter;
|
||||
use tower::ServiceExt as _;
|
||||
use tracing::{info, trace, warn};
|
||||
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
|
||||
@@ -203,14 +204,23 @@ where
|
||||
.map(tls::build_grpc_server_config)
|
||||
.transpose()
|
||||
.context("invalid gRPC TLS configuration")?;
|
||||
let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
|
||||
let (health_reporter, health_service) = health_reporter();
|
||||
let engine_health = state.engine_core_client().subscribe_health();
|
||||
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
|
||||
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
|
||||
let control_service =
|
||||
grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
|
||||
let generate_service =
|
||||
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
|
||||
let svc = TonicServer::builder()
|
||||
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
|
||||
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
|
||||
.layer(middleware::request_runtime_layer(state.clone()))
|
||||
.add_service(svc);
|
||||
.add_service(health_service)
|
||||
.add_service(control_service)
|
||||
.add_service(generate_service);
|
||||
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
|
||||
Some((grpc_listener, svc, grpc_tls))
|
||||
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -294,7 +304,8 @@ where
|
||||
let server_shutdown = server_shutdown.clone();
|
||||
let force_shutdown = force_shutdown.clone();
|
||||
async move {
|
||||
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
|
||||
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
|
||||
else {
|
||||
// No gRPC configured: just wait for shutdown so we do not race the
|
||||
// join! by resolving early and tripping the cancellation token.
|
||||
shutdown.cancelled().await;
|
||||
@@ -304,19 +315,26 @@ where
|
||||
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
|
||||
None => MaybeTlsListener::plain(grpc_listener),
|
||||
};
|
||||
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
|
||||
let server =
|
||||
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
|
||||
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
|
||||
|
||||
let result = tokio::select! {
|
||||
result = server => {
|
||||
result.context("gRPC server failed")
|
||||
}
|
||||
_ = force_shutdown.cancelled() => {
|
||||
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
||||
Ok(())
|
||||
}
|
||||
let server = async move {
|
||||
let result = tokio::select! {
|
||||
result = server => {
|
||||
result.context("gRPC server failed")
|
||||
}
|
||||
_ = force_shutdown.cancelled() => {
|
||||
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
|
||||
server_shutdown.cancel();
|
||||
result
|
||||
};
|
||||
|
||||
server_shutdown.cancel();
|
||||
let (result, ()) = tokio::join!(server, health_monitor);
|
||||
result
|
||||
}
|
||||
};
|
||||
|
||||
+108
-72
@@ -38,32 +38,22 @@ trait_set! {
|
||||
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
|
||||
}
|
||||
|
||||
/// Raw text facade above [`Llm`].
|
||||
///
|
||||
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
|
||||
/// in, decoded text deltas and terminal metadata flow out.
|
||||
pub struct TextLlm {
|
||||
/// Generate-only client owned by this text facade.
|
||||
llm: Llm,
|
||||
/// Text request preparation shared by inference and render-only frontends.
|
||||
pub struct TextRequestProcessor {
|
||||
/// Tokenizer/model metadata backend responsible for prompt encode/decode
|
||||
/// and sampling hints.
|
||||
backend: DynTextBackend,
|
||||
/// Runtime context window size reported by the engine startup handshake.
|
||||
/// Render-only frontends supply the downstream engine's effective value.
|
||||
max_model_len: u32,
|
||||
/// Maximum number of top log probabilities accepted by this text facade.
|
||||
max_logprobs: i32,
|
||||
}
|
||||
|
||||
impl TextLlm {
|
||||
/// Create a new text-generation facade from a shared LLM client plus a text
|
||||
/// backend.
|
||||
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
|
||||
// The engine-reported value reflects the post-profiling, auto-fitted
|
||||
// KV cache limit used at runtime.
|
||||
let max_model_len = llm.engine_core_client().max_model_len();
|
||||
|
||||
impl TextRequestProcessor {
|
||||
/// Create a processor with the effective model context length.
|
||||
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
|
||||
Self {
|
||||
llm,
|
||||
backend,
|
||||
max_model_len,
|
||||
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
|
||||
@@ -78,61 +68,18 @@ impl TextLlm {
|
||||
self
|
||||
}
|
||||
|
||||
/// Return the backend model ID.
|
||||
pub fn model_id(&self) -> &str {
|
||||
self.backend.model_id()
|
||||
}
|
||||
|
||||
/// Expose the underlying engine-core client for low-level utility/admin
|
||||
/// calls.
|
||||
pub fn engine_core_client(&self) -> &EngineCoreClient {
|
||||
self.llm.engine_core_client()
|
||||
}
|
||||
|
||||
/// Return the tokenizer used by this text backend.
|
||||
/// Return the tokenizer used by this processor.
|
||||
pub fn tokenizer(&self) -> DynTokenizer {
|
||||
self.backend.tokenizer()
|
||||
}
|
||||
|
||||
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
|
||||
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
|
||||
pub fn tokenizer_vocab_size(&self) -> usize {
|
||||
self.backend.tokenizer_vocab_size()
|
||||
/// Return the effective model context length.
|
||||
pub fn max_model_len(&self) -> u32 {
|
||||
self.max_model_len
|
||||
}
|
||||
|
||||
/// Model vocabulary size from the model config, used to bound generated
|
||||
/// token IDs and logits-domain sampling controls.
|
||||
pub fn model_vocab_size(&self) -> usize {
|
||||
self.backend.model_vocab_size()
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and return the raw
|
||||
/// token stream.
|
||||
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
|
||||
let (_, raw_stream) = self.generate_inner(request).await?;
|
||||
Ok(raw_stream)
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and stream
|
||||
/// incrementally decoded text.
|
||||
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
|
||||
let (text_request, raw_stream) = self.generate_inner(request).await?;
|
||||
let tokenizer = self.backend.tokenizer();
|
||||
let decoded_stream = output::decoded_text_event_stream(
|
||||
text_request.request_id,
|
||||
tokenizer,
|
||||
raw_stream,
|
||||
text_request.decode_options,
|
||||
text_request.intermediate,
|
||||
);
|
||||
|
||||
Ok(decoded_stream)
|
||||
}
|
||||
|
||||
async fn generate_inner(
|
||||
&self,
|
||||
mut request: TextRequest,
|
||||
) -> Result<(TextRequest, GenerateOutputStream)> {
|
||||
/// Tokenize and lower one request without submitting it to an engine.
|
||||
pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
|
||||
request.validate()?;
|
||||
|
||||
if request.arrival_time.is_none() {
|
||||
@@ -146,7 +93,6 @@ impl TextLlm {
|
||||
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
|
||||
Prompt::TokenIds(token_ids) => token_ids,
|
||||
};
|
||||
|
||||
let sampling_hints = self.backend.sampling_hints()?;
|
||||
let sampling_limits = SamplingLimits {
|
||||
max_model_len: self.max_model_len,
|
||||
@@ -155,16 +101,106 @@ impl TextLlm {
|
||||
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
|
||||
};
|
||||
|
||||
let PreparedTextRequest {
|
||||
text_request,
|
||||
generate_request,
|
||||
} = lower_text_request(
|
||||
lower_text_request(
|
||||
request,
|
||||
prompt_token_ids,
|
||||
sampling_hints,
|
||||
sampling_limits,
|
||||
&*tokenizer,
|
||||
)?;
|
||||
tokenizer.as_ref(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw text facade above [`Llm`].
|
||||
///
|
||||
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
|
||||
/// in, decoded text deltas and terminal metadata flow out.
|
||||
pub struct TextLlm {
|
||||
/// Generate-only client owned by this text facade.
|
||||
llm: Llm,
|
||||
/// Shared engine-free request preparation.
|
||||
processor: TextRequestProcessor,
|
||||
}
|
||||
|
||||
impl TextLlm {
|
||||
/// Create a new text-generation facade from a shared LLM client plus a text
|
||||
/// backend.
|
||||
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
|
||||
// The engine-reported value reflects the post-profiling, auto-fitted
|
||||
// KV cache limit used at runtime.
|
||||
let max_model_len = llm.engine_core_client().max_model_len();
|
||||
|
||||
Self {
|
||||
llm,
|
||||
processor: TextRequestProcessor::new(backend, max_model_len),
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the maximum accepted logprobs count.
|
||||
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
|
||||
self.processor = self.processor.with_max_logprobs(max_logprobs);
|
||||
self
|
||||
}
|
||||
|
||||
/// Return the backend model ID.
|
||||
pub fn model_id(&self) -> &str {
|
||||
self.processor.backend.model_id()
|
||||
}
|
||||
|
||||
/// Expose the underlying engine-core client for low-level utility/admin
|
||||
/// calls.
|
||||
pub fn engine_core_client(&self) -> &EngineCoreClient {
|
||||
self.llm.engine_core_client()
|
||||
}
|
||||
|
||||
/// Return the tokenizer used by this text backend.
|
||||
pub fn tokenizer(&self) -> DynTokenizer {
|
||||
self.processor.tokenizer()
|
||||
}
|
||||
|
||||
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
|
||||
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
|
||||
pub fn tokenizer_vocab_size(&self) -> usize {
|
||||
self.processor.backend.tokenizer_vocab_size()
|
||||
}
|
||||
|
||||
/// Model vocabulary size from the model config, used to bound generated
|
||||
/// token IDs and logits-domain sampling controls.
|
||||
pub fn model_vocab_size(&self) -> usize {
|
||||
self.processor.backend.model_vocab_size()
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and return the raw
|
||||
/// token stream.
|
||||
pub async fn generate_raw(&self, request: TextRequest) -> Result<GenerateOutputStream> {
|
||||
let (_, raw_stream) = self.generate_inner(request).await?;
|
||||
Ok(raw_stream)
|
||||
}
|
||||
|
||||
/// Tokenize if needed, lower to a generate request, and stream
|
||||
/// incrementally decoded text.
|
||||
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
|
||||
let (text_request, raw_stream) = self.generate_inner(request).await?;
|
||||
let tokenizer = self.processor.tokenizer();
|
||||
let decoded_stream = output::decoded_text_event_stream(
|
||||
text_request.request_id,
|
||||
tokenizer,
|
||||
raw_stream,
|
||||
text_request.decode_options,
|
||||
text_request.intermediate,
|
||||
);
|
||||
|
||||
Ok(decoded_stream)
|
||||
}
|
||||
|
||||
async fn generate_inner(
|
||||
&self,
|
||||
request: TextRequest,
|
||||
) -> Result<(TextRequest, GenerateOutputStream)> {
|
||||
let PreparedTextRequest {
|
||||
text_request,
|
||||
generate_request,
|
||||
} = self.processor.prepare(request)?;
|
||||
|
||||
let raw_stream = self.llm.generate(generate_request).await?;
|
||||
Ok((text_request, raw_stream))
|
||||
|
||||
@@ -214,3 +214,65 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
|
||||
base_hash = CacheConfig().compute_hash()
|
||||
assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
|
||||
assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
|
||||
|
||||
|
||||
def test_envs_compile_factors_relocation_invariant(tmp_path):
|
||||
"""Relocating HOME or the XDG roots must not change the compile-cache
|
||||
env hash.
|
||||
|
||||
Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME,
|
||||
VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about
|
||||
compiled artifacts, only about where directories live. When they leak
|
||||
into compile_factors(), a cache produced under one HOME/XDG layout
|
||||
silently misses under another - which defeats copying or pre-baking a
|
||||
compile cache into a container image.
|
||||
"""
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
code = """
|
||||
import sys
|
||||
import logging
|
||||
logging.disable(logging.CRITICAL)
|
||||
from vllm import envs
|
||||
from vllm.config.utils import hash_factors
|
||||
print(hash_factors(envs.compile_factors()))
|
||||
"""
|
||||
|
||||
def hash_with(extra_env):
|
||||
env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"}
|
||||
# Drop explicit overrides so the derived defaults are what is
|
||||
# exercised, then apply the relocation under test.
|
||||
for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"):
|
||||
env.pop(key, None)
|
||||
env.update(extra_env)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
env=env,
|
||||
)
|
||||
return result.stdout.strip()
|
||||
|
||||
xdg_cache = tmp_path / "relocated-xdg-cache"
|
||||
xdg_config = tmp_path / "relocated-xdg-config"
|
||||
new_home = tmp_path / "relocated-home"
|
||||
for d in (xdg_cache, xdg_config, new_home):
|
||||
d.mkdir()
|
||||
|
||||
base = hash_with({})
|
||||
relocated_xdg = hash_with(
|
||||
{"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)}
|
||||
)
|
||||
relocated_home = hash_with({"HOME": str(new_home)})
|
||||
|
||||
assert relocated_xdg == base, (
|
||||
"XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache "
|
||||
"env hash - a location-only derived var is leaking into the key"
|
||||
)
|
||||
assert relocated_home == base, (
|
||||
"HOME relocation changed the compile-cache env hash - a "
|
||||
"location-only derived var is leaking into the key"
|
||||
)
|
||||
|
||||
@@ -43,6 +43,15 @@ def test_language_model_only_affects_model_hash():
|
||||
assert base_hash != lm_only_hash
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"])
|
||||
def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str):
|
||||
config = MultiModalConfig(
|
||||
media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}}
|
||||
)
|
||||
|
||||
assert config.use_gpu_video_backend()
|
||||
|
||||
|
||||
def test_mm_encoder_fp8_scale_path_requires_fp8():
|
||||
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
|
||||
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
|
||||
|
||||
@@ -349,6 +349,20 @@ _T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature, dict)
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
def _fix_v4_tied_weights_keys(model_cls: type) -> None:
|
||||
"""Convert a v4 list-format _tied_weights_keys to the transformers v5 dict form."""
|
||||
tied = getattr(model_cls, "_tied_weights_keys", None)
|
||||
if not isinstance(tied, list) or not tied:
|
||||
return
|
||||
result = {
|
||||
k: "model.embed_tokens.weight"
|
||||
for k in tied
|
||||
if "lm_head" in k and k.endswith(".weight")
|
||||
}
|
||||
if result:
|
||||
setattr(model_cls, "_tied_weights_keys", result)
|
||||
|
||||
|
||||
class HfRunner:
|
||||
def get_default_device(self):
|
||||
from vllm.platforms import current_platform
|
||||
@@ -474,6 +488,22 @@ class HfRunner:
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
else:
|
||||
if trust_remote_code and hasattr(self.config, "auto_map"):
|
||||
cls_ref = self.config.auto_map.get(auto_cls.__name__)
|
||||
if cls_ref is not None:
|
||||
from vllm.transformers_utils.dynamic_module import (
|
||||
try_get_class_from_dynamic_module,
|
||||
)
|
||||
|
||||
model_cls = try_get_class_from_dynamic_module(
|
||||
cls_ref,
|
||||
model_name,
|
||||
trust_remote_code=trust_remote_code,
|
||||
warn_on_fail=False,
|
||||
)
|
||||
if model_cls is not None:
|
||||
_fix_v4_tied_weights_keys(model_cls)
|
||||
|
||||
model = cast(
|
||||
nn.Module,
|
||||
auto_cls.from_pretrained(
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
# 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 == []
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user