Compare commits

..
Author SHA1 Message Date
Kevin H. Luu 568afb3a13 [CI/Build] Refresh tags before building macOS wheel (#49901)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
(cherry picked from commit 0934b26790)
2026-07-26 17:57:50 -07:00
TJianandkhluu f2654939e6 [ROCm] [Release] [Bugfix] Fix the per commit wheel release pipeline. (#49245)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-07-25 00:23:46 -07:00
djramicandkhluu ffd46bfab2 [Bugfix] Register axk1 config to fix A.X-K1 init (#49727)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
(cherry picked from commit e222c33f2f)
2026-07-24 18:48:21 -07:00
Andrey Talmanandkhluu ffd6ee4bcc [CI] Bump PyTorch Compilation Unit Tests timeout to 150 min (#49606) 2026-07-23 21:20:25 -07:00
Nick Hillandkhluu bb26ce8e93 [CI] Increase timeout of pytorch-compilation-unit-tests (#49450)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-23 21:20:25 -07:00
Kevin H. LuuandOpenAI Codex 091db8b58f [CI] Increase timeouts for jobs exceeding current limits (#49374)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 21:20:25 -07:00
Nick Hillandkhluu ba694b86f2 [CI] Bump timeout of entrypoints-integration-api-server-openai-part-2 (#49359)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-23 21:20:25 -07:00
zhrrrandkhluu e5949f1000 [Bugfix] handle grammar compilation failures to avoid engine crash (#47312)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
(cherry picked from commit 12213c6795)
2026-07-23 11:13:25 -07:00
Michael Goinandkhluu 8b30569e83 [Bugfix] Fix DeepGEMM warmup when using FlashInferFp8DeepGEMMDynamicBlockScaledKernel (#49467)
Signed-off-by: mgoin <mgoin64@gmail.com>
(cherry picked from commit 917fdb5bf7)
2026-07-23 11:13:25 -07:00
Lucas Wilkinsonandkhluu 9d37a50c80 [Bugfix][Attention] Ignore empty MLA context chunks during merge (#49294)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
(cherry picked from commit 060b5f61dc)
2026-07-23 11:13:25 -07:00
aoshen02andkhluu 2dd1e7cd3b Update BGE-M3 token expectations for leading spaces (#49269)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
(cherry picked from commit d9aa35161d)
2026-07-23 11:13:25 -07:00
Alejandro Paredes La Torreandkhluu a54c93a146 [Bugfix] Fix WSL circular import from pin_memory warning_once (#48444)
Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
(cherry picked from commit 0a684ab0c0)
2026-07-23 11:13:25 -07:00
421 changed files with 4809 additions and 14411 deletions
+1 -5
View File
@@ -18,8 +18,6 @@ steps:
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
- tests/kernels/mamba/test_cpu_short_conv.py
- tests/kernels/mamba/test_causal_conv1d.py
- tests/kernels/mamba/test_mamba_ssm.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -30,9 +28,7 @@ steps:
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
# Note: SDE can't be downloaded from CI host because of AWS WAF
# - label: CPU-Compatibility Tests
+3
View File
@@ -7,6 +7,9 @@
set -euo pipefail
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
git fetch --tags --force origin
# The Rust frontend build needs protoc.
if ! command -v protoc >/dev/null 2>&1; then
brew install protobuf
+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
@@ -40,9 +40,7 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
@@ -99,4 +97,3 @@ function cpu_tests() {
# All of CPU tests are expected to be finished less than 40 mins.
export -f cpu_tests
timeout 2h bash -c cpu_tests
-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
-3
View File
@@ -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
@@ -10,9 +10,7 @@ steps:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
@@ -36,9 +34,7 @@ steps:
- vllm/engine/arg_utils.py
- vllm/config/model.py
- vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
-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:
-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(', ')}`);
+2 -2
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
+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:
-3
View File
@@ -430,7 +430,6 @@ set(VLLM_EXT_SRC
"csrc/cpu/layernorm.cpp"
"csrc/cpu/mla_decode.cpp"
"csrc/cpu/pos_encoding.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
@@ -490,7 +489,6 @@ if (ENABLE_X86_ISA)
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
"csrc/cpu/activation.cpp"
@@ -504,7 +502,6 @@ if (ENABLE_X86_ISA)
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
+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 ""
+5 -5
View File
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
else()
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _qutlass_fc_root)
@@ -125,6 +125,8 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
@@ -137,11 +139,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_MINIMAL_BUILD=1
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
TORCH_TARGET_VERSION=0x020B000000000000ULL
USE_CUDA)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+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) {
+7 -8
View File
@@ -336,14 +336,13 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
reg.val[1] = fp16_to_fp32_bits(raw_lo);
}
float reduce_sum() const {
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
// Step 1: pairwise sum of the two 4-wide halves
__vector float s = vec_add(reg.val[0], reg.val[1]);
// Step 2: rotate by 8 bytes (2 floats) and add
s = vec_add(s, vec_sld(s, s, 8));
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
s = vec_add(s, vec_sld(s, s, 4));
return vec_extract(s, 0);
AliasReg ar;
ar.reg = reg;
float result = 0;
unroll_loop<int, VEC_ELEM_NUM>(
[&result, &ar](int i) { result += ar.values[i]; });
return result;
}
FP32Vec8 exp() const {
f32x4x2_t out;
-285
View File
@@ -1,285 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
// mamba_kernels.hpp.
#include "cpu/mamba_kernels.hpp"
#include <ATen/ATen.h>
#include <torch/library.h>
#include <c10/util/Optional.h>
#include "cpu_types.hpp"
// ---------------------------------------------------------------------------
// causal_conv1d_update
// ---------------------------------------------------------------------------
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
bool do_silu = false;
if (activation.has_value()) {
const std::string& act = activation.value();
do_silu = (act == "silu" || act == "swish");
}
at::ScalarType dtype = x.scalar_type();
// Input x: contiguous in native dtype.
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
// conv_state: NEVER copy the full paged tensor just for layout reasons.
// If the dtype matches we work directly on conv_state (contiguous or not)
// by extracting strides and passing them to the kernel.
// Only a dtype-conversion copy is made when types differ (rare for BF16).
bool state_type_ok = (conv_state.scalar_type() == dtype);
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
// state_c and conv_state may be non-contiguous — that is intentional.
// Weight: coerce to same dtype if needed (should match in practice)
at::Tensor w_c =
(weight.scalar_type() != dtype)
? weight.to(dtype).contiguous()
: (weight.is_contiguous() ? weight : weight.contiguous());
// Bias stays float32 (small scalar, used only for fp32 accumulation)
at::Tensor bias_f32;
if (bias.has_value() && bias.value().defined())
bias_f32 = bias.value().to(at::kFloat).contiguous();
int64_t batch = x_c.size(0);
int64_t dim = x_c.size(1);
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
int64_t width = w_c.size(1);
int64_t state_len = state_c.size(2);
// Extract strides — works for contiguous AND non-contiguous (transposed)
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
// contiguous) stride(1): between conv channels (dim stride) stride(2):
// between state elements (=1 when contiguous, =dim when transposed)
int64_t stride_s_slot = state_c.stride(0);
int64_t stride_s_dim = state_c.stride(1);
int64_t stride_s_state = state_c.stride(2);
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
const int32_t* cache_idx_ptr = nullptr;
at::Tensor cache_idx_int;
if (conv_state_indices.has_value()) {
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
}
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<scalar_t>(), cache_idx_ptr,
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
do_silu);
});
// Write back only when a type-conversion copy was made.
// Layout-only non-contiguity is handled via strides above — no copy needed.
if (!state_type_ok) conv_state.copy_(state_c);
return out;
}
// ---------------------------------------------------------------------------
// selective_state_update
// ---------------------------------------------------------------------------
void selective_state_update_cpu_impl(
at::Tensor& state, // (nstates, nheads, dim, dstate)
const at::Tensor& x, // (N, nheads, dim)
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C, const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens) {
at::ScalarType state_type = state.scalar_type();
at::ScalarType input_type = x.scalar_type();
// x, B, C must be contiguous and match input_type
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_input(x);
at::Tensor B_in = ensure_input(B);
at::Tensor C_in = ensure_input(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
// A, D, dt_bias are float32 model parameters that arrive here as expanded
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
// We need just the scalar value per head as a (nheads,) 1-D array so that
// A_ptr[h] in the kernel correctly reads head h's value.
//
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
// → .select(2,0) → (nheads, head_dim) strides (1,0)
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
// No allocation, no type conversion (A is already float32).
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
at::Tensor D_f32, dt_bias_f32;
if (D.has_value() && D.value().defined())
D_f32 = to_per_head_1d_f32(D.value());
if (dt_bias.has_value() && dt_bias.value().defined())
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
// the type conversion so we convert head_dim x fewer elements.
at::Tensor dt_f32;
{
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
// take a zero-copy view of index 0 along that dim first.
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
}
int64_t nheads = state.size(1);
int64_t dim = state.size(2);
int64_t dstate = state.size(3);
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
? cu_seqlens.value().size(0) - 1
: x_in.size(0);
int64_t ngroups = B_in.size(1);
// Strides
int64_t stride_state_n = state.stride(0);
int64_t stride_state_h = state.stride(1);
int64_t stride_state_d = state.stride(2);
int64_t stride_x_n = x_in.stride(0);
int64_t stride_x_h = x_in.stride(1);
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
int64_t stride_BC_n = B_in.stride(0);
int64_t stride_BC_g = B_in.stride(1);
int64_t stride_out_n = out.stride(0);
int64_t stride_out_h = out.stride(1);
// Optional index pointers
auto get_int32_ptr =
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
return (opt.has_value() && opt.value().defined())
? opt.value().data_ptr<int32_t>()
: nullptr;
};
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
// Dispatch on (state_t, input_t, out_t): write directly into `out`
// without any intermediate float32 buffer.
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
using state_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
using input_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
using out_t = scalar_t;
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
nheads, ngroups, dim, dstate, dt_softplus);
});
});
});
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd_cpu
// ---------------------------------------------------------------------------
void mamba_chunk_scan_fwd_cpu_impl(
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
at::Tensor&
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
const at::Tensor& x, // [seqlen, nheads, headdim]
const at::Tensor&
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
const at::Tensor& A, // [nheads] float32
const at::Tensor& B, // [seqlen, ngroups, dstate]
const at::Tensor& C, // [seqlen, ngroups, dstate]
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
const at::Tensor& cu_seqlens // [batch+1] int32
) {
const at::ScalarType input_type = x.scalar_type();
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_contig(x);
at::Tensor B_in = ensure_contig(B);
at::Tensor C_in = ensure_contig(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
// A and D are float32 model parameters, potentially broadcast-expanded.
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_f32(A);
at::Tensor D_f32;
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
// Python.
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
const int64_t batch = final_states.size(0);
const int64_t nheads = final_states.size(1);
const int64_t headdim = final_states.size(2);
const int64_t dstate = final_states.size(3);
const int64_t ngroups = B_in.size(1);
TORCH_CHECK(final_states.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
TORCH_CHECK(out.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
"raw data_ptr)");
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
ngroups, headdim, dstate);
});
}
-382
View File
@@ -1,382 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Fused CPU vector kernels for Mamba decode-step hotspots:
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
// - selective_state_update (SSM recurrence, single-step)
#pragma once
#include "cpu_types.hpp"
#include <cmath>
#include <cstring>
#include <cstdint>
#include <algorithm>
namespace mamba_cpu {
// ---------------------------------------------------------------------------
// causal_conv1d_update — templated for native BF16/FP32
//
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
// Explicit strides are passed so the kernel writes directly into the
// correct memory locations without making a contiguous copy of the full
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
//
// stride_s_slot = state.stride(0) — between cache slots
// stride_s_dim = state.stride(1) — between conv_dim channels
// stride_s_state = state.stride(2) — between state elements
//
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
// ---------------------------------------------------------------------------
template <typename scalar_t>
inline void causal_conv1d_update_kernel(
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
int64_t width, int64_t state_len, bool do_silu) {
#pragma omp parallel for
for (int64_t b = 0; b < batch; ++b) {
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
if (cache_idx == pad_slot_id) continue;
for (int64_t t = 0; t < seqlen; ++t) {
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
// Base of this slot in the (possibly non-contiguous) paged state
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_b[d * seqlen]);
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
const scalar_t* w = weight_ptr + d * width;
// Accumulate in float32 for precision
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
for (int64_t k = 0; k < state_len; ++k) {
acc += static_cast<float>(w[k]) *
static_cast<float>(sd[k * stride_s_state]);
}
acc += static_cast<float>(w[state_len]) * x_val;
// Shift state left and append new input.
// Use memmove when contiguous (stride==1); element loop otherwise.
if (stride_s_state == 1) {
if (state_len > 1)
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
} else {
for (int64_t k = 0; k < state_len - 1; ++k)
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
if (state_len > 0)
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
}
if (do_silu) {
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
: std::exp(acc) / (1.0f + std::exp(acc));
acc *= sigmoid;
}
out_b[d * seqlen] = static_cast<scalar_t>(acc);
}
}
}
}
// ---------------------------------------------------------------------------
// selective_state_update
//
// Template parameters:
// state_t - dtype of ssm_state cache (typically BFloat16)
// input_t - dtype of x, B, C (typically BFloat16)
// out_t - dtype of output tensor (typically BFloat16)
// Write directly — no float32 intermediate buffer needed.
//
// A, D, dt_bias are accepted as const float* (they are always float32
// model parameters in Mamba2). This eliminates the per-call float32→BF16
// conversion and the .contiguous() materialisation of the broadcast-expand.
//
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
// ---------------------------------------------------------------------------
template <typename state_t, typename input_t, typename out_t = float>
inline void selective_state_update_kernel(
state_t* __restrict__ state_ptr, int64_t stride_state_n,
int64_t stride_state_h, int64_t stride_state_d,
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
// A: (nheads,) float32 — scalar per head
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
// D: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ D_ptr,
// z: same shape as x (optional)
const input_t* __restrict__ z_ptr,
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
int64_t stride_out_n, int64_t stride_out_h,
const int32_t* __restrict__ state_batch_indices,
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
const int32_t* __restrict__ num_accepted_tokens,
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
using state_vec_t = vec_op::vec_t<state_t>;
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
int64_t nheads_per_group = nheads / ngroups;
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
int64_t bos, seq_len;
if (cu_seqlens != nullptr) {
bos = cu_seqlens[seq_idx];
seq_len = cu_seqlens[seq_idx + 1] - bos;
} else {
bos = seq_idx;
seq_len = 1;
}
int64_t state_read_idx = (state_batch_indices != nullptr)
? state_batch_indices[seq_idx]
: seq_idx;
if (state_read_idx == null_block_id) continue;
int64_t state_write_idx = (num_accepted_tokens == nullptr)
? ((dst_state_batch_indices != nullptr)
? dst_state_batch_indices[seq_idx]
: state_read_idx)
: -1;
state_t* s = state_ptr + state_read_idx * stride_state_n;
for (int64_t t = 0; t < seq_len; ++t) {
int64_t token_idx = bos + t;
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
// dt: (N, nheads) — one float per head per token
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
out_t* out_tok = out_ptr + token_idx * stride_out_n;
#pragma omp parallel for
for (int64_t h = 0; h < nheads; ++h) {
int64_t g = h / nheads_per_group;
const input_t* x_h = x_tok + h * stride_x_h;
const input_t* B_g = B_tok + g * stride_BC_g;
const input_t* C_g = C_tok + g * stride_BC_g;
out_t* out_h = out_tok + h * stride_out_h;
state_t* s_h = s + h * stride_state_h;
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
float dt_val = dt_tok[h];
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
if (dt_softplus) {
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
}
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
const input_t* z_h =
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
: nullptr;
vec_op::FP32Vec8 dt_vec(dt_val);
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
// and broadcast. This saves 7 redundant std::exp() calls that
// FP32Vec8::exp() would otherwise make on the broadcast vector.
const float dA_scalar = std::exp(A_val * dt_val);
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_h[d]);
vec_op::FP32Vec8 out_vec(0.0f);
state_t* s_hd = s_h + d * stride_state_d;
const input_t* B_g_base = B_g;
const input_t* C_g_base = C_g;
vec_op::FP32Vec8 x_vec(x_val);
// dBx = B * x * dt — same dA for all dstate (A is scalar)
// s_new = s * dA + B * x * dt
int64_t n = 0;
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
state_vec_t(s_new).save(s_hd + n);
out_vec = out_vec + s_new * C_v;
}
float out_val = out_vec.reduce_sum();
for (; n < dstate; ++n) {
// Reuse dA_scalar computed once per head — no exp() re-call
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
s_hd[n] = static_cast<state_t>(s_new);
out_val += s_new * static_cast<float>(C_g[n]);
}
if (D_ptr != nullptr) out_val += x_val * D_val;
if (z_h != nullptr) {
float z_val = static_cast<float>(z_h[d]);
float sigmoid = (z_val >= 0)
? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
out_val *= z_val * sigmoid;
}
out_h[d] = static_cast<out_t>(out_val);
}
}
if (num_accepted_tokens != nullptr &&
dst_state_batch_indices != nullptr) {
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
state_write_idx != state_read_idx) {
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd
//
// Prefill SSM recurrence for Mamba2 / SSD models.
//
// Key difference from selective_state_update_kernel (decode path):
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
// Each thread owns a (batch, head) slice and runs the entire token
// sequence without any per-token OpenMP synchronisation overhead.
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
//
// `dt` arrives already processed (float32, after bias + softplus + clamp)
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
//
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
// tensor, pre-initialised by the caller (zero or from initial_states).
// Each (b, h) slice is private to exactly one thread via collapse(2), so
// there are no write conflicts.
//
// D is treated as a scalar per head ([nheads] float32).
// ---------------------------------------------------------------------------
template <typename input_t>
inline void mamba_chunk_scan_fwd_kernel(
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
const float* __restrict__ A_ptr, // [nheads] f32
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
int64_t dstate) {
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
const int64_t nheads_per_group = nheads / ngroups;
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
// guarantee)
const int64_t stride_s_b = nheads * headdim * dstate;
const int64_t stride_s_h = headdim * dstate;
// stride_s_d = dstate, stride_s_n = 1
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t b = 0; b < batch; ++b) {
for (int64_t h = 0; h < nheads; ++h) {
const int64_t seq_start = cu_seqlens[b];
const int64_t seq_end = cu_seqlens[b + 1];
const int64_t g = h / nheads_per_group;
const float A_val = A_ptr[h];
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
for (int64_t t = seq_start; t < seq_end; ++t) {
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
const float* dt_h = dt_ptr + t * nheads + h;
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
const input_t* z_h = (z_ptr != nullptr)
? z_ptr + t * nheads * headdim + h * headdim
: nullptr;
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
const float dt_val = *dt_h;
const float dA_val = std::exp(A_val * dt_val);
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
const vec_op::FP32Vec8 dt_vec(dt_val);
for (int64_t d = 0; d < headdim; ++d) {
const float x_val = static_cast<float>(x_h[d]);
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
// Vectorised SSM update + readout over dstate:
// s_new = s * dA + x * dt * B
// y += s_new * C
int64_t n = 0;
vec_op::FP32Vec8 y_vec(0.0f);
const vec_op::FP32Vec8 x_vec(x_val);
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
const vec_op::FP32Vec8 s_v(s_bhd + n);
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
s_new.save(s_bhd + n);
y_vec = y_vec + s_new * C_v;
}
float y_val = y_vec.reduce_sum();
// Scalar tail for remaining dstate elements
for (; n < dstate; ++n) {
const float B_n = static_cast<float>(B_g[n]);
const float C_n = static_cast<float>(C_g[n]);
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
s_bhd[n] = s_new;
y_val += s_new * C_n;
}
// D skip connection (scalar per head)
if (D_ptr != nullptr) y_val += x_val * D_val;
// z gating: out = y * z * sigmoid(z) (SiLU)
if (z_h != nullptr) {
const float z_val = static_cast<float>(z_h[d]);
const float sigmoid =
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
y_val *= z_val * sigmoid;
}
out_h[d] = static_cast<input_t>(y_val);
}
}
}
}
}
} // namespace mamba_cpu
-50
View File
@@ -213,32 +213,6 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
void selective_state_update_cpu_impl(
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens);
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C,
const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const at::Tensor& cu_seqlens);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
@@ -621,30 +595,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
// Mamba CPU kernels
ops.def(
"causal_conv1d_update_cpu_vec("
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
&causal_conv1d_update_cpu_impl);
ops.def(
"selective_state_update_cpu("
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
"SymInt null_block_id, Tensor(a13!) out, "
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
&selective_state_update_cpu_impl);
ops.def(
"mamba_chunk_scan_fwd_cpu("
"Tensor(a0!) out, Tensor(a1!) final_states, "
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
&mamba_chunk_scan_fwd_cpu_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
-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
@@ -39,15 +39,11 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
// AWQ zeros: (size_k // group_size, size_n // 8)
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
int32_t group_size) {
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
// same row (stride 1) instead of striding across rows (stride size_n/8).
int col = blockIdx.y * 32 + threadIdx.x;
if (col >= size_n / 8) return;
(void)size_k;
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
int32_t val =
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
int32_t zero =
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
blockIdx.y];
int32_t new_val = 0;
#pragma unroll
@@ -62,7 +58,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
zero >>= 4;
}
output[blockIdx.x * (size_n / 8) + col] = new_val;
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
}
torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -106,7 +102,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
"qweight.size(0) % qzeros.size(0) != 0");
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
dim3 blocks(size_k / 32, size_n / 8);
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
+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"
+2 -2
View File
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
+2 -2
View File
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
### SiLU+Mul + Quantization (`fuse_act_quant`)
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
### RMSNorm + Padding (`fuse_act_padding`)
+3 -4
View File
@@ -68,14 +68,13 @@ vllm serve <model> \
| --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
@@ -180,7 +179,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
+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]
@@ -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)}();
@@ -31,8 +31,10 @@
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
- `response_format`: Format of the response ("json", "text") (optional)
- `temperature`: Sampling temperature between 0 and 1 (optional)
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
**Response Format:**
-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
+11 -29
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",
]
@@ -5576,7 +5560,6 @@ dependencies = [
"tracing",
"tracing-subscriber",
"uuid",
"vllm-bench",
"vllm-chat",
"vllm-engine-core-client",
"vllm-managed-engine",
@@ -5746,7 +5729,6 @@ dependencies = [
"tokio-stream",
"tokio-util",
"tonic",
"tonic-health",
"tonic-prost",
"tonic-prost-build",
"tower",
@@ -6338,9 +6320,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
[[package]]
name = "xgrammar-structural-tag"
version = "0.2.0+xgrammar.0.2.4.dd729e7"
version = "0.1.0+xgrammar.0.2.2.4d145cc"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
dependencies = [
"auto_impl",
"serde",
+5 -7
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"] }
@@ -133,7 +132,6 @@ trait-set = "0.3.0"
url = "2.5.7"
uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
@@ -144,7 +142,7 @@ vllm-server = { path = "src/server" }
vllm-text = { path = "src/text" }
vllm-tokenizer = { path = "src/tokenizer" }
winnow = { version = "1.0.2", features = ["simd"] }
xgrammar-structural-tag = "0.2.0"
xgrammar-structural-tag = "0.1.0"
zeromq = { version = "0.6.0", default-features = false, features = [
"tokio-runtime",
"all-transport",
-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);
}
+11 -4
View File
@@ -3,6 +3,8 @@
use std::fmt;
use clap::Parser;
/// Backend type for the benchmark endpoint.
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendKind {
@@ -75,7 +77,7 @@ pub enum DatasetName {
ShareGpt,
#[value(name = "sonnet")]
Sonnet,
#[value(name = "speed-bench", alias = "speed_bench")]
#[value(name = "speed-bench")]
SpeedBench,
#[value(name = "hf")]
Hf,
@@ -142,8 +144,13 @@ impl fmt::Display for SpeedBenchConfig {
}
/// High-performance benchmark client for vLLM serving endpoints.
#[derive(clap::Args, Debug, Clone)]
pub struct BenchServeArgs {
#[derive(Parser, Debug, Clone)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
pub struct Cli {
/// The type of backend or endpoint to use for the benchmark.
#[arg(long, default_value = "openai")]
pub backend: BackendKind,
@@ -652,7 +659,7 @@ pub struct BenchServeArgs {
pub lora_assignment: LoraAssignment,
}
impl BenchServeArgs {
impl Cli {
/// Resolve the base URL from explicit --base-url or from --host/--port.
pub fn resolve_base_url(&self) -> String {
if let Some(ref base) = self.base_url {
+197 -230
View File
@@ -4,9 +4,7 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
use crate::error::{BenchError, Result};
@@ -217,63 +215,63 @@ pub struct BenchConfig {
}
impl BenchConfig {
pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
if args.burstiness <= 0.0 {
pub fn from_cli(cli: &Cli) -> Result<Self> {
if cli.burstiness <= 0.0 {
return Err(BenchError::Config("Burstiness must be positive".into()));
}
if args.num_prompts == 0 {
if cli.num_prompts == 0 {
return Err(BenchError::Config(
"--num-prompts must be at least 1".into(),
));
}
if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
return Err(BenchError::Config(
"--request-rate must be positive (or inf)".into(),
));
}
if args.max_model_len == Some(0) {
if cli.max_model_len == Some(0) {
return Err(BenchError::Config(
"--max-model-len must be at least 1".into(),
));
}
let base_url = args.resolve_base_url();
let api_url = args.resolve_api_url();
let base_url = cli.resolve_base_url();
let api_url = cli.resolve_api_url();
let extra_headers = args.parse_headers()?;
let mut extra_body = args.parse_extra_body()?;
let extra_headers = cli.parse_headers()?;
let mut extra_body = cli.parse_extra_body()?;
// Merge sampling parameters into extra_body (matches Python behavior).
// Python collects non-None sampling params and merges them UNDER extra_body,
// meaning extra_body keys take precedence over sampling params.
{
let mut sampling_params = serde_json::Map::new();
if let Some(v) = args.top_p {
if let Some(v) = cli.top_p {
sampling_params.insert("top_p".into(), serde_json::json!(v));
}
if let Some(v) = args.top_k {
if let Some(v) = cli.top_k {
sampling_params.insert("top_k".into(), serde_json::json!(v));
}
if let Some(v) = args.min_p {
if let Some(v) = cli.min_p {
sampling_params.insert("min_p".into(), serde_json::json!(v));
}
if let Some(v) = args.temperature {
if let Some(v) = cli.temperature {
sampling_params.insert("temperature".into(), serde_json::json!(v));
}
if let Some(v) = args.frequency_penalty {
if let Some(v) = cli.frequency_penalty {
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
}
if let Some(v) = args.presence_penalty {
if let Some(v) = cli.presence_penalty {
sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
}
if let Some(v) = args.repetition_penalty {
if let Some(v) = cli.repetition_penalty {
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
}
if !sampling_params.is_empty() {
if !args.backend.is_openai_compatible() {
if !cli.backend.is_openai_compatible() {
return Err(BenchError::Config(
"Sampling parameters are only supported by openai-compatible backends."
.into(),
@@ -288,18 +286,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,
@@ -309,7 +299,7 @@ impl BenchConfig {
}
// Parse metadata
let metadata = match &args.metadata {
let metadata = match &cli.metadata {
None => None,
Some(items) => {
let mut pairs = Vec::new();
@@ -324,24 +314,24 @@ impl BenchConfig {
};
// Parse goodput SLOs
let goodput = parse_goodput(&args.goodput)?;
let goodput = parse_goodput(&cli.goodput)?;
// Parse ramp-up config
let ramp_up = parse_ramp_up(args)?;
let ramp_up = parse_ramp_up(cli)?;
// Default percentile metrics based on backend type
let default_percentile_metrics = if args.backend.is_pooling() {
let default_percentile_metrics = if cli.backend.is_pooling() {
"e2el"
} else {
"ttft,tpot,itl,e2el"
};
let percentile_metrics_str =
args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
let selected_percentile_metrics: Vec<String> =
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
let sweep_summary_percentiles = args
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
let sweep_summary_percentiles = cli
.sweep_summary_percentiles
.as_deref()
.map(|raw| parse_percentiles(raw, true))
@@ -354,38 +344,38 @@ impl BenchConfig {
selected_percentiles.push(90.0);
}
let tokenizer_id = if args.skip_tokenizer_init {
let tokenizer_id = if cli.skip_tokenizer_init {
None
} else {
args.tokenizer.clone().or_else(|| args.model.clone())
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
};
// Resolve input/output lengths
let random_input_len = args.resolved_random_input_len();
let random_output_len = args.resolved_random_output_len();
let per_turn_input_len = args.resolved_per_turn_input_len();
let random_input_len = cli.resolved_random_input_len();
let random_output_len = cli.resolved_random_output_len();
let per_turn_input_len = cli.resolved_per_turn_input_len();
// Normalized multi-turn turn counts (computed in validation block below, defaults
// to num_turns if multi-turn mode is not active)
let mut multi_turn_min_turns = args.multi_turn_num_turns;
let mut multi_turn_max_turns = args.multi_turn_num_turns;
let mut multi_turn_min_turns = cli.multi_turn_num_turns;
let mut multi_turn_max_turns = cli.multi_turn_num_turns;
// For random datasets with openai-compatible backends, default to ignore_eos.
// Exception: multi-turn mode, where ignore_eos causes unbounded context growth
// across turns. Multi-turn uses min_tokens instead for output length control.
// Pooling backends don't generate tokens, so ignore_eos is irrelevant.
let ignore_eos = if args.backend.is_pooling() {
let ignore_eos = if cli.backend.is_pooling() {
false
} else {
args.ignore_eos
|| ((args.dataset_name == DatasetName::Random
|| args.dataset_name == DatasetName::RandomMm)
&& args.backend.is_openai_compatible()
&& !args.multi_turn)
cli.ignore_eos
|| ((cli.dataset_name == DatasetName::Random
|| cli.dataset_name == DatasetName::RandomMm)
&& cli.backend.is_openai_compatible()
&& !cli.multi_turn)
};
// Pooling backends don't support multi-turn
if args.backend.is_pooling() && args.multi_turn {
if cli.backend.is_pooling() && cli.multi_turn {
return Err(BenchError::Config(
"Pooling/embedding backends do not support --multi-turn".into(),
));
@@ -393,7 +383,7 @@ impl BenchConfig {
// LoRA validation. Adapter names must be non-empty after trim; pooling
// backends are out of scope (vLLM LoRA routing is for generative paths).
let lora_modules = match args.lora_modules.as_ref() {
let lora_modules = match cli.lora_modules.as_ref() {
None => None,
Some(names) => {
if names.is_empty() {
@@ -401,7 +391,7 @@ impl BenchConfig {
"--lora-modules requires at least one adapter name".into(),
));
}
if args.backend.is_pooling() {
if cli.backend.is_pooling() {
return Err(BenchError::Config(
"--lora-modules is not supported for pooling/embedding backends".into(),
));
@@ -421,18 +411,18 @@ impl BenchConfig {
};
// Random-MM validation and config parsing
let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
if args.backend != BackendKind::OpenaiChat {
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
if cli.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"Multi-modal content (images) is only supported on 'openai-chat' backend."
.into(),
));
}
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
&args.random_mm_limit_mm_per_prompt,
&cli.random_mm_limit_mm_per_prompt,
)?;
let buckets =
crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
(limit, buckets)
} else {
(MmLimitPerPrompt::default(), Vec::new())
@@ -442,18 +432,18 @@ impl BenchConfig {
// sonnet (uses built-in Shakespeare's sonnets).
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
// Batched inputs only make sense for pooling backends (the generation
// backends send one prompt per request).
if args.random_batch_size == 0 {
if cli.random_batch_size == 0 {
return Err(BenchError::Config(
"--random-batch-size must be at least 1".into(),
));
}
if args.random_batch_size > 1
&& !args.backend.is_pooling()
&& args.dataset_name != DatasetName::RandomRerank
if cli.random_batch_size > 1
&& !cli.backend.is_pooling()
&& cli.dataset_name != DatasetName::RandomRerank
{
return Err(BenchError::Config(
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -461,16 +451,16 @@ impl BenchConfig {
}
// random-rerank validation (mirrors Python RandomDatasetForReranking)
let is_reranker = !args.no_reranker;
if args.dataset_name == DatasetName::RandomRerank {
if !args.backend.is_pooling() {
let is_reranker = !cli.no_reranker;
if cli.dataset_name == DatasetName::RandomRerank {
if !cli.backend.is_pooling() {
return Err(BenchError::Config(
"--dataset-name random-rerank requires an embeddings/pooling backend \
(e.g. --backend vllm-rerank)"
.into(),
));
}
if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
return Err(BenchError::Config(
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
(the query is folded into the first batch slot)"
@@ -480,8 +470,8 @@ impl BenchConfig {
}
// Custom dataset validation
if args.dataset_name == DatasetName::Custom {
match args.dataset_path.as_deref() {
if cli.dataset_name == DatasetName::Custom {
match cli.dataset_path.as_deref() {
None => {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name custom \
@@ -496,38 +486,38 @@ impl BenchConfig {
}
_ => {}
}
if !args.skip_chat_template {
tracing::warn!(
dataset = "custom",
"client-side chat template rendering is unsupported; sending prompts raw"
if !cli.skip_chat_template {
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
);
}
}
// Prefix repetition validation
if args.dataset_name == DatasetName::PrefixRepetition {
if args.prefix_repetition_num_prefixes == 0 {
if cli.dataset_name == DatasetName::PrefixRepetition {
if cli.prefix_repetition_num_prefixes == 0 {
return Err(BenchError::Config(
"--prefix-repetition-num-prefixes must be at least 1".into(),
));
}
if args.num_prompts < args.prefix_repetition_num_prefixes {
if cli.num_prompts < cli.prefix_repetition_num_prefixes {
return Err(BenchError::Config(format!(
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
args.num_prompts, args.prefix_repetition_num_prefixes
cli.num_prompts, cli.prefix_repetition_num_prefixes
)));
}
}
// HF dataset validation
if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name hf \
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
.into(),
));
}
if let Some(len) = args.hf_output_len
if let Some(len) = cli.hf_output_len
&& len == 0
{
return Err(BenchError::Config(
@@ -536,13 +526,13 @@ impl BenchConfig {
}
// Multi-turn validation
if args.multi_turn {
if args.backend != BackendKind::OpenaiChat {
if cli.multi_turn {
if cli.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"--multi-turn requires --backend openai-chat".into(),
));
}
if args.multi_turn_num_turns == 0 {
if cli.multi_turn_num_turns == 0 {
return Err(BenchError::Config(
"--multi-turn-num-turns must be at least 1".into(),
));
@@ -551,18 +541,18 @@ impl BenchConfig {
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply.
if args.dataset_name == DatasetName::ShareGpt {
if args.multi_turn_max_turns == 1 {
if cli.dataset_name == DatasetName::ShareGpt {
if cli.multi_turn_max_turns == 1 {
return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
));
}
} else {
(multi_turn_min_turns, multi_turn_max_turns) =
match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
(0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
(m, 0) => (m, args.multi_turn_num_turns),
(0, x) => (args.multi_turn_num_turns, x),
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
(m, 0) => (m, cli.multi_turn_num_turns),
(0, x) => (cli.multi_turn_num_turns, x),
(m, x) => (m, x),
};
if multi_turn_min_turns < 1 {
@@ -578,16 +568,15 @@ 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."
);
}
// Validate prefix sharing ratios
let pg = args.multi_turn_prefix_global_ratio;
let pc = args.multi_turn_prefix_conversation_ratio;
let pg = cli.multi_turn_prefix_global_ratio;
let pc = cli.multi_turn_prefix_conversation_ratio;
if !(0.0..=1.0).contains(&pg) {
return Err(BenchError::Config(
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -603,20 +592,20 @@ impl BenchConfig {
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
));
}
if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
return Err(BenchError::Config(
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
));
}
}
if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
return Err(BenchError::Config(format!(
"--steady-state-threshold must be in (0.0, 1.0], got {}",
args.steady_state_threshold
cli.steady_state_threshold
)));
}
if let Some(mw) = args.steady_state_min_window
if let Some(mw) = cli.steady_state_min_window
&& mw < 0.0
{
return Err(BenchError::Config(format!(
@@ -624,122 +613,122 @@ impl BenchConfig {
)));
}
if args.profile_batch_threshold.is_some() && !args.profile {
if cli.profile_batch_threshold.is_some() && !cli.profile {
return Err(BenchError::Config(
"--profile-batch-threshold requires --profile".into(),
));
}
if args.profile_duration <= 0.0 {
if cli.profile_duration <= 0.0 {
return Err(BenchError::Config(
"--profile-duration must be positive".into(),
));
}
if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
return Err(BenchError::Config(
"--profile-duration requires --profile-batch-threshold".into(),
));
}
Ok(BenchConfig {
backend: args.backend,
backend: cli.backend,
base_url,
api_url,
model: args.model.clone(),
model_name: args.served_model_name.clone(),
model: cli.model.clone(),
model_name: cli.served_model_name.clone(),
tokenizer_id,
tokenizer_mode: args.tokenizer_mode.clone(),
trust_remote_code: args.trust_remote_code,
skip_tokenizer_init: args.skip_tokenizer_init,
dataset_name: args.dataset_name,
dataset_path: args.dataset_path.clone(),
max_model_len: args.max_model_len,
tokenizer_mode: cli.tokenizer_mode.clone(),
trust_remote_code: cli.trust_remote_code,
skip_tokenizer_init: cli.skip_tokenizer_init,
dataset_name: cli.dataset_name,
dataset_path: cli.dataset_path.clone(),
max_model_len: cli.max_model_len,
random_input_len,
random_output_len,
random_prefix_len: args.random_prefix_len,
random_prefix_len: cli.random_prefix_len,
random_range_ratio,
random_batch_size: args.random_batch_size,
random_batch_size: cli.random_batch_size,
is_reranker,
custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
prefix_repetition_output_len: args
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
prefix_repetition_output_len: cli
.output_len
.unwrap_or(args.prefix_repetition_output_len),
random_cache_hit_fraction: args.random_cache_hit_fraction,
random_cache_ratio: args.random_cache_ratio,
sharegpt_output_len: args.sharegpt_output_len,
sonnet_input_len: args.sonnet_input_len,
sonnet_output_len: args.sonnet_output_len,
sonnet_prefix_len: args.sonnet_prefix_len,
no_oversample: args.no_oversample,
disable_shuffle: args.disable_shuffle,
num_prompts: args.num_prompts,
request_rate: args.request_rate,
burstiness: args.burstiness,
max_concurrency: args.max_concurrency,
steady_state_threshold: args.steady_state_threshold,
steady_state_min_window: args.steady_state_min_window,
no_steady_state: args.no_steady_state,
disable_tqdm: args.disable_tqdm,
num_warmups: args.num_warmups,
profile: args.profile,
profile_batch_threshold: args.profile_batch_threshold,
profile_duration: args.profile_duration,
save_result: args.save_result,
save_detailed: args.save_detailed,
append_result: args.append_result,
result_dir: args.result_dir.clone(),
result_filename: args.result_filename.clone(),
seed: args.seed,
.unwrap_or(cli.prefix_repetition_output_len),
random_cache_hit_fraction: cli.random_cache_hit_fraction,
random_cache_ratio: cli.random_cache_ratio,
sharegpt_output_len: cli.sharegpt_output_len,
sonnet_input_len: cli.sonnet_input_len,
sonnet_output_len: cli.sonnet_output_len,
sonnet_prefix_len: cli.sonnet_prefix_len,
no_oversample: cli.no_oversample,
disable_shuffle: cli.disable_shuffle,
num_prompts: cli.num_prompts,
request_rate: cli.request_rate,
burstiness: cli.burstiness,
max_concurrency: cli.max_concurrency,
steady_state_threshold: cli.steady_state_threshold,
steady_state_min_window: cli.steady_state_min_window,
no_steady_state: cli.no_steady_state,
disable_tqdm: cli.disable_tqdm,
num_warmups: cli.num_warmups,
profile: cli.profile,
profile_batch_threshold: cli.profile_batch_threshold,
profile_duration: cli.profile_duration,
save_result: cli.save_result,
save_detailed: cli.save_detailed,
append_result: cli.append_result,
result_dir: cli.result_dir.clone(),
result_filename: cli.result_filename.clone(),
seed: cli.seed,
ignore_eos,
insecure: args.insecure,
insecure: cli.insecure,
selected_percentile_metrics,
selected_percentiles,
sweep_summary_percentiles,
label: args.label.clone(),
logprobs: args.logprobs,
request_id_prefix: args.get_request_id_prefix(),
ready_check_timeout_sec: args.ready_check_timeout_sec,
label: cli.label.clone(),
logprobs: cli.logprobs,
request_id_prefix: cli.get_request_id_prefix(),
ready_check_timeout_sec: cli.ready_check_timeout_sec,
extra_headers,
extra_body,
metadata,
dry_run: args.dry_run,
dry_run: cli.dry_run,
goodput,
ramp_up,
multi_turn: args.multi_turn,
multi_turn_num_turns: args.multi_turn_num_turns,
multi_turn: cli.multi_turn,
multi_turn_num_turns: cli.multi_turn_num_turns,
multi_turn_min_turns,
multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if args.multi_turn
&& args.dataset_name == DatasetName::ShareGpt
&& args.multi_turn_max_turns != 0
sharegpt_multi_turn_max_turns: if cli.multi_turn
&& cli.dataset_name == DatasetName::ShareGpt
&& cli.multi_turn_max_turns != 0
{
Some(args.multi_turn_max_turns)
Some(cli.multi_turn_max_turns)
} else {
None
},
per_turn_input_len,
multi_turn_concurrency: args.multi_turn_concurrency,
multi_turn_delay_ms: args.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
speed_bench_config: args.speed_bench_config,
speed_bench_category: args.speed_bench_category.clone(),
speed_bench_max_input_len: args.speed_bench_max_input_len,
hf_split: args.hf_split.clone(),
hf_subset: args.hf_subset.clone(),
hf_output_len: args.hf_output_len,
hf_text_column: args.hf_text_column.clone(),
reset_prefix_cache: args.reset_prefix_cache,
prompt_token_ids: args.prompt_token_ids,
random_mm_base_items_per_request: args.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
multi_turn_concurrency: cli.multi_turn_concurrency,
multi_turn_delay_ms: cli.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
speed_bench_config: cli.speed_bench_config,
speed_bench_category: cli.speed_bench_category.clone(),
speed_bench_max_input_len: cli.speed_bench_max_input_len,
hf_split: cli.hf_split.clone(),
hf_subset: cli.hf_subset.clone(),
hf_output_len: cli.hf_output_len,
hf_text_column: cli.hf_text_column.clone(),
reset_prefix_cache: cli.reset_prefix_cache,
prompt_token_ids: cli.prompt_token_ids,
random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
random_mm_limit,
random_mm_buckets,
enable_multimodal_chat: args.enable_multimodal_chat,
enable_multimodal_chat: cli.enable_multimodal_chat,
lora_modules,
lora_assignment: args.lora_assignment,
lora_assignment: cli.lora_assignment,
})
}
}
@@ -822,17 +811,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
Ok(config)
}
fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
let strategy = match args.ramp_up_strategy {
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
let strategy = match cli.ramp_up_strategy {
None => return Ok(None),
Some(s) => s,
};
let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
})?;
let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
})?;
@@ -854,21 +843,7 @@ mod tests {
use clap::Parser;
use super::*;
use crate::cli::BenchServeArgs;
#[derive(Parser)]
struct TestCli {
#[command(flatten)]
args: BenchServeArgs,
}
fn parse_args<I, T>(args: I) -> BenchServeArgs
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString> + Clone,
{
TestCli::parse_from(args).args
}
use crate::cli::Cli;
fn base_multi_turn_args() -> Vec<&'static str> {
vec![
@@ -884,8 +859,8 @@ mod tests {
#[test]
fn test_prefix_sharing_defaults_to_zero() {
let args = base_multi_turn_args();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
}
@@ -899,8 +874,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.8",
]);
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
}
@@ -914,8 +889,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.6",
]);
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
@@ -927,16 +902,16 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.5",
]);
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
fn test_prefix_sharing_out_of_range_fails() {
let mut args = base_multi_turn_args();
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
@@ -953,8 +928,8 @@ mod tests {
"--multi-turn-prefix-global-ratio",
"0.1",
];
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
@@ -969,8 +944,8 @@ mod tests {
"--dataset-name",
"sharegpt",
];
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.multi_turn_max_turns, 3);
assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -993,8 +968,8 @@ mod tests {
"--multi-turn-max-turns",
"2",
];
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
}
@@ -1012,8 +987,8 @@ mod tests {
"--multi-turn-max-turns",
"1",
];
let args = parse_args(args);
let err = BenchConfig::from_args(&args).unwrap_err().to_string();
let cli = Cli::parse_from(args);
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
assert!(
err.contains("at least 2 for ShareGPT"),
"expected ShareGPT-specific error, got: {err}"
@@ -1034,8 +1009,8 @@ mod tests {
"--multi-turn-max-turns",
"20",
];
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
}
@@ -1043,8 +1018,8 @@ mod tests {
#[test]
fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert!(config.sweep_summary_percentiles.is_empty());
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1059,8 +1034,8 @@ mod tests {
"--sweep-summary-percentiles",
"90,95,90",
]);
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1070,8 +1045,8 @@ mod tests {
fn test_invalid_sweep_summary_percentile_fails() {
let mut args = base_multi_turn_args();
args.extend(["--sweep-summary-percentiles", "101"]);
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
@@ -1083,20 +1058,12 @@ mod tests {
"--max-model-len",
"4096",
];
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
assert_eq!(config.max_model_len, Some(4096));
}
#[test]
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
let args = parse_args(["vllm-bench"]);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.tokenizer_id, None);
}
#[test]
fn test_zero_max_model_len_fails() {
let args = vec![
@@ -1106,9 +1073,9 @@ mod tests {
"--max-model-len",
"0",
];
let args = parse_args(args);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_args(&args).is_err());
assert!(BenchConfig::from_cli(&cli).is_err());
}
#[test]
fn test_range_ratio_parse_float() {
+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)
}
+32 -20
View File
@@ -1,39 +1,51 @@
// 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 api = hf_hub::api::tokio::Api::new()
.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())?
}
}
-86
View File
@@ -1,86 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
use anyhow::Context;
pub use cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use config::BenchConfig;
/// Prepare process-wide resources for a benchmark run.
pub fn prepare_process() {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
tracing::info!(soft_limit = new, "raised open-file limit");
}
}
/// Run the online serving benchmark.
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = args.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_args(&args).context("Configuration error")?;
async {
if config.multi_turn {
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
args.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if args.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, args.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
}
.await
.context("Benchmark failed")
}
+73 -25
View File
@@ -1,44 +1,92 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use anyhow::Context;
use clap::Parser;
#[derive(Parser)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
struct Cli {
#[command(flatten)]
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();
}
use cli::Cli;
use config::BenchConfig;
fn main() -> anyhow::Result<()> {
init_tracing();
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
let cli = Cli::parse();
vllm_bench::prepare_process();
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = cli.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.context("Failed to build tokio runtime")?;
.expect("Failed to build tokio runtime");
runtime.block_on(vllm_bench::run(cli.args))
runtime
.block_on(async {
if config.multi_turn {
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
cli.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
.await?;
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if cli.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, cli.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
})
.context("Benchmark failed")
}
+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 {
+4 -4
View File
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
let (value, field) = match spec.field_layout_for(key) {
Some(FieldLayout::Batched) => (
tensor.batched_wire_value_at(index)?,
tensor.batched_value_at(index)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
})?;
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
(
tensor.flat_wire_value_range(start, end)?,
tensor.flat_value_range(start, end)?,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
)
}
None => (
tensor.try_into()?,
tensor.clone(),
MmField::Shared(MmSharedField {
batch_size: len,
keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
data.insert(
key.clone(),
MmFieldElem {
data: Some(value),
data: Some(value.try_into()?),
field,
},
);
+85 -72
View File
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
use crate::error::{Error, Result, bail_multimodal, multimodal};
/// Representation for multimodal kwarg values for transformation.
#[derive(Debug)]
#[derive(Debug, Clone)]
pub(super) enum KwargValue {
/// Float tensor with row-major flat data and shape.
F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,19 +107,28 @@ impl KwargValue {
}
}
impl TryFrom<&KwargValue> for ProtocolKwargValue {
impl TryFrom<KwargValue> for ProtocolKwargValue {
type Error = Error;
fn try_from(value: &KwargValue) -> Result<Self> {
let tensor = match value {
KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
KwargValue::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
fn try_from(value: KwargValue) -> Result<Self> {
match value {
KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Passthrough(value) => Ok(value),
}
}
}
@@ -136,55 +145,63 @@ impl KwargValue {
}
}
/// Convert one media item from a batched tensor field to wire bytes.
/// Extract one media item from a batched tensor field.
///
/// Batched fields use their first axis as media-item index and drop that
/// axis in the per-feature value, matching vLLM's batched-field semantics.
pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
self.wire_value_range(index, index + 1, true)
}
/// Convert one media item's flat tensor range directly to wire bytes.
///
/// Flat fields keep the first axis as the sliced length for this item.
pub(super) fn flat_wire_value_range(
&self,
start: usize,
end: usize,
) -> Result<ProtocolKwargValue> {
self.wire_value_range(start, end, false)
}
fn wire_value_range(
&self,
start: usize,
end: usize,
drop_axis: bool,
) -> Result<ProtocolKwargValue> {
let tensor = match self {
pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f32(shape, data)
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F32Tensor { data, shape })
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f16(shape, data)
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F16Tensor { data, shape })
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_bf16(shape, data)
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::Bf16Tensor { data, shape })
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_i64(shape, data)
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::I64Tensor { data, shape })
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_u32(shape, data)
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::U32Tensor { data, shape })
}
Self::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
}
/// Extract one media item's variable-length range from a flat tensor field.
///
/// Flat fields keep the first axis as the sliced length for this item.
pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F32Tensor { data, shape })
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F16Tensor { data, shape })
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::Bf16Tensor { data, shape })
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::I64Tensor { data, shape })
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::U32Tensor { data, shape })
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
}
}
@@ -223,13 +240,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
}
/// Slice a flat row-major tensor along its first axis.
fn slice_first_axis_range<'a, T>(
fn slice_first_axis_range<T: Clone>(
shape: &[usize],
data: &'a [T],
data: &[T],
start: usize,
end: usize,
drop_axis: bool,
) -> Result<(Vec<usize>, &'a [T])> {
) -> Result<(Vec<usize>, Vec<T>)> {
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
if start > end || end > first_dim {
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -253,7 +270,7 @@ fn slice_first_axis_range<'a, T>(
shape[0] = end - start;
shape
};
Ok((out_shape, &data[data_start..data_end]))
Ok((out_shape, data[data_start..data_end].to_vec()))
}
#[cfg(test)]
@@ -261,39 +278,35 @@ mod tests {
use super::*;
#[test]
fn batched_wire_value_at_drops_first_axis() {
fn batched_value_at_drops_first_axis() {
let value = KwargValue::F32Tensor {
data: vec![1.0, 2.0, 3.0, 4.0],
shape: vec![2, 2],
};
let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
panic!("expected tensor");
};
let value = value.batched_value_at(1).unwrap();
assert_eq!(tensor.shape, vec![2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
);
assert!(matches!(
value,
KwargValue::F32Tensor { data, shape }
if shape == vec![2] && data == vec![3.0, 4.0]
));
}
#[test]
fn flat_wire_value_range_keeps_first_axis() {
fn flat_value_range_keeps_first_axis() {
let value = KwargValue::U32Tensor {
data: (0..10).collect(),
shape: vec![5, 2],
};
let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
panic!("expected tensor");
};
let value = value.flat_value_range(1, 3).unwrap();
assert_eq!(tensor.shape, vec![2, 2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
);
assert!(matches!(
value,
KwargValue::U32Tensor { data, shape }
if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
));
}
#[test]
@@ -323,7 +336,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
else {
panic!("expected tensor");
};
@@ -338,7 +351,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
else {
panic!("expected tensor");
};
+4 -4
View File
@@ -130,7 +130,7 @@ fn build_video_item(
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
let (value, field) = match support.spec.field_layout_for(&key) {
Some(FieldLayout::Batched) => (
tensor.batched_wire_value_at(0)?,
tensor.batched_value_at(0)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
.first_dim()
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
(
(&tensor).try_into()?,
tensor,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
)
}
None => (
(&tensor).try_into()?,
tensor,
MmField::Shared(MmSharedField {
batch_size: 1,
keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
data.insert(
key,
MmFieldElem {
data: Some(value),
data: Some(value.try_into()?),
field,
},
);
+1 -1
View File
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
};
apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
apply_structural_tag_constraint(request, parser.structural_tag_model())?;
if parser.preserve_special_tokens() {
request.decode_options.skip_special_tokens = false;
@@ -7,8 +7,7 @@ use thiserror_ext::AsReport;
use vllm_engine_core_client::protocol::structured_outputs::{
StructuredOutputBackend, StructuredOutputsParams,
};
use vllm_parser::tool::StructuralTagBuilder;
use xgrammar_structural_tag::builders::StructuralTagOptions;
use vllm_parser::tool::StructuralTagModel;
use xgrammar_structural_tag::{
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
build_structural_tag,
@@ -21,9 +20,9 @@ use crate::{Error, Result as ChatResult};
/// support and the request's tool choice.
pub(super) fn apply_structural_tag_constraint(
request: &mut ChatRequest,
builder: Option<&dyn StructuralTagBuilder>,
model: Option<StructuralTagModel>,
) -> ChatResult<()> {
let Some(builder) = builder else {
let Some(model) = model else {
return Ok(());
};
let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -43,16 +42,11 @@ pub(super) fn apply_structural_tag_constraint(
})
.collect::<Vec<_>>();
let structural_tag = build_structural_tag(
builder,
&tools,
tool_choice,
StructuralTagOptions::default().with_reasoning(false),
)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
// Overwrite any existing structured output settings with the structural tag constraint.
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -147,7 +141,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -160,7 +154,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -175,7 +169,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -190,7 +184,7 @@ mod tests {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request);
@@ -207,7 +201,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let params = structured_outputs(&request);
@@ -227,7 +221,7 @@ mod tests {
);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag should build");
let tag = structural_tag_value(&request).to_string();
@@ -240,7 +234,7 @@ mod tests {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none());
@@ -255,7 +249,7 @@ mod tests {
});
let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
.expect("structural tag decision should succeed");
let params = structured_outputs(&request);
+3 -14
View File
@@ -236,10 +236,9 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
loops.into_iter().any(|loop_ast| {
matches!(loop_ast.target, Expr::Var(_))
&& (is_var_access(&loop_ast.iter, "content")
|| message_varnames.iter().any(|varname| {
is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
}))
&& message_varnames
.iter()
.any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
})
}
@@ -316,16 +315,6 @@ mod tests {
);
}
#[test]
fn detects_openai_template_with_content_parameter_loop() {
assert_eq!(
detect(
"{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
),
ChatTemplateContentFormat::OpenAi
);
}
#[test]
fn detects_openai_template_with_messages_alias() {
assert_eq!(
-20
View File
@@ -1309,26 +1309,6 @@ mod tests {
.assert_eq(&rendered);
}
#[test]
fn qwen35_template_auto_detects_openai_multimodal_content() {
let mut request = image_request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_mm(
QWEN3_5_0_8B_TEMPLATE,
&request,
ChatTemplateContentFormatOption::Auto,
)
.unwrap();
expect![[r#"
Text(
"<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
)
"#]]
.assert_debug_eq(&rendered.prompt);
}
#[test]
fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
-1
View File
@@ -29,7 +29,6 @@ tokio-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
uuid.workspace = true
vllm-bench.workspace = true
vllm-chat.workspace = true
vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true
+1 -11
View File
@@ -79,23 +79,13 @@ impl Cli {
}
/// Supported top-level CLI commands.
#[derive(Debug, Subcommand)]
#[derive(Debug, Subcommand, PartialEq, Eq)]
pub enum Command {
/// Run the Rust OpenAI frontend as a Python-supervised worker.
Frontend(FrontendArgs),
/// Launch a managed Python headless engine, then run the Rust OpenAI
/// frontend.
Serve(ServeArgs),
/// Run vLLM benchmarks.
#[command(subcommand)]
Bench(BenchCommand),
}
/// Supported benchmark commands.
#[derive(Debug, Subcommand)]
pub enum BenchCommand {
/// Benchmark online serving throughput.
Serve(vllm_bench::BenchServeArgs),
}
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
+1 -21
View File
@@ -5,27 +5,7 @@ use expect_test::expect;
use vllm_engine_core_client::TransportMode;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{BenchCommand, Cli, Command};
#[test]
fn bench_serve_args_parse_without_managed_engine_repartition() {
let cli = Cli::try_parse_from([
"vllm-rs",
"bench",
"serve",
"--backend",
"openai-chat",
"--request-rate",
"inf",
])
.unwrap();
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
panic!("expected bench serve args");
};
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
assert!(args.request_rate.is_infinite());
}
use super::{Cli, Command};
#[test]
fn serve_args_forward_python_flags_with_separator() {
+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"));
+2 -14
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};
@@ -13,7 +12,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use vllm_managed_engine::ManagedEngineHandle;
use crate::cli::{BenchCommand, Cli, Command};
use crate::cli::{Cli, Command};
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -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();
@@ -108,10 +100,6 @@ fn main() -> Result<()> {
async fn async_main(cli: Cli) -> Result<()> {
match cli.command {
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
Command::Bench(BenchCommand::Serve(bench_args)) => {
vllm_bench::prepare_process();
vllm_bench::run(bench_args).await
}
Command::Serve(args) => {
let handshake_port = args.managed_engine.resolve_handshake_port()?;
@@ -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()

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