forked from Karylab-cklius/vllm
Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5ea7cac55b | ||
|
|
be3476447f | ||
|
|
f1473092c4 | ||
|
|
4dcbae8670 | ||
|
|
65dac3a770 | ||
|
|
0ba2500ef0 | ||
|
|
ef576befd2 | ||
|
|
35e4a36107 | ||
|
|
da5803d46e | ||
|
|
75ddfaf909 | ||
|
|
f36fe52add | ||
|
|
391d918d4d | ||
|
|
fca040885b | ||
|
|
a7fd4c7482 | ||
|
|
5131691063 | ||
|
|
80e00e5ac6 |
@@ -813,8 +813,8 @@ steps:
|
||||
|
||||
# Download artifacts from current build
|
||||
echo "Downloading artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
# buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
# buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
|
||||
# # Run upload script
|
||||
bash .buildkite/scripts/upload-rocm-wheels.sh
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 53
|
||||
timeout_in_minutes: 35
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
|
||||
@@ -39,7 +39,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -59,7 +59,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -78,7 +78,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 83
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -156,7 +156,7 @@ steps:
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 83
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
@@ -90,7 +90,6 @@ steps:
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/simple_kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/streaming_input
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/ec_connector/unit
|
||||
- tests/v1/metrics
|
||||
@@ -104,7 +103,6 @@ steps:
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s v1/streaming_input
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -145,7 +143,6 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
|
||||
@@ -5,14 +5,12 @@ steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- 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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 98
|
||||
timeout_in_minutes: 65
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
device: h200_35gb
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 110
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(', ')}`);
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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 ""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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()),
|
||||
|
||||
+1
-1
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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"]
|
||||
}
|
||||
|
||||
@@ -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" {
|
||||
|
||||
@@ -68,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.15.post1"
|
||||
"default": "0.6.14"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -75,7 +75,7 @@ vllm serve <model> \
|
||||
| `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, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)}();
|
||||
@@ -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:**
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
|
||||
# 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"
|
||||
|
||||
@@ -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
|
||||
|
||||
Generated
+31
-68
@@ -3446,6 +3446,7 @@ dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
@@ -4252,6 +4253,18 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "subenum"
|
||||
version = "1.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec3d08fe7078c57309d5c3d938e50eba95ba1d33b9c3a101a8465fc6861a5416"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "subtle"
|
||||
version = "2.6.1"
|
||||
@@ -4863,7 +4876,6 @@ dependencies = [
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4925,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",
|
||||
@@ -4954,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",
|
||||
@@ -4964,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",
|
||||
@@ -4990,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",
|
||||
@@ -5485,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",
|
||||
]
|
||||
@@ -5504,11 +5500,15 @@ dependencies = [
|
||||
"asynk-strim-attr",
|
||||
"bytes",
|
||||
"clap",
|
||||
"easy-ext",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"half",
|
||||
"indexmap 2.13.0",
|
||||
"itertools 0.14.0",
|
||||
"llm-multimodal",
|
||||
"minijinja",
|
||||
"minijinja-contrib",
|
||||
"ndarray 0.17.2",
|
||||
"oss-harmony",
|
||||
"paste",
|
||||
@@ -5519,16 +5519,17 @@ dependencies = [
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"serial_test",
|
||||
"strum",
|
||||
"subenum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"trait-set",
|
||||
"uuid",
|
||||
"vllm-chat-renderer",
|
||||
"vllm-chat-types",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-parser",
|
||||
@@ -5538,42 +5539,6 @@ dependencies = [
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-chat-renderer"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"enum-as-inner",
|
||||
"expect-test",
|
||||
"indexmap 2.13.0",
|
||||
"itertools 0.14.0",
|
||||
"minijinja",
|
||||
"minijinja-contrib",
|
||||
"oss-harmony",
|
||||
"serde",
|
||||
"serde-json-fmt",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"strum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"time",
|
||||
"tracing",
|
||||
"vllm-chat-types",
|
||||
"vllm-tokenizer",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-chat-types"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"easy-ext",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-cmd"
|
||||
version = "0.1.0"
|
||||
@@ -5714,11 +5679,11 @@ dependencies = [
|
||||
"expect-test",
|
||||
"futures",
|
||||
"openai-protocol",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tool-parser",
|
||||
"vllm-chat-types",
|
||||
"vllm-tokenizer",
|
||||
"winnow",
|
||||
"xgrammar-structural-tag",
|
||||
@@ -5765,7 +5730,6 @@ dependencies = [
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tonic",
|
||||
"tonic-health",
|
||||
"tonic-prost",
|
||||
"tonic-prost-build",
|
||||
"tower",
|
||||
@@ -5825,7 +5789,6 @@ dependencies = [
|
||||
"rustc-hash 1.1.0",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"tekken-rs",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
@@ -6358,9 +6321,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
-10
@@ -2,8 +2,6 @@
|
||||
members = [
|
||||
"src/bench",
|
||||
"src/chat",
|
||||
"src/chat-renderer",
|
||||
"src/chat-types",
|
||||
"src/cmd",
|
||||
"src/engine-core-client",
|
||||
"src/llm",
|
||||
@@ -120,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"] }
|
||||
@@ -137,8 +134,6 @@ 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-chat-renderer = { path = "src/chat-renderer" }
|
||||
vllm-chat-types = { path = "src/chat-types" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
vllm-managed-engine = { path = "src/managed-engine" }
|
||||
@@ -148,7 +143,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",
|
||||
|
||||
@@ -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 {}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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);
|
||||
}
|
||||
|
||||
@@ -288,18 +288,10 @@ impl BenchConfig {
|
||||
}
|
||||
Some(other) => {
|
||||
// extra_body was not an object — just use sampling params
|
||||
let value_type = match &other {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "boolean",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => unreachable!(),
|
||||
};
|
||||
tracing::warn!(
|
||||
value_type,
|
||||
"sampling parameters may be lost because --extra-body is not a JSON object"
|
||||
eprintln!(
|
||||
"Warning: --extra-body is not a JSON object, sampling params may be lost"
|
||||
);
|
||||
let _ = other;
|
||||
sampling_params
|
||||
}
|
||||
None => sampling_params,
|
||||
@@ -497,9 +489,9 @@ impl BenchConfig {
|
||||
_ => {}
|
||||
}
|
||||
if !args.skip_chat_template {
|
||||
tracing::warn!(
|
||||
dataset = "custom",
|
||||
"client-side chat template rendering is unsupported; sending prompts raw"
|
||||
eprintln!(
|
||||
"NOTE: client-side chat template rendering is not supported; custom \
|
||||
dataset prompts are sent raw (equivalent to --skip-chat-template)."
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -578,10 +570,9 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
if ignore_eos {
|
||||
tracing::warn!(
|
||||
ignore_eos,
|
||||
multi_turn = true,
|
||||
"output length limits may be ignored, causing unbounded context growth"
|
||||
eprintln!(
|
||||
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
|
||||
respect output length limits, causing unbounded context growth."
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)));
|
||||
}
|
||||
}
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,6 @@ use rand::seq::SliceRandom;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use super::SampleRequest;
|
||||
use super::progress::RowDownloadReporter;
|
||||
use crate::cli::SpeedBenchConfig;
|
||||
use crate::error::{BenchError, Result};
|
||||
use crate::tokenizer::TokenizerKind;
|
||||
@@ -26,7 +25,7 @@ fn cache_dir() -> std::path::PathBuf {
|
||||
|
||||
/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
|
||||
/// Results are cached as JSON locally for subsequent runs.
|
||||
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let config_name = config.as_str();
|
||||
|
||||
let dir = cache_dir();
|
||||
@@ -36,13 +35,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
// Return cached file if it exists
|
||||
if cache_path.exists() {
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
|
||||
println!("SPEED-Bench ({config_name}) cached: {path_str}");
|
||||
return Ok(path_str);
|
||||
}
|
||||
|
||||
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
|
||||
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
|
||||
@@ -50,7 +49,6 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let mut all_rows: Vec<serde_json::Value> = Vec::new();
|
||||
let mut offset = 0usize;
|
||||
let page_size = 100usize;
|
||||
let mut progress = RowDownloadReporter::new();
|
||||
|
||||
loop {
|
||||
let url = format!(
|
||||
@@ -66,14 +64,13 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let max_retries = 3;
|
||||
let mut data: Option<serde_json::Value> = None;
|
||||
for attempt in 0..=max_retries {
|
||||
let resp = match client.get(&url).send().await {
|
||||
let resp = match client.get(&url).send() {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(
|
||||
std::thread::sleep(std::time::Duration::from_secs(
|
||||
2 * (attempt as u64 + 1),
|
||||
))
|
||||
.await;
|
||||
));
|
||||
continue;
|
||||
}
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -83,7 +80,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
};
|
||||
|
||||
if resp.status().is_server_error() && attempt < max_retries {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
|
||||
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -94,7 +91,7 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
)));
|
||||
}
|
||||
|
||||
data = Some(resp.json().await.map_err(|e| {
|
||||
data = Some(resp.json().map_err(|e| {
|
||||
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
|
||||
})?);
|
||||
break;
|
||||
@@ -119,14 +116,15 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
let fetched = rows.len();
|
||||
offset += fetched;
|
||||
|
||||
// Print progress
|
||||
let total = data["num_rows_total"].as_u64().unwrap_or(0);
|
||||
progress.update(offset, total);
|
||||
eprint!("\r Fetched {offset}/{total} rows...");
|
||||
|
||||
if fetched < page_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
progress.finish();
|
||||
eprintln!(); // newline after progress
|
||||
|
||||
if all_rows.is_empty() {
|
||||
return Err(BenchError::Config(
|
||||
@@ -139,11 +137,9 @@ pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
|
||||
std::fs::write(&cache_path, &json_str)?;
|
||||
|
||||
let path_str = cache_path.to_string_lossy().to_string();
|
||||
tracing::info!(
|
||||
config = config_name,
|
||||
rows = all_rows.len(),
|
||||
path = %path_str,
|
||||
"saved SPEED-Bench dataset"
|
||||
println!(
|
||||
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
|
||||
all_rows.len()
|
||||
);
|
||||
Ok(path_str)
|
||||
}
|
||||
@@ -267,11 +263,9 @@ pub fn load_speed_bench_dataset(
|
||||
// Oversample if needed
|
||||
if samples.len() < num_requests {
|
||||
if no_oversample {
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
samples = samples.len(),
|
||||
requested = num_requests,
|
||||
"skipping dataset oversampling"
|
||||
println!(
|
||||
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
|
||||
samples.len()
|
||||
);
|
||||
} else if !samples.is_empty() {
|
||||
let original_len = samples.len();
|
||||
@@ -281,11 +275,9 @@ pub fn load_speed_bench_dataset(
|
||||
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
|
||||
samples.push(req);
|
||||
}
|
||||
tracing::info!(
|
||||
dataset = "speed-bench",
|
||||
original_samples = original_len,
|
||||
samples = samples.len(),
|
||||
"oversampled dataset"
|
||||
println!(
|
||||
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
|
||||
samples.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -296,6 +288,7 @@ pub fn load_speed_bench_dataset(
|
||||
));
|
||||
}
|
||||
|
||||
// Print category distribution
|
||||
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
|
||||
for entry in &filtered[..filtered.len().min(samples.len())] {
|
||||
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
|
||||
@@ -304,7 +297,7 @@ pub fn load_speed_bench_dataset(
|
||||
let mut cats: Vec<_> = cat_counts.into_iter().collect();
|
||||
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
|
||||
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
|
||||
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
|
||||
println!("SPEED-Bench categories: {}", cat_str.join(", "));
|
||||
|
||||
Ok(samples)
|
||||
}
|
||||
|
||||
+35
-20
@@ -1,39 +1,54 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::path::PathBuf;
|
||||
//! Sync facade over the async `hf_hub` API.
|
||||
//!
|
||||
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
|
||||
//! backend unconditionally pulls ureq's default rustls feature. So we use the
|
||||
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
|
||||
//! loaders, tokenizer fallback in rayon threads) by running each download on a
|
||||
//! dedicated thread with its own single-threaded runtime.
|
||||
|
||||
use hf_hub::Repo;
|
||||
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
|
||||
pub struct HubRepo {
|
||||
repo: ApiRepo,
|
||||
repo: hf_hub::Repo,
|
||||
}
|
||||
|
||||
impl HubRepo {
|
||||
pub fn model(model_id: String) -> Result<Self, String> {
|
||||
Self::new(Repo::model(model_id))
|
||||
pub fn model(model_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::model(model_id),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dataset(repo_id: String) -> Result<Self, String> {
|
||||
Self::new(Repo::dataset(repo_id))
|
||||
}
|
||||
|
||||
fn new(repo: Repo) -> Result<Self, String> {
|
||||
let mut builder = ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
pub fn dataset(repo_id: String) -> Self {
|
||||
Self {
|
||||
repo: hf_hub::Repo::dataset(repo_id),
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
Ok(Self {
|
||||
repo: api.repo(repo),
|
||||
})
|
||||
}
|
||||
|
||||
/// Download (or fetch from cache) a single file from the repo.
|
||||
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
|
||||
pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
self.repo.get(filename).await.map_err(|e| format!("{e}"))
|
||||
pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
|
||||
let repo = self.repo.clone();
|
||||
let filename = filename.to_string();
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
|
||||
rt.block_on(async move {
|
||||
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
|
||||
})
|
||||
})
|
||||
.join()
|
||||
.map_err(|_| "HF Hub download thread panicked".to_string())?
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ pub fn prepare_process() {
|
||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||
&& new > 1024
|
||||
{
|
||||
tracing::info!(soft_limit = new, "raised open-file limit");
|
||||
eprintln!("Open-file limit: {new}");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,19 +19,7 @@ struct Cli {
|
||||
args: vllm_bench::BenchServeArgs,
|
||||
}
|
||||
|
||||
// TODO: unify the tracing subscriber used by different binaries.
|
||||
fn init_tracing() {
|
||||
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(filter)
|
||||
.with_writer(std::io::stderr)
|
||||
.try_init();
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
init_tracing();
|
||||
|
||||
let cli = Cli::parse();
|
||||
vllm_bench::prepare_process();
|
||||
|
||||
|
||||
@@ -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> =
|
||||
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
[package]
|
||||
name = "vllm-chat-renderer"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
enum-as-inner.workspace = true
|
||||
indexmap.workspace = true
|
||||
itertools.workspace = true
|
||||
minijinja.workspace = true
|
||||
minijinja-contrib.workspace = true
|
||||
openai-harmony.workspace = true
|
||||
serde.workspace = true
|
||||
serde-json-fmt.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_with.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
time.workspace = true
|
||||
tracing.workspace = true
|
||||
vllm-chat-types.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
expect-test.workspace = true
|
||||
tempfile.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,54 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
type BoxedError = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
/// Error returned while constructing or applying a chat renderer.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum Error {
|
||||
/// Rendering requires at least one chat message.
|
||||
#[error("chat request must contain at least one message")]
|
||||
EmptyMessages,
|
||||
/// Continuation mode requires the final message to be an assistant turn.
|
||||
#[error("cannot continue the final message when the last message is not from the assistant")]
|
||||
ContinueFinalAssistantWithoutFinalAssistant,
|
||||
/// The selected renderer requires a chat template and none was configured.
|
||||
#[error("chat template is required but none was configured")]
|
||||
MissingChatTemplate,
|
||||
/// A chat template could not be compiled or applied.
|
||||
#[error("chat template error: {0}")]
|
||||
ChatTemplate(String),
|
||||
/// The selected renderer cannot represent the given multimodal part.
|
||||
#[error("unsupported multimodal content: {0}")]
|
||||
UnsupportedMultimodalContent(&'static str),
|
||||
/// The process-wide GPT-OSS Harmony encoding could not be initialized.
|
||||
#[error("failed to initialize the Harmony encoding")]
|
||||
HarmonyEncoding {
|
||||
/// Underlying Harmony initialization failure.
|
||||
#[source]
|
||||
error: BoxedError,
|
||||
},
|
||||
/// Tokenizer construction or encoding failed.
|
||||
#[error(transparent)]
|
||||
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
|
||||
}
|
||||
|
||||
impl Error {
|
||||
/// Whether this error should be reported as invalid request input when
|
||||
/// raised while rendering.
|
||||
pub fn is_request_validation_error(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::EmptyMessages
|
||||
| Self::ContinueFinalAssistantWithoutFinalAssistant
|
||||
| Self::MissingChatTemplate
|
||||
| Self::ChatTemplate(_)
|
||||
| Self::UnsupportedMultimodalContent(_)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Result returned by chat renderer operations.
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
@@ -1,172 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use serde_json::Value;
|
||||
use vllm_chat_types::{
|
||||
ChatMessage, ChatOptions, ChatRole, ChatToolChoice, GenerationPromptMode, Tool,
|
||||
};
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Borrowed input consumed by one chat renderer.
|
||||
///
|
||||
/// This view contains only chat-domain values that affect prompt construction.
|
||||
/// Serving, sampling, decoding, scheduling, and engine metadata stay with the
|
||||
/// caller.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct RenderRequest<'a> {
|
||||
/// Ordered chat history to render.
|
||||
pub messages: &'a [ChatMessage],
|
||||
/// Chat-template and generation-prompt controls.
|
||||
pub chat_options: &'a ChatOptions,
|
||||
/// Request-level tools available to the model.
|
||||
pub tools: &'a [Tool],
|
||||
/// Tool-choice behavior used to decide whether tools are exposed.
|
||||
pub tool_choice: &'a ChatToolChoice,
|
||||
/// Optional retrieval documents exposed to HF chat templates.
|
||||
pub documents: Option<&'a [Value]>,
|
||||
}
|
||||
|
||||
impl RenderRequest<'_> {
|
||||
/// Validate renderer-owned request invariants.
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.messages.is_empty() {
|
||||
return Err(Error::EmptyMessages);
|
||||
}
|
||||
match (
|
||||
self.chat_options.generation_prompt_mode,
|
||||
self.messages.last().map(ChatMessage::role),
|
||||
) {
|
||||
(GenerationPromptMode::ContinueFinalAssistant, Some(ChatRole::Assistant)) => {}
|
||||
(GenerationPromptMode::ContinueFinalAssistant, _) => {
|
||||
return Err(Error::ContinueFinalAssistantWithoutFinalAssistant);
|
||||
}
|
||||
(GenerationPromptMode::NoGenerationPrompt, _)
|
||||
| (GenerationPromptMode::StartNewAssistant, _) => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Return whether any message contains multimodal content.
|
||||
pub fn has_multimodal(&self) -> bool {
|
||||
self.messages.iter().any(ChatMessage::has_multimodal)
|
||||
}
|
||||
|
||||
/// Return whether request-level tools should be exposed to the renderer.
|
||||
pub fn tool_parsing_enabled(&self) -> bool {
|
||||
!matches!(self.tool_choice, ChatToolChoice::None) && !self.tools.is_empty()
|
||||
}
|
||||
|
||||
/// Return the request-level thinking toggle when explicitly requested.
|
||||
///
|
||||
/// The `thinking` and `enable_thinking` kwargs must be booleans when
|
||||
/// present and must carry the same value when both are set.
|
||||
pub fn enable_thinking(&self) -> Result<Option<bool>> {
|
||||
let thinking = self.parse_template_bool("thinking")?;
|
||||
let enable_thinking = self.parse_template_bool("enable_thinking")?;
|
||||
|
||||
match (thinking, enable_thinking) {
|
||||
(None, None) => Ok(None),
|
||||
(Some(thinking), Some(enable_thinking)) if thinking != enable_thinking => {
|
||||
Err(Error::ChatTemplate(
|
||||
"template kwargs `thinking` and `enable_thinking` must match when both are set"
|
||||
.to_string(),
|
||||
))
|
||||
}
|
||||
(Some(thinking), _) => Ok(Some(thinking)),
|
||||
(None, Some(enable_thinking)) => Ok(Some(enable_thinking)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse one optional boolean chat-template kwarg.
|
||||
pub fn parse_template_bool(&self, key: &str) -> Result<Option<bool>> {
|
||||
match self.chat_options.template_kwargs.get(key) {
|
||||
None => Ok(None),
|
||||
Some(Value::Bool(value)) => Ok(Some(*value)),
|
||||
Some(other) => Err(Error::ChatTemplate(format!(
|
||||
"template kwarg `{key}` must be a boolean, got {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::RenderRequest;
|
||||
use crate::{ChatMessage, ChatToolChoice, Error, GenerationPromptMode, TestRenderRequest};
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_message_history() {
|
||||
let mut request = TestRenderRequest::for_test();
|
||||
request.messages.clear();
|
||||
|
||||
assert!(matches!(
|
||||
request.as_request().validate(),
|
||||
Err(Error::EmptyMessages)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn continuation_requires_a_final_assistant_message() {
|
||||
let mut request = TestRenderRequest::for_test();
|
||||
request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
|
||||
|
||||
assert!(matches!(
|
||||
request.as_request().validate(),
|
||||
Err(Error::ContinueFinalAssistantWithoutFinalAssistant)
|
||||
));
|
||||
|
||||
request.messages.push(ChatMessage::assistant_text("partial"));
|
||||
request.as_request().validate().unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thinking_aliases_must_match() {
|
||||
let mut request = TestRenderRequest::for_test();
|
||||
request
|
||||
.chat_options
|
||||
.template_kwargs
|
||||
.insert("thinking".to_string(), json!(false));
|
||||
request
|
||||
.chat_options
|
||||
.template_kwargs
|
||||
.insert("enable_thinking".to_string(), json!(true));
|
||||
|
||||
assert!(matches!(
|
||||
request.as_request().enable_thinking(),
|
||||
Err(Error::ChatTemplate(message))
|
||||
if message.contains("`thinking` and `enable_thinking` must match")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_exposure_requires_tools_and_an_enabled_choice() {
|
||||
let mut request = TestRenderRequest::for_test();
|
||||
assert!(!request.as_request().tool_parsing_enabled());
|
||||
|
||||
request.tools.push(crate::Tool {
|
||||
name: "lookup".to_string(),
|
||||
description: None,
|
||||
parameters: json!({"type": "object"}),
|
||||
strict: None,
|
||||
});
|
||||
assert!(!request.as_request().tool_parsing_enabled());
|
||||
|
||||
request.tool_choice = ChatToolChoice::Auto;
|
||||
assert!(request.as_request().tool_parsing_enabled());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn borrowed_request_carries_optional_documents() {
|
||||
let mut request = TestRenderRequest::for_test();
|
||||
request.documents = Some(vec![json!({"title": "doc"})]);
|
||||
|
||||
let borrowed: RenderRequest<'_> = request.as_request();
|
||||
assert_eq!(
|
||||
borrowed.documents.unwrap(),
|
||||
request.documents.as_deref().unwrap()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
# vLLM Example Templates
|
||||
|
||||
These fixtures are copied from `vllm/examples/`.
|
||||
|
||||
They are used by `src/chat-renderer/src/hf/format.rs` tests to keep
|
||||
chat-template content-format detection aligned with Python vLLM behavior.
|
||||
@@ -1,14 +0,0 @@
|
||||
[package]
|
||||
name = "vllm-chat-types"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
easy-ext.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_with.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,155 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::ops::Deref;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// One finalized assistant tool call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AssistantToolCall {
|
||||
/// Stable tool-call identifier.
|
||||
pub id: String,
|
||||
/// Function name selected by the assistant.
|
||||
pub name: String,
|
||||
/// Serialized function arguments.
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// Semantic kind of one assistant output block.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AssistantBlockKind {
|
||||
/// Visible final-answer text.
|
||||
Text,
|
||||
/// Extracted reasoning content.
|
||||
Reasoning,
|
||||
/// One finalized tool call.
|
||||
ToolCall,
|
||||
}
|
||||
|
||||
/// One structured assistant output block.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AssistantContentBlock {
|
||||
/// Visible final-answer text.
|
||||
Text {
|
||||
/// Visible text.
|
||||
text: String,
|
||||
},
|
||||
/// Extracted reasoning content.
|
||||
Reasoning {
|
||||
/// Reasoning text.
|
||||
text: String,
|
||||
},
|
||||
/// One finalized tool call.
|
||||
ToolCall(AssistantToolCall),
|
||||
}
|
||||
|
||||
impl AssistantContentBlock {
|
||||
/// Return the semantic kind of this block.
|
||||
pub fn kind(&self) -> AssistantBlockKind {
|
||||
match self {
|
||||
Self::Text { .. } => AssistantBlockKind::Text,
|
||||
Self::Reasoning { .. } => AssistantBlockKind::Reasoning,
|
||||
Self::ToolCall(..) => AssistantBlockKind::ToolCall,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return this block as one finalized tool call when applicable.
|
||||
pub fn as_tool_call(&self) -> Option<&AssistantToolCall> {
|
||||
match self {
|
||||
Self::ToolCall(call) => Some(call),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Trim whitespace from text and tool arguments.
|
||||
///
|
||||
/// Returns `None` when trimming makes a text or reasoning block empty.
|
||||
pub fn trim(mut self) -> Option<Self> {
|
||||
match &mut self {
|
||||
Self::Text { text } | Self::Reasoning { text } => {
|
||||
let trimmed_text = text.trim();
|
||||
if trimmed_text.is_empty() {
|
||||
return None;
|
||||
}
|
||||
*text = trimmed_text.to_string();
|
||||
}
|
||||
Self::ToolCall(call) => {
|
||||
call.arguments = call.arguments.trim().to_string();
|
||||
}
|
||||
}
|
||||
Some(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[easy_ext::ext(AssistantMessageExt)]
|
||||
impl [AssistantContentBlock] {
|
||||
/// Concatenate all visible final-answer text blocks.
|
||||
pub fn text(&self) -> String {
|
||||
self.iter()
|
||||
.filter_map(|block| match block {
|
||||
AssistantContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Concatenate all extracted reasoning blocks.
|
||||
pub fn reasoning(&self) -> Option<String> {
|
||||
Some(
|
||||
self.iter()
|
||||
.filter_map(|block| match block {
|
||||
AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.filter(|text: &String| !text.is_empty())
|
||||
}
|
||||
|
||||
/// Return whether this assistant message contains reasoning text.
|
||||
pub fn has_reasoning(&self) -> bool {
|
||||
self.iter().any(|block| match block {
|
||||
AssistantContentBlock::Reasoning { text } => !text.is_empty(),
|
||||
_ => false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Iterate over finalized assistant tool calls in encounter order.
|
||||
pub fn tool_calls(&self) -> impl Iterator<Item = &AssistantToolCall> {
|
||||
self.iter().filter_map(AssistantContentBlock::as_tool_call)
|
||||
}
|
||||
|
||||
/// Return whether this assistant message contains any tool-call blocks.
|
||||
pub fn has_tool_calls(&self) -> bool {
|
||||
self.iter().any(|block| matches!(block, AssistantContentBlock::ToolCall(_)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Final structured assistant message assembled from parsed output.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AssistantMessage {
|
||||
/// Assistant content blocks in emission order.
|
||||
pub content: Vec<AssistantContentBlock>,
|
||||
}
|
||||
|
||||
impl Deref for AssistantMessage {
|
||||
type Target = [AssistantContentBlock];
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.content
|
||||
}
|
||||
}
|
||||
|
||||
impl AssistantMessage {
|
||||
/// Push one new block to the end of the message content.
|
||||
pub fn push_block(&mut self, block: AssistantContentBlock) {
|
||||
self.content.push(block);
|
||||
}
|
||||
|
||||
/// Trim all blocks and remove text blocks that become empty.
|
||||
pub fn trim(mut self) -> Self {
|
||||
self.content = self.content.into_iter().filter_map(AssistantContentBlock::trim).collect();
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -1,191 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Detail level requested for an OpenAI-style image input.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ImageDetail {
|
||||
/// Let the model-specific multimodal processor select the detail level.
|
||||
#[default]
|
||||
Auto,
|
||||
/// Request low-detail image processing.
|
||||
Low,
|
||||
/// Request high-detail image processing.
|
||||
High,
|
||||
}
|
||||
|
||||
/// One chat content part in OpenAI-style block format.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ChatContentPart {
|
||||
/// One plain-text content block.
|
||||
Text {
|
||||
/// Plain-text content.
|
||||
text: String,
|
||||
},
|
||||
/// One image URL or data URL content block.
|
||||
ImageUrl {
|
||||
/// Image URL or data URL.
|
||||
image_url: String,
|
||||
/// Requested image detail level.
|
||||
detail: Option<ImageDetail>,
|
||||
/// Optional caller-provided media identifier.
|
||||
uuid: Option<String>,
|
||||
},
|
||||
/// One video URL or data URL content block.
|
||||
VideoUrl {
|
||||
/// Video URL or data URL.
|
||||
video_url: String,
|
||||
/// Optional caller-provided media identifier.
|
||||
uuid: Option<String>,
|
||||
},
|
||||
/// One `input_audio` content block carrying base64-encoded audio bytes.
|
||||
InputAudio {
|
||||
/// Base64-encoded audio bytes.
|
||||
data: String,
|
||||
/// Optional audio format such as `wav` or `mp3`.
|
||||
format: Option<String>,
|
||||
/// Optional caller-provided media identifier.
|
||||
uuid: Option<String>,
|
||||
},
|
||||
/// One audio URL or data URL content block.
|
||||
AudioUrl {
|
||||
/// Audio URL or data URL.
|
||||
audio_url: String,
|
||||
/// Optional caller-provided media identifier.
|
||||
uuid: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ChatContentPart {
|
||||
/// Construct one text content part with plain string content.
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self::Text { text: text.into() }
|
||||
}
|
||||
|
||||
/// Construct one image URL content part with the given URL string.
|
||||
pub fn image_url(image_url: impl Into<String>) -> Self {
|
||||
Self::ImageUrl {
|
||||
image_url: image_url.into(),
|
||||
detail: None,
|
||||
uuid: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one video URL content part with the given URL string.
|
||||
pub fn video_url(video_url: impl Into<String>) -> Self {
|
||||
Self::VideoUrl {
|
||||
video_url: video_url.into(),
|
||||
uuid: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one base64-encoded input-audio content part.
|
||||
pub fn input_audio(data: impl Into<String>, format: Option<String>) -> Self {
|
||||
Self::InputAudio {
|
||||
data: data.into(),
|
||||
format,
|
||||
uuid: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one audio URL content part with the given URL string.
|
||||
pub fn audio_url(audio_url: impl Into<String>) -> Self {
|
||||
Self::AudioUrl {
|
||||
audio_url: audio_url.into(),
|
||||
uuid: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the text content of this part.
|
||||
///
|
||||
/// Returns the static content-part type for multimodal content.
|
||||
pub fn as_text(&self) -> Result<&str, &'static str> {
|
||||
match self {
|
||||
Self::Text { text } => Ok(text),
|
||||
Self::ImageUrl { .. } => Err("image_url"),
|
||||
Self::VideoUrl { .. } => Err("video_url"),
|
||||
Self::InputAudio { .. } => Err("input_audio"),
|
||||
Self::AudioUrl { .. } => Err("audio_url"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return whether this part is a text block with empty content.
|
||||
fn is_empty_text(&self) -> bool {
|
||||
matches!(self, Self::Text { text } if text.is_empty())
|
||||
}
|
||||
|
||||
/// Return whether this part contains any multimodal content.
|
||||
fn is_multimodal(&self) -> bool {
|
||||
match self {
|
||||
Self::Text { .. } => false,
|
||||
Self::ImageUrl { .. }
|
||||
| Self::VideoUrl { .. }
|
||||
| Self::InputAudio { .. }
|
||||
| Self::AudioUrl { .. } => true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Chat content represented as a string or OpenAI-style content parts.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatContent {
|
||||
/// Simple text content.
|
||||
Text(String),
|
||||
/// OpenAI-style content parts.
|
||||
Parts(Vec<ChatContentPart>),
|
||||
}
|
||||
|
||||
impl ChatContent {
|
||||
/// Flatten text parts into one string without adding separators.
|
||||
///
|
||||
/// Returns the static content-part type when the content is multimodal.
|
||||
pub fn try_flatten_to_text(&self) -> Result<String, &'static str> {
|
||||
Ok(match self {
|
||||
Self::Text(text) => text.clone(),
|
||||
Self::Parts(parts) => parts
|
||||
.iter()
|
||||
.map(ChatContentPart::as_text)
|
||||
.collect::<Result<Vec<_>, _>>()?
|
||||
.concat(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Return whether the content has no text or only empty text blocks.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
match self {
|
||||
Self::Text(text) => text.is_empty(),
|
||||
Self::Parts(parts) => parts.iter().all(ChatContentPart::is_empty_text),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return whether this content contains any multimodal parts.
|
||||
pub fn has_multimodal(&self) -> bool {
|
||||
match self {
|
||||
Self::Text(_) => false,
|
||||
Self::Parts(parts) => parts.iter().any(ChatContentPart::is_multimodal),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ChatContent {
|
||||
fn from(value: String) -> Self {
|
||||
Self::Text(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for ChatContent {
|
||||
fn from(value: &str) -> Self {
|
||||
Self::Text(value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Vec<ChatContentPart>> for ChatContent {
|
||||
fn from(value: Vec<ChatContentPart>) -> Self {
|
||||
Self::Parts(value)
|
||||
}
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
//! Engine-independent data types shared by chat renderers and output parsers.
|
||||
//!
|
||||
//! This crate defines chat history, rendering options, tool descriptions, and
|
||||
//! structured assistant payloads. Serving requests, streamed events, renderer
|
||||
//! implementations, parser state, and engine metadata live in their owning
|
||||
//! crates.
|
||||
|
||||
mod assistant;
|
||||
mod content;
|
||||
mod message;
|
||||
mod options;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
mod tool;
|
||||
|
||||
pub use assistant::{
|
||||
AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
|
||||
AssistantToolCall,
|
||||
};
|
||||
pub use content::{ChatContent, ChatContentPart, ImageDetail};
|
||||
pub use message::{ChatMessage, ChatRole};
|
||||
pub use options::{ChatOptions, ChatToolChoice, GenerationPromptMode, ReasoningEffort};
|
||||
pub use tool::Tool;
|
||||
@@ -1,195 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, ChatContent, Tool};
|
||||
|
||||
/// Role label for one chat message.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ChatRole {
|
||||
/// System instructions.
|
||||
System,
|
||||
/// Developer instructions.
|
||||
Developer,
|
||||
/// User input.
|
||||
User,
|
||||
/// Assistant history.
|
||||
Assistant,
|
||||
/// Result of an assistant tool call.
|
||||
ToolResponse,
|
||||
}
|
||||
|
||||
impl ChatRole {
|
||||
/// Return the role string exposed to chat templates.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::System => "system",
|
||||
Self::Developer => "developer",
|
||||
Self::User => "user",
|
||||
Self::Assistant => "assistant",
|
||||
Self::ToolResponse => "tool_response",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One chat message.
|
||||
///
|
||||
/// Original Python API reference:
|
||||
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/chat_utils.py#L309-L333>
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "role", rename_all = "snake_case")]
|
||||
pub enum ChatMessage {
|
||||
/// System message.
|
||||
System {
|
||||
/// Message content.
|
||||
content: ChatContent,
|
||||
},
|
||||
/// Developer message with optional message-local tools.
|
||||
Developer {
|
||||
/// Message content.
|
||||
content: ChatContent,
|
||||
/// Tools introduced by this developer message.
|
||||
tools: Option<Vec<Tool>>,
|
||||
},
|
||||
/// User message.
|
||||
User {
|
||||
/// Message content.
|
||||
content: ChatContent,
|
||||
},
|
||||
/// Assistant history assembled from structured blocks.
|
||||
Assistant {
|
||||
/// Structured assistant content.
|
||||
content: Vec<AssistantContentBlock>,
|
||||
},
|
||||
/// Tool response associated with one prior assistant tool call.
|
||||
ToolResponse {
|
||||
/// Tool response content.
|
||||
content: ChatContent,
|
||||
/// Identifier of the assistant tool call being answered.
|
||||
tool_call_id: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
/// Construct one chat message with plain string content.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics for [`ChatRole::ToolResponse`], which requires a tool-call ID.
|
||||
/// Use [`Self::tool_response`] for tool responses.
|
||||
pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
|
||||
let content: String = text.into();
|
||||
|
||||
match role {
|
||||
ChatRole::System => Self::system(content),
|
||||
ChatRole::Developer => Self::developer(content, None),
|
||||
ChatRole::User => Self::user(content),
|
||||
ChatRole::Assistant => Self::assistant_text(content),
|
||||
ChatRole::ToolResponse => {
|
||||
panic!(
|
||||
"tool response messages require a tool_call_id; \
|
||||
use ChatMessage::tool_response() instead"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one system message.
|
||||
pub fn system(content: impl Into<ChatContent>) -> Self {
|
||||
Self::System {
|
||||
content: content.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one developer message.
|
||||
pub fn developer(content: impl Into<ChatContent>, tools: Option<Vec<Tool>>) -> Self {
|
||||
Self::Developer {
|
||||
content: content.into(),
|
||||
tools,
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one user message.
|
||||
pub fn user(content: impl Into<ChatContent>) -> Self {
|
||||
Self::User {
|
||||
content: content.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one assistant message with plain string content.
|
||||
pub fn assistant_text(text: impl Into<String>) -> Self {
|
||||
Self::Assistant {
|
||||
content: vec![AssistantContentBlock::Text { text: text.into() }],
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct one assistant message with structured content blocks.
|
||||
pub fn assistant_blocks(content: Vec<AssistantContentBlock>) -> Self {
|
||||
Self::Assistant { content }
|
||||
}
|
||||
|
||||
/// Construct one tool-response message.
|
||||
pub fn tool_response(content: impl Into<ChatContent>, tool_call_id: impl Into<String>) -> Self {
|
||||
Self::ToolResponse {
|
||||
content: content.into(),
|
||||
tool_call_id: tool_call_id.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the role of this message.
|
||||
pub fn role(&self) -> ChatRole {
|
||||
match self {
|
||||
Self::System { .. } => ChatRole::System,
|
||||
Self::Developer { .. } => ChatRole::Developer,
|
||||
Self::User { .. } => ChatRole::User,
|
||||
Self::Assistant { .. } => ChatRole::Assistant,
|
||||
Self::ToolResponse { .. } => ChatRole::ToolResponse,
|
||||
}
|
||||
}
|
||||
|
||||
/// Concatenate the visible text carried by this message.
|
||||
///
|
||||
/// Returns the static content-part type when a non-assistant message
|
||||
/// contains multimodal content.
|
||||
pub fn text_content(&self) -> Result<String, &'static str> {
|
||||
match self {
|
||||
Self::System { content }
|
||||
| Self::Developer { content, .. }
|
||||
| Self::User { content }
|
||||
| Self::ToolResponse { content, .. } => content.try_flatten_to_text(),
|
||||
Self::Assistant { content } => Ok(content.text()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Concatenate assistant reasoning text when present.
|
||||
pub fn reasoning_content(&self) -> Option<String> {
|
||||
match self {
|
||||
Self::Assistant { content } => content.reasoning(),
|
||||
Self::System { .. }
|
||||
| Self::Developer { .. }
|
||||
| Self::User { .. }
|
||||
| Self::ToolResponse { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return whether this message contains multimodal content.
|
||||
pub fn has_multimodal(&self) -> bool {
|
||||
match self {
|
||||
Self::System { content }
|
||||
| Self::Developer { content, .. }
|
||||
| Self::User { content }
|
||||
| Self::ToolResponse { content, .. } => content.has_multimodal(),
|
||||
Self::Assistant { .. } => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AssistantMessage> for ChatMessage {
|
||||
fn from(value: AssistantMessage) -> Self {
|
||||
Self::Assistant {
|
||||
content: value.content,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,134 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Controls how prompt rendering should end after the existing chat history.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum GenerationPromptMode {
|
||||
/// Append a generation prompt for a new assistant turn.
|
||||
///
|
||||
/// Equivalent to `add_generation_prompt = true` and
|
||||
/// `continue_final_message = false`.
|
||||
#[default]
|
||||
StartNewAssistant,
|
||||
/// Leave the final assistant message open so generation continues it.
|
||||
///
|
||||
/// Equivalent to `add_generation_prompt = false` and
|
||||
/// `continue_final_message = true`.
|
||||
ContinueFinalAssistant,
|
||||
/// Render the existing chat history without adding any trailing generation
|
||||
/// prompt.
|
||||
///
|
||||
/// Equivalent to `add_generation_prompt = false` and
|
||||
/// `continue_final_message = false`.
|
||||
NoGenerationPrompt,
|
||||
}
|
||||
|
||||
/// Effort level for reasoning models.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ReasoningEffort {
|
||||
/// Disable reasoning.
|
||||
None,
|
||||
/// Use the smallest available reasoning effort.
|
||||
Minimal,
|
||||
/// Use low reasoning effort.
|
||||
Low,
|
||||
/// Use medium reasoning effort.
|
||||
Medium,
|
||||
/// Use high reasoning effort.
|
||||
High,
|
||||
/// Use extra-high reasoning effort.
|
||||
XHigh,
|
||||
/// Use the largest available reasoning effort.
|
||||
Max,
|
||||
}
|
||||
|
||||
impl ReasoningEffort {
|
||||
/// Return the lowercase value exposed to chat templates.
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::None => "none",
|
||||
Self::Minimal => "minimal",
|
||||
Self::Low => "low",
|
||||
Self::Medium => "medium",
|
||||
Self::High => "high",
|
||||
Self::XHigh => "xhigh",
|
||||
Self::Max => "max",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Chat-template-related request options.
|
||||
///
|
||||
/// These are the chat controls that currently affect prompt rendering.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatOptions {
|
||||
/// Controls whether rendering starts a new assistant turn, continues the
|
||||
/// final assistant message, or emits no trailing generation prompt.
|
||||
pub generation_prompt_mode: GenerationPromptMode,
|
||||
|
||||
/// Per-request Jinja chat template override.
|
||||
///
|
||||
/// The renderer uses this template in place of the model's default chat
|
||||
/// template when it is present.
|
||||
pub chat_template: Option<String>,
|
||||
|
||||
/// Effort level exposed to chat templates for reasoning models.
|
||||
pub reasoning_effort: Option<ReasoningEffort>,
|
||||
|
||||
/// Additional keyword arguments exposed to the chat template.
|
||||
pub template_kwargs: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
impl Default for ChatOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
generation_prompt_mode: GenerationPromptMode::StartNewAssistant,
|
||||
chat_template: None,
|
||||
reasoning_effort: None,
|
||||
template_kwargs: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatOptions {
|
||||
/// Return whether rendering adds a prompt for a new assistant turn.
|
||||
pub fn add_generation_prompt(&self) -> bool {
|
||||
matches!(
|
||||
self.generation_prompt_mode,
|
||||
GenerationPromptMode::StartNewAssistant
|
||||
)
|
||||
}
|
||||
|
||||
/// Return whether rendering continues the final assistant message.
|
||||
pub fn continue_final_message(&self) -> bool {
|
||||
matches!(
|
||||
self.generation_prompt_mode,
|
||||
GenerationPromptMode::ContinueFinalAssistant
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool-choice semantics supported by the shared chat types.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ChatToolChoice {
|
||||
/// Disable tool calling.
|
||||
#[default]
|
||||
None,
|
||||
/// Let the model choose whether to call a tool.
|
||||
Auto,
|
||||
/// Require the model to call a tool.
|
||||
Required,
|
||||
/// Require one named function.
|
||||
Function {
|
||||
/// Required function name.
|
||||
name: String,
|
||||
},
|
||||
}
|
||||
@@ -1,117 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use serde_json::{json, to_value};
|
||||
|
||||
use crate::{AssistantContentBlock, ChatContent, ChatContentPart, ChatMessage, ChatRole, Tool};
|
||||
|
||||
#[test]
|
||||
fn chat_content_deserializes_from_raw_string() {
|
||||
let content: ChatContent = serde_json::from_value(json!("hello")).unwrap();
|
||||
assert_eq!(content, ChatContent::Text("hello".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_content_video_url_part_round_trips_through_serde() {
|
||||
let content = ChatContent::Parts(vec![ChatContentPart::VideoUrl {
|
||||
video_url: "https://example.com/demo.mp4".to_string(),
|
||||
uuid: Some("video-1".to_string()),
|
||||
}]);
|
||||
|
||||
let value = to_value(&content).unwrap();
|
||||
assert_eq!(
|
||||
value,
|
||||
json!([{
|
||||
"type": "video_url",
|
||||
"video_url": "https://example.com/demo.mp4",
|
||||
"uuid": "video-1",
|
||||
}])
|
||||
);
|
||||
let decoded: ChatContent = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(decoded, content);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_content_deserializes_from_openai_text_blocks() {
|
||||
let content: ChatContent =
|
||||
serde_json::from_value(json!([{ "type": "text", "text": "hello" }])).unwrap();
|
||||
assert_eq!(
|
||||
content,
|
||||
ChatContent::Parts(vec![ChatContentPart::text("hello")])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_content_from_string_like_values_builds_text() {
|
||||
assert_eq!(
|
||||
ChatContent::from("hello"),
|
||||
ChatContent::Text("hello".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
ChatContent::from("hello".to_string()),
|
||||
ChatContent::Text("hello".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_content_try_flattens_text_parts_without_separators() {
|
||||
let content = ChatContent::Parts(vec![
|
||||
ChatContentPart::text("hello"),
|
||||
ChatContentPart::text(" world"),
|
||||
]);
|
||||
assert_eq!(content.try_flatten_to_text().unwrap(), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multimodal_content_parts_return_static_type_names() {
|
||||
let parts = [
|
||||
(ChatContentPart::image_url("image"), "image_url"),
|
||||
(ChatContentPart::video_url("video"), "video_url"),
|
||||
(ChatContentPart::input_audio("audio", None), "input_audio"),
|
||||
(ChatContentPart::audio_url("audio"), "audio_url"),
|
||||
];
|
||||
|
||||
for (part, expected) in parts {
|
||||
assert_eq!(part.as_text(), Err(expected));
|
||||
assert_eq!(
|
||||
ChatContent::Parts(vec![part]).try_flatten_to_text(),
|
||||
Err(expected)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_message_collects_visible_and_reasoning_text() {
|
||||
let message = ChatMessage::assistant_blocks(vec![
|
||||
AssistantContentBlock::Reasoning {
|
||||
text: "inner".to_string(),
|
||||
},
|
||||
AssistantContentBlock::Text {
|
||||
text: "outer".to_string(),
|
||||
},
|
||||
]);
|
||||
|
||||
assert_eq!(message.role(), ChatRole::Assistant);
|
||||
assert_eq!(message.text_content().unwrap(), "outer");
|
||||
assert_eq!(message.reasoning_content().as_deref(), Some("inner"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn developer_message_round_trips_through_serde() {
|
||||
let message = ChatMessage::developer(
|
||||
"hello",
|
||||
Some(vec![Tool {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
}),
|
||||
strict: Some(true),
|
||||
}]),
|
||||
);
|
||||
|
||||
let value = to_value(&message).unwrap();
|
||||
let decoded: ChatMessage = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(decoded, message);
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// One function-style tool made available to the model.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Tool {
|
||||
/// Function name exposed to the model.
|
||||
pub name: String,
|
||||
/// Optional human-readable function description.
|
||||
pub description: Option<String>,
|
||||
/// JSON Schema describing the function parameters.
|
||||
pub parameters: Value,
|
||||
/// Optional strict-schema enforcement request.
|
||||
pub strict: Option<bool>,
|
||||
}
|
||||
@@ -7,24 +7,29 @@ license.workspace = true
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
asynk-strim-attr.workspace = true
|
||||
easy-ext.workspace = true
|
||||
futures.workspace = true
|
||||
half.workspace = true
|
||||
indexmap.workspace = true
|
||||
itertools.workspace = true
|
||||
llm-multimodal.workspace = true
|
||||
minijinja.workspace = true
|
||||
minijinja-contrib.workspace = true
|
||||
openai-harmony.workspace = true
|
||||
reqwest-0-13.workspace = true
|
||||
serde.workspace = true
|
||||
serde-json-fmt.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_with.workspace = true
|
||||
strum.workspace = true
|
||||
subenum.workspace = true
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
time.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
trait-set.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-chat-renderer.workspace = true
|
||||
vllm-chat-types.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
|
||||
@@ -4,13 +4,6 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use tracing::info;
|
||||
use vllm_chat_renderer::hf::{
|
||||
HfChatRenderer, HfRendererConfig, HfRendererFiles, MultimodalRenderInfo,
|
||||
};
|
||||
use vllm_chat_renderer::{
|
||||
DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
|
||||
InklingChatRenderer,
|
||||
};
|
||||
use vllm_text::backend::hf::{HfTextBackend, ResolvedModelFiles, load_model_config};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::{DynTextBackend, TextBackend as _};
|
||||
@@ -24,6 +17,11 @@ use crate::multimodal::{MultimodalConfigFiles, MultimodalModelInfo};
|
||||
use crate::output::{
|
||||
DefaultChatOutputProcessor, HarmonyChatOutputProcessor, validate_harmony_parser_overrides,
|
||||
};
|
||||
use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo};
|
||||
use crate::renderer::{
|
||||
DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer,
|
||||
InklingChatRenderer,
|
||||
};
|
||||
use crate::request::ChatRequest;
|
||||
use crate::{DynChatOutputProcessor, RendererSelection};
|
||||
|
||||
@@ -67,16 +65,9 @@ impl HfChatBackend {
|
||||
let chat_renderer: DynChatRenderer = match renderer {
|
||||
RendererSelection::Auto => unreachable!("renderer auto should be resolved above"),
|
||||
RendererSelection::Hf => Arc::new(HfChatRenderer::load(
|
||||
HfRendererFiles {
|
||||
tokenizer_config: files.tokenizer_config_path.as_deref(),
|
||||
chat_template: files.chat_template_path.as_deref(),
|
||||
},
|
||||
HfRendererConfig {
|
||||
chat_template: options.chat_template,
|
||||
default_template_kwargs: options.default_chat_template_kwargs,
|
||||
content_format: options.chat_template_content_format,
|
||||
multimodal: multimodal_render_info,
|
||||
},
|
||||
&files,
|
||||
options,
|
||||
multimodal_render_info,
|
||||
)?),
|
||||
RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()),
|
||||
RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()),
|
||||
@@ -174,6 +165,7 @@ mod tests {
|
||||
|
||||
use tempfile::tempdir;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::TokenizerSource;
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
@@ -181,7 +173,7 @@ mod tests {
|
||||
use super::HfChatBackend;
|
||||
use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
|
||||
use crate::request::{ChatContent, ChatMessage, ChatRequest};
|
||||
use crate::{ParserSelection, RenderedPromptContent, RendererSelection};
|
||||
use crate::{ParserSelection, RendererSelection};
|
||||
|
||||
fn request_with_user_text(text: &str) -> ChatRequest {
|
||||
ChatRequest {
|
||||
@@ -249,13 +241,11 @@ mod tests {
|
||||
config_json: &str,
|
||||
tokenizer_config_json: &str,
|
||||
) -> String {
|
||||
let backend = backend_for_selection(renderer, config_json, tokenizer_config_json);
|
||||
let request = request_with_user_text("hello");
|
||||
backend
|
||||
backend_for_selection(renderer, config_json, tokenizer_config_json)
|
||||
.chat_renderer()
|
||||
.render(request.as_render_request())
|
||||
.render(&request_with_user_text("hello"))
|
||||
.unwrap()
|
||||
.content
|
||||
.prompt
|
||||
.into_text()
|
||||
.expect("renderer should return text prompt")
|
||||
}
|
||||
@@ -293,13 +283,9 @@ mod tests {
|
||||
r#"{"chat_template":"{{ messages[0].content }}"}"#,
|
||||
);
|
||||
|
||||
let render_request = request_with_user_text("hello");
|
||||
let prompt = backend
|
||||
.chat_renderer()
|
||||
.render(render_request.as_render_request())
|
||||
.unwrap()
|
||||
.content;
|
||||
assert!(matches!(prompt, RenderedPromptContent::TokenIds(_)));
|
||||
let prompt =
|
||||
backend.chat_renderer().render(&request_with_user_text("hello")).unwrap().prompt;
|
||||
assert!(matches!(prompt, Prompt::TokenIds(_)));
|
||||
|
||||
let mut request = request_with_user_text("hello");
|
||||
let error = match backend.new_chat_output_processor(
|
||||
|
||||
@@ -10,8 +10,9 @@ use vllm_text::{DynTextBackend, TextBackend};
|
||||
use crate::error::Result;
|
||||
use crate::multimodal::MultimodalModelInfo;
|
||||
use crate::output::DynChatOutputProcessor;
|
||||
use crate::renderer::DynChatRenderer;
|
||||
use crate::request::ChatRequest;
|
||||
use crate::{ChatTemplateContentFormatOption, DynChatRenderer, ParserSelection, RendererSelection};
|
||||
use crate::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
|
||||
|
||||
pub mod hf;
|
||||
|
||||
|
||||
@@ -9,8 +9,18 @@ type BoxedError = Box<dyn std::error::Error + Send + Sync>;
|
||||
#[derive(Debug, Error, Macro)]
|
||||
#[thiserror_ext(macro(path = "crate::error"))]
|
||||
pub enum Error {
|
||||
#[error("chat request must contain at least one message")]
|
||||
EmptyMessages,
|
||||
#[error("cannot continue the final message when the last message is not from the assistant")]
|
||||
ContinueFinalAssistantWithoutFinalAssistant,
|
||||
#[error("chat template is required but none was configured")]
|
||||
MissingChatTemplate,
|
||||
#[error("chat template error: {0}")]
|
||||
ChatTemplate(String),
|
||||
#[error("multimodal input is not supported by this chat renderer")]
|
||||
UnsupportedMultimodalRenderer,
|
||||
#[error("unsupported multimodal content: {0}")]
|
||||
UnsupportedMultimodalContent(&'static str),
|
||||
#[error("`{modality}` input is not supported by this model")]
|
||||
UnsupportedModality { modality: String },
|
||||
#[error("multimodal preprocessing error: {0}")]
|
||||
@@ -64,8 +74,6 @@ pub enum Error {
|
||||
#[error(transparent)]
|
||||
Text(#[from] vllm_text::Error),
|
||||
#[error(transparent)]
|
||||
Renderer(#[from] vllm_chat_renderer::Error),
|
||||
#[error(transparent)]
|
||||
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
|
||||
}
|
||||
|
||||
@@ -77,8 +85,10 @@ impl Error {
|
||||
match self {
|
||||
Self::PromptTooLong { .. } => true,
|
||||
Self::Text(error) => error.is_request_validation_error(),
|
||||
Self::Renderer(error) => error.is_request_validation_error(),
|
||||
Self::UnsupportedMultimodalRenderer | Self::UnsupportedModality { .. } => true,
|
||||
Self::UnsupportedMultimodalRenderer
|
||||
| Self::UnsupportedMultimodalContent(_)
|
||||
| Self::UnsupportedModality { .. } => true,
|
||||
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
+146
-29
@@ -1,17 +1,155 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use vllm_llm::TokenUsage;
|
||||
use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
|
||||
|
||||
use crate::FinishReason;
|
||||
|
||||
pub use vllm_chat_types::{
|
||||
AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
|
||||
AssistantToolCall,
|
||||
};
|
||||
/// One finalized assistant tool call.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AssistantToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
/// Semantic kind of one assistant output block.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AssistantBlockKind {
|
||||
/// Visible final-answer text.
|
||||
Text,
|
||||
/// Extracted reasoning content.
|
||||
Reasoning,
|
||||
/// One finalized tool call.
|
||||
ToolCall,
|
||||
}
|
||||
|
||||
/// One structured assistant output block.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AssistantContentBlock {
|
||||
/// Visible final-answer text.
|
||||
Text { text: String },
|
||||
/// Extracted reasoning content.
|
||||
Reasoning { text: String },
|
||||
/// One finalized tool call.
|
||||
ToolCall(AssistantToolCall),
|
||||
}
|
||||
|
||||
impl AssistantContentBlock {
|
||||
/// Return the semantic kind of this block.
|
||||
pub fn kind(&self) -> AssistantBlockKind {
|
||||
match self {
|
||||
Self::Text { .. } => AssistantBlockKind::Text,
|
||||
Self::Reasoning { .. } => AssistantBlockKind::Reasoning,
|
||||
Self::ToolCall(..) => AssistantBlockKind::ToolCall,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return this block as one finalized tool call, if applicable.
|
||||
pub fn as_tool_call(&self) -> Option<&AssistantToolCall> {
|
||||
match self {
|
||||
Self::ToolCall(call) => Some(call),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Return a copy of this block with leading and trailing whitespace trimmed from all text
|
||||
/// fields and tool call arguments, or `None` if the resulting text would be empty.
|
||||
pub fn trim(mut self) -> Option<Self> {
|
||||
match &mut self {
|
||||
Self::Text { text } | Self::Reasoning { text } => {
|
||||
let trimmed_text = text.trim();
|
||||
if trimmed_text.is_empty() {
|
||||
return None;
|
||||
} else {
|
||||
*text = trimmed_text.to_string();
|
||||
}
|
||||
}
|
||||
Self::ToolCall(call) => {
|
||||
call.arguments = call.arguments.trim().to_string();
|
||||
}
|
||||
}
|
||||
Some(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[easy_ext::ext(AssistantMessageExt)]
|
||||
impl [AssistantContentBlock] {
|
||||
/// Concatenate all visible final-answer text blocks.
|
||||
pub fn text(&self) -> String {
|
||||
self.iter()
|
||||
.filter_map(|block| match block {
|
||||
AssistantContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Concatenate all extracted reasoning blocks, if any.
|
||||
pub fn reasoning(&self) -> Option<String> {
|
||||
Some(
|
||||
self.iter()
|
||||
.filter_map(|block| match block {
|
||||
AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
.filter(|s: &String| !s.is_empty())
|
||||
}
|
||||
|
||||
/// Return whether this assistant message contains any non-empty reasoning
|
||||
/// text blocks.
|
||||
pub fn has_reasoning(&self) -> bool {
|
||||
self.iter().any(|block| match block {
|
||||
AssistantContentBlock::Reasoning { text } => !text.is_empty(),
|
||||
_ => false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Return finalized assistant tool calls in encounter order.
|
||||
pub fn tool_calls(&self) -> impl Iterator<Item = &AssistantToolCall> {
|
||||
self.iter().filter_map(AssistantContentBlock::as_tool_call)
|
||||
}
|
||||
|
||||
/// Return whether this assistant message contains any tool-call blocks.
|
||||
pub fn has_tool_calls(&self) -> bool {
|
||||
self.iter().any(|block| matches!(block, AssistantContentBlock::ToolCall(_)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Final structured assistant message assembled from the event stream.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AssistantMessage {
|
||||
pub content: Vec<AssistantContentBlock>,
|
||||
}
|
||||
|
||||
impl Deref for AssistantMessage {
|
||||
type Target = [AssistantContentBlock];
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.content
|
||||
}
|
||||
}
|
||||
|
||||
impl AssistantMessage {
|
||||
/// Push one new block to the end of the message content.
|
||||
pub(crate) fn push_block(&mut self, block: AssistantContentBlock) {
|
||||
self.content.push(block);
|
||||
}
|
||||
|
||||
/// Return a copy of this message with leading and trailing whitespace trimmed from all text
|
||||
/// fields and tool call arguments, and with any blocks that are empty after trimming removed.
|
||||
pub fn trim(mut self) -> Self {
|
||||
self.content = self.content.into_iter().filter_map(|block| block.trim()).collect();
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Streamed chat event emitted by [`crate::ChatEventStream`].
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
@@ -26,65 +164,44 @@ pub enum ChatEvent {
|
||||
},
|
||||
/// A new assistant output block has started.
|
||||
BlockStart {
|
||||
/// Stable block index within the assistant message.
|
||||
index: usize,
|
||||
/// Semantic kind of the opened block.
|
||||
kind: AssistantBlockKind,
|
||||
},
|
||||
/// A newly observed delta for one open assistant output block.
|
||||
BlockDelta {
|
||||
/// Stable block index within the assistant message.
|
||||
index: usize,
|
||||
/// Semantic kind of the open block.
|
||||
kind: AssistantBlockKind,
|
||||
/// Newly emitted text.
|
||||
delta: String,
|
||||
},
|
||||
/// Per-decoded-update sample metadata.
|
||||
/// Per-decoded-update sample metadata: logprobs and/or output token IDs.
|
||||
LogprobsDelta {
|
||||
/// Decoded output logprobs, when requested.
|
||||
logprobs: Option<DecodedLogprobs>,
|
||||
/// Output token IDs emitted by this update.
|
||||
token_ids: Vec<u32>,
|
||||
},
|
||||
/// One assistant output block has ended.
|
||||
BlockEnd {
|
||||
/// Stable block index within the assistant message.
|
||||
index: usize,
|
||||
/// Finalized block.
|
||||
block: AssistantContentBlock,
|
||||
},
|
||||
/// One tool call has started.
|
||||
ToolCallStart {
|
||||
/// Stable tool-call index within the assistant message.
|
||||
index: usize,
|
||||
/// Stable tool-call identifier.
|
||||
id: String,
|
||||
/// Function name selected by the assistant.
|
||||
name: String,
|
||||
},
|
||||
/// One incremental tool-call arguments delta.
|
||||
ToolCallArgumentsDelta {
|
||||
/// Stable tool-call index within the assistant message.
|
||||
index: usize,
|
||||
/// Newly emitted arguments text.
|
||||
delta: String,
|
||||
},
|
||||
/// One incremental tool-call arguments delta for the currently open tool
|
||||
/// call.
|
||||
ToolCallArgumentsDelta { index: usize, delta: String },
|
||||
/// One tool call has ended.
|
||||
ToolCallEnd {
|
||||
/// Stable tool-call index within the assistant message.
|
||||
index: usize,
|
||||
/// Finalized tool call.
|
||||
call: AssistantToolCall,
|
||||
},
|
||||
/// Terminal event carrying the final assembled assistant message and finish
|
||||
/// metadata.
|
||||
Done {
|
||||
/// Final structured assistant message.
|
||||
message: AssistantMessage,
|
||||
/// Final token usage.
|
||||
usage: TokenUsage,
|
||||
/// Reason generation stopped.
|
||||
finish_reason: FinishReason,
|
||||
/// Connector-specific KV transfer parameters for disaggregated serving.
|
||||
kv_transfer_params: Option<serde_json::Value>,
|
||||
|
||||
+71
-119
@@ -30,17 +30,16 @@ pub use parser::reasoning::{
|
||||
ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory,
|
||||
};
|
||||
pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory};
|
||||
pub use renderer::hf::ChatTemplateContentFormatOption;
|
||||
pub use renderer::{
|
||||
ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer,
|
||||
HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection,
|
||||
};
|
||||
pub use request::{
|
||||
ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
|
||||
ChatToolChoice, GenerationPromptMode, ImageDetail, ReasoningEffort, SamplingParams,
|
||||
ChatToolChoice, GenerationPromptMode, ReasoningEffort, SamplingParams,
|
||||
};
|
||||
pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage};
|
||||
pub use vllm_chat_renderer::hf::ChatTemplateContentFormatOption;
|
||||
pub use vllm_chat_renderer::{
|
||||
ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer,
|
||||
Error as RendererError, HarmonyChatRenderer, InklingChatRenderer, RenderRequest,
|
||||
RenderedPrompt, RenderedPromptContent, RendererSelection, Result as RendererResult,
|
||||
};
|
||||
pub use vllm_llm::FinishReason;
|
||||
|
||||
mod backend;
|
||||
@@ -49,12 +48,12 @@ mod event;
|
||||
pub mod multimodal;
|
||||
mod output;
|
||||
mod parser;
|
||||
mod renderer;
|
||||
mod request;
|
||||
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,94 +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((multimodal::to_text_prompt(rendered.content), 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.as_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
|
||||
@@ -184,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.
|
||||
@@ -199,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,
|
||||
}
|
||||
@@ -226,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
|
||||
}
|
||||
|
||||
@@ -258,37 +172,75 @@ 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.
|
||||
///
|
||||
/// Uses the same render, multimodal finalization, and encoding pipeline as
|
||||
/// [`Self::chat`], but stops after token IDs so `/tokenize` counts match
|
||||
/// what generation would see. Used by `POST /tokenize` (chat form).
|
||||
/// Same render → [`multimodal::finalize_rendered_prompt`] → encode pipeline as
|
||||
/// [`Self::chat`], but stops after token IDs so `/tokenize` counts match what
|
||||
/// generation would see. Used by `POST /tokenize` (chat form).
|
||||
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
|
||||
request.validate()?;
|
||||
|
||||
let rendered =
|
||||
self.processor.backend.chat_renderer().render(request.as_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 {
|
||||
|
||||
@@ -32,8 +32,8 @@ use vllm_text::Prompt;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
|
||||
use crate::error::{Error, Result, bail_multimodal, multimodal};
|
||||
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ImageDetail};
|
||||
use crate::{RenderedPrompt, RenderedPromptContent};
|
||||
use crate::renderer::RenderedPrompt;
|
||||
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
|
||||
|
||||
mod audio;
|
||||
mod expand;
|
||||
@@ -486,16 +486,16 @@ pub(crate) async fn finalize_rendered_prompt(
|
||||
model_dtype: ModelDtype,
|
||||
) -> Result<(Prompt, Option<MmFeatures>)> {
|
||||
if !request.has_multimodal() {
|
||||
return Ok((to_text_prompt(rendered.content), None));
|
||||
return Ok((rendered.prompt, None));
|
||||
}
|
||||
let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?;
|
||||
let mut prompt_token_ids = match rendered.content {
|
||||
RenderedPromptContent::Text(prompt) => info
|
||||
let mut prompt_token_ids = match rendered.prompt {
|
||||
Prompt::Text(prompt) => info
|
||||
.context
|
||||
.tokenizer()
|
||||
.encode(&prompt, request.add_special_tokens)
|
||||
.map_err(|error| multimodal!("{error}"))?,
|
||||
RenderedPromptContent::TokenIds(token_ids) => token_ids,
|
||||
Prompt::TokenIds(token_ids) => token_ids,
|
||||
};
|
||||
let media_parts = extract_media_parts(request)?;
|
||||
let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?;
|
||||
@@ -503,14 +503,6 @@ pub(crate) async fn finalize_rendered_prompt(
|
||||
Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared)))
|
||||
}
|
||||
|
||||
/// Lower an engine-independent renderer artifact into a text-backend prompt.
|
||||
pub(crate) fn to_text_prompt(content: RenderedPromptContent) -> Prompt {
|
||||
match content {
|
||||
RenderedPromptContent::Text(text) => Prompt::Text(text),
|
||||
RenderedPromptContent::TokenIds(token_ids) => Prompt::TokenIds(token_ids),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract media parts from chat messages in message/content order.
|
||||
///
|
||||
/// Assistant history is skipped because generated assistant blocks are already
|
||||
@@ -537,7 +529,7 @@ fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
|
||||
uuid,
|
||||
} => all_parts.push(MediaContentPart::ImageUrl {
|
||||
url: image_url.clone(),
|
||||
detail: detail.map(to_multimodal_image_detail),
|
||||
detail: *detail,
|
||||
uuid: uuid.clone(),
|
||||
}),
|
||||
ChatContentPart::VideoUrl { video_url, uuid } => {
|
||||
@@ -564,15 +556,6 @@ fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
|
||||
Ok(all_parts)
|
||||
}
|
||||
|
||||
/// Convert the protocol-level image detail into the multimodal processor type.
|
||||
fn to_multimodal_image_detail(detail: ImageDetail) -> llm_multimodal::ImageDetail {
|
||||
match detail {
|
||||
ImageDetail::Auto => llm_multimodal::ImageDetail::Auto,
|
||||
ImageDetail::Low => llm_multimodal::ImageDetail::Low,
|
||||
ImageDetail::High => llm_multimodal::ImageDetail::High,
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`.
|
||||
fn input_audio_data_url(data: &str, format: Option<&str>) -> Result<String> {
|
||||
let mime_type = match format {
|
||||
|
||||
@@ -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;
|
||||
@@ -178,8 +178,8 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
|
||||
/// events through two sequential stages once text decoding has
|
||||
/// already happened:
|
||||
///
|
||||
/// 1. `unified_event_stream` — reasoning and tool-call parsing
|
||||
/// 2. `structured_chat_event_stream` — final block assembly
|
||||
/// 1. [`unified_event_stream`] — reasoning and tool-call parsing
|
||||
/// 2. [`structured_chat_event_stream`] — final block assembly
|
||||
fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
|
||||
let parsed = unified_event_stream(decoded, self.parser);
|
||||
let structured = structured_chat_event_stream(parsed, self.parallel_tool_calls);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -21,8 +21,8 @@ use crate::output::{
|
||||
generate_tool_call_id,
|
||||
};
|
||||
use crate::parser::ParserSelection;
|
||||
use crate::renderer::harmony::encoding::harmony_encoding;
|
||||
use crate::request::ChatRequest;
|
||||
use vllm_chat_renderer::harmony::harmony_encoding;
|
||||
|
||||
/// Request-scoped Harmony output processor used for `model_type == "gpt_oss"`.
|
||||
///
|
||||
|
||||
+13
-14
@@ -11,10 +11,8 @@ use serde_json::Value;
|
||||
use serde_json_fmt::JsonFormat;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::{
|
||||
AssistantContentBlock, AssistantMessageExt, AssistantToolCall, ChatContent, ChatMessage,
|
||||
ChatRole, RenderRequest, Tool,
|
||||
};
|
||||
use crate::request::{ChatContent, ChatMessage, ChatRequest, ChatRole, ChatTool};
|
||||
use crate::{AssistantContentBlock, AssistantMessageExt, AssistantToolCall};
|
||||
|
||||
const BOS_TOKEN: &str = "<|begin▁of▁sentence|>";
|
||||
const EOS_TOKEN: &str = "<|end▁of▁sentence|>";
|
||||
@@ -41,7 +39,7 @@ struct RenderedToolSchema<'a> {
|
||||
}
|
||||
|
||||
/// Render one chat request into the final prompt string.
|
||||
pub(super) fn render_request(request: &RenderRequest<'_>) -> Result<String> {
|
||||
pub(super) fn render_request(request: &ChatRequest) -> Result<String> {
|
||||
let thinking_mode = match request.enable_thinking()?.unwrap_or(false) {
|
||||
true => ThinkingMode::Thinking,
|
||||
false => ThinkingMode::Chat,
|
||||
@@ -51,19 +49,20 @@ pub(super) fn render_request(request: &RenderRequest<'_>) -> Result<String> {
|
||||
Some(ChatRole::User | ChatRole::Developer)
|
||||
);
|
||||
let render_offset = isize::from(request.tool_parsing_enabled());
|
||||
let last_user_render_index = find_last_user_render_index(request.messages, render_offset);
|
||||
let last_user_actual_index = find_last_user_actual_index(request.messages);
|
||||
let last_user_render_index =
|
||||
find_last_user_render_index(request.messages.as_slice(), render_offset);
|
||||
let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice());
|
||||
let continue_final_message = request.chat_options.continue_final_message();
|
||||
let mut prompt = String::from(BOS_TOKEN);
|
||||
|
||||
if request.tool_parsing_enabled() {
|
||||
render_system_message(&mut prompt, None, request.tools)?;
|
||||
render_system_message(&mut prompt, None, &request.tools)?;
|
||||
}
|
||||
|
||||
for (message_index, message) in request.messages.iter().enumerate() {
|
||||
render_message(
|
||||
&mut prompt,
|
||||
request.messages,
|
||||
request.messages.as_slice(),
|
||||
message_index,
|
||||
message,
|
||||
render_offset,
|
||||
@@ -169,7 +168,7 @@ fn find_last_user_actual_index(messages: &[ChatMessage]) -> usize {
|
||||
fn render_system_message(
|
||||
out: &mut String,
|
||||
content: Option<&ChatContent>,
|
||||
tools: &[Tool],
|
||||
tools: &[ChatTool],
|
||||
) -> Result<()> {
|
||||
if let Some(content) = content {
|
||||
write_chat_content(out, content)?;
|
||||
@@ -186,7 +185,7 @@ fn render_system_message(
|
||||
fn render_developer_message(
|
||||
out: &mut String,
|
||||
content: &ChatContent,
|
||||
tools: &[Tool],
|
||||
tools: &[ChatTool],
|
||||
opens_thinking: bool,
|
||||
) -> Result<()> {
|
||||
if content.is_empty() {
|
||||
@@ -453,7 +452,7 @@ fn encode_arguments_to_dsml(out: &mut String, tool_call: &AssistantToolCall) ->
|
||||
}
|
||||
|
||||
/// Render the full tool preamble shown to the model.
|
||||
fn render_tools(out: &mut String, tools: &[Tool]) -> Result<()> {
|
||||
fn render_tools(out: &mut String, tools: &[ChatTool]) -> Result<()> {
|
||||
out.push_str(
|
||||
r#"## Tools
|
||||
|
||||
@@ -501,7 +500,7 @@ Here are the functions available in JSONSchema format:
|
||||
|
||||
/// Serialize one typed tool schema into the JSON shape embedded inside
|
||||
/// `<functions>`.
|
||||
fn render_tool_schema(out: &mut String, tool: &Tool) -> Result<()> {
|
||||
fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> {
|
||||
out.push_str(&json_dumps(&RenderedToolSchema {
|
||||
name: &tool.name,
|
||||
description: tool.description.as_deref(),
|
||||
@@ -518,7 +517,7 @@ fn write_chat_content(out: &mut String, content: &ChatContent) -> Result<()> {
|
||||
ChatContent::Text(text) => out.push_str(text),
|
||||
ChatContent::Parts(parts) => {
|
||||
for part in parts {
|
||||
out.push_str(part.as_text().map_err(Error::UnsupportedMultimodalContent)?);
|
||||
out.push_str(part.as_text()?);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user