Compare commits

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

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

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

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

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

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

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

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

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

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

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

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

Signed-off-by: Shengqi Chen <harry-chen@outlook.com>
2026-06-30 19:27:42 +08:00
315 changed files with 3695 additions and 8094 deletions
+28 -28
View File
@@ -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_nixl 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_rixl 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.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.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.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.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_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 NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
for var in RIXL_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 nixl_branch=""
local rixl_branch=""
local ucx_branch=""
local rocshmem_branch=""
local deepep_branch=""
local nixl_material=""
local rixl_material=""
local rocshmem_material=""
local deepep_material=""
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_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 "${nixl_branch}" && -n "${ucx_branch}" ]]; then
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
NIXL_CACHE_KEY=$(
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
RIXL_CACHE_KEY=$(
compose_dependency_cache_key \
"${nixl_branch}-ucx-${ucx_branch}" \
"${nixl_material}"
"${rixl_branch}-ucx-${ucx_branch}" \
"${rixl_material}"
)
export NIXL_CACHE_KEY
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
export RIXL_CACHE_KEY
echo "RIXL dependency cache key: ${RIXL_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
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:-}"
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:-}"
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 nixl_ref=""
local rixl_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 nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
for target in rixl-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 "${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}"
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}"
else
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
add_dependency_cache_target "nixl-rocm-ci"
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
add_dependency_cache_target "rixl-rocm-ci"
fi
fi
@@ -1,11 +1,10 @@
#!/bin/bash
set -euox pipefail
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_KVCACHE_SPACE=1
export VLLM_CPU_CI_ENV=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.
# Reduce sub-processes for acceleration
export TORCH_COMPILE_DISABLE=1
export VLLM_ENABLE_V1_MULTIPROCESSING=0
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
@@ -50,15 +49,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 --enforce-eager > test_0.log 2>&1 &
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > 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 --enforce-eager > test_1.log 2>&1 &
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > 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 --enforce-eager > test_2.log 2>&1 &
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
PID_TEST_2=$!
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 45
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
-3
View File
@@ -26,10 +26,7 @@ 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
+1 -1
View File
@@ -51,7 +51,7 @@ steps:
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 53
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- vllm/v1/
+4 -4
View File
@@ -39,7 +39,7 @@ steps:
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
timeout_in_minutes: 75
timeout_in_minutes: 50
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: 68
timeout_in_minutes: 45
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: 83
timeout_in_minutes: 45
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: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
+1 -4
View File
@@ -31,7 +31,7 @@ steps:
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 83
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/config/
@@ -90,7 +90,6 @@ 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
@@ -104,7 +103,6 @@ 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
@@ -145,7 +143,6 @@ 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
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 60
timeout_in_minutes: 45
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
-13
View File
@@ -46,19 +46,6 @@ 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:
+1 -1
View File
@@ -137,7 +137,7 @@ steps:
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
optional: true
source_file_dependencies:
+4 -4
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 68
timeout_in_minutes: 45
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: 75
timeout_in_minutes: 50
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: 75
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 98
timeout_in_minutes: 65
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 150
timeout_in_minutes: 90
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
-16
View File
@@ -170,19 +170,3 @@ 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
-1
View File
@@ -47,7 +47,6 @@
# Rust Frontend
/rust/ @BugenZhao @njhill
/rust/src/bench @esmeetu
/build_rust.sh @BugenZhao @njhill
/rust-toolchain.toml @BugenZhao @njhill
/.buildkite/test_areas/rust* @BugenZhao @njhill
-12
View File
@@ -181,18 +181,6 @@ 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:
+1 -42
View File
@@ -130,47 +130,6 @@ 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: [...],
@@ -532,4 +491,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(', ')}`);
+88 -44
View File
@@ -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.13.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.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,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
cuda_archs_loose_intersection(CUDA_ARCHS
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
set(VLLM_COMPILED_CUDA_ARCHS)
foreach(_ARCH ${CUDA_ARCHS})
set(_COMPILED_ARCH "${_ARCH}")
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
if(_ARCH MATCHES "^(10|11|12)\\.0$")
set(_COMPILED_ARCH "${_ARCH}f")
endif()
elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
set(_COMPILED_ARCH "${_ARCH}a")
endif()
endif()
list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
endforeach()
list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
else()
#
# For other GPU targets override the GPU architectures detected by cmake/torch
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
#
# Shared entry sources are part of the base extension source list, but some
# optional kernel families below append source-local feature macros to them.
set(VLLM_STABLE_TORCH_BINDINGS_SRC
"csrc/libtorch_stable/torch_bindings.cpp")
set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
set(SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
set(NVFP4_QUANT_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
set(NVFP4_SCALED_MM_ENTRY_SRC
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
"csrc/libtorch_stable/activation_kernels.cu"
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/sampler.cu"
"csrc/libtorch_stable/topk.cu"
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"${CACHE_KERNELS_SRC}"
"csrc/libtorch_stable/cache_kernels_fused.cu"
"csrc/libtorch_stable/custom_all_reduce.cu"
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
endif()
if(COOPERATIVE_TOPK_ARCHS)
list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
endif()
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
"${SCALED_MM_ENTRY_SRC}"
"${NVFP4_QUANT_ENTRY_SRC}"
"${NVFP4_SCALED_MM_ENTRY_SRC}"
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")
#
# Machete kernels
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
CUDA_ARCHS "${CUDA_ARCHS}")
if(COOPERATIVE_TOPK_ARCHS)
set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cooperative_topk.cu")
"${COOPERATIVE_TOPK_SRC}")
set_gencode_flags_for_srcs(
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
SRCS "${COOPERATIVE_TOPK_SRC}"
CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Only build Marlin kernels if we are building for at least some compatible archs.
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM120_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
# Let scaled_mm_c2x know it doesn't need to build these arches
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${SCALED_MM_C2X_SRCS}"
CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
DEFINITIONS ENABLE_SCALED_MM_C2X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
else()
if (SCALED_MM_3X_ARCHS)
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM90_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${CUTLASS_MOE_SM100_SRCS}"
CUDA_ARCHS "${SCALED_MM_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
# The implementation is named sm100 historically, but it is built for the
# SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
# SM110. Keep the compile-time macro aligned with runtime dispatch.
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
else()
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
#
# SM12x FP4 kernels. These share some generic NVFP4 quantization entry
# sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
# per-source flags, so shared files accumulate both SM12x and SM10x/11x
# gencodes when both families are requested.
# Shared FP4 implementation sources live next to the FP4 arch logic because
# both SM12x and SM10x/11x append family-specific gencodes and feature macros
# to them.
set(FP4_SHARED_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
else()
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
set(FP4_SM120_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
)
set_gencode_flags_for_srcs(
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM120=1)
set_compile_definitions_for_srcs(
SRCS "${SCALED_MM_ENTRY_SRC}"
DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
else()
message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
set(FP4_SM100_SRCS
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
${FP4_SHARED_SRCS}
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
message(STATUS
"Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set_gencode_flags_for_srcs(
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
set_compile_definitions_for_srcs(
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
DEFINITIONS ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
else()
message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${CUTLASS_MLA_SRCS}"
CUDA_ARCHS "${MLA_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
# Add MLA-specific include directories only to MLA source files
set_source_files_properties(${CUTLASS_MLA_SRCS}
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
# Needed to use cuda/hip APIs from C-shim
if(VLLM_GPU_LANG STREQUAL "CUDA")
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
if(COOPERATIVE_TOPK_ARCHS)
target_compile_definitions(_C_stable_libtorch PRIVATE
VLLM_ENABLE_COOPERATIVE_TOPK=1)
endif()
# Needed by CUTLASS kernels
target_compile_definitions(_C_stable_libtorch PRIVATE
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+1 -1
View File
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
- Tool calling and reasoning parsers
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
- Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
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 ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+22
View File
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
endif()
endmacro()
#
# For a list of source files append preprocessor definitions to file-specific
# compile options. Use this for optional kernel feature macros so toggling one
# kernel family does not perturb the compile command for every source in the
# extension target.
#
macro(set_compile_definitions_for_srcs)
set(options)
set(oneValueArgs)
set(multiValueArgs SRCS DEFINITIONS)
cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
"${multiValueArgs}" ${ARGN})
foreach(_DEF ${arg_DEFINITIONS})
set_property(
SOURCE ${arg_SRCS}
APPEND PROPERTY
COMPILE_DEFINITIONS "${_DEF}"
)
endforeach()
endmacro()
#
# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
# `<major>.<minor>[letter]` compute the "loose intersection" with the
+1 -3
View File
@@ -102,9 +102,7 @@ class TileGemm82 {
kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto fp32_b_regs = load_b_pair_vec(curr_b);
auto fp32_b_0_reg = fp32_b_regs.first;
auto fp32_b_1_reg = fp32_b_regs.second;
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-3
View File
@@ -1025,9 +1025,6 @@ __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 =
+3 -6
View File
@@ -9,16 +9,14 @@ 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> is_padding);
std::optional<torch::stable::Tensor> bias);
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,
std::optional<torch::stable::Tensor> is_padding);
double routed_scaling_factor);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -27,8 +25,7 @@ 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>& is_padding);
const std::optional<torch::stable::Tensor>& tid2eid);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
@@ -174,8 +174,7 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int end_expert,
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
const bool* is_padding)
const double routed_scaling_factor)
{
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -229,14 +228,12 @@ __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] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : num_experts);
assert(is_pad_row || indices[idx] >= 0);
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -280,7 +277,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 bool* is_padding)
const float* bias, const double routed_scaling_factor)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -548,14 +545,12 @@ __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] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -610,7 +605,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 bool* is_padding)
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
{
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>;
@@ -621,7 +616,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, is_padding);
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
}
#ifndef USE_ROCM
@@ -632,7 +627,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, is_padding);
bias, routed_scaling_factor, stream);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -640,13 +635,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, is_padding); \
bias, routed_scaling_factor, stream); \
} 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, is_padding); \
bias, routed_scaling_factor, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -666,8 +661,7 @@ void topkGatingKernelLauncher(
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
cudaStream_t stream,
const bool* is_padding) {
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
@@ -742,7 +736,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, is_padding);
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
}
}
}
@@ -761,8 +755,7 @@ 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,
std::optional<torch::stable::Tensor> is_padding)
cudaStream_t stream)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
@@ -776,18 +769,6 @@ 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()),
@@ -796,7 +777,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, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
} 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()),
@@ -805,7 +786,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, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
@@ -815,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, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
}
}
@@ -825,8 +806,7 @@ 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> is_padding)
std::optional<torch::stable::Tensor> bias)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -845,15 +825,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, is_padding);
bias, 1.0, stream);
} 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, is_padding);
bias, 1.0, stream);
} 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, is_padding);
bias, 1.0, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -866,8 +846,7 @@ void topk_sigmoid(
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding)
double routed_scaling_factor)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -886,15 +865,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, is_padding);
bias, routed_scaling_factor, stream);
} 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, is_padding);
bias, routed_scaling_factor, stream);
} 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, is_padding);
bias, routed_scaling_factor, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -80,27 +80,22 @@ __launch_bounds__(128) __global__
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid,
const bool* is_padding) {
const HashIndType* tid2eid) {
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 && !is_pad_row) {
if (lane < 6) {
// 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
@@ -116,8 +111,7 @@ __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] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-1);
indices[offset] = static_cast<OutIndType>(expert);
}
}
@@ -126,8 +120,7 @@ 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 bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -141,7 +134,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, is_padding);
num_experts, scale, input_ids, tid2eid);
}
#endif
@@ -173,8 +166,7 @@ __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 bool* is_padding) {
const HashIndType* input_ids, const HashIndType* tid2eid) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -239,7 +231,6 @@ __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.
@@ -258,12 +249,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
cudaGridDependencySynchronize();
#endif
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>) {
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
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 =
@@ -327,22 +315,12 @@ __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;
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;
}
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);
}
float selected_sum = 0.f;
#pragma unroll
@@ -357,8 +335,7 @@ __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] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-1);
indices[idx] = static_cast<IndType>(expert);
selected_sum += row_chunk[ii];
break;
}
@@ -402,31 +379,23 @@ __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);
// 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;
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];
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
@@ -481,19 +450,18 @@ __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 && !is_pad_row;
const bool should_process_row = row_is_active && node_uses_expert;
// 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 && should_process_row) {
if (correction_bias != nullptr) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] =
!is_pad_row ? expert - start_expert : 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;
@@ -576,7 +544,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 bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -605,12 +573,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, is_padding);
input_ids, tid2eid);
#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, is_padding);
input_ids, tid2eid);
#endif
})
}
@@ -624,7 +592,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, is_padding);
stream);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -633,14 +601,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, is_padding); \
tid2eid, stream); \
} 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, is_padding); \
tid2eid, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -654,14 +622,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, const bool* is_padding) {
cudaStream_t stream) {
#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, is_padding);
routed_scaling_factor, input_ids, tid2eid, stream);
return;
}
}
@@ -760,8 +728,7 @@ 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>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -770,22 +737,6 @@ 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");
@@ -800,7 +751,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, is_padding_ptr);
tid2eid.value().const_data_ptr<int64_t>(), stream);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
@@ -810,7 +761,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, is_padding_ptr);
tid2eid.value().const_data_ptr<int>(), stream);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -818,7 +769,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, is_padding_ptr);
static_cast<const OutIndType*>(nullptr), stream);
}
};
@@ -842,8 +793,7 @@ 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>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
@@ -856,22 +806,21 @@ 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,
is_padding);
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
} 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, is_padding);
input_ids, tid2eid, stream);
} 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, is_padding);
input_ids, tid2eid, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
+3 -3
View File
@@ -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, Tensor? is_padding) -> ()");
"bias) -> ()");
// 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? is_padding) -> ()");
"Tensor? bias, float routed_scaling_factor) -> ()");
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, Tensor? is_padding) -> ()");
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
// Calculate the result of moe by summing up the partial results
// from all selected experts. topk_ids/expert_map are optional and, when
+2
View File
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
torch::stable::Tensor const& perm);
#ifndef USE_ROCM
std::string get_compiled_cuda_archs();
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
#endif
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& a_tensors,
torch::stable::Tensor const& b_tensors,
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
std::optional<torch::stable::Tensor> const& bias);
#endif
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
(defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
void get_cutlass_moe_mm_data_caller(
const torch::stable::Tensor& topk_ids,
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
// CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
// actual cutlass_moe_mm dispatch compiled into this file.
// or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
// SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
// implementation. Only report archs that have an actual cutlass_moe_mm
// dispatch compiled into this file.
#if defined CUDA_VERSION
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
return CUDA_VERSION >= 12080;
}
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
torch::stable::Tensor const& c_strides, bool per_act_token,
bool per_out_ch) {
int32_t version_num = get_sm_version_num();
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
if (version_num >= 100 && version_num < 110) {
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
// Keep runtime dispatch aligned with the CMake arch list and support query:
// CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
if (version_num >= 100 && version_num < 120) {
cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
expert_offsets, problem_sizes, a_strides, b_strides,
c_strides, per_act_token, per_out_ch);
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
problem_sizes2, input_permutation,
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
const bool swap_ab) {
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
// This function currently gets compiled only if we have a valid cutlass moe
// mm to run it for.
int32_t version_num = get_sm_version_num();
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
problem_sizes2, expert_num_tokens,
+6
View File
@@ -4,6 +4,10 @@
#include <torch/csrc/stable/library.h>
#ifndef USE_ROCM
std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
#endif
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
// Note: We register under namespace "_C" so ops are accessible as
// torch.ops._C.<op_name> for compatibility with existing code.
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.def("get_compiled_cuda_archs() -> str");
#ifndef USE_ROCM
@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
#ifndef USE_ROCM
ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
ops.impl("cutlass_scaled_mm_supports_fp8",
TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
ops.impl("cutlass_group_gemm_supported",
+2 -2
View File
@@ -22,7 +22,7 @@
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
# =============================================================================
ARG CUDA_VERSION=13.0.3
ARG CUDA_VERSION=13.0.2
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.15.post1
ARG FLASHINFER_VERSION=0.6.14
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 '.')
+31 -53
View File
@@ -339,17 +339,18 @@ 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
# 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"
# 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"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV NIXL_HOME=/usr/local/nixl
ENV RIXL_HOME=/usr/local/rixl
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
# NIXL build system dependencies and RDMA support
# RIXL build system dependences and RDMA support
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
libgrpc-dev \
libgrpc++-dev \
@@ -367,8 +368,7 @@ 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 meson-python pybind11 pyyaml types-PyYAML \
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
uv pip install --system meson auditwheel patchelf tomlkit
RUN --mount=type=cache,target=/root/.cache/ccache \
cd /usr/local/src && \
@@ -396,50 +396,30 @@ 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 ${NIXL_REPO} /opt/nixl && \
cd /opt/nixl && \
git checkout ${NIXL_BRANCH} && \
git clone ${RIXL_REPO} /opt/rixl && \
cd /opt/rixl && \
git checkout ${RIXL_BRANCH} && \
CC="ccache gcc" CXX="ccache g++" \
meson setup build --prefix=${NIXL_HOME} \
meson setup build --prefix=${RIXL_HOME} \
-Ducx_path=${UCX_HOME} \
-Dwheel_variant=rocm \
-Dbuild_tests=false \
-Dbuild_examples=false && \
-Drocm_path=${ROCM_PATH} && \
cd build && \
ninja -j$(nproc) && \
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
ninja install
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
# so configure the ROCm wheel variant directly through Meson.
# Generate RIXL wheel
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
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 \
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} \
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
/tmp/nixl_wheels/repaired/*.whl && \
cp /tmp/nixl_wheels/repaired/*.whl /app/install
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
# invalidate the slow ROCShmem build.
@@ -680,10 +660,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install NIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
# Install RIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
# Copy ROCShmem runtime libraries.
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -744,7 +724,6 @@ 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
@@ -817,9 +796,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
&& pip uninstall -y vllm \
&& uv pip install --system *.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
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
ARG COMMON_WORKDIR
ARG BASE_IMAGE
@@ -834,7 +813,6 @@ 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
+1 -1
View File
@@ -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.post5"
ARG AITER_BRANCH="v0.1.16.post3"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
+1 -25
View File
@@ -86,29 +86,6 @@ 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
@@ -119,13 +96,11 @@ 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 \
@@ -183,6 +158,7 @@ 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
View File
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
}
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# rebuilding RIXL/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 "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
default = ""
}
variable "NIXL_CACHE_KEY" {
variable "RIXL_CACHE_KEY" {
default = ""
}
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
])
}
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
# Cache functions for upstream dependency stages (RIXL/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([
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}" : ""),
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}" : ""),
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_nixl" {
function "get_cache_to_rocm_rixl" {
params = []
result = compact([
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" : ""),
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" : ""),
])
}
@@ -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 "nixl-rocm-ci" {
target "rixl-rocm-ci" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "build_nixl"
target = "build_rixl"
cache-from = get_cache_from_rocm_deps()
cache-to = get_cache_to_rocm_nixl()
cache-to = get_cache_to_rocm_rixl()
output = ["type=cacheonly"]
}
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
output = ["type=cacheonly"]
}
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
# Builds only the ci_base stage (RIXL, 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 NIXL/ROCShmem/DeepEP stages
# Import upstream dependency caches so RIXL/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 = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
}
+2 -2
View File
@@ -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 "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -106,7 +106,7 @@ target "test-rocm" {
output = ["type=docker"]
}
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
# CI base image target - builds only the ci_base stage (RIXL, 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" {
+4 -4
View File
@@ -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.3"
"default": "13.0.2"
},
"PYTHON_VERSION": {
"default": "3.12"
@@ -11,10 +11,10 @@
"default": "22.04"
},
"BUILD_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
},
"FINAL_BASE_IMAGE": {
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
},
"BUILD_OS": {
"default": "ubuntu"
@@ -68,7 +68,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.15.post1"
"default": "0.6.14"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+5 -1
View File
@@ -13,7 +13,11 @@ 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, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
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 non-cuda platform, please install nixl with ucx build from source, instructed as below.
+1 -1
View File
@@ -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 only)
- compressed-tensor INT8 W8A8 (x86, s390x)
### 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 from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
--8<-- [end:requirements]
--8<-- [start:set-up-using-python]
@@ -28,24 +28,13 @@ Install the following packages from the package manager before building the vLLM
```bash
dnf install -y \
which procps findutils tar vim git patch xz ninja-build \
gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
which procps findutils tar vim git 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
```
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.
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
```bash
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
@@ -55,79 +44,26 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
Execute the following commands to build and install vLLM from source.
!!! tip
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
```
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
```bash
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
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
```
??? console "pip"
```bash
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
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
```
--8<-- [end:build-wheel-from-source]
@@ -144,20 +80,19 @@ docker build -f docker/Dockerfile.s390x \
# Launch OpenAI server
docker run --rm \
--security-opt seccomp=unconfined \
--cap-add SYS_NICE \
--privileged true \
--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 bfloat16 \
--dtype float \
other vLLM OpenAI server arguments
```
!!! tip
Alternatively, `--privileged=true` also works but is broader and not generally recommended.
An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
--8<-- [end:build-image-from-source]
--8<-- [start:extra-information]
@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
```
!!! warning "CUDA architecture coverage"
Pre-built CUDA wheels are compiled for the CUDA architectures selected by
vLLM's release and build pipelines. This list is intentionally smaller than
every architecture CMake can build, because each additional architecture
increases wheel size.
In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
architecture-specific Blackwell/Thor targets, such as `sm_103` or
`sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
cover the corresponding major-version GPU family.
If vLLM logs a warning that your visible CUDA device is not covered by the
wheel's compiled CUDA architectures, or if you see a CUDA error such as
`no kernel image is available for execution on the device`, install a CUDA
13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
that includes your GPU.
#### Install the latest code
LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
@@ -58,40 +58,7 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
--8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images]
vLLM offers official Docker images for deployment.
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
```bash
docker run --rm \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
vllm/vllm-openai-xpu:<tag> \
--model Qwen/Qwen3-0.6B
```
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
???+ console "Commands"
```bash
docker run --rm -it \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
--entrypoint /bin/bash \
vllm/vllm-openai-xpu:<tag>
```
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
--8<-- [end:pre-built-images]
--8<-- [start:build-image-from-source]
-9
View File
@@ -65,15 +65,6 @@ This guide will help you quickly get started with vLLM to perform:
!!! tip
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
=== "Intel GPU"
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
!!! tip
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
=== "Google TPU"
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
-3
View File
@@ -1,3 +0,0 @@
// 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)}();
+27
View File
@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc
If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.
## CUDA architecture not covered by the wheel
If vLLM logs a warning that the current wheel was built for a set of CUDA
architectures but one of your visible CUDA devices is not covered, the installed
wheel may not contain native CUDA kernels for that GPU. The same issue may also
surface later as a CUDA runtime error such as `no kernel image is available for
execution on the device`.
The architecture list in a pre-built wheel is determined by the vLLM release and
build pipelines, then filtered by CMake before individual kernels choose their
own per-kernel architectures. It is not the same as the full set of
architectures that vLLM can build from source.
CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
them from source. CUDA 13 wheels can use family-specific targets such as
`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
major-version GPU family.
If you see this warning or error, install a CUDA 13 wheel when possible.
Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
your GPU's compute capability. See the CUDA installation guide's
[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
sections for details.
## Python multiprocessing
### `RuntimeError` Exception
-1
View File
@@ -160,4 +160,3 @@ 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
View File
@@ -7,7 +7,7 @@ requires = [
"setuptools>=77.0.3,<81.0.0",
"setuptools-scm>=8.0",
"setuptools-rust>=1.9.0",
"torch == 2.13.0",
"torch == 2.11.0",
"wheel",
"jinja2",
]
+2 -2
View File
@@ -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.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"
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"
wheel
jinja2>=3.1.6
regex
+1 -1
View File
@@ -5,7 +5,7 @@ packaging>=24.2
setuptools>=77.0.3,<81.0.0
setuptools-scm>=8
setuptools-rust>=1.9.0
torch==2.13.0
torch==2.11.0
wheel
jinja2>=3.1.6
regex
+2 -2
View File
@@ -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.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"
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"
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
+5 -5
View File
@@ -4,18 +4,18 @@
numba == 0.65.0 # Required for N-gram speculative decoding
# Dependencies for NVIDIA GPUs
torch==2.13.0
torch==2.11.0
torchaudio==2.11.0
# These must be updated alongside torch
torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
torchvision==0.26.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.15.post1
flashinfer-cubin==0.6.15.post1
flashinfer-python==0.6.14
flashinfer-cubin==0.6.14
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.6.1 # Required for CUTLASS DSL 4.6 by MSA
quack-kernels>=0.4.0 # Required for tml-fa4
# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.1.8; platform_system == "Linux"
+2 -2
View File
@@ -1107,7 +1107,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.13.0+cpu
torch==2.11.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.28.0+cpu
torchvision==0.26.0+cpu
# via
# -r requirements/test/cuda.in
# open-clip-torch
+2 -2
View File
@@ -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.13.0
torch==2.11.0
torchaudio==2.11.0
torchvision==0.28.0
torchvision==0.26.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
+9 -11
View File
@@ -159,7 +159,7 @@ cuda-bindings==13.0.3
# via torch
cuda-pathfinder==1.3.3
# via cuda-bindings
cuda-toolkit==13.0.3.0
cuda-toolkit==13.0.2
# via torch
cupy-cuda12x==13.6.0
# via ray
@@ -599,7 +599,7 @@ numpy==2.2.6
# tritonclient
# vocos
# xgrammar
nvidia-cublas==13.1.1.3
nvidia-cublas==13.1.0.3
# via
# cuda-toolkit
# nvidia-cudnn-cu13
@@ -607,12 +607,10 @@ nvidia-cublas==13.1.1.3
nvidia-cuda-cupti==13.0.85
# via cuda-toolkit
nvidia-cuda-nvrtc==13.0.88
# via
# cuda-toolkit
# nvidia-cublas
# via cuda-toolkit
nvidia-cuda-runtime==13.0.96
# via cuda-toolkit
nvidia-cudnn-cu13==9.20.0.48
nvidia-cudnn-cu13==9.19.0.56
# via torch
nvidia-cufft==12.0.0.61
# via cuda-toolkit
@@ -626,9 +624,9 @@ nvidia-cusparse==12.6.3.3
# via
# cuda-toolkit
# nvidia-cusolver
nvidia-cusparselt-cu13==0.8.1
nvidia-cusparselt-cu13==0.8.0
# via torch
nvidia-nccl-cu13==2.29.7
nvidia-nccl-cu13==2.28.9
# via torch
nvidia-nvjitlink==13.0.88
# via
@@ -1204,7 +1202,7 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/cuda.in
# transformers
torch==2.13.0+cu130
torch==2.11.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1235,7 +1233,7 @@ torchcodec==0.14.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
torchvision==0.28.0+cu130
torchvision==0.26.0+cu130
# via
# -c requirements/cuda.txt
# -r requirements/test/cuda.in
@@ -1272,7 +1270,7 @@ transformers==5.13.1
# xgrammar
transformers-stream-generator==0.0.5
# via -r requirements/test/cuda.in
triton==3.7.1
triton==3.6.0
# via
# torch
# xgrammar
+1 -1
View File
@@ -12,4 +12,4 @@ ray[data]
setuptools==78.1.0
setuptools-rust>=1.9.0
nixl==0.3.0
tpu-inference==0.25.0
tpu-inference==0.24.0
+9 -26
View File
@@ -3446,6 +3446,7 @@ dependencies = [
"base64 0.22.1",
"bytes",
"encoding_rs",
"futures-channel",
"futures-core",
"futures-util",
"h2",
@@ -4875,7 +4876,6 @@ dependencies = [
"futures-core",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
@@ -4937,9 +4937,9 @@ dependencies = [
[[package]]
name = "tonic"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
dependencies = [
"async-trait",
"axum",
@@ -4966,9 +4966,9 @@ dependencies = [
[[package]]
name = "tonic-build"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -4976,24 +4976,11 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "tonic-health"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
dependencies = [
"prost",
"tokio",
"tokio-stream",
"tonic",
"tonic-prost",
]
[[package]]
name = "tonic-prost"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
dependencies = [
"bytes",
"prost",
@@ -5002,9 +4989,9 @@ dependencies = [
[[package]]
name = "tonic-prost-build"
version = "0.14.6"
version = "0.14.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
dependencies = [
"prettyplease",
"proc-macro2",
@@ -5497,13 +5484,10 @@ dependencies = [
"serde",
"serde_json",
"thiserror 2.0.18",
"thiserror-ext",
"tiktoken-rs 0.9.1",
"tokenizers",
"tokio",
"tokio-stream",
"tracing",
"tracing-subscriber",
"url",
"uuid",
]
@@ -5746,7 +5730,6 @@ dependencies = [
"tokio-stream",
"tokio-util",
"tonic",
"tonic-health",
"tonic-prost",
"tonic-prost-build",
"tower",
+4 -5
View File
@@ -118,11 +118,10 @@ tokio = { version = "1.47.1", features = [
tokio-openssl = "0.6"
tokio-stream = "0.1"
tokio-util = { version = "0.7.18", features = ["rt"] }
tonic = "0.14.6"
tonic-build = "0.14.6"
tonic-health = "0.14.6"
tonic-prost = "0.14.6"
tonic-prost-build = "0.14.6"
tonic = "0.14.5"
tonic-build = "0.14.5"
tonic-prost = "0.14.5"
tonic-prost-build = "0.14.5"
tool-parser = "1.2.0"
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
-13
View File
@@ -14,10 +14,6 @@ service Generate {
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
}
service Control {
rpc Abort (AbortRequest) returns (AbortResponse) {}
}
// ======================================================================================
// Generate Request
// ======================================================================================
@@ -205,12 +201,3 @@ message TokenIds {
repeated uint32 ids = 1;
}
// ======================================================================================
// Control
// ======================================================================================
message AbortRequest {
repeated string request_ids = 1;
}
message AbortResponse {}
+1 -4
View File
@@ -20,19 +20,16 @@ mimalloc.workspace = true
rand.workspace = true
rand_distr.workspace = true
rayon.workspace = true
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
rlimit.workspace = true
rustc-hash.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json = { workspace = true, features = ["raw_value"] }
thiserror.workspace = true
thiserror-ext.workspace = true
tiktoken-rs.workspace = true
tokenizers.workspace = true
tokio.workspace = true
tokio-stream.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true
uuid.workspace = true
+7 -8
View File
@@ -158,10 +158,9 @@ impl PoolingBackend {
// (mirrors Python async_request_vllm_rerank).
if let Some(ref list) = input.prompt_list {
if list.len() < 2 {
tracing::warn!(
backend = "vllm-rerank",
inputs = list.len(),
"rerank request has no documents"
eprintln!(
"WARNING: vllm-rerank request has no documents \
(prompt_list needs [query, doc, ...])"
);
}
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -176,10 +175,10 @@ impl PoolingBackend {
// Legacy path: text prompt as query, documents via --extra-body.
let query = input.prompt.as_ref();
if query.is_empty() && input.prompt_token_ids.is_some() {
tracing::warn!(
backend = "vllm-rerank",
dataset = "random",
"rerank request has an empty query; use the random-rerank dataset"
eprintln!(
"WARNING: vllm-rerank received empty query (random dataset uses \
token IDs only). Use --dataset-name random-rerank for meaningful \
rerank benchmarks."
);
}
serde_json::json!({
+84 -114
View File
@@ -6,7 +6,6 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,12 +72,12 @@ pub fn pre_resolve_dns(
v4.extend(v6);
if !v4.is_empty() {
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
println!("Pre-resolved {host} -> {ips:?}");
builder = builder.resolve_to_addrs(host, &v4);
}
}
Err(e) => {
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
}
}
@@ -347,14 +346,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
tracing::info!(base_url = %config.base_url, "fetching first model from server");
println!("Model not specified, fetching first model from server...");
let (name, id) =
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
println!("First model name: {name}, first model id: {id}");
(id, Some(name))
};
@@ -363,10 +358,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
tracing::info!(tokenizer = tid, "loading tokenizer");
println!("Loading tokenizer: {tid}");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
let has_tokenizer = tokenizer.is_some();
@@ -426,12 +421,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.num_prompts, config.random_batch_size, config.is_reranker,
),
};
tracing::info!(
dataset = ?config.dataset_name,
prompts = config.num_prompts,
description = %dataset_label,
"generating benchmark dataset"
);
println!("Generating {dataset_label}...");
let gen_start = Instant::now();
let mut input_requests = match config.dataset_name {
@@ -482,7 +472,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded.as_str()
}
};
@@ -522,8 +512,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None => {
downloaded = crate::datasets::speed_bench::download_speed_bench(
config.speed_bench_config,
)
.await?;
)?;
downloaded.as_str()
}
};
@@ -554,8 +543,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.hf_subset.as_deref(),
config.hf_split.as_deref(),
config.num_prompts,
)
.await?;
)?;
crate::datasets::hf_dataset::load_hf_dataset(
tok,
&downloaded_path,
@@ -620,19 +608,18 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
};
let gen_elapsed = gen_start.elapsed();
tracing::info!(
prompts = input_requests.len(),
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated benchmark dataset"
println!(
"Generated {} prompts in {:.2}s",
input_requests.len(),
gen_elapsed.as_secs_f64()
);
let filtered_count =
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
if filtered_count > 0 {
tracing::info!(
filtered_prompts = filtered_count,
max_model_len = config.max_model_len.unwrap(),
"filtered prompts above maximum model length"
println!(
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
config.max_model_len.unwrap()
);
}
if input_requests.is_empty() {
@@ -683,7 +670,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Ready check
if config.ready_check_timeout_sec > 0 {
tracing::info!("starting initial single-prompt test run");
println!("Starting initial single prompt test run...");
let test_output = wait_for_endpoint(
config.backend,
&client,
@@ -698,7 +685,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
test_output.error
)));
}
tracing::info!("initial single-prompt test run completed");
println!("Initial test run completed.");
}
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -716,15 +703,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
DatasetName::Random | DatasetName::PrefixRepetition
);
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
tracing::info!(
reason = "prompt_token_ids",
"skipping server tokenizer verification"
);
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
}
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
if is_tokenizer_verified(&cache_key) {
tracing::info!(reason = "cached", "skipping server tokenizer verification");
println!("Tokenizer verified in previous run (cached), skipping verification.");
} else {
let num_special =
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -739,17 +723,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await?
{
SampleVerifyOutcome::Passed => {
tracing::info!("tokenizer sample verification passed");
println!("Sample verification passed, skipping full verification.");
mark_tokenizer_verified(&cache_key);
}
SampleVerifyOutcome::Skipped(reason) => {
tracing::warn!(
reason = %reason,
"server tokenizer unavailable; skipping prompt verification"
);
println!("Server /tokenize unavailable ({reason}), skipping verification.");
}
SampleVerifyOutcome::Mismatch => {
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
println!("Sample verification found mismatch, running full verify+fix...");
match verify_and_fix_prompt_lengths(
&client,
&config.base_url,
@@ -761,16 +742,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await
{
Ok(()) => {
tracing::info!(
prompts = input_requests.len(),
"verified exact prompt token lengths"
println!(
"All {} prompts verified: exact token length match.",
input_requests.len()
);
mark_tokenizer_verified(&cache_key);
}
Err(BenchError::TokenizeUnavailable(reason)) => {
tracing::warn!(
reason = %reason,
"server tokenizer became unavailable; using client token counts"
println!(
"Server /tokenize became unavailable during verification \
({reason}); proceeding with client-side token counts."
);
}
Err(e) => return Err(e),
@@ -782,7 +763,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Warmup
if config.num_warmups > 0 {
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
println!("Warming up with {} requests...", config.num_warmups);
run_warmup(
config.backend,
&client,
@@ -795,7 +776,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.disable_tqdm,
)
.await;
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
println!("Warmup run completed.");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -833,22 +814,28 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
tracing::info!("detected speculative decoding; collecting metrics");
println!("Speculative decoding detected, will collect metrics.");
}
// Main benchmark
println!("Starting main benchmark run...");
let distribution = if config.burstiness == 1.0 {
"Poisson process"
} else {
"Gamma distribution"
};
tracing::info!(
request_rate = config.request_rate,
burstiness = config.burstiness,
distribution,
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
prompts = config.num_prompts,
"starting main benchmark run"
println!(
"Traffic request rate: {}",
if config.request_rate.is_infinite() {
"inf".to_string()
} else {
format!("{}", config.request_rate)
}
);
println!("Burstiness factor: {} ({distribution})", config.burstiness);
println!(
"Maximum request concurrency: {}",
config.max_concurrency.unwrap_or(config.num_prompts)
);
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -860,11 +847,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
);
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
"assigned LoRA adapters"
println!(
"LoRA adapters ({}): {:?} [assignment={:?}]",
modules.len(),
names,
config.lora_assignment
);
}
@@ -1138,7 +1125,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
tracing::error!(error = %e.as_report(), "profiler background task failed");
eprintln!("WARNING: Profile background task failed: {e}");
}
}
@@ -1302,14 +1289,12 @@ pub(crate) async fn start_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Starting profiler...");
let profile_url = format!("{base_url}/start_profile");
tracing::info!(url = %profile_url, "starting profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
}
Ok(true) => println!("Profiler started"),
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
}
}
@@ -1319,14 +1304,12 @@ pub(crate) async fn stop_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Stopping profiler...");
let profile_url = format!("{base_url}/stop_profile");
tracing::info!(url = %profile_url, "stopping profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
}
Ok(true) => println!("Profiler stopped"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
}
}
@@ -1388,30 +1371,25 @@ pub(crate) async fn profile_on_batch_threshold(
duration_secs: f64,
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
) {
tracing::info!(
threshold,
duration_seconds = duration_secs,
"waiting for profiler batch threshold"
println!(
"Waiting for batch size >= {threshold} before starting profiler \
(will capture {duration_secs}s)..."
);
loop {
if let Some(running) = fetch_num_requests_running(client, base_url).await
&& running >= threshold
{
tracing::info!(
running_requests = running,
threshold,
"profiler batch threshold reached"
);
println!("Batch size {running} >= {threshold}, starting profiler...");
break;
}
// Wait 500ms or until the benchmark signals cancellation
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = &mut cancel_rx => {
tracing::warn!(
threshold,
"benchmark finished before profiler batch threshold; skipping profiling"
eprintln!(
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
profiling skipped."
);
return;
}
@@ -1420,13 +1398,13 @@ pub(crate) async fn profile_on_batch_threshold(
let start_url = format!("{base_url}/start_profile");
match send_profile_request(client, &start_url, extra_headers).await {
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
Ok(true) => println!("Profiler started"),
Ok(false) => {
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
eprintln!("WARNING: Profiler start request returned non-success");
return;
}
Err(e) => {
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
eprintln!("WARNING: Failed to start profiler: {e}");
return;
}
}
@@ -1435,17 +1413,15 @@ pub(crate) async fn profile_on_batch_threshold(
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
_ = &mut cancel_rx => {
tracing::info!("benchmark finished; stopping profiler early");
println!("Benchmark finished, stopping profiler early...");
}
}
let stop_url = format!("{base_url}/stop_profile");
match send_profile_request(client, &stop_url, extra_headers).await {
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
}
Ok(true) => println!("Profiler stopped after capturing"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
}
}
@@ -1526,11 +1502,10 @@ async fn verify_and_fix_prompt_lengths(
let excess = tokens.len().saturating_sub(expected_input_len);
let compensate = if excess > 0 && last_excess == Some(excess) {
if _iter == 1 {
tracing::warn!(
prompt_index = i,
extra_tokens = excess,
adjusted_target = expected_input_len.saturating_sub(excess),
"server consistently adds prompt tokens; compensating verification target"
eprintln!(
"Prompt {i}: server consistently adds {excess} extra token(s) \
(likely BOS), compensating target to {}.",
expected_input_len.saturating_sub(excess),
);
}
excess
@@ -1588,10 +1563,7 @@ async fn verify_and_fix_prompt_lengths(
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
if fc > 0 {
tracing::info!(
fixed_prompts = fc,
"fixed prompt lengths using server tokenizer"
);
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
}
Ok(())
@@ -1846,7 +1818,7 @@ async fn sample_verify_prompts(
let tokenize_url = format!("{base_url}/tokenize");
let api_key = std::env::var("OPENAI_API_KEY").ok();
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
println!("Sampling {sample_size} prompts for verification...");
for (i, request) in requests.iter().enumerate().take(sample_size) {
let tokens = match server_tokenize(
@@ -1869,11 +1841,9 @@ async fn sample_verify_prompts(
let expected = request.prompt_len + num_special;
if tokens.len() != expected {
tracing::warn!(
prompt_index = i,
expected_tokens = expected,
actual_tokens = tokens.len(),
"tokenizer verification sample mismatch"
println!(
"Prompt {i}: expected {expected} tokens, server returned {}",
tokens.len()
);
return Ok(SampleVerifyOutcome::Mismatch);
}
+9 -18
View File
@@ -288,18 +288,10 @@ impl BenchConfig {
}
Some(other) => {
// extra_body was not an object — just use sampling params
let value_type = match &other {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => unreachable!(),
};
tracing::warn!(
value_type,
"sampling parameters may be lost because --extra-body is not a JSON object"
eprintln!(
"Warning: --extra-body is not a JSON object, sampling params may be lost"
);
let _ = other;
sampling_params
}
None => sampling_params,
@@ -497,9 +489,9 @@ impl BenchConfig {
_ => {}
}
if !args.skip_chat_template {
tracing::warn!(
dataset = "custom",
"client-side chat template rendering is unsupported; sending prompts raw"
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
);
}
}
@@ -578,10 +570,9 @@ impl BenchConfig {
}
if ignore_eos {
tracing::warn!(
ignore_eos,
multi_turn = true,
"output length limits may be ignored, causing unbounded context growth"
eprintln!(
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
respect output length limits, causing unbounded context growth."
);
}
+2 -4
View File
@@ -128,10 +128,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
#[test]
+37 -74
View File
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -51,19 +50,18 @@ enum ColumnFormat {
/// Make a GET request with retry logic (3 retries with exponential backoff).
/// Returns the parsed JSON response.
async fn get_with_retry(
client: &reqwest::Client,
fn get_with_retry(
client: &reqwest::blocking::Client,
url: &str,
label: &str,
) -> Result<serde_json::Value> {
let max_retries = 3;
for attempt in 0..=max_retries {
let resp = match client.get(url).send().await {
let resp = match client.get(url).send() {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
.await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
return Err(BenchError::Config(format!(
@@ -82,7 +80,7 @@ async fn get_with_retry(
}
if status.is_server_error() && attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
@@ -94,7 +92,6 @@ async fn get_with_retry(
let data: serde_json::Value = resp
.json()
.await
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
return Ok(data);
}
@@ -108,7 +105,7 @@ async fn get_with_retry(
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
/// or the dataset is exhausted.
pub async fn download_hf_dataset(
pub fn download_hf_dataset(
dataset: &str,
subset: Option<&str>,
split: Option<&str>,
@@ -118,7 +115,7 @@ pub async fn download_hf_dataset(
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
let mut client_builder =
reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
// Add HF_TOKEN auth header if available
if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -141,7 +138,7 @@ pub async fn download_hf_dataset(
// Call /info to discover available configs and splits
let info_url =
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
let dataset_info =
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -204,12 +201,7 @@ pub async fn download_hf_dataset(
(resolved_config, resolved_split)
};
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"resolved Hugging Face dataset"
);
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
// Check cache
let dir = cache_dir();
@@ -223,16 +215,11 @@ pub async fn download_hf_dataset(
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
println!("HF dataset cached: {path_str}");
return Ok((path_str, resolved_config, resolved_split));
}
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"downloading Hugging Face dataset"
);
println!("Downloading HF dataset '{dataset}' from datasets-server...");
let encoded_config: String =
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -242,7 +229,6 @@ pub async fn download_hf_dataset(
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -254,7 +240,7 @@ pub async fn download_hf_dataset(
&length={page_size}"
);
let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
let data = get_with_retry(&client, &url, "HF dataset /rows")?;
let rows = data["rows"]
.as_array()
@@ -274,14 +260,14 @@ pub async fn download_hf_dataset(
offset += fetched;
let total = data["num_rows_total"].as_u64().unwrap_or(0);
progress.update(offset, total);
eprint!("\r Fetched {offset}/{total} rows...");
// Stop if we have enough rows or reached end of dataset
if all_rows.len() >= num_rows_needed || fetched < page_size {
break;
}
}
progress.finish();
eprintln!(); // newline after progress
if all_rows.is_empty() {
return Err(BenchError::Config(format!(
@@ -294,12 +280,7 @@ pub async fn download_hf_dataset(
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(
dataset,
rows = all_rows.len(),
path = %path_str,
"saved Hugging Face dataset"
);
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
Ok((path_str, resolved_config, resolved_split))
}
@@ -502,31 +483,21 @@ pub fn load_hf_dataset(
// Detect column format from first row
let format = detect_column_format(&entries[0], text_column_override)?;
// Print detected format
match &format {
ColumnFormat::Chat(col) => {
tracing::info!(
format = "chat",
column = col,
"detected Hugging Face dataset format"
);
}
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
ColumnFormat::Text {
prompt_col,
output_col,
} => {
tracing::info!(
format = "text",
prompt_column = prompt_col,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
let out_msg = output_col.as_deref().unwrap_or("none");
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
}
ColumnFormat::Combined { cols, output_col } => {
tracing::info!(
format = "combined",
prompt_columns = ?cols,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
let out_msg = output_col.as_deref().unwrap_or("none");
println!(
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
cols
);
}
}
@@ -637,10 +608,9 @@ pub fn load_hf_dataset(
if len == 0 { 128 } else { len }
} else {
if !warned_no_output {
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
);
warned_no_output = true;
}
@@ -648,10 +618,9 @@ pub fn load_hf_dataset(
}
} else {
if !warned_no_output {
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
);
warned_no_output = true;
}
@@ -671,11 +640,9 @@ pub fn load_hf_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "hf",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -685,11 +652,9 @@ pub fn load_hf_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "hf",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled HF dataset from {original_len} to {} total samples.",
samples.len()
);
}
}
@@ -1037,10 +1002,8 @@ mod tests {
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
crate::tokenizer::TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
/// Write JSON data to a unique temp file and return the path string.
+6 -9
View File
@@ -5,7 +5,6 @@ pub mod custom;
pub mod hf_dataset;
pub mod multi_turn;
pub mod prefix_repetition;
mod progress;
pub mod random;
pub mod random_mm;
pub mod random_rerank;
@@ -91,10 +90,9 @@ pub fn oversample_requests(
return;
}
if no_oversample {
tracing::info!(
samples = requests.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
requests.len()
);
return;
}
@@ -105,10 +103,9 @@ pub fn oversample_requests(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
requests.push(req);
}
tracing::info!(
original_samples = original_len,
samples = requests.len(),
"oversampled dataset"
println!(
"Oversampled requests from {original_len} to {} total samples.",
requests.len()
);
}
+18 -29
View File
@@ -445,10 +445,9 @@ pub fn load_sharegpt_multi_turn(
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
conversations.push(conv);
}
tracing::info!(
original_conversations = original_len,
conversations = conversations.len(),
"oversampled multi-turn conversations"
println!(
"Oversampled multi-turn conversations from {original_len} to {} total.",
conversations.len()
);
}
@@ -526,12 +525,10 @@ mod tests {
len
}
#[tokio::test]
#[test]
#[ignore]
async fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 5,
@@ -613,12 +610,10 @@ mod tests {
println!("All prefix sharing checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 4,
@@ -655,12 +650,10 @@ mod tests {
println!("per_turn_input_len default-mode checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 50,
@@ -691,12 +684,10 @@ mod tests {
println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
let cfg = MultiTurnRandomConfig {
num_conversations: 10,
@@ -718,12 +709,10 @@ mod tests {
println!("variable_turns_fixed checks passed!");
}
#[tokio::test]
#[test]
#[ignore]
async fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
@@ -41,13 +41,11 @@ pub fn generate_prefix_repetition_dataset(
}
let total = prompts_per_prefix * num_prefixes;
if total != num_requests {
tracing::info!(
requested = num_requests,
generated = total,
prefixes = num_prefixes,
prompts_per_prefix,
dropped = num_requests - total,
"adjusted prefix-repetition request count"
println!(
"prefix_repetition: generating {total} requests \
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
{} dropped to divide evenly)",
num_requests - total
);
}
@@ -111,10 +109,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
#[test]
-77
View File
@@ -1,77 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
/// Reports row download progress to an interactive progress bar, or through
/// periodic tracing events when the progress bar is hidden on a non-TTY.
pub(super) struct RowDownloadReporter {
progress: ProgressBar,
next_report: Instant,
}
impl RowDownloadReporter {
/// Creates a reporter that emits non-TTY updates every 10 seconds.
pub fn new() -> Self {
let progress = ProgressBar::new(0);
progress.set_style(
ProgressStyle::with_template(
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
)
.unwrap()
.progress_chars("#>-"),
);
Self {
progress,
next_report: Instant::now() + REPORT_INTERVAL,
}
}
/// Updates the current row count and reports progress when due.
pub fn update(&mut self, rows: usize, total: u64) {
let rows = rows as u64;
let total = total.max(rows);
self.progress.set_length(total);
self.progress.set_position(rows);
if self.should_report(Instant::now()) {
tracing::info!(rows, total, "fetching dataset rows");
}
}
/// Clears the interactive progress bar after the download completes.
pub fn finish(self) {
self.progress.finish_and_clear();
}
fn should_report(&mut self, now: Instant) -> bool {
if !self.progress.is_hidden() || now < self.next_report {
return false;
}
self.next_report = now + REPORT_INTERVAL;
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hidden_reporter_uses_ten_second_deadline() {
let start = Instant::now();
let mut reporter = RowDownloadReporter {
progress: ProgressBar::hidden(),
next_report: start + REPORT_INTERVAL,
};
assert!(!reporter.should_report(start + Duration::from_secs(9)));
assert!(reporter.should_report(start + Duration::from_secs(10)));
assert!(!reporter.should_report(start + Duration::from_secs(19)));
assert!(reporter.should_report(start + Duration::from_secs(20)));
}
}
+12 -18
View File
@@ -49,12 +49,9 @@ pub fn generate_random_dataset(
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
let (output_low, output_high) = range_ratio.output_bounds(output_len);
if !range_ratio.is_fixed() {
tracing::info!(
input_low,
input_high,
output_low,
output_high,
"sampling random request lengths"
println!(
"Sampling input_len from [{input_low}, {input_high}] and \
output_len from [{output_low}, {output_high}]"
);
}
@@ -308,8 +305,7 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_token_ids() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -341,8 +337,7 @@ mod tests {
#[test]
#[ignore]
fn test_generate_random_dataset_text() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let requests = generate_random_dataset(
&tokenizer,
10, // num_requests
@@ -376,8 +371,7 @@ mod tests {
#[test]
#[ignore]
fn test_token_length_exact_local() {
let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
let target_len = 512;
let requests = generate_random_dataset(
&tokenizer,
@@ -411,11 +405,11 @@ mod tests {
}
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
#[tokio::test]
#[test]
#[ignore]
async fn test_token_length_exact_tiktoken() {
fn test_token_length_exact_tiktoken() {
// Use Qwen2.5 which has a tiktoken-format tokenizer
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
@@ -459,10 +453,10 @@ mod tests {
/// Test encode/decode roundtrip stability for tiktoken.
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
#[tokio::test]
#[test]
#[ignore]
async fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
let tokenizer = match tokenizer {
Ok(t) => t,
Err(e) => {
+2 -4
View File
@@ -139,10 +139,8 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind {
TokenizerKind::Tiktoken(
crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
crate::tokenizer::load_tokenizer("gpt2", false, None)
.expect("gpt2 built-in tiktoken should always load without network")
}
fn fixed_ratio() -> RangeRatio {
+12 -19
View File
@@ -22,21 +22,18 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
/// Download the default ShareGPT dataset from HuggingFace Hub.
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
pub async fn download_sharegpt_dataset() -> Result<String> {
tracing::info!(
repository = DEFAULT_SHAREGPT_REPO,
file = DEFAULT_SHAREGPT_FILE,
"downloading ShareGPT dataset"
pub fn download_sharegpt_dataset() -> Result<String> {
println!(
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
);
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
.map_err(BenchError::Config)?;
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
BenchError::Config(format!(
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
))
})?;
let path_str = path.to_string_lossy().to_string();
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
println!("ShareGPT dataset ready: {path_str}");
Ok(path_str)
}
@@ -138,11 +135,9 @@ pub fn load_sharegpt_dataset(
// Oversample if dataset is smaller than requested
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "sharegpt",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let needed = num_requests - samples.len();
@@ -152,11 +147,9 @@ pub fn load_sharegpt_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "sharegpt",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled requests from {original_len} to {} total samples.",
samples.len()
);
}
}
+23 -30
View File
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::cli::SpeedBenchConfig;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -26,7 +25,7 @@ fn cache_dir() -> std::path::PathBuf {
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
/// Results are cached as JSON locally for subsequent runs.
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let config_name = config.as_str();
let dir = cache_dir();
@@ -36,13 +35,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
// Return cached file if it exists
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
println!("SPEED-Bench ({config_name}) cached: {path_str}");
return Ok(path_str);
}
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
let client = reqwest::Client::builder()
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -50,7 +49,6 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -66,14 +64,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let max_retries = 3;
let mut data: Option<serde_json::Value> = None;
for attempt in 0..=max_retries {
let resp = match client.get(&url).send().await {
let resp = match client.get(&url).send() {
Ok(r) => r,
Err(e) => {
if attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(
std::thread::sleep(std::time::Duration::from_secs(
2 * (attempt as u64 + 1),
))
.await;
));
continue;
}
return Err(BenchError::Config(format!(
@@ -83,7 +80,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
};
if resp.status().is_server_error() && attempt < max_retries {
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
continue;
}
@@ -94,7 +91,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
)));
}
data = Some(resp.json().await.map_err(|e| {
data = Some(resp.json().map_err(|e| {
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
})?);
break;
@@ -119,14 +116,15 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let fetched = rows.len();
offset += fetched;
// Print progress
let total = data["num_rows_total"].as_u64().unwrap_or(0);
progress.update(offset, total);
eprint!("\r Fetched {offset}/{total} rows...");
if fetched < page_size {
break;
}
}
progress.finish();
eprintln!(); // newline after progress
if all_rows.is_empty() {
return Err(BenchError::Config(
@@ -139,11 +137,9 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
tracing::info!(
config = config_name,
rows = all_rows.len(),
path = %path_str,
"saved SPEED-Bench dataset"
println!(
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
all_rows.len()
);
Ok(path_str)
}
@@ -267,11 +263,9 @@ pub fn load_speed_bench_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
tracing::info!(
dataset = "speed-bench",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -281,11 +275,9 @@ pub fn load_speed_bench_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
tracing::info!(
dataset = "speed-bench",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
println!(
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
samples.len()
);
}
}
@@ -296,6 +288,7 @@ pub fn load_speed_bench_dataset(
));
}
// Print category distribution
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
for entry in &filtered[..filtered.len().min(samples.len())] {
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -304,7 +297,7 @@ pub fn load_speed_bench_dataset(
let mut cats: Vec<_> = cat_counts.into_iter().collect();
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
println!("SPEED-Bench categories: {}", cat_str.join(", "));
Ok(samples)
}
+35 -20
View File
@@ -1,39 +1,54 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::path::PathBuf;
//! Sync facade over the async `hf_hub` API.
//!
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
//! backend unconditionally pulls ureq's default rustls feature. So we use the
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
//! loaders, tokenizer fallback in rayon threads) by running each download on a
//! dedicated thread with its own single-threaded runtime.
use hf_hub::Repo;
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
use std::path::PathBuf;
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
pub struct HubRepo {
repo: ApiRepo,
repo: hf_hub::Repo,
}
impl HubRepo {
pub fn model(model_id: String) -> Result<Self, String> {
Self::new(Repo::model(model_id))
pub fn model(model_id: String) -> Self {
Self {
repo: hf_hub::Repo::model(model_id),
}
}
pub fn dataset(repo_id: String) -> Result<Self, String> {
Self::new(Repo::dataset(repo_id))
}
fn new(repo: Repo) -> Result<Self, String> {
let mut builder = ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
pub fn dataset(repo_id: String) -> Self {
Self {
repo: hf_hub::Repo::dataset(repo_id),
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
Ok(Self {
repo: api.repo(repo),
})
}
/// Download (or fetch from cache) a single file from the repo.
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
self.repo.get(filename).await.map_err(|e| format!("{e}"))
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
let repo = self.repo.clone();
let filename = filename.to_string();
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
.join()
.map_err(|_| "HF Hub download thread panicked".to_string())?
}
}
+1 -1
View File
@@ -33,7 +33,7 @@ pub fn prepare_process() {
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
tracing::info!(soft_limit = new, "raised open-file limit");
eprintln!("Open-file limit: {new}");
}
}
-12
View File
@@ -19,19 +19,7 @@ struct Cli {
args: vllm_bench::BenchServeArgs,
}
// TODO: unify the tracing subscriber used by different binaries.
fn init_tracing() {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
let _ = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_writer(std::io::stderr)
.try_init();
}
fn main() -> anyhow::Result<()> {
init_tracing();
let cli = Cli::parse();
vllm_bench::prepare_process();
+16 -18
View File
@@ -7,22 +7,6 @@ use crate::datasets::SampleRequest;
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
use crate::multi_turn::ConversationOutput;
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
if failed_outputs.is_empty() {
return;
}
tracing::warn!(
failed_requests = failed_outputs.len(),
displayed_errors = failed_outputs.len().min(10),
"benchmark requests failed"
);
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
tracing::warn!(index, error = %output.error, "benchmark request failed");
}
}
/// Calculate benchmark metrics from request outputs.
///
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
@@ -79,7 +63,14 @@ pub fn calculate_metrics(
let failed = outputs.len() - completed;
log_failed_requests(outputs);
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Calculate max output tokens per second and max concurrent requests
let mut max_output_tokens_per_s = 0.0_f64;
@@ -304,7 +295,14 @@ pub fn calculate_embedding_metrics(
let failed = outputs.len() - completed;
log_failed_requests(outputs);
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Compute peak concurrent requests from start_time + latency windows
let successful_outputs: Vec<&RequestFuncOutput> =
+39 -53
View File
@@ -6,7 +6,6 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
@@ -73,13 +72,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
tracing::info!(base_url = %config.base_url, "fetching first model from server");
println!("Model not specified, fetching first model from server...");
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
println!("First model name: {name}, first model id: {id}");
(id, Some(name))
};
@@ -88,19 +83,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
tracing::info!(tokenizer = tid, "loading tokenizer");
println!("Loading tokenizer: {tid}");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t =
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
// Generate/load conversations
tracing::info!(
dataset = ?config.dataset_name,
conversations = config.num_prompts,
"generating multi-turn conversations"
);
println!("Generating multi-turn conversations...");
let gen_start = Instant::now();
let mut conversations = match config.dataset_name {
@@ -140,7 +131,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let path = match config.dataset_path.as_deref() {
Some(p) => p,
None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
downloaded.as_str()
}
};
@@ -188,11 +179,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (filtered_conversations, filtered_turns) =
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
if filtered_turns > 0 || filtered_conversations > 0 {
tracing::info!(
filtered_turns,
filtered_conversations,
max_model_len,
"filtered conversations above maximum model length"
println!(
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
);
}
if conversations.is_empty() {
@@ -204,11 +192,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let gen_elapsed = gen_start.elapsed();
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
tracing::info!(
conversations = conversations.len(),
println!(
"Generated {} conversations ({} total turns) in {:.2}s",
conversations.len(),
total_turns,
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated multi-turn conversations"
gen_elapsed.as_secs_f64()
);
// Log prefix sharing info
@@ -220,15 +208,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let conv_tokens =
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
tracing::info!(
global_ratio = config.multi_turn_prefix_global_ratio,
println!(
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
config.multi_turn_prefix_global_ratio * 100.0,
global_tokens,
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
conversation_tokens = conv_tokens,
config.multi_turn_prefix_conversation_ratio * 100.0,
conv_tokens,
(1.0 - config.multi_turn_prefix_global_ratio
- config.multi_turn_prefix_conversation_ratio)
* 100.0,
unique_tokens,
history_accumulation = false,
"configured multi-turn prefix sharing"
);
println!("No history accumulation: each turn sends fixed-length prompt only.");
}
if config.dry_run {
@@ -262,7 +253,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
..Default::default()
};
tracing::info!("starting initial single-prompt test run");
println!("Starting initial single prompt test run...");
let test_output = crate::ready_checker::wait_for_endpoint(
config.backend,
&client,
@@ -277,7 +268,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
test_output.error
)));
}
tracing::info!("initial single-prompt test run completed");
println!("Initial test run completed.");
}
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -292,10 +283,9 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
"min_tokens".to_string(),
serde_json::json!(config.random_output_len),
);
tracing::info!(
min_tokens = config.random_output_len,
dataset = "random",
"set minimum output tokens for multi-turn dataset"
println!(
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
config.random_output_len
);
}
Some(body)
@@ -307,7 +297,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
tracing::info!("detected speculative decoding; collecting metrics");
println!("Speculative decoding detected, will collect metrics.");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -340,13 +330,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
};
// Main benchmark
tracing::info!(
conversations = conversations.len(),
total_turns,
concurrency,
inter_turn_delay_ms = config.multi_turn_delay_ms,
"starting multi-turn benchmark"
);
println!("Starting multi-turn benchmark...");
println!("Conversations: {}", conversations.len());
println!("Concurrency: {concurrency}");
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
@@ -377,12 +364,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
);
if let Some(modules) = config.lora_modules.as_ref() {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
scope = "conversation",
"assigned LoRA adapters"
println!(
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
modules.len(),
names,
config.lora_assignment
);
}
@@ -447,7 +433,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
match handle.await {
Ok(output) => all_outputs.push(output),
Err(e) => {
tracing::error!(error = %e.as_report(), "conversation task panicked");
eprintln!("Conversation task panicked: {e}");
}
}
}
@@ -467,7 +453,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
tracing::error!(error = %e.as_report(), "profiler background task failed");
eprintln!("WARNING: Profile background task failed: {e}");
}
}
+2 -2
View File
@@ -603,7 +603,7 @@ fn add_metric_stats(
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
let content = serde_json::to_string(json)?;
std::fs::write(file_path, content)?;
tracing::info!(path = file_path, "saved benchmark results");
println!("Results saved to {file_path}");
Ok(())
}
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
file.write_all(b"\n")?;
}
file.write_all(content.as_bytes())?;
tracing::info!(path = file_path, "appended benchmark results");
println!("Results appended to {file_path}");
Ok(())
}
+2 -8
View File
@@ -23,11 +23,7 @@ pub async fn wait_for_endpoint(
let backend = get_backend(backend)?;
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
tracing::info!(
timeout_seconds,
retry_interval,
"waiting for endpoint readiness"
);
println!("Waiting for endpoint to become up in {timeout_seconds}s");
let pb = ProgressBar::new(timeout_seconds);
pb.set_style(
@@ -57,9 +53,7 @@ pub async fn wait_for_endpoint(
Ok(output) => {
let err = output.error.clone();
let err_last_line = err.lines().last().unwrap_or(&err);
pb.suspend(|| {
tracing::warn!(error = err_last_line, "endpoint is not ready");
});
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
last_error = err;
}
Err(e) => {
+1 -1
View File
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
.await
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
if resp.status().is_success() {
tracing::info!(url = %url, "reset prefix cache");
println!("Prefix cache reset successfully.");
} else {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
+24 -36
View File
@@ -199,17 +199,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
}
};
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
tracing::info!(
encoding,
kind = "built-in-tiktoken",
vocab_size,
"loaded tokenizer"
);
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
}
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
// Phase 1: If model_id is a local directory, look for tiktoken files there
let local_dir = Path::new(model_id);
if local_dir.is_dir() {
@@ -217,7 +212,7 @@ pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
}
// Phase 2: Fall back to HuggingFace Hub download
try_load_tiktoken_from_hf(model_id).await
try_load_tiktoken_from_hf(model_id)
}
/// Common tiktoken model filenames to search for.
@@ -252,28 +247,25 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
}
/// Load a tiktoken tokenizer from a HuggingFace model repo.
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string());
let mut model_path = None;
for filename in TIKTOKEN_MODEL_FILENAMES {
if let Ok(path) = repo.get(filename).await {
model_path = Some(path);
break;
}
}
let model_path = model_path.ok_or_else(|| {
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
})?;
let model_path = repo
.get("tiktoken.model")
.or_else(|_| repo.get("qwen.tiktoken"))
.or_else(|_| repo.get("vocab.tiktoken"))
.map_err(|_| {
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
})?;
let num_base_tokens = count_base_tokens(&model_path)?;
let config = match repo.get("tokenizer_config.json").await {
let config = match repo.get("tokenizer_config.json") {
Ok(config_path) => read_tokenizer_config(&config_path),
Err(_) => None,
};
let pattern = extract_pat_str_from_repo(&repo).await;
let pattern = extract_pat_str_from_repo(&repo);
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
}
@@ -317,16 +309,15 @@ fn build_tiktoken(
}
}
tracing::info!(
model = model_id,
base_tokens = num_base_tokens,
special_tokens = all_special_tokens.len(),
pattern = if pattern.is_some() {
println!(
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
num_base_tokens,
all_special_tokens.len(),
if pattern.is_some() {
"custom"
} else {
"default"
},
"loading tiktoken model"
);
TiktokenTokenizer::from_file(
@@ -406,12 +397,9 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
/// Try to download the Python tokenizer source file and extract pat_str via regex.
/// Returns None if unavailable or unparsable.
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
// Try common Python tokenizer filenames
let py_path = match repo.get("tokenization_kimi.py").await {
Ok(path) => path,
Err(_) => repo.get("tokenizer.py").await.ok()?,
};
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
let source = std::fs::read_to_string(&py_path).ok()?;
@@ -450,9 +438,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
if !fragments.is_empty() {
let pattern = fragments.join("|");
tracing::debug!(
fragments = fragments.len(),
"extracted tiktoken pattern from Python source"
println!(
"Extracted pat_str from Python source: {} fragments",
fragments.len()
);
return Some(pattern);
}
+21 -98
View File
@@ -2,10 +2,8 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::HashSet;
use std::future::Future;
use std::path::Path;
use thiserror_ext::AsReport as _;
use tokenizers::Tokenizer;
use crate::error::{BenchError, Result};
@@ -20,8 +18,7 @@ pub enum TokenizerKind {
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
pub struct ServerTokenizer {
client: reqwest::Client,
runtime: tokio::runtime::Handle,
client: reqwest::blocking::Client,
tokenize_url: String,
detokenize_url: String,
model: String,
@@ -30,8 +27,8 @@ pub struct ServerTokenizer {
impl ServerTokenizer {
/// Create a new server tokenizer and verify connectivity.
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::Client::builder()
pub fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -41,7 +38,6 @@ impl ServerTokenizer {
let st = Self {
client,
runtime: tokio::runtime::Handle::current(),
tokenize_url,
detokenize_url,
model: model.to_string(),
@@ -49,7 +45,7 @@ impl ServerTokenizer {
};
// Probe the endpoint to verify it works and discover vocab size
let test_tokens = st.encode_async("test").await?;
let test_tokens = st.encode_inner("test")?;
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
let estimated_vocab = (max_id * 2).max(131072);
@@ -60,10 +56,6 @@ impl ServerTokenizer {
}
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
self.block_on(self.encode_async(text))
}
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
let payload = serde_json::json!({
"model": self.model,
"prompt": text,
@@ -74,7 +66,6 @@ impl ServerTokenizer {
.post(&self.tokenize_url)
.json(&payload)
.send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
if !resp.status().is_success() {
@@ -84,7 +75,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().await.map_err(|e| {
let data: serde_json::Value = resp.json().map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
})?;
@@ -104,10 +95,6 @@ impl ServerTokenizer {
}
fn decode_inner(&self, ids: &[u32]) -> Result<String> {
self.block_on(self.decode_async(ids))
}
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
let payload = serde_json::json!({
"model": self.model,
"tokens": ids,
@@ -118,7 +105,6 @@ impl ServerTokenizer {
.post(&self.detokenize_url)
.json(&payload)
.send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
if !resp.status().is_success() {
@@ -128,7 +114,7 @@ impl ServerTokenizer {
)));
}
let data: serde_json::Value = resp.json().await.map_err(|e| {
let data: serde_json::Value = resp.json().map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
})?;
@@ -137,26 +123,6 @@ impl ServerTokenizer {
.map(|s| s.to_string())
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
}
fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
if matches!(
self.runtime.runtime_flavor(),
tokio::runtime::RuntimeFlavor::CurrentThread
) {
return Err(BenchError::Tokenizer(
"Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
));
}
// Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
// Tokio workers must enter a blocking region before re-entering the runtime;
// Rayon workers can drive the future directly with the saved runtime handle.
if tokio::runtime::Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
}
// --- TokenizerKind methods ---
@@ -226,7 +192,7 @@ impl TokenizerKind {
/// 3. Server-side /tokenize + /detokenize endpoints
///
/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
pub async fn load_tokenizer(
pub fn load_tokenizer(
model_id: &str,
_trust_remote_code: bool,
server_info: Option<(&str, &str)>,
@@ -246,48 +212,31 @@ pub async fn load_tokenizer(
}
// 1. Try local HuggingFace tokenizer (tokenizer.json)
match try_load_local(model_id).await {
match try_load_local(model_id) {
Ok(tok) => {
tracing::info!(
model = model_id,
kind = "local",
vocab_size = tok.get_vocab_size(true),
"loaded tokenizer"
);
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
Ok(TokenizerKind::Local(Box::new(tok)))
}
Err(local_err) => {
// 2. Try tiktoken format
tracing::info!(
model = model_id,
error = %local_err.as_report(),
"local tokenizer unavailable; trying tiktoken"
);
match crate::tiktoken::try_load_tiktoken(model_id).await {
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
match crate::tiktoken::try_load_tiktoken(model_id) {
Ok(tok) => {
tracing::info!(
model = model_id,
kind = "tiktoken",
vocab_size = tok.vocab_size(),
"loaded tokenizer"
);
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
Ok(TokenizerKind::Tiktoken(tok))
}
Err(tiktoken_err) => {
// 3. Try server-side fallback
if let Some((base_url, model)) = server_info {
tracing::info!(
model = model_id,
error = %tiktoken_err.as_report(),
"tiktoken unavailable; trying server-side tokenization"
println!(
"Tiktoken also not available ({tiktoken_err}), \
trying server-side tokenization..."
);
match ServerTokenizer::new(base_url, model).await {
match ServerTokenizer::new(base_url, model) {
Ok(srv) => {
tracing::info!(
model = model_id,
kind = "server",
vocab_size = srv.cached_vocab_size,
"loaded tokenizer"
println!(
"Tokenizer: Server (vocab_size≈{})",
srv.cached_vocab_size
);
return Ok(TokenizerKind::Server(srv));
}
@@ -315,7 +264,7 @@ pub async fn load_tokenizer(
}
/// Try loading tokenizer.json from local path or HuggingFace Hub.
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
fn try_load_local(model_id: &str) -> Result<Tokenizer> {
// 1. Try local directory with tokenizer.json
let local_path = Path::new(model_id).join("tokenizer.json");
if local_path.exists() {
@@ -341,37 +290,11 @@ async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
}
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
let repo = crate::hub::HubRepo::model(model_id.to_string());
let tokenizer_path = repo
.get("tokenizer.json")
.await
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
Tokenizer::from_file(&tokenizer_path)
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_tokenizer_sync_bridge() {
let tokenizer = std::sync::Arc::new(ServerTokenizer {
client: reqwest::Client::new(),
runtime: tokio::runtime::Handle::current(),
tokenize_url: String::new(),
detokenize_url: String::new(),
model: String::new(),
cached_vocab_size: 0,
});
assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
let (tx, rx) = tokio::sync::oneshot::channel();
rayon::spawn(move || {
let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
});
assert_eq!(rx.await.unwrap().unwrap(), 2);
}
}
+61 -106
View File
@@ -54,7 +54,6 @@ 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};
@@ -89,92 +88,6 @@ 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
@@ -182,7 +95,9 @@ impl ChatRequestProcessor {
/// request semantics such as tool calls.
pub struct ChatLlm {
text: TextLlm,
processor: ChatRequestProcessor,
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
model_dtype: ModelDtype,
/// Tool-call parser selection.
tool_call_parser: ParserSelection,
/// Reasoning parser selection.
@@ -197,7 +112,8 @@ impl ChatLlm {
Self {
text,
processor: ChatRequestProcessor::new(backend, model_dtype),
backend,
model_dtype,
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
}
@@ -224,7 +140,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.processor.model_dtype = Some(model_dtype);
self.model_dtype = model_dtype;
self
}
@@ -256,23 +172,57 @@ impl ChatLlm {
}
/// Render, tokenize, and submit one chat request.
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();
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),
};
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_id, structured_stream))
Ok(ChatEventStream::new(request.request_id, structured_stream))
}
/// Render through the chat template and tokenize, without submitting to the engine.
@@ -283,9 +233,14 @@ impl ChatLlm {
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
request.validate()?;
let rendered = self.processor.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) =
self.processor.finalize_rendered_prompt(&request, rendered).await?;
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 tokenizer = self.text.tokenizer();
let token_ids = match prompt {
+7 -17
View File
@@ -26,21 +26,18 @@ const RESET: &str = "\x1b[0m";
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
format_description!("[month]-[day] [hour]:[minute]:[second]");
const PROCESS_LABEL: &str = "RustFrontend";
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
pub(crate) fn init_tracing(process_label: &str) {
pub(crate) fn init_tracing() {
let filter = build_targets_filter(
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
env::var("RUST_LOG").ok().as_deref(),
);
let formatter = VllmEventFormatter::new(process_label);
let formatter = VllmEventFormatter::new();
let _ = tracing_subscriber::registry()
.with(
tracing_subscriber::fmt::layer()
.event_format(formatter)
.with_writer(std::io::stderr)
.with_filter(filter),
)
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
.try_init();
}
@@ -97,9 +94,9 @@ struct VllmEventFormatter {
}
impl VllmEventFormatter {
fn new(process_label: &str) -> Self {
fn new() -> Self {
Self {
prefix: format!("({process_label} pid={})", process::id()),
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
timer: VllmLocalTimer::default(),
}
}
@@ -294,13 +291,6 @@ fn map_python_log_level(level: &str) -> LevelFilter {
mod tests {
use super::*;
#[test]
fn formatter_prefix_uses_process_label() {
let formatter = VllmEventFormatter::new("Bench");
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
}
#[test]
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
+1 -9
View File
@@ -5,7 +5,6 @@ mod cli;
mod logging;
use std::env;
use std::ffi::OsStr;
use std::process::ExitStatus;
use anyhow::{Context, Result, anyhow, bail};
@@ -83,14 +82,7 @@ fn shutdown_signal() -> CancellationToken {
}
fn main() -> Result<()> {
let process_label =
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
"bench" => "Bench",
"serve" | "frontend" => "RustFrontend",
_ => "Rust",
};
logging::init_tracing(process_label);
logging::init_tracing();
let cli = Cli::parse();
let mut runtime = tokio::runtime::Builder::new_multi_thread();
@@ -460,12 +460,6 @@ impl EngineCoreClient {
self.inner.is_healthy()
}
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
self.inner.subscribe_health()
}
/// Return the first persistent health error observed by the client, if any.
pub fn health_error(&self) -> Option<Arc<Error>> {
self.inner.health_error()
+1 -20
View File
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
use parking_lot::Mutex;
use thiserror_ext::AsReport as _;
use tokio::runtime::Handle;
use tokio::sync::{mpsc, watch};
use tokio::sync::mpsc;
use tracing::{debug, info, trace, warn};
use vllm_metrics::METRICS;
use zeromq::RouterSendHalf;
@@ -36,7 +36,6 @@ pub(crate) struct ClientInner {
request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>,
health_tx: watch::Sender<bool>,
}
impl ClientInner {
@@ -58,7 +57,6 @@ impl ClientInner {
request_reg: Mutex::new(RequestRegistry::new(engines)),
utility_reg: Mutex::new(UtilityRegistry::default()),
health_error: ArcSwapOption::empty(),
health_tx: watch::Sender::new(true),
}
}
@@ -171,7 +169,6 @@ impl ClientInner {
/// persistent health error.
pub fn close_registries(&self, error: Arc<Error>) {
let persistent_error = self.record_health_error(error);
self.publish_unhealthy();
let request_senders = self.request_reg.lock().close();
let utility_senders = self.utility_reg.lock().close();
@@ -194,12 +191,6 @@ impl ClientInner {
self.health_error.load().is_none()
}
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> watch::Receiver<bool> {
self.health_tx.subscribe()
}
/// Resolve one utility output to the waiting caller. Returns `true` if a
/// waiting caller existed.
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
@@ -289,11 +280,6 @@ impl ClientInner {
.expect("health error must be recorded before registries close")
}
/// Publish the sticky healthy-to-unhealthy transition.
fn publish_unhealthy(&self) {
self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
}
/// Assert there is a recorded health error and return a `Shared` variant
/// wrapping it for error returns when the client is already closed.
fn closed_error(&self) -> Error {
@@ -475,18 +461,13 @@ mod tests {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn close_registries_records_first_health_error_only() {
let inner = test_inner().await;
let mut health = inner.subscribe_health();
assert!(*health.borrow());
inner.close_registries(Arc::new(Error::EngineCoreDead));
health.changed().await.expect("health sender remains open");
assert!(!inner.is_healthy());
assert!(!*health.borrow());
assert!(matches!(
inner.health_error().as_deref(),
Some(Error::EngineCoreDead)
));
assert!(!*inner.subscribe_health().borrow());
inner.close_registries(Arc::new(client_closed!("shutdown")));
assert!(matches!(
@@ -269,19 +269,6 @@ impl WireLogprobs {
);
}
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
if token_ids.rows == 0 {
return Ok(Logprobs {
positions: Vec::new(),
});
}
if token_ids.cols == 0 {
bail_ext_value_decode!(
"{field_prefix}: zero-column logprobs payload with {} rows",
token_ids.rows
);
}
let mut positions = Vec::with_capacity(token_ids.rows);
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
.data
@@ -303,49 +303,3 @@ fn rejects_non_none_cu_num_generated_tokens() {
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
);
}
#[test]
fn decodes_zero_row_logprobs_as_empty() {
for shape in [[0usize, 0], [0, 3]] {
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
None,
Some(Value::Array(vec![
ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
Value::Nil,
])),
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0]
.new_prompt_logprobs_tensors
.clone()
.unwrap()
.into_direct()
.unwrap();
assert!(logprobs.is_empty());
}
}
#[test]
fn rejects_zero_column_logprobs_with_rows() {
let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
Some(Value::Array(vec![
ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[2], ranks),
Value::Nil,
])),
None,
)))];
let error = decode_engine_core_outputs(&frames).unwrap_err();
let crate::error::Error::ExtValueDecode { message } = &error else {
panic!("expected ExtValueDecode");
};
assert_eq!(
message,
"new_logprobs: zero-column logprobs payload with 2 rows"
);
}
-1
View File
@@ -35,7 +35,6 @@ tokio-openssl.workspace = true
tokio-stream.workspace = true
tokio-util.workspace = true
tonic.workspace = true
tonic-health.workspace = true
tonic-prost.workspace = true
tower.workspace = true
tower-http.workspace = true
-65
View File
@@ -1,65 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tonic::server::NamedService;
use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::{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;
}
-36
View File
@@ -4,7 +4,6 @@
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
mod convert;
mod health;
use std::pin::Pin;
use std::sync::Arc;
@@ -25,13 +24,8 @@ 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;
@@ -46,36 +40,6 @@ 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 =
+21 -288
View File
@@ -16,10 +16,6 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use tokio::net::TcpStream;
use tokio_openssl::SslStream;
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
use tonic_health::pb::HealthCheckRequest;
use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
use tonic_health::pb::health_client::HealthClient;
use tonic_health::server::health_reporter;
use tower::service_fn;
use vllm_chat::{
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
@@ -38,9 +34,8 @@ 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::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
use super::{GenerateServer, GenerateServiceImpl, pb};
use crate::listener::{Listener, MaybeTlsListener};
use crate::state::AppState;
use crate::tls;
@@ -154,6 +149,10 @@ 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;
@@ -201,12 +200,7 @@ impl ChatRenderer for FakeTextBackend {
async fn setup_grpc_service(
engine_id: impl Into<EngineId>,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (
GenerateServer<GenerateServiceImpl>,
ControlServer<ControlServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {
let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint();
let engine_id = engine_id.into();
@@ -238,17 +232,14 @@ async fn setup_grpc_service(
)
.await
.expect("connect client");
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
Llm::new(client),
test_llm(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.clone())),
ControlServer::new(ControlServiceImpl::new(state)),
engine_health,
GenerateServer::new(GenerateServiceImpl::new(state)),
engine_task,
)
}
@@ -263,56 +254,25 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
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 (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
let server = TonicServer::builder()
.add_service(health_service)
.add_service(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");
TonicServer::builder()
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
});
let channel = Endpoint::from_shared(format!("http://{addr}"))
.expect("grpc endpoint")
.connect()
let grpc_client = GenerateClient::connect(format!("http://{addr}"))
.await
.expect("connect grpc channel");
.expect("connect grpc client");
(channel, server_task)
(grpc_client, server_task, engine_task)
}
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -323,8 +283,7 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -334,8 +293,7 @@ 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(control_service)
.add_service(generate_service)
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc tls server");
@@ -415,8 +373,7 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr").to_string();
@@ -431,8 +388,7 @@ async fn grpc_server_with_keepalive(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
builder
.add_service(control_service)
.add_service(generate_service)
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
@@ -1079,226 +1035,3 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn 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");
}
+14 -32
View File
@@ -39,7 +39,6 @@ use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
use tokio_util::sync::CancellationToken;
use tonic::transport::Server as TonicServer;
use tonic_health::server::health_reporter;
use tower::ServiceExt as _;
use tracing::{info, trace, warn};
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
@@ -204,23 +203,14 @@ where
.map(tls::build_grpc_server_config)
.transpose()
.context("invalid gRPC TLS configuration")?;
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
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 = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service);
.add_service(svc);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
Some((grpc_listener, svc, grpc_tls))
} else {
None
};
@@ -304,8 +294,7 @@ where
let server_shutdown = server_shutdown.clone();
let force_shutdown = force_shutdown.clone();
async move {
let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
else {
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {
// No gRPC configured: just wait for shutdown so we do not race the
// join! by resolving early and tripping the cancellation token.
shutdown.cancelled().await;
@@ -315,26 +304,19 @@ where
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
None => MaybeTlsListener::plain(grpc_listener),
};
let server =
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
let server = async move {
let result = tokio::select! {
result = server => {
result.context("gRPC server failed")
}
_ = force_shutdown.cancelled() => {
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
Ok(())
}
};
server_shutdown.cancel();
result
let result = tokio::select! {
result = server => {
result.context("gRPC server failed")
}
_ = force_shutdown.cancelled() => {
warn!("gRPC graceful shutdown deadline elapsed; aborting server");
Ok(())
}
};
let (result, ()) = tokio::join!(server, health_monitor);
server_shutdown.cancel();
result
}
};
+51 -87
View File
@@ -38,22 +38,32 @@ trait_set! {
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
}
/// Text request preparation shared by inference and render-only frontends.
pub struct TextRequestProcessor {
/// 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,
/// 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 TextRequestProcessor {
/// Create a processor with the effective model context length.
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
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,
backend,
max_model_len,
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
@@ -68,83 +78,9 @@ impl TextRequestProcessor {
self
}
/// Return the tokenizer used by this processor.
pub fn tokenizer(&self) -> DynTokenizer {
self.backend.tokenizer()
}
/// Return the effective model context length.
pub fn max_model_len(&self) -> u32 {
self.max_model_len
}
/// 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() {
request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
}
let tokenizer = self.backend.tokenizer();
let prompt_token_ids = match take(&mut request.prompt) {
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
// Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
// 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,
max_logprobs: self.max_logprobs,
model_vocab_size: self.backend.model_vocab_size(),
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
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()
self.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
@@ -155,19 +91,19 @@ impl TextLlm {
/// Return the tokenizer used by this text backend.
pub fn tokenizer(&self) -> DynTokenizer {
self.processor.tokenizer()
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.processor.backend.tokenizer_vocab_size()
self.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()
self.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
@@ -181,7 +117,7 @@ impl TextLlm {
/// 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 tokenizer = self.backend.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
@@ -195,12 +131,40 @@ impl TextLlm {
async fn generate_inner(
&self,
request: TextRequest,
mut request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
request.validate()?;
if request.arrival_time.is_none() {
request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
}
let tokenizer = self.backend.tokenizer();
let prompt_token_ids = match take(&mut request.prompt) {
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
// Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
// 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,
max_logprobs: self.max_logprobs,
model_vocab_size: self.backend.model_vocab_size(),
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
let PreparedTextRequest {
text_request,
generate_request,
} = self.processor.prepare(request)?;
} = lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
&*tokenizer,
)?;
let raw_stream = self.llm.generate(generate_request).await?;
Ok((text_request, raw_stream))
-62
View File
@@ -214,65 +214,3 @@ 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"
)
-9
View File
@@ -43,15 +43,6 @@ 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")
-30
View File
@@ -349,20 +349,6 @@ _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
@@ -488,22 +474,6 @@ 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(
+131
View File
@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from types import ModuleType, SimpleNamespace
from typing import Any
import pytest
from vllm.platforms.interface import DeviceCapability
@pytest.fixture
def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
monkeypatch.setitem(
sys.modules,
"vllm._C_stable_libtorch",
stable_libtorch_module,
)
from vllm.platforms.cuda import CudaPlatformBase
return CudaPlatformBase
def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
assert cuda_platform_base._compiled_arch_covers_device(
"12.1a", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"12.0a", DeviceCapability(12, 1)
)
assert cuda_platform_base._compiled_arch_covers_device(
"12.0f", DeviceCapability(12, 1)
)
assert not cuda_platform_base._compiled_arch_covers_device(
"10.0f", DeviceCapability(12, 1)
)
def test_warn_if_device_arch_not_compiled(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 2
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
capabilities = {
0: DeviceCapability(12, 1),
1: DeviceCapability(10, 3),
}
return capabilities[device_id]
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
return f"GPU {device_id}"
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(
_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert len(warnings) == 1
assert warnings[0][1] == "12.0f, 10.0a"
assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
def test_warn_if_device_arch_not_compiled_no_warning(
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
) -> None:
def device_count(cls: type[Any]) -> int:
return 1
def get_device_capability(
cls: type[Any], device_id: int = 0
) -> DeviceCapability | None:
return DeviceCapability(10, 3)
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
raise AssertionError("get_device_name should not be called for covered devices")
warnings: list[tuple[str, str, str]] = []
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
warnings.append((message, compiled_archs, devices))
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
monkeypatch.setattr(
cuda_platform_base,
"get_device_capability",
classmethod(get_device_capability),
)
monkeypatch.setattr(
cuda_platform_base, "get_device_name", classmethod(get_device_name)
)
monkeypatch.setattr(
"vllm.platforms.cuda.torch.ops",
SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
)
monkeypatch.setattr(
"vllm.platforms.cuda.logger",
SimpleNamespace(warning_once=warning_once),
)
cuda_platform_base._warn_if_device_arch_not_compiled()
assert warnings == []

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