forked from Karylab-cklius/vllm
Compare commits
100
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
efa494e397 | ||
|
|
74b77a593a | ||
|
|
ec53889a3b | ||
|
|
ede3d4ddf6 | ||
|
|
7301a834fa | ||
|
|
7104ac6b5d | ||
|
|
9a98eccb49 | ||
|
|
2ce1e2ba71 | ||
|
|
2496f66f7f | ||
|
|
d18b7f2723 | ||
|
|
80b633f076 | ||
|
|
28f589ebb1 | ||
|
|
98998279d4 | ||
|
|
71d208089e | ||
|
|
0ac6fc8352 | ||
|
|
fdfb2566c0 | ||
|
|
8a5cf1ccd6 | ||
|
|
fe1d923afc | ||
|
|
32daf56b42 | ||
|
|
82a42234be | ||
|
|
af9f583344 | ||
|
|
bd2d83ff31 | ||
|
|
bb78168b21 | ||
|
|
89c6a41001 | ||
|
|
7fdfa6441d | ||
|
|
e9b728de8a | ||
|
|
5828a205ef | ||
|
|
7a74f31d2e | ||
|
|
47930b59ca | ||
|
|
6aec99f030 | ||
|
|
f4966f8b3d | ||
|
|
2c9c07c85e | ||
|
|
320c52b134 | ||
|
|
6deb05e0e4 | ||
|
|
d82ac00923 | ||
|
|
dac9e9a640 | ||
|
|
d7607ad273 | ||
|
|
d955745d58 | ||
|
|
e1ed89dbee | ||
|
|
1c2ffc6f88 | ||
|
|
ca4cfd8731 | ||
|
|
c9c1540e61 | ||
|
|
c1d754d681 | ||
|
|
01d8cd92dd | ||
|
|
a4b14b98c6 | ||
|
|
cf1c906724 | ||
|
|
766ce2bb6b | ||
|
|
3d119f78f7 | ||
|
|
1b1359c332 | ||
|
|
cad4ca12b8 | ||
|
|
b697119800 | ||
|
|
b4c6dc6454 | ||
|
|
2ee5106372 | ||
|
|
7a89b72564 | ||
|
|
dc10e467a9 | ||
|
|
ee4d7df2b5 | ||
|
|
3e8afdf785 | ||
|
|
6690a0c4de | ||
|
|
1c23c42030 | ||
|
|
b12e42d132 | ||
|
|
69fdaffbcd | ||
|
|
80e2c4462d | ||
|
|
5b3807e862 | ||
|
|
59401ac9f1 | ||
|
|
d841386d27 | ||
|
|
fff9210b2a | ||
|
|
70db1488c5 | ||
|
|
2385e140d6 | ||
|
|
dab60fc658 | ||
|
|
996222f4bf | ||
|
|
e6fc848d4f | ||
|
|
f843ac1a1c | ||
|
|
7c2aa3108a | ||
|
|
ebf53ba373 | ||
|
|
baacbfcebf | ||
|
|
d8218b1ee7 | ||
|
|
9f153aa781 | ||
|
|
d3de61502f | ||
|
|
540aaf2140 | ||
|
|
4128605ad4 | ||
|
|
e2f993dc41 | ||
|
|
05cb606cad | ||
|
|
3f627ebef7 | ||
|
|
bc941f375d | ||
|
|
6afa25000c | ||
|
|
823a0ab754 | ||
|
|
2c27c294c0 | ||
|
|
ba94a3b998 | ||
|
|
dc68bd8c41 | ||
|
|
753e9d55e6 | ||
|
|
ac3409d162 | ||
|
|
93ee4cd47f | ||
|
|
980796cd07 | ||
|
|
5add018beb | ||
|
|
d5fe994e79 | ||
|
|
fa662b1a8b | ||
|
|
3c0b4432be | ||
|
|
469f3dcf1d | ||
|
|
94fcdd007f | ||
|
|
d9ff7e4e9a |
@@ -0,0 +1,23 @@
|
||||
group: Expert Parallelism
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: EPLB Algorithm
|
||||
key: eplb-algorithm
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/eplb
|
||||
- tests/distributed/test_eplb_algo.py
|
||||
- tests/distributed/test_eplb_utils.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s distributed/test_eplb_algo.py'
|
||||
@@ -38,7 +38,17 @@ steps:
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/logger.py
|
||||
- vllm/model_executor/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1/sample
|
||||
- tests/v1/logits_processors
|
||||
- tests/v1/test_oracle.py
|
||||
@@ -52,7 +62,8 @@ steps:
|
||||
pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
|
||||
pytest -v -s v1/test_oracle.py &&
|
||||
pytest -v -s v1/test_request.py &&
|
||||
pytest -v -s v1/test_outputs.py'
|
||||
pytest -v -s v1/test_outputs.py &&
|
||||
pytest -v -s v1/sample/test_topk_topp_sampler.py'
|
||||
|
||||
- label: XPU CPU Offload
|
||||
timeout_in_minutes: 60
|
||||
@@ -76,3 +87,101 @@ steps:
|
||||
cd tests &&
|
||||
pytest -v -s v1/kv_offload &&
|
||||
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/test_regression
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install modelscope &&
|
||||
cd tests &&
|
||||
pytest -v -s test_regression.py'
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs)
|
||||
key: metrics-tracing-2-gpus
|
||||
timeout_in_minutes: 30
|
||||
num_devices: 2
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/tracing/
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/v1/tracing
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/tracing'
|
||||
|
||||
- label: Async Engine, Inputs, Utils, Worker
|
||||
key: async-engine-inputs-utils-worker
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/assets/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/inputs/
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/tokenizers/
|
||||
- vllm/transformers_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pip install av &&
|
||||
pytest -v -s detokenizer &&
|
||||
pytest -v -s -m "not cpu_test" ./multimodal &&
|
||||
pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
group: Models - Multimodal
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
|
||||
pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
|
||||
key: multi-modal-models-standard-3-llava-qwen2-vl
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
pip install open-clip-torch --no-deps &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
|
||||
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
|
||||
--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
|
||||
parallelism: 4
|
||||
@@ -85,5 +85,5 @@ steps:
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av &&
|
||||
cd tests &&
|
||||
pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
|
||||
pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
|
||||
pytest -v -s benchmarks/test_serve_cli.py'
|
||||
|
||||
@@ -324,23 +324,6 @@ IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
|
||||
|
||||
echo "Using image: ${IMAGE}"
|
||||
|
||||
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo "Image already exists locally, skipping pull"
|
||||
else
|
||||
echo "Image not found locally, waiting for lock..."
|
||||
|
||||
flock /tmp/docker-pull.lock bash -c "
|
||||
if docker image inspect '${IMAGE}' >/dev/null 2>&1; then
|
||||
echo 'Image already pulled by another runner'
|
||||
else
|
||||
echo 'Pulling image...'
|
||||
timeout 900 docker pull '${IMAGE}'
|
||||
fi
|
||||
"
|
||||
|
||||
echo "Pull step completed"
|
||||
fi
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
}
|
||||
@@ -357,9 +340,12 @@ export HF_TOKEN ZE_AFFINITY_MASK
|
||||
|
||||
{
|
||||
flock 9
|
||||
if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo 'Image missing before container creation, pulling again...'
|
||||
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
|
||||
echo "Image already exists locally, skipping pull"
|
||||
else
|
||||
echo "Image not found locally, pulling image..."
|
||||
timeout 900 docker pull "${IMAGE}"
|
||||
echo "Pull step completed"
|
||||
fi
|
||||
|
||||
docker create \
|
||||
@@ -372,6 +358,8 @@ export HF_TOKEN ZE_AFFINITY_MASK
|
||||
--entrypoint='' \
|
||||
-e HF_TOKEN \
|
||||
-e ZE_AFFINITY_MASK \
|
||||
-e BUILDKITE_PARALLEL_JOB \
|
||||
-e BUILDKITE_PARALLEL_JOB_COUNT \
|
||||
-e CMDS \
|
||||
--name "${container_name}" \
|
||||
"${IMAGE}" \
|
||||
|
||||
@@ -110,6 +110,36 @@ install_uv() {
|
||||
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
|
||||
}
|
||||
|
||||
setup_pyo3_python() {
|
||||
local python_version="${PYO3_PYTHON_VERSION:-3.12}"
|
||||
|
||||
log_section "Installing Python ${python_version} for PyO3 tests"
|
||||
uv python install "$python_version"
|
||||
PYO3_PYTHON="$(uv python find \
|
||||
--managed-python \
|
||||
--no-project \
|
||||
--resolve-links \
|
||||
"$python_version")"
|
||||
export PYO3_PYTHON
|
||||
|
||||
local python_libdir
|
||||
python_libdir="$("$PYO3_PYTHON" - <<'PY'
|
||||
import pathlib
|
||||
import sysconfig
|
||||
|
||||
libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
|
||||
ldlibrary = sysconfig.get_config_var("LDLIBRARY")
|
||||
assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
|
||||
assert ldlibrary
|
||||
assert (libdir / ldlibrary).exists(), libdir / ldlibrary
|
||||
print(libdir)
|
||||
PY
|
||||
)"
|
||||
|
||||
export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
|
||||
export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
|
||||
}
|
||||
|
||||
run_style_clippy() {
|
||||
install_cargo_sort
|
||||
|
||||
@@ -132,6 +162,7 @@ run_style_clippy() {
|
||||
|
||||
run_tests() {
|
||||
install_uv
|
||||
setup_pyo3_python
|
||||
install_cargo_nextest
|
||||
|
||||
log_section "Running cargo nextest"
|
||||
|
||||
+146
-115
@@ -88,16 +88,16 @@
|
||||
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
|
||||
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
|
||||
# * [Plugin Tests (2 GPUs)]: #
|
||||
# - {`pytest -v -s entrypoints/openai/test_oot_registration.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s models/test_oot_registration.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
|
||||
# - {`pytest -v -s plugins_tests/test_oot_registration_online.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins_tests/test_oot_registration_offline.py`}: It needs a clean process #
|
||||
# - {`pytest -v -s plugins_tests/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
|
||||
# * [LoRA TP (Distributed)]: #
|
||||
# - There is some Tensor Parallelism related processing logic in LoRA that requires multi-GPU testing for validation. #
|
||||
# - {`pytest -v -s -x lora/test_gptoss_tp.py`}: Disabled for now because MXFP4 backend on non-cuda platform doesn't support #
|
||||
# LoRA yet. #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Don't test llama model here, it seems hf implementation is buggy. See: #
|
||||
# https://github.com/vllm-project/vllm/pull/5689 #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
|
||||
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
|
||||
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in #
|
||||
# this file as it's deprecated. #
|
||||
# #
|
||||
@@ -617,9 +617,9 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s models/test_oot_registration.py
|
||||
- pytest -v -s plugins/lora_resolvers
|
||||
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------#
|
||||
|
||||
@@ -804,7 +804,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: V1 e2e (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -830,7 +830,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
|
||||
|
||||
@@ -1187,88 +1187,20 @@ steps:
|
||||
|
||||
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
- label: Entrypoints Unit Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/unit_tests
|
||||
- tests/entrypoints/weight_transfer
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
- pytest -v -s entrypoints/unit_tests
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1287,19 +1219,75 @@ steps:
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
- label: Entrypoints Integration (API Server) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
|
||||
- label: Entrypoints Integration (Responses API) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1315,19 +1303,47 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Entrypoints Unit Tests # TBD
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: OpenAI API correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1338,7 +1354,6 @@ steps:
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
@@ -2391,7 +2406,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2427,7 +2442,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2441,7 +2456,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2455,7 +2470,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: V1 e2e (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2783,7 +2798,7 @@ steps:
|
||||
|
||||
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) # TBD
|
||||
- label: Entrypoints Integration (API Server) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2799,7 +2814,7 @@ steps:
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2813,9 +2828,9 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2827,29 +2842,31 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
- label: Entrypoints Integration (API Server Generate) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2865,6 +2882,20 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3042,7 +3073,7 @@ steps:
|
||||
- vllm/_aiter_ops.py
|
||||
commands:
|
||||
- rocm-smi
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
@@ -3406,7 +3437,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -3421,7 +3452,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -3436,7 +3467,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -8,10 +8,11 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/
|
||||
- tests/entrypoints/unit_tests
|
||||
- tests/entrypoints/weight_transfer
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- pytest -v -s entrypoints/unit_tests
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
@@ -31,68 +32,9 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
key: entrypoints-integration-api-server-openai-part-3
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-2
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -108,6 +50,75 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate)
|
||||
key: entrypoints-integration-api-server-generate
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/tool_use
|
||||
- tests/entrypoints/tool_parsers
|
||||
- tests/entrypoints/anthropic
|
||||
- tests/entrypoints/generate
|
||||
commands:
|
||||
- pytest -v -s tool_use
|
||||
- pytest -v -s entrypoints/tool_parsers
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s entrypoints/anthropic
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-speech_to_text
|
||||
@@ -120,6 +131,18 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-multimodal
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
@@ -131,16 +154,6 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: OpenAI API Correctness
|
||||
key: openai-api-correctness
|
||||
timeout_in_minutes: 30
|
||||
@@ -158,7 +171,6 @@ steps:
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
|
||||
@@ -21,8 +21,9 @@ steps:
|
||||
- csrc/
|
||||
- tests/kernels/core
|
||||
- tests/kernels/test_concat_mla_q.py
|
||||
- tests/kernels/test_fused_qk_norm_rope_gate.py
|
||||
commands:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py
|
||||
|
||||
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
|
||||
key: kernels-minimax-reduce-rms-test-2-gpus
|
||||
@@ -299,3 +300,4 @@ steps:
|
||||
- vllm/config
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_moe_layer.py
|
||||
- pytest -v -s kernels/moe/test_deepep_v2_moe.py
|
||||
|
||||
@@ -300,9 +300,9 @@ steps:
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
- tests/tokenizers_
|
||||
- tests/reasoning
|
||||
- tests/tool_parsers
|
||||
- tests/tokenizers_
|
||||
- tests/parser
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
@@ -315,9 +315,9 @@ steps:
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s parser
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
@@ -37,6 +37,6 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
|
||||
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
+6
-1
@@ -23,9 +23,14 @@
|
||||
|
||||
# Any change to the VllmConfig changes can have a large user-facing impact,
|
||||
# so spam a lot of people
|
||||
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
|
||||
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
|
||||
/vllm/config/cache.py @heheda12345
|
||||
|
||||
# Config utils
|
||||
/vllm/config/utils.py @hmellor
|
||||
/vllm/engine/arg_utils.py @hmellor
|
||||
/vllm/utils/argparse_utils.py
|
||||
|
||||
# Entrypoints
|
||||
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
|
||||
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
|
||||
|
||||
@@ -4,6 +4,7 @@ include requirements/cuda.txt
|
||||
include requirements/rocm.txt
|
||||
include requirements/cpu.txt
|
||||
include CMakeLists.txt
|
||||
include tools/build_rust.py
|
||||
|
||||
recursive-include cmake *
|
||||
recursive-include csrc *
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# benchmark the overhead of disaggregated prefill.
|
||||
# methodology:
|
||||
# - send all request to prefill vLLM instance. It will buffer KV cache.
|
||||
# - then send all request to decode instance.
|
||||
# - The TTFT of decode instance is the overhead.
|
||||
|
||||
set -ex
|
||||
|
||||
kill_gpu_processes() {
|
||||
# kill all processes on GPU.
|
||||
pgrep pt_main_thread | xargs -r kill -9
|
||||
pgrep python3 | xargs -r kill -9
|
||||
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
|
||||
pgrep VLLM | xargs -r kill -9
|
||||
sleep 10
|
||||
|
||||
# remove vllm config file
|
||||
rm -rf ~/.config/vllm
|
||||
|
||||
# Print the GPU memory usage
|
||||
# so that we know if all GPU processes are killed.
|
||||
gpu_memory_usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0)
|
||||
# The memory usage should be 0 MB.
|
||||
echo "GPU 0 Memory Usage: $gpu_memory_usage MB"
|
||||
}
|
||||
|
||||
wait_for_server() {
|
||||
# wait for vllm server to start
|
||||
# return 1 if vllm server crashes
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -s localhost:${port}/v1/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
benchmark() {
|
||||
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
|
||||
|
||||
# compare chunked prefill with disaggregated prefill
|
||||
|
||||
results_folder="./results"
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
dataset_name="sonnet"
|
||||
dataset_path="../sonnet_4x.txt"
|
||||
num_prompts=10
|
||||
qps=$1
|
||||
prefix_len=50
|
||||
input_len=2048
|
||||
output_len=$2
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
|
||||
# let the prefill instance finish prefill
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8100 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename disagg_prefill_tp1.json \
|
||||
--request-rate "inf"
|
||||
|
||||
|
||||
# send the request to decode.
|
||||
# The TTFT of this command will be the overhead of disagg prefill impl.
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8200 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename disagg_prefill_tp1_overhead.json \
|
||||
--request-rate "$qps"
|
||||
kill_gpu_processes
|
||||
|
||||
}
|
||||
|
||||
|
||||
main() {
|
||||
|
||||
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
|
||||
(which jq) || (apt-get -y install jq)
|
||||
(which socat) || (apt-get -y install socat)
|
||||
|
||||
pip install quart httpx datasets
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
cd ..
|
||||
# create sonnet-4x.txt
|
||||
echo "" > sonnet_4x.txt
|
||||
for _ in {1..4}
|
||||
do
|
||||
cat sonnet.txt >> sonnet_4x.txt
|
||||
done
|
||||
cd disagg_benchmarks
|
||||
|
||||
rm -rf results
|
||||
mkdir results
|
||||
|
||||
default_qps=1
|
||||
default_output_len=1
|
||||
benchmark $default_qps $default_output_len
|
||||
|
||||
}
|
||||
|
||||
|
||||
main "$@"
|
||||
@@ -1,157 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Requirement: 2x GPUs.
|
||||
|
||||
|
||||
# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
|
||||
# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
|
||||
# Resource: 2x GPU
|
||||
# Approaches:
|
||||
# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
|
||||
# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
|
||||
# Prefilling instance: max_output_token=1
|
||||
# Decoding instance: force the input tokens be the same across requests to bypass prefilling
|
||||
|
||||
set -ex
|
||||
|
||||
kill_gpu_processes() {
|
||||
# kill all processes on GPU.
|
||||
pgrep pt_main_thread | xargs -r kill -9
|
||||
pgrep python3 | xargs -r kill -9
|
||||
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
|
||||
pgrep VLLM | xargs -r kill -9
|
||||
for port in 8000 8100 8200; do lsof -t -i:$port | xargs -r kill -9; done
|
||||
sleep 1
|
||||
}
|
||||
|
||||
wait_for_server() {
|
||||
# wait for vllm server to start
|
||||
# return 1 if vllm server crashes
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -s localhost:${port}/v1/completions > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
launch_chunked_prefill() {
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
# disagg prefill
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--enable-chunked-prefill \
|
||||
--gpu-memory-utilization 0.6 &
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--enable-chunked-prefill \
|
||||
--gpu-memory-utilization 0.6 &
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
python3 round_robin_proxy.py &
|
||||
sleep 1
|
||||
}
|
||||
|
||||
|
||||
launch_disagg_prefill() {
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
# disagg prefill
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
|
||||
--port 8100 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
|
||||
--port 8200 \
|
||||
--max-model-len 10000 \
|
||||
--gpu-memory-utilization 0.6 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
python3 disagg_prefill_proxy_server.py &
|
||||
sleep 1
|
||||
}
|
||||
|
||||
|
||||
benchmark() {
|
||||
results_folder="./results"
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
|
||||
dataset_name="sonnet"
|
||||
dataset_path="../sonnet_4x.txt"
|
||||
num_prompts=100
|
||||
qps=$1
|
||||
prefix_len=50
|
||||
input_len=1024
|
||||
output_len=$2
|
||||
tag=$3
|
||||
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $model \
|
||||
--dataset-name $dataset_name \
|
||||
--dataset-path $dataset_path \
|
||||
--sonnet-input-len $input_len \
|
||||
--sonnet-output-len "$output_len" \
|
||||
--sonnet-prefix-len $prefix_len \
|
||||
--num-prompts $num_prompts \
|
||||
--port 8000 \
|
||||
--save-result \
|
||||
--result-dir $results_folder \
|
||||
--result-filename "$tag"-qps-"$qps".json \
|
||||
--request-rate "$qps"
|
||||
|
||||
sleep 2
|
||||
}
|
||||
|
||||
|
||||
main() {
|
||||
|
||||
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
|
||||
(which jq) || (apt-get -y install jq)
|
||||
(which socat) || (apt-get -y install socat)
|
||||
(which lsof) || (apt-get -y install lsof)
|
||||
|
||||
pip install quart httpx matplotlib aiohttp datasets
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
cd ..
|
||||
# create sonnet-4x.txt so that we can sample 2048 tokens for input
|
||||
echo "" > sonnet_4x.txt
|
||||
for _ in {1..4}
|
||||
do
|
||||
cat sonnet.txt >> sonnet_4x.txt
|
||||
done
|
||||
cd disagg_benchmarks
|
||||
|
||||
rm -rf results
|
||||
mkdir results
|
||||
|
||||
default_output_len=6
|
||||
|
||||
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
|
||||
|
||||
launch_chunked_prefill
|
||||
for qps in 2 4 6 8; do
|
||||
benchmark $qps $default_output_len chunked_prefill
|
||||
done
|
||||
kill_gpu_processes
|
||||
|
||||
launch_disagg_prefill
|
||||
for qps in 2 4 6 8; do
|
||||
benchmark $qps $default_output_len disagg_prefill
|
||||
done
|
||||
kill_gpu_processes
|
||||
|
||||
python3 visualize_benchmark_results.py
|
||||
|
||||
}
|
||||
|
||||
|
||||
main "$@"
|
||||
@@ -1,260 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import aiohttp
|
||||
from quart import Quart, Response, make_response, request
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def parse_args():
|
||||
"""parse command line arguments"""
|
||||
parser = argparse.ArgumentParser(description="vLLM P/D disaggregation proxy server")
|
||||
|
||||
# Add args
|
||||
parser.add_argument(
|
||||
"--timeout",
|
||||
type=float,
|
||||
default=6 * 60 * 60,
|
||||
help="Timeout for backend service requests in seconds (default: 21600)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8000,
|
||||
help="Port to run the server on (default: 8000)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefill-url",
|
||||
type=str,
|
||||
default="http://localhost:8100",
|
||||
help="Prefill service base URL (protocol + host[:port])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--decode-url",
|
||||
type=str,
|
||||
default="http://localhost:8200",
|
||||
help="Decode service base URL (protocol + host[:port])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kv-host",
|
||||
type=str,
|
||||
default="localhost",
|
||||
help="Hostname or IP used by KV transfer (default: localhost)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefill-kv-port",
|
||||
type=int,
|
||||
default=14579,
|
||||
help="Prefill KV port (default: 14579)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--decode-kv-port",
|
||||
type=int,
|
||||
default=14580,
|
||||
help="Decode KV port (default: 14580)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
"""parse command line arguments"""
|
||||
args = parse_args()
|
||||
|
||||
# Initialize configuration using command line parameters
|
||||
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
|
||||
PREFILL_SERVICE_URL = args.prefill_url
|
||||
DECODE_SERVICE_URL = args.decode_url
|
||||
PORT = args.port
|
||||
|
||||
PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
|
||||
DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
|
||||
|
||||
logger.info(
|
||||
"Proxy resolved KV addresses -> prefill: %s, decode: %s",
|
||||
PREFILL_KV_ADDR,
|
||||
DECODE_KV_ADDR,
|
||||
)
|
||||
|
||||
app = Quart(__name__)
|
||||
|
||||
# Attach the configuration object to the application instance so helper
|
||||
# coroutines can read the resolved backend URLs and timeouts without using
|
||||
# globals.
|
||||
app.config.update(
|
||||
{
|
||||
"AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
|
||||
"PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
|
||||
"DECODE_SERVICE_URL": DECODE_SERVICE_URL,
|
||||
"PREFILL_KV_ADDR": PREFILL_KV_ADDR,
|
||||
"DECODE_KV_ADDR": DECODE_KV_ADDR,
|
||||
}
|
||||
)
|
||||
|
||||
def _normalize_base_url(url: str) -> str:
|
||||
"""Remove any trailing slash so path joins behave predictably."""
|
||||
return url.rstrip("/")
|
||||
|
||||
def _get_host_port(url: str) -> str:
|
||||
"""Return the hostname:port portion for logging and KV headers."""
|
||||
parsed = urlparse(url)
|
||||
host = parsed.hostname or "localhost"
|
||||
port = parsed.port
|
||||
if port is None:
|
||||
port = 80 if parsed.scheme == "http" else 443
|
||||
return f"{host}:{port}"
|
||||
|
||||
PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
|
||||
DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
|
||||
KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
|
||||
|
||||
def _build_headers(request_id: str) -> dict[str, str]:
|
||||
"""Construct the headers expected by vLLM's P2P disagg connector."""
|
||||
headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
async def _run_prefill(
|
||||
request_path: str,
|
||||
payload: dict,
|
||||
headers: dict[str, str],
|
||||
request_id: str,
|
||||
):
|
||||
url = f"{PREFILL_BASE}{request_path}"
|
||||
start_ts = time.perf_counter()
|
||||
logger.info("[prefill] start request_id=%s url=%s", request_id, url)
|
||||
try:
|
||||
async with (
|
||||
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
|
||||
session.post(url=url, json=payload, headers=headers) as resp,
|
||||
):
|
||||
if resp.status != 200:
|
||||
error_text = await resp.text()
|
||||
raise RuntimeError(
|
||||
f"Prefill backend error {resp.status}: {error_text}"
|
||||
)
|
||||
await resp.read()
|
||||
logger.info(
|
||||
"[prefill] done request_id=%s status=%s elapsed=%.2fs",
|
||||
request_id,
|
||||
resp.status,
|
||||
time.perf_counter() - start_ts,
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise RuntimeError(f"Prefill service timeout at {url}") from exc
|
||||
except aiohttp.ClientError as exc:
|
||||
raise RuntimeError(f"Prefill service unavailable at {url}") from exc
|
||||
|
||||
async def _stream_decode(
|
||||
request_path: str,
|
||||
payload: dict,
|
||||
headers: dict[str, str],
|
||||
request_id: str,
|
||||
):
|
||||
url = f"{DECODE_BASE}{request_path}"
|
||||
# Stream tokens from the decode service once the prefill stage has
|
||||
# materialized KV caches on the target workers.
|
||||
logger.info("[decode] start request_id=%s url=%s", request_id, url)
|
||||
try:
|
||||
async with (
|
||||
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
|
||||
session.post(url=url, json=payload, headers=headers) as resp,
|
||||
):
|
||||
if resp.status != 200:
|
||||
error_text = await resp.text()
|
||||
logger.error(
|
||||
"Decode backend error %s - %s", resp.status, error_text
|
||||
)
|
||||
err_msg = (
|
||||
'{"error": "Decode backend error ' + str(resp.status) + '"}'
|
||||
)
|
||||
yield err_msg.encode()
|
||||
return
|
||||
logger.info(
|
||||
"[decode] streaming response request_id=%s status=%s",
|
||||
request_id,
|
||||
resp.status,
|
||||
)
|
||||
async for chunk_bytes in resp.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
logger.info("[decode] finished streaming request_id=%s", request_id)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error("Decode service timeout at %s", url)
|
||||
yield b'{"error": "Decode service timeout"}'
|
||||
except aiohttp.ClientError as exc:
|
||||
logger.error("Decode service error at %s: %s", url, exc)
|
||||
yield b'{"error": "Decode service unavailable"}'
|
||||
|
||||
async def process_request():
|
||||
"""Process a single request through prefill and decode stages"""
|
||||
try:
|
||||
original_request_data = await request.get_json()
|
||||
|
||||
# Create prefill request (max_tokens=1)
|
||||
prefill_request = original_request_data.copy()
|
||||
prefill_request["max_tokens"] = 1
|
||||
if "max_completion_tokens" in prefill_request:
|
||||
prefill_request["max_completion_tokens"] = 1
|
||||
|
||||
# Execute prefill stage
|
||||
# The request id encodes both KV socket addresses so the backend can
|
||||
# shuttle tensors directly via NCCL once the prefill response
|
||||
# completes.
|
||||
request_id = (
|
||||
f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
|
||||
f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
|
||||
)
|
||||
|
||||
headers = _build_headers(request_id)
|
||||
await _run_prefill(request.path, prefill_request, headers, request_id)
|
||||
|
||||
# Execute decode stage and stream response
|
||||
# Pass the unmodified user request so the decode phase can continue
|
||||
# sampling with the already-populated KV cache.
|
||||
generator = _stream_decode(
|
||||
request.path, original_request_data, headers, request_id
|
||||
)
|
||||
response = await make_response(generator)
|
||||
response.timeout = None # Disable timeout for streaming response
|
||||
return response
|
||||
|
||||
except Exception:
|
||||
logger.exception("Error processing request")
|
||||
return Response(
|
||||
response=b'{"error": "Internal server error"}',
|
||||
status=500,
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
async def handle_request():
|
||||
"""Handle incoming API requests with concurrency and rate limiting"""
|
||||
try:
|
||||
return await process_request()
|
||||
except asyncio.CancelledError:
|
||||
logger.warning("Request cancelled")
|
||||
return Response(
|
||||
response=b'{"error": "Request cancelled"}',
|
||||
status=503,
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
# Start the Quart server with host can be set to 0.0.0.0
|
||||
app.run(port=PORT)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,63 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
class RoundRobinProxy:
|
||||
def __init__(self, target_ports):
|
||||
self.target_ports = target_ports
|
||||
self.port_cycle = itertools.cycle(self.target_ports)
|
||||
|
||||
async def handle_request(self, request):
|
||||
target_port = next(self.port_cycle)
|
||||
target_url = f"http://localhost:{target_port}{request.path_qs}"
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
# Forward the request
|
||||
async with session.request(
|
||||
method=request.method,
|
||||
url=target_url,
|
||||
headers=request.headers,
|
||||
data=request.content,
|
||||
) as response:
|
||||
# Start sending the response
|
||||
resp = web.StreamResponse(
|
||||
status=response.status, headers=response.headers
|
||||
)
|
||||
await resp.prepare(request)
|
||||
|
||||
# Stream the response content
|
||||
async for chunk in response.content.iter_any():
|
||||
await resp.write(chunk)
|
||||
|
||||
await resp.write_eof()
|
||||
return resp
|
||||
|
||||
except Exception as e:
|
||||
return web.Response(text=f"Error: {str(e)}", status=500)
|
||||
|
||||
|
||||
async def main():
|
||||
proxy = RoundRobinProxy([8100, 8200])
|
||||
app = web.Application()
|
||||
app.router.add_route("*", "/{path:.*}", proxy.handle_request)
|
||||
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "localhost", 8000)
|
||||
await site.start()
|
||||
|
||||
print("Proxy server started on http://localhost:8000")
|
||||
|
||||
# Keep the server running
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -1,47 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
|
||||
if __name__ == "__main__":
|
||||
data = []
|
||||
for name in ["disagg_prefill", "chunked_prefill"]:
|
||||
for qps in [2, 4, 6, 8]:
|
||||
with open(f"results/{name}-qps-{qps}.json") as f:
|
||||
x = json.load(f)
|
||||
x["name"] = name
|
||||
x["qps"] = qps
|
||||
data.append(x)
|
||||
|
||||
df = pd.DataFrame.from_dict(data)
|
||||
dis_df = df[df["name"] == "disagg_prefill"]
|
||||
chu_df = df[df["name"] == "chunked_prefill"]
|
||||
|
||||
plt.style.use("bmh")
|
||||
plt.rcParams["font.size"] = 20
|
||||
|
||||
for key in [
|
||||
"mean_ttft_ms",
|
||||
"median_ttft_ms",
|
||||
"p99_ttft_ms",
|
||||
"mean_itl_ms",
|
||||
"median_itl_ms",
|
||||
"p99_itl_ms",
|
||||
]:
|
||||
fig, ax = plt.subplots(figsize=(11, 7))
|
||||
plt.plot(
|
||||
dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4
|
||||
)
|
||||
plt.plot(
|
||||
chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4
|
||||
)
|
||||
ax.legend()
|
||||
|
||||
ax.set_xlabel("QPS")
|
||||
ax.set_ylabel(key)
|
||||
ax.set_ylim(bottom=0)
|
||||
fig.savefig(f"results/{key}.png")
|
||||
plt.close(fig)
|
||||
@@ -250,7 +250,7 @@ def benchmark_config(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=topk,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size_per_partition=shard_intermediate_size,
|
||||
intermediate_size=shard_intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
activation=MoEActivation.SILU,
|
||||
|
||||
@@ -65,6 +65,32 @@ class RequestArgs(NamedTuple):
|
||||
limit_min_tokens: int # Use negative value for no limit
|
||||
limit_max_tokens: int # Use negative value for no limit
|
||||
timeout_sec: int
|
||||
send_conversation_id: bool
|
||||
headers: dict[str, str]
|
||||
|
||||
|
||||
def parse_custom_header(header: str) -> tuple[str, str]:
|
||||
separators = (":", "=")
|
||||
for separator in separators:
|
||||
if separator in header:
|
||||
key, value = header.split(separator, 1)
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
if key:
|
||||
return key, value
|
||||
break
|
||||
raise argparse.ArgumentTypeError(
|
||||
"Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
|
||||
)
|
||||
|
||||
|
||||
def build_request_headers(
|
||||
api_key: str | None, custom_headers: list[tuple[str, str]] | None
|
||||
) -> dict[str, str]:
|
||||
headers = dict(custom_headers or [])
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
|
||||
class BenchmarkArgs(NamedTuple):
|
||||
@@ -218,12 +244,11 @@ async def send_request(
|
||||
max_tokens: int | None = None,
|
||||
timeout_sec: int = 120,
|
||||
conversation_id: str | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> ServerResponse:
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"seed": 0,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
if conversation_id is not None:
|
||||
@@ -233,13 +258,17 @@ async def send_request(
|
||||
payload["stream"] = True
|
||||
payload["stream_options"] = {"include_usage": False}
|
||||
|
||||
if min_tokens is not None:
|
||||
payload["min_tokens"] = min_tokens
|
||||
# if min_tokens is not None:
|
||||
# payload["min_tokens"] = min_tokens
|
||||
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
headers = {"Content-Type": "application/json"}
|
||||
request_headers = {"Content-Type": "application/json"}
|
||||
if conversation_id is not None:
|
||||
request_headers["X-Session-ID"] = str(conversation_id)
|
||||
if headers is not None:
|
||||
request_headers.update(headers)
|
||||
|
||||
# Calculate the timeout for the request
|
||||
if max_tokens is not None:
|
||||
@@ -265,7 +294,7 @@ async def send_request(
|
||||
most_recent_timestamp: int = start_time
|
||||
|
||||
async with session.post(
|
||||
url=chat_url, json=payload, headers=headers, timeout=timeout
|
||||
url=chat_url, json=payload, headers=request_headers, timeout=timeout
|
||||
) as response:
|
||||
http_status = HTTPStatus(response.status)
|
||||
if http_status == HTTPStatus.OK:
|
||||
@@ -317,6 +346,8 @@ async def send_request(
|
||||
latency = time.perf_counter_ns() - start_time
|
||||
|
||||
if ttft is None:
|
||||
if stream:
|
||||
valid_response = False
|
||||
# The response was a single chunk
|
||||
ttft = latency
|
||||
|
||||
@@ -423,7 +454,8 @@ async def send_turn(
|
||||
min_tokens,
|
||||
max_tokens,
|
||||
req_args.timeout_sec,
|
||||
conversation_id=conv_id,
|
||||
conversation_id=conv_id if req_args.send_conversation_id else None,
|
||||
headers=req_args.headers,
|
||||
)
|
||||
|
||||
if response.valid is False:
|
||||
@@ -872,6 +904,7 @@ def get_client_config(
|
||||
# Arguments for API requests
|
||||
chat_url = f"{args.url}/v1/chat/completions"
|
||||
model_name = args.served_model_name if args.served_model_name else args.model
|
||||
headers = build_request_headers(args.api_key, args.header)
|
||||
|
||||
req_args = RequestArgs(
|
||||
chat_url=chat_url,
|
||||
@@ -880,6 +913,8 @@ def get_client_config(
|
||||
limit_min_tokens=args.limit_min_tokens,
|
||||
limit_max_tokens=args.limit_max_tokens,
|
||||
timeout_sec=args.request_timeout_sec,
|
||||
send_conversation_id=args.send_conversation_id,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
return client_args, req_args
|
||||
@@ -1245,19 +1280,19 @@ def process_statistics(
|
||||
)
|
||||
|
||||
|
||||
async def get_server_info(url: str) -> None:
|
||||
async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
|
||||
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Get server version (not mandatory, "version" endpoint may not exist)
|
||||
url_version = f"{url}/version"
|
||||
async with session.get(url_version) as response:
|
||||
async with session.get(url_version, headers=headers) as response:
|
||||
if HTTPStatus(response.status) == HTTPStatus.OK:
|
||||
text = await response.text()
|
||||
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
|
||||
|
||||
# Get available models
|
||||
url_models = f"{url}/v1/models"
|
||||
async with session.get(url_models) as response:
|
||||
async with session.get(url_models, headers=headers) as response:
|
||||
if HTTPStatus(response.status) == HTTPStatus.OK:
|
||||
text = await response.text()
|
||||
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
|
||||
@@ -1323,6 +1358,22 @@ async def main() -> None:
|
||||
help="Base URL for the LLM API server",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--api-key",
|
||||
type=str,
|
||||
default=None,
|
||||
help="API key to send as an Authorization bearer token",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--header",
|
||||
action="append",
|
||||
type=parse_custom_header,
|
||||
default=None,
|
||||
metavar="KEY=VALUE",
|
||||
help="Custom request header. Can be specified multiple times. "
|
||||
"Accepts 'Header-Name: value' or 'Header-Name=value'.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--num-clients",
|
||||
@@ -1437,6 +1488,22 @@ async def main() -> None:
|
||||
help="Disable stream/streaming mode (set 'stream' to False in the API request)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--send-conversation-id",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help=(
|
||||
"Inject a `conversation_id` field into each Chat Completions "
|
||||
"payload. This is a non-standard OpenAI extension consumed by "
|
||||
"vLLM's disaggregated multi-turn proxy "
|
||||
"(examples/disaggregated/disaggregated_serving/"
|
||||
"disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
|
||||
"Leave disabled (default) when targeting strict "
|
||||
"OpenAI-compatible endpoints; enable when benchmarking the "
|
||||
"disaggregated proxy."
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-e",
|
||||
"--excel-output",
|
||||
@@ -1525,7 +1592,8 @@ async def main() -> None:
|
||||
args.model, trust_remote_code=args.trust_remote_code
|
||||
)
|
||||
|
||||
await get_server_info(args.url)
|
||||
headers = build_request_headers(args.api_key, args.header)
|
||||
await get_server_info(args.url, headers=headers)
|
||||
|
||||
# Load the input file (either conversations of configuration file)
|
||||
logger.info(f"Reading input file: {args.input_file}")
|
||||
|
||||
+4
-15
@@ -1,5 +1,5 @@
|
||||
#!/bin/bash
|
||||
# Build the vllm-rs Rust frontend binary and install it into the vllm package.
|
||||
# Build vLLM Rust artifacts and install them into the vllm package.
|
||||
# Usage: ./build_rust.sh [--debug]
|
||||
#
|
||||
# By default builds in release mode. Pass --debug for faster compile times
|
||||
@@ -8,8 +8,6 @@
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
|
||||
RUST_DIR="$REPO_ROOT/rust"
|
||||
TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
|
||||
|
||||
# Read the required toolchain from rust-toolchain.toml.
|
||||
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
|
||||
@@ -27,18 +25,9 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
|
||||
fi
|
||||
|
||||
if [[ "${1:-}" == "--debug" ]]; then
|
||||
PROFILE_ARGS=()
|
||||
PROFILE_DIR="debug"
|
||||
PROFILE_ARG="--debug"
|
||||
else
|
||||
PROFILE_ARGS=(--release)
|
||||
PROFILE_DIR="release"
|
||||
PROFILE_ARG="--release"
|
||||
fi
|
||||
|
||||
cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
|
||||
--manifest-path "$RUST_DIR/Cargo.toml" \
|
||||
--bin vllm-rs \
|
||||
--features native-tls-vendored
|
||||
|
||||
mkdir -p "$(dirname "$TARGET_PATH")"
|
||||
cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
|
||||
echo "Installed vllm-rs to $TARGET_PATH"
|
||||
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
static const char* PYARGS_PARSE = "KKKK";
|
||||
#else
|
||||
#include <cstdlib>
|
||||
#include <cstdint>
|
||||
#include <cerrno>
|
||||
#include <climits>
|
||||
|
||||
@@ -46,6 +47,29 @@ static inline unsigned long long my_min(unsigned long long a,
|
||||
return a < b ? a : b;
|
||||
}
|
||||
|
||||
static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
|
||||
size_t alignment) {
|
||||
CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);
|
||||
if (status == CUresult(0) || alignment == 0) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Some ROCm stacks can report OOM while reserving VA with an explicit
|
||||
// alignment even when physical VRAM is free. Let HIP choose the default
|
||||
// alignment, then verify that the returned address still satisfies the
|
||||
// requested alignment before accepting it.
|
||||
status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
|
||||
if (status != CUresult(0)) {
|
||||
return status;
|
||||
}
|
||||
if (((std::uintptr_t)(*d_mem) % alignment) == 0) {
|
||||
return status;
|
||||
}
|
||||
|
||||
(void)cuMemAddressFree(*d_mem, size);
|
||||
return hipErrorNotSupported;
|
||||
}
|
||||
|
||||
static const char* PYARGS_PARSE = "KKKO";
|
||||
#endif
|
||||
|
||||
@@ -325,7 +349,7 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
|
||||
return nullptr;
|
||||
}
|
||||
#else
|
||||
CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
|
||||
CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity));
|
||||
if (error_code != 0) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
+22
-22
@@ -255,51 +255,49 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
|
||||
# build stage doesn't need the rust toolchain, protoc, or the rust source.
|
||||
# This stage runs in parallel with csrc-build/extensions-build.
|
||||
FROM ${BUILD_BASE_IMAGE} AS rust-build
|
||||
# This stage reuses the Python environment from base and runs in parallel with
|
||||
# csrc-build/extensions-build.
|
||||
FROM base AS rust-build
|
||||
ARG BUILD_OS
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install a basic C toolchain (some rust crates compile C in their build.rs
|
||||
# scripts) and unzip (used to extract the pinned protoc release below).
|
||||
# Install native tools needed only for Rust/protoc builds.
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
dnf install -y --setopt=install_weak_deps=False \
|
||||
ca-certificates curl git gcc gcc-c++ make unzip \
|
||||
make unzip \
|
||||
&& dnf clean all && rm -rf /var/cache/dnf; \
|
||||
else \
|
||||
apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
make unzip \
|
||||
&& rm -rf /var/lib/apt/lists/*; \
|
||||
fi
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Copy only the rust workspace — the binary is the sole artifact we need.
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
# Build the release binary. Cache cargo registry/git and target/, but copy the
|
||||
# binary out of the target/ cache mount so it persists into the image layer
|
||||
# for later COPY --from=rust-build.
|
||||
# Build the release artifacts. Cache cargo registry/git, but not target/,
|
||||
# because stale target metadata can outlive source updates across BuildKit
|
||||
# cache reuse.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
|
||||
#################### CSRC BUILD IMAGE ####################
|
||||
@@ -342,6 +340,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY pyproject.toml setup.py CMakeLists.txt ./
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY cmake cmake/
|
||||
COPY csrc csrc/
|
||||
COPY vllm/envs.py vllm/envs.py
|
||||
@@ -506,9 +505,10 @@ WORKDIR /workspace
|
||||
COPY --from=csrc-build /workspace/dist /precompiled-wheels
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
|
||||
+17
-17
@@ -93,35 +93,34 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Copy only the rust workspace — the binary is the sole artifact we need.
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
# Build the release binary. Cache cargo registry/git and target/, but copy the
|
||||
# binary out of the target/ cache mount so it persists into the image layer
|
||||
# for later COPY --from=rust-build.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
# Build the release artifacts. Cache cargo registry/git, but not target/,
|
||||
# because stale target metadata can outlive source updates across BuildKit
|
||||
# cache reuse.
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
|
||||
bash build_rust.sh
|
||||
|
||||
######################### BUILD IMAGE #########################
|
||||
FROM base AS vllm-build
|
||||
@@ -154,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
|
||||
@@ -102,21 +102,22 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
@@ -125,8 +126,7 @@ ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
@@ -139,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
|
||||
+17
-16
@@ -138,27 +138,25 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
ENV CARGO_NET_RETRY=10
|
||||
ENV RUSTUP_MAX_RETRIES=10
|
||||
|
||||
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
cd ${COMMON_WORKDIR}/vllm \
|
||||
&& uv pip install --system -r requirements/build/rust.txt
|
||||
|
||||
# Build the release binary. Cargo's registry/git caches can be written by
|
||||
# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
|
||||
# keeping the cache benefit. Copy the binary out so it persists into the image
|
||||
# layer for later COPY --from=rust-build.
|
||||
# keeping the cache benefit. Do not cache target/, because stale target metadata
|
||||
# can outlive source updates across BuildKit cache reuse.
|
||||
RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
|
||||
--mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
|
||||
cd ${COMMON_WORKDIR}/vllm \
|
||||
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
|
||||
&& test -x /tmp/vllm-rs
|
||||
&& bash build_rust.sh \
|
||||
&& test -x vllm/vllm-rs
|
||||
|
||||
# -----------------------
|
||||
# vLLM native build stages
|
||||
@@ -178,6 +176,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
|
||||
# not invalidate the expensive native build layer.
|
||||
COPY setup.py CMakeLists.txt ./
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY cmake cmake/
|
||||
COPY csrc csrc/
|
||||
COPY vllm/envs.py vllm/envs.py
|
||||
@@ -209,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm
|
||||
|
||||
COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
|
||||
|
||||
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
|
||||
cd vllm \
|
||||
@@ -418,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
|
||||
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
|
||||
|
||||
# Create /install directory for custom wheels
|
||||
RUN mkdir -p /install
|
||||
|
||||
+104
-90
@@ -6,21 +6,22 @@ ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
|
||||
sh -s -- -y --profile minimal --default-toolchain none
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
@@ -29,8 +30,7 @@ ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry \
|
||||
--mount=type=cache,target=/root/.cargo/git \
|
||||
--mount=type=cache,target=/workspace/rust/target \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
bash build_rust.sh
|
||||
|
||||
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
|
||||
|
||||
@@ -89,8 +89,8 @@ RUN mkdir neo && \
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
|
||||
&& uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
@@ -99,8 +99,8 @@ RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
|
||||
RUN rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
|
||||
rm -f /opt/intel/oneapi/ccl/latest && \
|
||||
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
|
||||
|
||||
SHELL ["bash", "-c"]
|
||||
@@ -119,24 +119,104 @@ ENV UV_LINK_MODE="copy"
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
|
||||
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
|
||||
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt
|
||||
uv pip install --upgrade pip
|
||||
|
||||
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
######################### UCX + NIXL BUILD STAGE #########################
|
||||
# Build UCX and NIXL in a dedicated stage so compiler/autotools layers are
|
||||
# never included in the final runtime image (mirrors ROCm's build_rixl stage).
|
||||
FROM vllm-base AS ucx-nixl-build
|
||||
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
# Build-time only: compiler, autotools, and verbs dev headers
|
||||
RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
pkg-config \
|
||||
libibverbs-dev \
|
||||
librdmacm-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Build UCX and produce a NIXL wheel so the final image needs no compiler.
|
||||
# patchelf (installed via uv) is used by the NIXL wheel build to rewrite
|
||||
# RPATH entries, making the wheel portable across stages.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
|
||||
LD_LIBRARY_PATH=/tmp/ucx_install/lib \
|
||||
python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
|
||||
find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
|
||||
rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
# Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
|
||||
# No compiler or autotools are installed in this stage.
|
||||
COPY --from=ucx-nixl-build /tmp/ucx_install /tmp/ucx_install
|
||||
COPY --from=ucx-nixl-build /tmp/nixl_wheels /tmp/nixl_wheels
|
||||
|
||||
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
# Install RDMA runtime libraries (no build tools) and the pre-built NIXL wheel.
|
||||
# Do not uninstall/reinstall large Python packages here to avoid extra layer
|
||||
# churn; final package resolution remains in the later app install step.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
rdma-core \
|
||||
libibverbs1 \
|
||||
librdmacm1 \
|
||||
libibumad3 \
|
||||
libibmad5 \
|
||||
libmlx5-1 \
|
||||
libmlx4-1 \
|
||||
ibverbs-providers \
|
||||
librdmacm1t64 \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& uv pip install --no-deps /tmp/nixl_wheels/nixl-*.whl \
|
||||
&& rm -rf /tmp/nixl_wheels
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
|
||||
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
uv pip install -r /workspace/vllm/requirements/xpu.txt && \
|
||||
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.1 && \
|
||||
uv pip uninstall oneccl oneccl-devel && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
# Keep source-dependent layers near the end so frequent code-only changes
|
||||
# don't invalidate heavy dependency and UCX/NIXL layers.
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built rust frontend binary into the source tree. setup.py
|
||||
# detects it and ships it as-is, skipping the local cargo build.
|
||||
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
@@ -147,75 +227,9 @@ ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
uv pip install --no-build-isolation .
|
||||
uv pip install --no-build-isolation --no-deps .
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# install NIXL and UCX from source code
|
||||
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
|
||||
ARG NIXL_VERSION=0.7.0
|
||||
|
||||
RUN apt-get update && apt-get install -y \
|
||||
pciutils \
|
||||
net-tools \
|
||||
iproute2 \
|
||||
hwloc \
|
||||
numactl \
|
||||
wget \
|
||||
curl \
|
||||
git \
|
||||
build-essential \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
pkg-config \
|
||||
rdma-core \
|
||||
libibverbs-dev \
|
||||
ibverbs-utils \
|
||||
libibverbs1 \
|
||||
librdmacm-dev \
|
||||
librdmacm1 \
|
||||
libibumad-dev \
|
||||
libibumad3 \
|
||||
libibmad-dev \
|
||||
libibmad5 \
|
||||
infiniband-diags \
|
||||
perftest \
|
||||
ibutils \
|
||||
libmlx5-1 \
|
||||
libmlx4-1 \
|
||||
ibverbs-providers \
|
||||
librdmacm1t64
|
||||
|
||||
ENV PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig:${PKG_CONFIG_PATH}
|
||||
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
cd /tmp/nixl_source && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
uv pip install . && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall oneccl oneccl-devel
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 388 KiB After Width: | Height: | Size: 373 KiB |
@@ -405,6 +405,47 @@ vllm bench serve \
|
||||
|
||||
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
|
||||
|
||||
#### BFCL (Tool-Calling) Benchmark
|
||||
|
||||
The Berkeley Function Calling Leaderboard (BFCL) dataset measures serving
|
||||
latency and throughput on realistic tool-calling traffic. Each request
|
||||
carries a per-sample `tools` schema and chat history, so the server must
|
||||
expose `/v1/chat/completions` with an auto-tool-choice parser enabled.
|
||||
The benchmark client always uses the `openai-chat` backend.
|
||||
|
||||
Start a tool-parser-enabled server, then run the bench. For example, with
|
||||
`gpt-oss-20b`:
|
||||
|
||||
```bash
|
||||
# Server
|
||||
vllm serve openai/gpt-oss-20b \
|
||||
--enable-auto-tool-choice \
|
||||
--tool-call-parser openai \
|
||||
--reasoning-parser openai_gptoss
|
||||
|
||||
# Client
|
||||
vllm bench serve \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--model openai/gpt-oss-20b \
|
||||
--dataset-name hf \
|
||||
--dataset-path gorilla-llm/Berkeley-Function-Calling-Leaderboard \
|
||||
--bfcl-categories simple,live_simple,multiple \
|
||||
--num-prompts 200
|
||||
```
|
||||
|
||||
`--bfcl-categories` is a comma-separated list of BFCL v3 category names
|
||||
(without the `BFCL_v3_` prefix or `.json` suffix). Defaults to
|
||||
`simple,live_simple,multiple`. Other supported non-multi-turn categories
|
||||
include `parallel`, `live_parallel`, `parallel_multiple`,
|
||||
`live_parallel_multiple`, `irrelevance`, `live_irrelevance`,
|
||||
`live_relevance`, `java`, `javascript`, and `rest`. Multi-turn categories
|
||||
are not yet supported.
|
||||
|
||||
The dataset class normalizes BFCL's loose schema dialect (`dict` →
|
||||
`object`, `float` → `number`, `tuple` → `array`, `any` → `string`) so
|
||||
modern grammar backends accept the translated tool definitions.
|
||||
|
||||
#### Other HuggingFaceDataset Examples
|
||||
|
||||
```bash
|
||||
|
||||
@@ -88,6 +88,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
|
||||
| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
|
||||
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
|
||||
|
||||
@@ -60,7 +60,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
|
||||
- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
|
||||
- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
|
||||
- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod]
|
||||
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
|
||||
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.UnquantizedFusedMoEMethod]
|
||||
|
||||
## Fused Experts Kernels
|
||||
|
||||
|
||||
@@ -1,319 +0,0 @@
|
||||
# P2P NCCL Connector
|
||||
|
||||
An implementation of xPyD with dynamic scaling based on point-to-point communication, partly inspired by Dynamo.
|
||||
|
||||
## Detailed Design
|
||||
|
||||
### Overall Process
|
||||
|
||||
As shown in Figure 1, the overall process of this **PD disaggregation** solution is described through a request flow:
|
||||
|
||||
1. The client sends an HTTP request to the Proxy/Router's `/v1/completions` interface.
|
||||
2. The Proxy/Router selects a **1P1D (1 Prefill instance + 1 Decode instance)** through either through round-robin or random selection, generates a `request_id` (rules to be introduced later), modifies the `max_tokens` in the HTTP request message to **1**, and then forwards the request to the **P instance**.
|
||||
3. Immediately afterward, the Proxy/Router forwards the **original HTTP request** to the **D instance**.
|
||||
4. The **P instance** performs **Prefill** and then **actively sends the generated KV cache** to the D instance (using **PUT_ASYNC** mode). The D instance's `zmq_addr` can be resolved through the `request_id`.
|
||||
5. The **D instance** has a **dedicated thread** for receiving the KV cache (to avoid blocking the main process). The received KV cache is saved into the **GPU memory buffer**, the size of which is determined by the vLLM startup parameter `kv_buffer_size`. When the GPU buffer is full, the KV cache is stored in the **local Tensor memory pool**.
|
||||
6. During the **Decode**, the D instance's main process retrieves the KV cache (transmitted by the P instance) from either the **GPU buffer** or the **memory pool**, thereby **skipping Prefill**.
|
||||
7. After completing **Decode**, the D instance returns the result to the **Proxy/Router**, which then forwards it to the **client**.
|
||||
|
||||

|
||||
|
||||
### Proxy/Router (Demo)
|
||||
|
||||
A simple HTTP service acts as the entry point for client requests and starts a background thread to listen for P/D instances reporting their HTTP IP and PORT, as well as ZMQ IP and PORT. It maintains a dictionary of `http_addr -> zmq_addr`. The `http_addr` is the IP:PORT for the vLLM instance's request, while the `zmq_addr` is the address for KV cache handshake and metadata reception.
|
||||
|
||||
The Proxy/Router is responsible for selecting 1P1D based on the characteristics of the client request, such as the prompt, and generating a corresponding `request_id`, for example:
|
||||
|
||||
```text
|
||||
cmpl-___prefill_addr_10.0.1.2:21001___decode_addr_10.0.1.3:22001_93923d63113b4b338973f24d19d4bf11-0
|
||||
```
|
||||
|
||||
Currently, to quickly verify whether xPyD can work, a round-robin selection of 1P1D is used. In the future, it is planned to use a trie combined with the load status of instances to select appropriate P and D.
|
||||
|
||||
Each P/D instance periodically sends a heartbeat packet to the Proxy/Router (currently every 3 seconds) to register (i.e., report `http_addr -> zmq_addr`) and keep the connection alive. If an instance crashes and fails to send a ping for a certain period of time, the Proxy/Router will remove the timed-out instance (this feature has not yet been developed).
|
||||
|
||||
### KV Cache Transfer Methods
|
||||
|
||||
There are three methods for KVCache transfer: PUT, GET, and PUT_ASYNC. These methods can be specified using the `--kv-transfer-config` and `kv_connector_extra_config` parameters, specifically through the `send_type` field. Both PUT and PUT_ASYNC involve the P instance actively sending KVCache to the D instance. The difference is that PUT is a synchronous transfer method that blocks the main process, while PUT_ASYNC is an asynchronous transfer method. PUT_ASYNC uses a dedicated thread for sending KVCache, which means it does not block the main process. In contrast, the GET method involves the P instance saving the KVCache to the memory buffer after computing the prefill. The D instance then actively retrieves the computed KVCache from the P instance once it has allocated space for the KVCache.
|
||||
|
||||
Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
|
||||
|
||||
### P2P Communication via ZMQ & NCCL
|
||||
|
||||
As long as the address of the counterpart is known, point-to-point KV cache transfer (using NCCL) can be performed, without being constrained by rank and world size. To support dynamic scaling (expansion and contraction) of instances with PD disaggregation. This means that adding or removing P/D instances does not require a full system restart.
|
||||
|
||||
Each P/D instance only needs to create a single `P2pNcclEngine` instance. This instance maintains a ZMQ Server, which runs a dedicated thread to listen on the `zmq_addr` address and receive control flow requests from other instances. These requests include requests to establish an NCCL connection and requests to send KVCache metadata (such as tensor shapes and data types). However, it does not actually transmit the KVCache data itself.
|
||||
|
||||
When a P instance and a D instance transmit KVCache for the first time, they need to establish a ZMQ connection and an NCCL group. For subsequent KVCache transmissions, this ZMQ connection and NCCL group are reused. The NCCL group consists of only two ranks, meaning the world size is equal to 2. This design is intended to support dynamic scaling, which means that adding or removing P/D instances does not require a full system restart. As long as the address of the counterpart is known, point-to-point KVCache transmission can be performed, without being restricted by rank or world size.
|
||||
|
||||
### NCCL Group Topology
|
||||
|
||||
Currently, only symmetric TP (Tensor Parallelism) methods are supported for KVCache transmission. Asymmetric TP and PP (Pipeline Parallelism) methods will be supported in the future. Figure 2 illustrates the 1P2D setup, where each instance has a TP (Tensor Parallelism) degree of 2. There are a total of 7 NCCL groups: three vLLM instances each have one NCCL group with TP=2. Additionally, the 0th GPU card of the P instance establishes an NCCL group with the 0th GPU card of each D instance. Similarly, the 1st GPU card of the P instance establishes an NCCL group with the 1st GPU card of each D instance.
|
||||
|
||||

|
||||
|
||||
Each NCCL group occupies a certain amount of GPU memory buffer for communication, the size of which is primarily influenced by the `NCCL_MAX_NCHANNELS` environment variable. When `NCCL_MAX_NCHANNELS=16`, an NCCL group typically occupies 100MB, while when `NCCL_MAX_NCHANNELS=8`, it usually takes up 52MB. For large-scale xPyD configurations—such as DeepSeek's 96P144D—this implementation is currently not feasible. Moving forward, we are considering using RDMA for point-to-point communication and are also keeping an eye on UCCL.
|
||||
|
||||
### GPU Memory Buffer and Tensor Memory Pool
|
||||
|
||||
The trade-off in the size of the memory buffer is as follows: For P instances, the memory buffer is not required in PUT and PUT_ASYNC modes, but it is necessary in GET mode. For D instances, a memory buffer is needed in all three modes. The memory buffer for D instances should not be too large. Similarly, for P instances in GET mode, the memory buffer should also not be too large. The memory buffer of D instances is used to temporarily store KVCache sent by P instances. If it is too large, it will reduce the KVCache space available for normal inference by D instances, thereby decreasing the inference batch size and ultimately leading to a reduction in output throughput. The size of the memory buffer is configured by the parameter `kv_buffer_size`, measured in bytes, and is typically set to 5%~10% of the memory size.
|
||||
|
||||
If the `--max-num-seqs` parameter for P instances is set to a large value, due to the large batch size, P instances will generate a large amount of KVCache simultaneously. This may exceed the capacity of the memory buffer of D instances, resulting in KVCache loss. Once KVCache is lost, D instances need to recompute Prefill, which is equivalent to performing Prefill twice. Consequently, the time-to-first-token (TTFT) will significantly increase, leading to degraded performance.
|
||||
|
||||
To address the above issues, I have designed and developed a local Tensor memory pool for storing KVCache, inspired by the buddy system used in Linux memory modules. Since the memory is sufficiently large, typically in the TB range on servers, there is no need to consider prefix caching or using block-based designs to reuse memory, thereby saving space. When the memory buffer is insufficient, KVCache can be directly stored in the Tensor memory pool, and D instances can subsequently retrieve KVCache from it. The read and write speed is that of PCIe, with PCIe 4.0 having a speed of approximately 21 GB/s, which is usually faster than the Prefill speed. Otherwise, solutions like Mooncake and lmcache would not be necessary. The Tensor memory pool acts as a flood diversion area, typically unused except during sudden traffic surges. In the worst-case scenario, my solution performs no worse than the normal situation with a Cache store.
|
||||
|
||||
## Install vLLM
|
||||
|
||||
```shell
|
||||
pip install "vllm>=0.9.2"
|
||||
```
|
||||
|
||||
## Run xPyD
|
||||
|
||||
### Instructions
|
||||
|
||||
- The following examples are run on an A800 (80GB) device, using the Meta-Llama-3.1-8B-Instruct model.
|
||||
- Pay attention to the setting of the `kv_buffer_size` (in bytes). The empirical value is 10% of the GPU memory size. This is related to the kvcache size. If it is too small, the GPU memory buffer for temporarily storing the received kvcache will overflow, causing the kvcache to be stored in the tensor memory pool, which increases latency. If it is too large, the kvcache available for inference will be reduced, leading to a smaller batch size and decreased throughput.
|
||||
- For Prefill instances, when using non-GET mode, the `kv_buffer_size` can be set to 1, as Prefill currently does not need to receive kvcache. However, when using GET mode, a larger `kv_buffer_size` is required because it needs to store the kvcache sent to the D instance.
|
||||
- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
|
||||
- `PUT_ASYNC` offers the best performance and should be prioritized.
|
||||
- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
|
||||
- The `disagg_proxy_p2p_nccl_xpyd.py` script will use port 10001 (for receiving client requests) and port 30001 (for receiving service discovery from P and D instances).
|
||||
- The node running the proxy must have `quart` installed.
|
||||
- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
|
||||
- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
|
||||
|
||||
### Run 1P3D
|
||||
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20001 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode2 (e.g. 10.0.1.4 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20003 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20004 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
### Run 3P1D
|
||||
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20001 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20002 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20003 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
#### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
|
||||
--host 0.0.0.0 \
|
||||
--port 20004 \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--served-model-name base_model \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
|
||||
```
|
||||
|
||||
## Single request
|
||||
|
||||
```shell
|
||||
curl -X POST -s http://10.0.1.1:10001/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "base_model",
|
||||
"prompt": "San Francisco is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}'
|
||||
```
|
||||
|
||||
## Benchmark
|
||||
|
||||
??? console "Command"
|
||||
|
||||
```shell
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model base_model \
|
||||
--tokenizer meta-llama/Llama-3.1-8B-Instruct \
|
||||
--dataset-name "random" \
|
||||
--host 10.0.1.1 \
|
||||
--port 10001 \
|
||||
--random-input-len 1024 \
|
||||
--random-output-len 1024 \
|
||||
--ignore-eos \
|
||||
--burstiness 100 \
|
||||
--percentile-metrics "ttft,tpot,itl,e2el" \
|
||||
--metric-percentiles "90,95,99" \
|
||||
--seed $(date +%s) \
|
||||
--trust-remote-code \
|
||||
--request-rate 3 \
|
||||
--num-prompts 1000
|
||||
```
|
||||
|
||||
## Shut down
|
||||
|
||||
```shell
|
||||
pgrep python | xargs kill -9 && pkill -f python
|
||||
```
|
||||
|
||||
## Test data
|
||||
|
||||
### **Scenario**: 1K input & 200 output tokens, E2E P99 latency ~2s
|
||||
|
||||

|
||||
@@ -17,8 +17,6 @@ Two main reasons:
|
||||
|
||||
## Usage example
|
||||
|
||||
Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
|
||||
|
||||
Now supports 9 types of connectors:
|
||||
|
||||
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
|
||||
@@ -29,7 +27,6 @@ Now supports 9 types of connectors:
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
|
||||
```
|
||||
|
||||
- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
|
||||
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
|
||||
- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
|
||||
- **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
|
||||
@@ -44,16 +41,14 @@ Now supports 9 types of connectors:
|
||||
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
|
||||
```
|
||||
|
||||
For multi-tier offloading (e.g., CPU + filesystem tier) and the full configuration reference, see the [KV Offloading Usage Guide](kv_offloading_usage.md).
|
||||
|
||||
- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
|
||||
```
|
||||
|
||||
## Benchmarks
|
||||
|
||||
Please refer to [benchmarks/disagg_benchmarks](../../benchmarks/disagg_benchmarks) for disaggregated prefilling benchmarks.
|
||||
|
||||
## Development
|
||||
|
||||
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
# KV Offloading Usage Guide
|
||||
|
||||
This guide covers configuration of the [`OffloadingConnector`](disagg_prefill.md), which extends the prefix cache by offloading completed KV blocks to slower but larger tiers (CPU host memory, plus optional secondary tiers) as they are produced. Hits in the offload tiers are promoted back to GPU on demand. Transfers between GPU and CPU use DMA (`cudaMemcpyAsync`) and run asynchronously alongside model computation, so offloading adds minimal CPU- and GPU-core overhead.
|
||||
|
||||
!!! note
|
||||
The `OffloadingConnector` currently supports CUDA, ROCm, and XPU only.
|
||||
|
||||
## Overview
|
||||
|
||||
Two specs are available, selected by the `spec_name` key in `kv_connector_extra_config`:
|
||||
|
||||
- `CPUOffloadingSpec` (default): single CPU tier. Completed GPU blocks are copied into pinned host memory.
|
||||
- `TieringOffloadingSpec`: multi-tier. A CPU primary tier plus one or more secondary tiers.
|
||||
|
||||
Only the CPU primary tier has direct GPU access. Secondary tiers cannot read from or write to GPU memory; all GPU↔secondary transfers are staged through the CPU primary tier.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
GPU <--> CPU["CPU primary tier"]
|
||||
CPU <--> S0["Secondary tier 0"]
|
||||
CPU <--> S1["Secondary tier 1"]
|
||||
CPU <--> SN["..."]
|
||||
```
|
||||
|
||||
## Single-Tier Setup (CPU Only)
|
||||
|
||||
```bash
|
||||
vllm serve <model> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "OffloadingConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {
|
||||
"block_size": 64,
|
||||
"cpu_bytes_to_use": 1000000000
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## Multi-Tier Setup
|
||||
|
||||
Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields. The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
|
||||
|
||||
```bash
|
||||
vllm serve <model> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "OffloadingConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {
|
||||
"spec_name": "TieringOffloadingSpec",
|
||||
"cpu_bytes_to_use": 10737418240,
|
||||
"block_size": 16,
|
||||
"eviction_policy": "lru",
|
||||
"secondary_tiers": [
|
||||
{
|
||||
"type": "fs",
|
||||
"root_dir": "/mnt/kv_cache",
|
||||
"n_read_threads": 32,
|
||||
"n_write_threads": 16
|
||||
}
|
||||
]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## `kv_connector_extra_config` Reference
|
||||
|
||||
| Key | Required | Default | Scope | Notes |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
|
||||
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
|
||||
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
|
||||
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `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
|
||||
|
||||
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
|
||||
|
||||
### Filesystem (FS)
|
||||
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
| `type` | yes | — | Must be `fs`. |
|
||||
| `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). |
|
||||
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
|
||||
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
|
||||
|
||||
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
|
||||
|
||||
#### On-Disk Layout
|
||||
|
||||
Under `root_dir`, vLLM creates a subdirectory `<model>_<digest>`, where `<model>` is the model name with `/` replaced by `_` (so HuggingFace IDs like `meta-llama/Llama-3-8B` don't nest), and `<digest>` is a short SHA256 prefix derived from the run configuration (model, block size, parallelism, dtype, etc.). Runs with the same configuration share the same subdirectory; runs with different configurations live side-by-side under the same `root_dir` without colliding.
|
||||
|
||||
Inside that subdirectory, blocks are sharded across hash-prefix subdirectories to limit directory fan-out:
|
||||
|
||||
```text
|
||||
<root_dir>/
|
||||
<model>_<digest>/
|
||||
config.json
|
||||
<model>_<digest>_r<rank>/
|
||||
<hhh>/ # first 3 hex chars of the block hash
|
||||
<hh>_g<group_idx>/ # next 2 hex chars + KV cache group index
|
||||
<hash_hex>.bin # full block hash (in hex)
|
||||
```
|
||||
|
||||
`config.json` records the run (block size, number of KV groups, etc.) and is written on first start. Each rank writes blocks under its own `_r<rank>` sibling directory, so multiple ranks can safely share the same `root_dir`.
|
||||
|
||||
#### Cross-Process Sharing
|
||||
|
||||
To enable KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC), the `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g., `"0"`) on every instance. Without this, each process initializes `NONE_HASH` (the chain-hash seed for block content hashes) with random bytes, producing different block filenames for identical token content.
|
||||
|
||||
```bash
|
||||
PYTHONHASHSEED=0 vllm serve ...
|
||||
```
|
||||
|
||||
## Tuning Tips
|
||||
|
||||
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
|
||||
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
|
||||
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
|
||||
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
|
||||
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
|
||||
|
||||
## Further Reading
|
||||
|
||||
- [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).
|
||||
@@ -294,6 +294,21 @@ curl http://localhost:8000/v1/chat/completions \
|
||||
!!! note
|
||||
The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine.
|
||||
|
||||
### Benchmarking the multi-turn proxy
|
||||
|
||||
[`benchmarks/multi_turn/benchmark_serving_multi_turn.py`](../../benchmarks/multi_turn/benchmark_serving_multi_turn.py) supports targeting the disaggregated multi-turn proxy with the `--send-conversation-id` flag, which injects a per-conversation `conversation_id` into every request payload so the proxy can key cross-turn KV cache reuse.
|
||||
|
||||
The flag is **off by default** so the benchmark is compatible with strict OpenAI-compatible frontends that reject unknown top-level fields. When benchmarking the multi-turn proxy you must pass it explicitly — otherwise every turn lands as a cache MISS and the bidirectional KV transfer path is never exercised.
|
||||
|
||||
```bash
|
||||
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \
|
||||
--model <MODEL> --served-model-name <NAME> \
|
||||
--url http://<proxy_host>:8000 \
|
||||
--input-file benchmarks/multi_turn/generate_multi_turn.json \
|
||||
--num-clients 2 --max-active-conversations 6 \
|
||||
--send-conversation-id
|
||||
```
|
||||
|
||||
### Limitations
|
||||
|
||||
- Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns.
|
||||
|
||||
@@ -504,6 +504,13 @@ Flags: `--tool-call-parser pythonic --chat-template {see_above}`
|
||||
!!! warning
|
||||
Llama's smaller models frequently fail to emit tool calls in the correct format. Results may vary depending on the model.
|
||||
|
||||
## Benchmarking Tool-Calling Performance
|
||||
|
||||
To measure serving latency and throughput on realistic tool-calling traffic,
|
||||
use the BFCL (Berkeley Function Calling Leaderboard) dataset with
|
||||
`vllm bench serve`. See the [BFCL benchmark example](../benchmarking/cli.md#bfcl-tool-calling-benchmark)
|
||||
for the full server + client commands.
|
||||
|
||||
## How to Write a Tool Parser Plugin
|
||||
|
||||
A tool parser plugin is a Python file containing one or more ToolParser implementations. You can write a ToolParser similar to the `Hermes2ProToolParser` in [vllm/tool_parsers/hermes_tool_parser.py](../../vllm/tool_parsers/hermes_tool_parser.py).
|
||||
|
||||
@@ -41,12 +41,12 @@ pip install -v -r requirements/xpu.txt
|
||||
|
||||
```bash
|
||||
pip uninstall -y triton triton-xpu
|
||||
pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu
|
||||
pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
|
||||
```
|
||||
|
||||
!!! note
|
||||
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
|
||||
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
- For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
|
||||
- Finally, build and install vLLM XPU backend:
|
||||
|
||||
|
||||
@@ -1,127 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This file demonstrates the example usage of disaggregated prefilling
|
||||
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
|
||||
and then transfer the KV cache between them.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from multiprocessing import Event, Process
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import KVTransferConfig
|
||||
|
||||
|
||||
def run_prefill(prefill_done):
|
||||
# We use GPU 0 for prefill node.
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
|
||||
# The prefill node receives two requests, while the decode node receives
|
||||
# three requests. So the decode node will only receive the KV Cache for
|
||||
# requests 1 and 3. The decode node will use the KV Cache of requests 1
|
||||
# and 3 and do prefilling on request 2.
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"Hi, your name is",
|
||||
# The decode node will actually "prefill" this request.
|
||||
"Tell me a very long story",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
|
||||
|
||||
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
|
||||
# This instance is the prefill node (kv_producer, rank 0).
|
||||
# The number of parallel instances for KV cache transfer is set to 2,
|
||||
# as required for P2pNcclConnector.
|
||||
ktc = KVTransferConfig(
|
||||
kv_connector="P2pNcclConnector",
|
||||
kv_role="kv_producer",
|
||||
kv_rank=0,
|
||||
kv_parallel_size=2,
|
||||
)
|
||||
|
||||
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
|
||||
# memory. You may need to adjust the value to fit your GPU.
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
llm.generate(prompts, sampling_params)
|
||||
print("Prefill node is finished.")
|
||||
prefill_done.set()
|
||||
|
||||
# To keep the prefill node running in case the decode node is not done;
|
||||
# otherwise, the script might exit prematurely, causing incomplete decoding.
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
print("Script stopped by user.")
|
||||
|
||||
|
||||
def run_decode(prefill_done):
|
||||
# We use GPU 1 for decode node.
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"Hi, your name is",
|
||||
"Tell me a very long story",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95)
|
||||
|
||||
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
|
||||
# This instance is the decode node (kv_consumer, rank 1).
|
||||
# The number of parallel instances for KV cache transfer is set to 2,
|
||||
# as required for P2pNcclConnector.
|
||||
ktc = KVTransferConfig(
|
||||
kv_connector="P2pNcclConnector",
|
||||
kv_role="kv_consumer",
|
||||
kv_rank=1,
|
||||
kv_parallel_size=2,
|
||||
)
|
||||
|
||||
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
|
||||
# memory. You may need to adjust the value to fit your GPU.
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
kv_transfer_config=ktc,
|
||||
max_model_len=2000,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
# Wait for the producer to start the pipe
|
||||
print("Waiting for prefill node to finish...")
|
||||
prefill_done.wait()
|
||||
|
||||
# At this point when the prefill_done is set, the kv-cache should have been
|
||||
# transferred to this decode node, so we can start decoding.
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
|
||||
def main():
|
||||
prefill_done = Event()
|
||||
prefill_process = Process(target=run_prefill, args=(prefill_done,))
|
||||
decode_process = Process(target=run_decode, args=(prefill_done,))
|
||||
|
||||
# Start prefill node
|
||||
prefill_process.start()
|
||||
|
||||
# Start decode node
|
||||
decode_process.start()
|
||||
|
||||
# Terminate the prefill node when decode is finished
|
||||
decode_process.join()
|
||||
prefill_process.terminate()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,125 +0,0 @@
|
||||
#!/bin/bash
|
||||
# This file demonstrates the example usage of disaggregated prefilling
|
||||
# We will launch 2 vllm instances (1 for prefill and 1 for decode),
|
||||
# and then transfer the KV cache between them.
|
||||
|
||||
set -xe
|
||||
|
||||
echo "🚧🚧 Warning: The usage of disaggregated prefill is experimental and subject to change 🚧🚧"
|
||||
sleep 1
|
||||
|
||||
# meta-llama/Meta-Llama-3.1-8B-Instruct or deepseek-ai/DeepSeek-V2-Lite
|
||||
MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
|
||||
|
||||
# Trap the SIGINT signal (triggered by Ctrl+C)
|
||||
trap 'cleanup' INT
|
||||
|
||||
# Cleanup function
|
||||
cleanup() {
|
||||
echo "Caught Ctrl+C, cleaning up..."
|
||||
# Cleanup commands
|
||||
pgrep python | xargs kill -9
|
||||
pkill -f python
|
||||
echo "Cleanup complete. Exiting."
|
||||
exit 0
|
||||
}
|
||||
|
||||
|
||||
if [[ -z "${VLLM_HOST_IP:-}" ]]; then
|
||||
export VLLM_HOST_IP=127.0.0.1
|
||||
echo "Using default VLLM_HOST_IP=127.0.0.1 (override by exporting VLLM_HOST_IP before running this script)"
|
||||
else
|
||||
echo "Using provided VLLM_HOST_IP=${VLLM_HOST_IP}"
|
||||
fi
|
||||
|
||||
|
||||
# install quart first -- required for disagg prefill proxy serve
|
||||
if python3 -c "import quart" &> /dev/null; then
|
||||
echo "Quart is already installed."
|
||||
else
|
||||
echo "Quart is not installed. Installing..."
|
||||
python3 -m pip install quart
|
||||
fi
|
||||
|
||||
# a function that waits vLLM server to start
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
timeout 1200 bash -c "
|
||||
until curl -i localhost:${port}/v1/models > /dev/null; do
|
||||
sleep 1
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
|
||||
# You can also adjust --kv-ip and --kv-port for distributed inference.
|
||||
|
||||
# prefilling instance, which is the KV producer
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# decoding instance, which is the KV consumer
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8200 \
|
||||
--max-model-len 100 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--trust-remote-code \
|
||||
--kv-transfer-config \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":"1e10","kv_port":"14580","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8200","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# wait until prefill and decode instances are ready
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
|
||||
# launch a proxy server that opens the service at port 8000
|
||||
# the workflow of this proxy:
|
||||
# - send the request to prefill vLLM instance (port 8100), change max_tokens
|
||||
# to 1
|
||||
# - after the prefill vLLM finishes prefill, send the request to decode vLLM
|
||||
# instance
|
||||
# NOTE: the usage of this API is subject to change --- in the future we will
|
||||
# introduce "vllm connect" to connect between prefill and decode instances
|
||||
python3 ../../benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py &
|
||||
sleep 1
|
||||
|
||||
# serve two example requests
|
||||
output1=$(curl -X POST -s http://localhost:8000/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'"$MODEL_NAME"'",
|
||||
"prompt": "San Francisco is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}')
|
||||
|
||||
output2=$(curl -X POST -s http://localhost:8000/v1/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'"$MODEL_NAME"'",
|
||||
"prompt": "Santa Clara is a",
|
||||
"max_tokens": 10,
|
||||
"temperature": 0
|
||||
}')
|
||||
|
||||
|
||||
# Cleanup commands
|
||||
pgrep python | xargs kill -9
|
||||
pkill -f python
|
||||
|
||||
echo ""
|
||||
|
||||
sleep 1
|
||||
|
||||
# Print the outputs of the curl requests
|
||||
echo ""
|
||||
echo "Output of first request: $output1"
|
||||
echo "Output of second request: $output2"
|
||||
|
||||
echo "🎉🎉 Successfully finished 2 test requests! 🎉🎉"
|
||||
echo ""
|
||||
@@ -35,12 +35,36 @@ Conversation isolation:
|
||||
the JSON body) to scope the KV cache across turns. Without it, the
|
||||
proxy cannot link turns and falls back to no-cache behavior.
|
||||
|
||||
``conversation_id`` is a non-standard extension to the OpenAI Chat
|
||||
Completions schema, consumed by this proxy and not forwarded to the
|
||||
vLLM engine. Strict OpenAI-compatible frontends reject unknown
|
||||
fields, so clients must opt in only when targeting this proxy.
|
||||
|
||||
Usage:
|
||||
python disagg_proxy_multiturn.py \\
|
||||
--host 0.0.0.0 --port 8000 \\
|
||||
--prefiller-host 10.0.0.1 --prefiller-port 8100 \\
|
||||
--decoder-host 10.0.0.2 --decoder-port 8200
|
||||
|
||||
Benchmarking:
|
||||
Use ``benchmarks/multi_turn/benchmark_serving_multi_turn.py`` with
|
||||
the ``--send-conversation-id`` flag to inject a per-conversation
|
||||
``conversation_id`` into every request so this proxy can key
|
||||
cross-turn KV cache reuse. The flag is *off by default*: without
|
||||
it the benchmark sends OpenAI-schema-compliant payloads and every
|
||||
turn lands as a cache MISS in this proxy.
|
||||
|
||||
Example:
|
||||
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\
|
||||
--model <MODEL> --served-model-name <NAME> \\
|
||||
--url http://<proxy_host>:8000 \\
|
||||
--input-file generate_multi_turn.json \\
|
||||
--num-clients 2 --max-active-conversations 6 \\
|
||||
--send-conversation-id
|
||||
|
||||
See ``docs/features/nixl_connector_usage.md`` for the broader
|
||||
bidirectional-KV-transfer setup these benchmarks exercise.
|
||||
|
||||
Dependencies:
|
||||
pip install fastapi uvicorn httpx
|
||||
"""
|
||||
@@ -373,7 +397,9 @@ async def _handle_request(api_path: str, request: Request):
|
||||
logger.warning(
|
||||
"[%s] No conversation_id provided — KV cache reuse disabled "
|
||||
"for this request. Add a 'conversation_id' field to enable "
|
||||
"cross-turn KV sharing.",
|
||||
"cross-turn KV sharing. When using "
|
||||
"benchmarks/multi_turn/benchmark_serving_multi_turn.py, pass "
|
||||
"--send-conversation-id (off by default).",
|
||||
request_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,245 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# =============================================================================
|
||||
# vLLM Disaggregated Serving Script - P2P NCCL XpYd Architecture
|
||||
# =============================================================================
|
||||
# This script demonstrates disaggregated prefill and decode serving using
|
||||
# P2P NCCL communication. The architecture supports various XpYd configurations:
|
||||
#
|
||||
# - 1P3D: 1 Prefill server + 3 Decode servers (current default)
|
||||
# - 3P1D: 3 Prefill servers + 1 Decode server
|
||||
# - etc.
|
||||
#
|
||||
# Configuration can be customized via environment variables:
|
||||
# MODEL: Model to serve
|
||||
# PREFILL_GPUS: Comma-separated GPU IDs for prefill servers
|
||||
# DECODE_GPUS: Comma-separated GPU IDs for decode servers
|
||||
# PREFILL_PORTS: Comma-separated ports for prefill servers
|
||||
# DECODE_PORTS: Comma-separated ports for decode servers
|
||||
# PROXY_PORT: Proxy server port used to setup XpYd connection.
|
||||
# TIMEOUT_SECONDS: Server startup timeout
|
||||
# =============================================================================
|
||||
|
||||
# Configuration - can be overridden via environment variables
|
||||
MODEL=${MODEL:-meta-llama/Llama-3.1-8B-Instruct}
|
||||
TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200}
|
||||
PROXY_PORT=${PROXY_PORT:-30001}
|
||||
|
||||
# Default 1P3D configuration (1 Prefill + 3 Decode)
|
||||
PREFILL_GPUS=${PREFILL_GPUS:-0}
|
||||
DECODE_GPUS=${DECODE_GPUS:-1,2,3}
|
||||
PREFILL_PORTS=${PREFILL_PORTS:-20003}
|
||||
DECODE_PORTS=${DECODE_PORTS:-20005,20007,20009}
|
||||
|
||||
echo "Warning: P2P NCCL disaggregated prefill XpYd support for vLLM v1 is experimental and subject to change."
|
||||
echo ""
|
||||
echo "Architecture Configuration:"
|
||||
echo " Model: $MODEL"
|
||||
echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS"
|
||||
echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS"
|
||||
echo " Proxy Port: $PROXY_PORT"
|
||||
echo " Timeout: ${TIMEOUT_SECONDS}s"
|
||||
echo ""
|
||||
|
||||
PIDS=()
|
||||
|
||||
# Switch to the directory of the current script
|
||||
cd "$(dirname "${BASH_SOURCE[0]}")"
|
||||
|
||||
check_required_files() {
|
||||
local files=("disagg_proxy_p2p_nccl_xpyd.py")
|
||||
for file in "${files[@]}"; do
|
||||
if [[ ! -f "$file" ]]; then
|
||||
echo "Required file $file not found in $(pwd)"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
check_hf_token() {
|
||||
if [ -z "$HF_TOKEN" ]; then
|
||||
echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
|
||||
echo "Example: export HF_TOKEN=your_token_here"
|
||||
exit 1
|
||||
fi
|
||||
if [[ "$HF_TOKEN" != hf_* ]]; then
|
||||
echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
|
||||
exit 1
|
||||
fi
|
||||
echo "HF_TOKEN is set and valid."
|
||||
}
|
||||
|
||||
check_num_gpus() {
|
||||
# Check if the number of GPUs are >=2 via nvidia-smi
|
||||
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
|
||||
if [ "$num_gpus" -lt 2 ]; then
|
||||
echo "You need at least 2 GPUs to run disaggregated prefill."
|
||||
exit 1
|
||||
else
|
||||
echo "Found $num_gpus GPUs."
|
||||
fi
|
||||
}
|
||||
|
||||
ensure_python_library_installed() {
|
||||
echo "Checking if $1 is installed..."
|
||||
if ! python3 -c "import $1" > /dev/null 2>&1; then
|
||||
echo "$1 is not installed. Please install it via pip install $1."
|
||||
exit 1
|
||||
else
|
||||
echo "$1 is installed."
|
||||
fi
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
echo "Stopping everything…"
|
||||
trap - INT TERM # prevent re-entrancy
|
||||
pkill -9 -f "disagg_proxy_p2p_nccl_xpyd.py"
|
||||
kill -- -$$ # negative PID == "this whole process-group"
|
||||
wait # reap children so we don't leave zombies
|
||||
exit 0
|
||||
}
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local timeout_seconds=$TIMEOUT_SECONDS
|
||||
local start_time=$(date +%s)
|
||||
|
||||
echo "Waiting for server on port $port..."
|
||||
|
||||
while true; do
|
||||
if curl -s "localhost:${port}/v1/completions" > /dev/null; then
|
||||
echo "Server on port $port is ready."
|
||||
return 0
|
||||
fi
|
||||
|
||||
local now=$(date +%s)
|
||||
if (( now - start_time >= timeout_seconds )); then
|
||||
echo "Timeout waiting for server on port $port"
|
||||
return 1
|
||||
fi
|
||||
|
||||
sleep 1
|
||||
done
|
||||
}
|
||||
|
||||
main() {
|
||||
check_required_files
|
||||
check_hf_token
|
||||
check_num_gpus
|
||||
ensure_python_library_installed pandas
|
||||
ensure_python_library_installed datasets
|
||||
ensure_python_library_installed vllm
|
||||
ensure_python_library_installed quart
|
||||
|
||||
trap cleanup INT
|
||||
trap cleanup USR1
|
||||
trap cleanup TERM
|
||||
|
||||
echo "Launching disaggregated serving components..."
|
||||
echo "Please check the log files for detailed output:"
|
||||
echo " - prefill*.log: Prefill server logs"
|
||||
echo " - decode*.log: Decode server logs"
|
||||
echo " - proxy.log: Proxy server log"
|
||||
|
||||
# =============================================================================
|
||||
# Launch Proxy Server
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Starting proxy server on port $PROXY_PORT..."
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
PIDS+=($!)
|
||||
|
||||
# Parse GPU and port arrays
|
||||
IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS"
|
||||
IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS"
|
||||
IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS"
|
||||
IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
|
||||
|
||||
# =============================================================================
|
||||
# Launch Prefill Servers (X Producers)
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..."
|
||||
for i in "${!PREFILL_GPU_ARRAY[@]}"; do
|
||||
local gpu_id=${PREFILL_GPU_ARRAY[$i]}
|
||||
local port=${PREFILL_PORT_ARRAY[$i]}
|
||||
local kv_port=$((21001 + i))
|
||||
|
||||
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"1e1\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > prefill$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
# Launch Decode Servers (Y Decoders)
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..."
|
||||
for i in "${!DECODE_GPU_ARRAY[@]}"; do
|
||||
local gpu_id=${DECODE_GPU_ARRAY[$i]}
|
||||
local port=${DECODE_PORT_ARRAY[$i]}
|
||||
local kv_port=$((22001 + i))
|
||||
|
||||
echo " Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
--max-model-len 10000 \
|
||||
--max-num-batched-tokens 10000 \
|
||||
--max-num-seqs 256 \
|
||||
--trust-remote-code \
|
||||
--gpu-memory-utilization 0.7 \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > decode$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
# Wait for All Servers to Start
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Waiting for all servers to start..."
|
||||
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
|
||||
if ! wait_for_server "$port"; then
|
||||
echo "Failed to start server on port $port"
|
||||
cleanup
|
||||
# shellcheck disable=SC2317
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
echo ""
|
||||
echo "All servers are up. Starting benchmark..."
|
||||
|
||||
# =============================================================================
|
||||
# Run Benchmark
|
||||
# =============================================================================
|
||||
cd ../../../benchmarks/
|
||||
vllm bench serve --port 10001 --seed "$(date +%s)" \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
|
||||
|
||||
echo "Benchmarking done. Cleaning up..."
|
||||
|
||||
cleanup
|
||||
}
|
||||
|
||||
main
|
||||
@@ -1,190 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
import msgpack
|
||||
import zmq
|
||||
from quart import Quart, make_response, request
|
||||
|
||||
count = 0
|
||||
prefill_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
|
||||
decode_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
|
||||
|
||||
prefill_cv = threading.Condition()
|
||||
decode_cv = threading.Condition()
|
||||
|
||||
DEFAULT_PING_SECONDS = 5
|
||||
|
||||
|
||||
def _remove_oldest_instances(instances: dict[str, Any]) -> None:
|
||||
oldest_key = next(iter(instances), None)
|
||||
while oldest_key is not None:
|
||||
value = instances[oldest_key]
|
||||
if value[1] > time.time():
|
||||
break
|
||||
print(f"🔴Remove [HTTP:{oldest_key}, ZMQ:{value[0]}, stamp:{value[1]}]")
|
||||
instances.pop(oldest_key, None)
|
||||
oldest_key = next(iter(instances), None)
|
||||
|
||||
|
||||
def _listen_for_register(poller, router_socket):
|
||||
while True:
|
||||
socks = dict(poller.poll())
|
||||
if router_socket in socks:
|
||||
remote_address, message = router_socket.recv_multipart()
|
||||
# data: {"type": "P", "http_address": "ip:port",
|
||||
# "zmq_address": "ip:port"}
|
||||
data = msgpack.loads(message)
|
||||
if data["type"] == "P":
|
||||
global prefill_instances
|
||||
global prefill_cv
|
||||
with prefill_cv:
|
||||
node = prefill_instances.get(data["http_address"], None)
|
||||
prefill_instances[data["http_address"]] = (
|
||||
data["zmq_address"],
|
||||
time.time() + DEFAULT_PING_SECONDS,
|
||||
)
|
||||
_remove_oldest_instances(prefill_instances)
|
||||
|
||||
elif data["type"] == "D":
|
||||
global decode_instances
|
||||
global decode_cv
|
||||
with decode_cv:
|
||||
node = decode_instances.get(data["http_address"], None)
|
||||
decode_instances[data["http_address"]] = (
|
||||
data["zmq_address"],
|
||||
time.time() + DEFAULT_PING_SECONDS,
|
||||
)
|
||||
_remove_oldest_instances(decode_instances)
|
||||
else:
|
||||
print(
|
||||
"Unexpected, Received message from %s, data: %s",
|
||||
remote_address,
|
||||
data,
|
||||
)
|
||||
return
|
||||
|
||||
if node is None:
|
||||
print(f"🔵Add [HTTP:{data['http_address']}, ZMQ:{data['zmq_address']}]")
|
||||
|
||||
|
||||
def start_service_discovery(hostname, port):
|
||||
if not hostname:
|
||||
hostname = socket.gethostname()
|
||||
if port == 0:
|
||||
raise ValueError("Port cannot be 0")
|
||||
|
||||
context = zmq.Context()
|
||||
router_socket = context.socket(zmq.ROUTER)
|
||||
router_socket.bind(f"tcp://{hostname}:{port}")
|
||||
|
||||
poller = zmq.Poller()
|
||||
poller.register(router_socket, zmq.POLLIN)
|
||||
|
||||
_listener_thread = threading.Thread(
|
||||
target=_listen_for_register, args=[poller, router_socket], daemon=True
|
||||
)
|
||||
_listener_thread.start()
|
||||
return _listener_thread
|
||||
|
||||
|
||||
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
|
||||
|
||||
app = Quart(__name__)
|
||||
|
||||
|
||||
def random_uuid() -> str:
|
||||
return str(uuid.uuid4().hex)
|
||||
|
||||
|
||||
async def forward_request(url, data, request_id):
|
||||
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
|
||||
"X-Request-Id": request_id,
|
||||
}
|
||||
async with session.post(url=url, json=data, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
if True:
|
||||
async for chunk_bytes in response.content.iter_chunked(1024):
|
||||
yield chunk_bytes
|
||||
else:
|
||||
content = await response.read()
|
||||
yield content
|
||||
|
||||
|
||||
@app.route("/v1/completions", methods=["POST"])
|
||||
@app.route("/v1/chat/completions", methods=["POST"])
|
||||
async def handle_request():
|
||||
try:
|
||||
original_request_data = await request.get_json()
|
||||
|
||||
prefill_request = original_request_data.copy()
|
||||
# change max_tokens = 1 to let it only do prefill
|
||||
prefill_request["max_tokens"] = 1
|
||||
if "max_completion_tokens" in prefill_request:
|
||||
prefill_request["max_completion_tokens"] = 1
|
||||
|
||||
global count
|
||||
global prefill_instances
|
||||
global prefill_cv
|
||||
with prefill_cv:
|
||||
prefill_list = list(prefill_instances.items())
|
||||
prefill_addr, prefill_zmq_addr = prefill_list[count % len(prefill_list)]
|
||||
prefill_zmq_addr = prefill_zmq_addr[0]
|
||||
|
||||
global decode_instances
|
||||
global decode_cv
|
||||
with decode_cv:
|
||||
decode_list = list(decode_instances.items())
|
||||
decode_addr, decode_zmq_addr = decode_list[count % len(decode_list)]
|
||||
decode_zmq_addr = decode_zmq_addr[0]
|
||||
|
||||
print(
|
||||
f"handle_request count: {count}, [HTTP:{prefill_addr}, "
|
||||
f"ZMQ:{prefill_zmq_addr}] 👉 [HTTP:{decode_addr}, "
|
||||
f"ZMQ:{decode_zmq_addr}]"
|
||||
)
|
||||
count += 1
|
||||
|
||||
request_id = (
|
||||
f"___prefill_addr_{prefill_zmq_addr}___decode_addr_"
|
||||
f"{decode_zmq_addr}_{random_uuid()}"
|
||||
)
|
||||
|
||||
# finish prefill
|
||||
async for _ in forward_request(
|
||||
f"http://{prefill_addr}{request.path}", prefill_request, request_id
|
||||
):
|
||||
continue
|
||||
|
||||
# return decode
|
||||
generator = forward_request(
|
||||
f"http://{decode_addr}{request.path}", original_request_data, request_id
|
||||
)
|
||||
response = await make_response(generator)
|
||||
response.timeout = None
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
exc_info = sys.exc_info()
|
||||
print("Error occurred in disagg prefill proxy server")
|
||||
print(e)
|
||||
print("".join(traceback.format_exception(*exc_info)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
t = start_service_discovery("0.0.0.0", 30001)
|
||||
app.run(host="0.0.0.0", port=10001)
|
||||
t.join()
|
||||
@@ -17,9 +17,7 @@ class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
|
||||
events: list[Any]
|
||||
|
||||
|
||||
class KVCacheEvent(
|
||||
msgspec.Struct, array_like=True, omit_defaults=True, gc=False, tag=True
|
||||
):
|
||||
class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
|
||||
"""Base class for all KV cache-related events"""
|
||||
|
||||
|
||||
|
||||
@@ -2562,6 +2562,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
"stepvl",
|
||||
"glm4_1v",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
# Dependencies for building Rust artifacts through setuptools-rust.
|
||||
setuptools>=77.0.3,<81.0.0
|
||||
setuptools-rust>=1.9.0
|
||||
wheel
|
||||
@@ -31,7 +31,7 @@ schemathesis
|
||||
jiwer
|
||||
bm25s
|
||||
pystemmer
|
||||
mteb[bm25s]
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
num2words
|
||||
pqdm
|
||||
|
||||
|
||||
@@ -646,7 +646,7 @@ tokenizers==0.22.2
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# transformers
|
||||
torch==2.11.0+xpu
|
||||
torch==2.12.0+xpu
|
||||
# via
|
||||
# -c requirements/xpu.txt
|
||||
# accelerate
|
||||
@@ -654,7 +654,7 @@ torch==2.11.0+xpu
|
||||
# sentence-transformers
|
||||
# timm
|
||||
# torchvision
|
||||
torchvision==0.26.0+xpu
|
||||
torchvision==0.27.0+xpu
|
||||
# via timm
|
||||
tqdm==4.67.3
|
||||
# via
|
||||
@@ -672,7 +672,7 @@ transformers==5.5.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# sentence-transformers
|
||||
triton-xpu==3.7.0
|
||||
triton-xpu==3.7.1
|
||||
# via torch
|
||||
typepy==1.3.4
|
||||
# via
|
||||
|
||||
@@ -12,9 +12,9 @@ jinja2>=3.1.6
|
||||
datasets # for benchmark scripts
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
--extra-index-url=https://download.pytorch.org/whl/xpu
|
||||
torch==2.11.0
|
||||
torch==2.12.0
|
||||
torchaudio
|
||||
torchvision
|
||||
|
||||
auto_round_lib>=0.13.0
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9/vllm_xpu_kernels-0.1.9-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
|
||||
Generated
+89
@@ -3458,6 +3458,75 @@ version = "0.1.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
|
||||
|
||||
[[package]]
|
||||
name = "pyo3"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"once_cell",
|
||||
"portable-atomic",
|
||||
"pyo3-build-config",
|
||||
"pyo3-ffi",
|
||||
"pyo3-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
|
||||
dependencies = [
|
||||
"target-lexicon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-ffi"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"pyo3-build-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"pyo3-macros-backend",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros-backend"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"pyo3-build-config",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pythonize"
|
||||
version = "0.28.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
|
||||
dependencies = [
|
||||
"pyo3",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "qoi"
|
||||
version = "0.4.1"
|
||||
@@ -4669,6 +4738,12 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "target-lexicon"
|
||||
version = "0.13.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
|
||||
|
||||
[[package]]
|
||||
name = "task-local"
|
||||
version = "0.1.1"
|
||||
@@ -5799,6 +5874,7 @@ dependencies = [
|
||||
"axum",
|
||||
"bytes",
|
||||
"clap",
|
||||
"educe",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"http-body",
|
||||
@@ -5813,7 +5889,9 @@ dependencies = [
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"serial_test",
|
||||
"sha2",
|
||||
"socket2",
|
||||
"subtle",
|
||||
"thiserror-ext",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
@@ -5901,6 +5979,17 @@ dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-tool-parser-py"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"pyo3",
|
||||
"pythonize",
|
||||
"serde_json",
|
||||
"thiserror-ext",
|
||||
"vllm-tool-parser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.5.0"
|
||||
|
||||
@@ -12,6 +12,7 @@ members = [
|
||||
"src/text",
|
||||
"src/tokenizer",
|
||||
"src/tool-parser",
|
||||
"src/tool-parser/python",
|
||||
]
|
||||
resolver = "3"
|
||||
|
||||
@@ -60,6 +61,8 @@ prometheus-client = "0.24.0"
|
||||
prometheus-client-derive-encode = "0.5.0"
|
||||
prost = "0.14.3"
|
||||
prost-types = "0.14.3"
|
||||
pyo3 = "0.28.3"
|
||||
pythonize = "0.28.0"
|
||||
rand = "0.9.2"
|
||||
reasoning-parser = "1.2.2"
|
||||
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
|
||||
@@ -75,8 +78,10 @@ serde_repr = "0.1.20"
|
||||
serde_tuple = "1.1.3"
|
||||
serde_with = "3.18.0"
|
||||
serial_test = { version = "3.2.0", features = ["file_locks"] }
|
||||
sha2 = "0.10.9"
|
||||
socket2 = "0.6.3"
|
||||
subenum = "1.1.3"
|
||||
subtle = "2.6"
|
||||
task-local = "0.1.1"
|
||||
tekken = { package = "tekken-rs", version = "0.1.1", default-features = false }
|
||||
tempfile = "3.23.0"
|
||||
|
||||
+1
-1
@@ -71,7 +71,7 @@ To build the `vllm-rs` in isolation:
|
||||
|
||||
```bash
|
||||
# from the local checkout
|
||||
cargo install --path src/cmd --bin vllm-rs
|
||||
./build_rust.sh
|
||||
```
|
||||
|
||||
### Example Request
|
||||
|
||||
@@ -66,6 +66,8 @@ pub enum Error {
|
||||
ToolCallStreamInvariant { message: String },
|
||||
#[error(transparent)]
|
||||
Text(#[from] vllm_text::Error),
|
||||
#[error(transparent)]
|
||||
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
@@ -52,7 +52,7 @@ mod stream;
|
||||
use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::ModelDtype;
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{TextLlm, TextRequest};
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
|
||||
/// Validate explicit parser override names without starting request processing.
|
||||
pub fn validate_parser_overrides(
|
||||
@@ -198,6 +198,33 @@ impl ChatLlm {
|
||||
Ok(ChatEventStream::new(request.request_id, structured_stream))
|
||||
}
|
||||
|
||||
/// Render through the chat template and tokenize, without submitting to the engine.
|
||||
///
|
||||
/// 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.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 {
|
||||
// Rendered string from the template (usual chat path).
|
||||
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
|
||||
// Already tokenized (e.g. multimodal path); pass through unchanged.
|
||||
Prompt::TokenIds(ids) => ids,
|
||||
};
|
||||
Ok(token_ids)
|
||||
}
|
||||
|
||||
/// Shut down the underlying LLM client and its background tasks.
|
||||
pub async fn shutdown(self) -> Result<()> {
|
||||
self.text.shutdown().await?;
|
||||
@@ -245,6 +272,6 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, step3)"].assert_eq(&error.to_report_string());
|
||||
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,7 +121,11 @@ impl ToolState {
|
||||
None => true,
|
||||
};
|
||||
if is_new_tool {
|
||||
let id = generate_tool_call_id();
|
||||
let id = self
|
||||
.parser
|
||||
.tool_call_id(item.tool_index)
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(generate_tool_call_id);
|
||||
self.open_call_index = Some(item.tool_index);
|
||||
events.push(AssistantEvent::ToolCallStart { id, name });
|
||||
}
|
||||
@@ -291,6 +295,11 @@ mod tests {
|
||||
buffered: String,
|
||||
}
|
||||
|
||||
struct IdScriptedParser {
|
||||
output: ToolParserOutput,
|
||||
tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
impl ToolParser for FailingParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
@@ -351,6 +360,35 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for IdScriptedParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self {
|
||||
output: ToolParserOutput::default(),
|
||||
tool_call_id: None,
|
||||
}))
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
(tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten()
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
output.append(std::mem::take(&mut self.output));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ToolParserOutput> {
|
||||
Ok(ToolParserOutput::default())
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolParser for PartialThenFailParser {
|
||||
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
|
||||
where
|
||||
@@ -505,6 +543,70 @@ mod tests {
|
||||
assert!(matches!(events[3], AssistantEvent::Done { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_preserves_parser_provided_tool_call_id() {
|
||||
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
})]);
|
||||
let parser = IdScriptedParser {
|
||||
output: ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: "{}".to_string(),
|
||||
}],
|
||||
},
|
||||
tool_call_id: Some("functions.get_weather:0".to_string()),
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { id, name }
|
||||
if id == "functions.get_weather:0" && name == "get_weather"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_stream_generates_tool_call_id_when_parser_omits_one() {
|
||||
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "ignored".to_string(),
|
||||
})]);
|
||||
let parser = IdScriptedParser {
|
||||
output: ToolParserOutput {
|
||||
normal_text: String::new(),
|
||||
calls: vec![crate::parser::tool::ToolCallDelta {
|
||||
tool_index: 0,
|
||||
name: Some("get_weather".to_string()),
|
||||
arguments: "{}".to_string(),
|
||||
}],
|
||||
},
|
||||
tool_call_id: None,
|
||||
};
|
||||
|
||||
let events = tool_event_stream(events, Some(Box::new(parser)))
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<crate::Result<Vec<_>>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
AssistantEvent::ToolCallStart { id, name }
|
||||
if id.starts_with("call_") && name == "get_weather"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn real_buffered_parser_error_matches_streaming_and_non_streaming() {
|
||||
let prefix = "I will check both.\n";
|
||||
|
||||
@@ -128,8 +128,6 @@ trait_set! {
|
||||
|
||||
/// Generate the northbound tool-call ID using the OpenAI-style `call_<id>`
|
||||
/// format.
|
||||
// TODO: support other ID scheme like Kimi-K2's
|
||||
// `functions.{name}:{global_index}`.
|
||||
pub(crate) fn generate_tool_call_id() -> String {
|
||||
format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])
|
||||
}
|
||||
|
||||
@@ -6,7 +6,8 @@ pub use vllm_reasoning_parser::{
|
||||
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
|
||||
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
|
||||
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
|
||||
ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
|
||||
ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
|
||||
Step3p5ReasoningParser,
|
||||
};
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
@@ -25,7 +26,9 @@ pub mod names {
|
||||
pub const MINIMAX_M2: &str = "minimax_m2";
|
||||
pub const NEMOTRON_V3: &str = "nemotron_v3";
|
||||
pub const QWEN3: &str = "qwen3";
|
||||
pub const SEED_OSS: &str = "seed_oss";
|
||||
pub const STEP3: &str = "step3";
|
||||
pub const STEP3P5: &str = "step3p5";
|
||||
}
|
||||
|
||||
/// Constructor signature for one registered reasoning parser implementation.
|
||||
@@ -61,7 +64,9 @@ impl ReasoningParserFactory {
|
||||
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
|
||||
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
|
||||
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
|
||||
.register_parser::<Step3ReasoningParser>(names::STEP3);
|
||||
.register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
|
||||
.register_parser::<Step3ReasoningParser>(names::STEP3)
|
||||
.register_parser::<Step3p5ReasoningParser>(names::STEP3P5);
|
||||
|
||||
factory
|
||||
.register_pattern("deepseek-r1", names::DEEPSEEK_R1)
|
||||
@@ -77,7 +82,14 @@ impl ReasoningParserFactory {
|
||||
.register_pattern("glm-4.5", names::GLM45)
|
||||
.register_pattern("kimi-k2", names::KIMI_K2)
|
||||
.register_pattern("kimi", names::KIMI)
|
||||
// step3p5 patterns must precede `step3`: substring matching would
|
||||
// otherwise route step3p5 IDs to step3.
|
||||
.register_pattern("step-3p5", names::STEP3P5)
|
||||
.register_pattern("step3p5", names::STEP3P5)
|
||||
.register_pattern("step-3.5", names::STEP3P5)
|
||||
.register_pattern("step3", names::STEP3)
|
||||
.register_pattern("seed-oss", names::SEED_OSS)
|
||||
.register_pattern("seedoss", names::SEED_OSS)
|
||||
.register_pattern("minimax", names::MINIMAX_M2)
|
||||
.register_pattern("mm-m2", names::MINIMAX_M2)
|
||||
.register_pattern("cohere", names::COHERE_CMD)
|
||||
|
||||
@@ -32,8 +32,12 @@ fn factory_contains_and_lists_registered_parsers() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
assert!(factory.contains(names::QWEN3));
|
||||
assert!(factory.contains(names::DEEPSEEK_V4));
|
||||
assert!(factory.contains(names::SEED_OSS));
|
||||
assert!(factory.contains(names::STEP3P5));
|
||||
assert!(factory.list().contains(&names::QWEN3.to_string()));
|
||||
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
|
||||
assert!(factory.list().contains(&names::SEED_OSS.to_string()));
|
||||
assert!(factory.list().contains(&names::STEP3P5.to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -49,6 +53,41 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_routes_step3p5_models_to_dedicated_parser() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
// step3p5 patterns must beat the bare `step3` substring.
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("step-3p5-instruct"),
|
||||
Some(names::STEP3P5)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("step3p5"),
|
||||
Some(names::STEP3P5)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("step-3.5-base"),
|
||||
Some(names::STEP3P5)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("step3-base"),
|
||||
Some(names::STEP3)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_routes_seed_oss_models() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("ByteDance-Seed/Seed-OSS-36B-Instruct"),
|
||||
Some(names::SEED_OSS)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("seedoss-7b"),
|
||||
Some(names::SEED_OSS)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_rejects_unknown_parser_names() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
|
||||
@@ -144,6 +144,30 @@ impl RoundtripCase {
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
|
||||
fn seed_oss() -> Self {
|
||||
Self {
|
||||
model_id: "ByteDance-Seed/Seed-OSS-36B-Instruct",
|
||||
assistant_stop_suffix: "<seed:eos>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Always { value: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Step-3.5 with `<think>` / `</think>` reasoning tags and newline trimming.
|
||||
fn step3p5() -> Self {
|
||||
Self {
|
||||
model_id: "stepfun-ai/Step-3.5-Flash",
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Always { value: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! roundtrip_tests {
|
||||
@@ -168,10 +192,11 @@ roundtrip_tests! {
|
||||
minimax_m25 => [reasoning_and_content, tool_call_mix],
|
||||
deepseek_v4 => [reasoning_and_content, tool_call_mix],
|
||||
glm47 => [reasoning_and_content, tool_call_mix],
|
||||
seed_oss => [reasoning_and_content],
|
||||
step3p5 => [reasoning_and_content],
|
||||
|
||||
// Note: Kimi K2.5 strips the reasoning content in history.
|
||||
// TODO: we don't respect model-generated tool call id now so `tool_call_mix` cannot pass.
|
||||
// kimi_k25 => [tool_call_mix],
|
||||
kimi_k25 => [tool_call_mix],
|
||||
}
|
||||
|
||||
/// Run the fixed reasoning+content fixture for one model/parser case.
|
||||
|
||||
+25
-1
@@ -16,6 +16,7 @@ use educe::Educe;
|
||||
use serde::Deserialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde_json::Value;
|
||||
use serde_with::{DefaultOnNull, OneOrMany, serde_as};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use uuid::Uuid;
|
||||
use vllm_engine_core_client::TransportMode;
|
||||
@@ -84,6 +85,7 @@ pub enum Command {
|
||||
}
|
||||
|
||||
/// Runtime arguments shared by the external-engine and managed-engine paths.
|
||||
#[serde_as]
|
||||
#[derive(Educe, Clone, Args, PartialEq, Eq, Deserialize)]
|
||||
#[educe(Debug)]
|
||||
pub struct SharedRuntimeArgs {
|
||||
@@ -178,6 +180,14 @@ pub struct SharedRuntimeArgs {
|
||||
#[serde(default)]
|
||||
pub enable_request_id_headers: bool,
|
||||
|
||||
/// If provided, the server will require one of these keys to be presented
|
||||
/// in the Authorization header.
|
||||
#[educe(Debug(ignore))]
|
||||
#[arg(long, env = "VLLM_API_KEY", value_delimiter = ' ')]
|
||||
#[serde_as(as = "DefaultOnNull<OneOrMany<_>>")]
|
||||
#[serde(default)]
|
||||
pub api_key: Vec<String>,
|
||||
|
||||
/// Disable periodic logging of engine statistics (throughput, queue depth,
|
||||
/// cache usage).
|
||||
#[arg(long)]
|
||||
@@ -215,6 +225,15 @@ impl SharedRuntimeArgs {
|
||||
Duration::from_secs(self.shutdown_timeout)
|
||||
}
|
||||
|
||||
/// Apply fallback logic for API key configuration from env variables.
|
||||
fn apply_env_api_key_fallback(&mut self) {
|
||||
if self.api_key.is_empty()
|
||||
&& let Ok(api_key) = std::env::var("VLLM_API_KEY")
|
||||
{
|
||||
self.api_key.push(api_key);
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the OpenAI-server config for the Python-bootstrap worker contract.
|
||||
///
|
||||
/// The resulting config binds the Python-supplied transport addresses and
|
||||
@@ -253,6 +272,7 @@ impl SharedRuntimeArgs {
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
api_keys: self.api_key,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
@@ -295,6 +315,7 @@ impl SharedRuntimeArgs {
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
api_keys: self.api_key,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
@@ -311,8 +332,11 @@ fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
|
||||
}
|
||||
|
||||
fn parse_runtime_args_json(value: &str) -> Result<SharedRuntimeArgs, String> {
|
||||
let args: SharedRuntimeArgs = serde_json::from_str(value)
|
||||
let mut args: SharedRuntimeArgs = serde_json::from_str(value)
|
||||
.map_err(|e| format!("invalid JSON arguments: {}", e.as_report()))?;
|
||||
// --args-json is parsed with serde, so clap's env support does not run for
|
||||
// the Python-supervised frontend path.
|
||||
args.apply_env_api_key_fallback();
|
||||
args.unsupported.check()?;
|
||||
Ok(args)
|
||||
}
|
||||
|
||||
@@ -169,6 +169,76 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() {
|
||||
assert!(config.enable_request_id_headers);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_passes_api_keys_into_config() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"serve",
|
||||
"Qwen/Qwen3-0.6B",
|
||||
"--api-key",
|
||||
"secret-a",
|
||||
"--api-key",
|
||||
"secret-b",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Serve(args) = cli.command else {
|
||||
panic!("expected serve args");
|
||||
};
|
||||
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
|
||||
assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
|
||||
let debug = format!("{config:#?}");
|
||||
assert!(debug.contains("api_keys: [<redacted>; 2]"));
|
||||
assert!(!debug.contains("secret-a"));
|
||||
assert!(!debug.contains("secret-b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_accepts_api_key_string() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":"secret"}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
let config = args.into_config();
|
||||
assert_eq!(config.api_keys, vec!["secret"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_accepts_api_key_list() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":["secret-a","secret-b"]}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
let config = args.into_config();
|
||||
assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_reject_unknown_renderer_value() {
|
||||
let error = Cli::try_parse_from([
|
||||
@@ -446,20 +516,21 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() {
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}"#,
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}"#,
|
||||
])
|
||||
.unwrap_err();
|
||||
|
||||
let actual = error.to_string().replace(": \n", ":\n");
|
||||
expect![[r#"
|
||||
error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}' for '--args-json <JSON>':
|
||||
error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>':
|
||||
The following arguments are not implemented in Rust frontend yet:
|
||||
- allow_credentials
|
||||
- api_key
|
||||
- ssl_keyfile
|
||||
|
||||
Remove these arguments to continue.
|
||||
|
||||
For more information, try '--help'.
|
||||
"#]].assert_eq(&error.to_string());
|
||||
"#]].assert_eq(&actual);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -793,6 +864,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
api_keys: [],
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -857,6 +929,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
api_keys: [],
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -936,6 +1009,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
api_keys: [],
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
|
||||
@@ -564,11 +564,6 @@ pub struct ServerUnsupportedArgs {
|
||||
#[arg(long)]
|
||||
pub allowed_headers: Option<Unsupported>,
|
||||
|
||||
/// If provided, the server will require one of these keys to be presented
|
||||
/// in the header.
|
||||
#[arg(long)]
|
||||
pub api_key: Option<Unsupported>,
|
||||
|
||||
/// The file path to the SSL key file.
|
||||
#[arg(long)]
|
||||
pub ssl_keyfile: Option<Unsupported>,
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::time::Duration;
|
||||
use anyhow::{Context, Result, bail};
|
||||
use clap::Parser;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use vllm_engine_core_client::protocol::utility::PauseMode;
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
@@ -32,8 +33,8 @@ struct Args {
|
||||
reset_external: bool,
|
||||
#[arg(long, default_value_t = 1)]
|
||||
sleep_level: u32,
|
||||
#[arg(long, default_value = "abort")]
|
||||
sleep_mode: String,
|
||||
#[arg(long, default_value_t = PauseMode::Abort)]
|
||||
sleep_mode: PauseMode,
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = false,
|
||||
@@ -106,7 +107,7 @@ async fn main() -> Result<()> {
|
||||
if args.skip_sleep_wake {
|
||||
println!("sleep_wake=skipped");
|
||||
} else {
|
||||
client.sleep(args.sleep_level, &args.sleep_mode).await.with_context(|| {
|
||||
client.sleep(args.sleep_level, args.sleep_mode).await.with_context(|| {
|
||||
format!(
|
||||
"failed to call sleep utility with level={} mode={}",
|
||||
args.sleep_level, args.sleep_mode
|
||||
|
||||
@@ -2,6 +2,7 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::future::{join_all, try_join_all};
|
||||
use itertools::Itertools;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
@@ -12,7 +13,7 @@ use crate::coordinator::CoordinatorHandle;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::handshake::EngineCoreReadyResponse;
|
||||
use crate::protocol::lora::LoraRequest;
|
||||
use crate::protocol::utility::EngineCoreUtilityRequest;
|
||||
use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
|
||||
use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
|
||||
use crate::transport::{self, ConnectedEngine};
|
||||
|
||||
@@ -571,6 +572,26 @@ impl EngineCoreClient {
|
||||
try_join_all(futures).await
|
||||
}
|
||||
|
||||
/// Call a utility method on all connected engines and return the shared
|
||||
/// result if every engine agrees.
|
||||
pub async fn call_utility_consensus<T, A>(&self, method: &str, args: A) -> Result<T>
|
||||
where
|
||||
T: serde::de::DeserializeOwned + std::fmt::Debug + PartialEq,
|
||||
A: serde::Serialize + std::fmt::Debug,
|
||||
{
|
||||
let results: Vec<T> = self.call_utility(method, args).await?;
|
||||
|
||||
if results.iter().all_equal() {
|
||||
// `engine_count >= 1` is enforced during startup handshake so `results` must be non-empty.
|
||||
Ok(results.into_iter().next().unwrap())
|
||||
} else {
|
||||
Err(Error::InconsistentUtilityResults {
|
||||
method: method.to_string(),
|
||||
values: format!("{results:?}"),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute `collective_rpc` on all engines and flatten all engine results
|
||||
/// into one list.
|
||||
pub async fn collective_rpc<A, K>(
|
||||
@@ -599,27 +620,8 @@ impl EngineCoreClient {
|
||||
}
|
||||
|
||||
/// Return whether the engine is currently sleeping at any level.
|
||||
///
|
||||
/// Under data parallel, all engines should agree on the sleep state: a
|
||||
/// divergence signals a control-plane bug. Returns
|
||||
/// `Error::InconsistentUtilityResults` if engines disagree.
|
||||
pub async fn is_sleeping(&self) -> Result<bool> {
|
||||
let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
|
||||
// `engine_count >= 1` is enforced during startup handshake, so `results`
|
||||
// is normally non-empty; fall back to a fail-loud error rather than
|
||||
// indexing in case that invariant is ever bypassed.
|
||||
let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
|
||||
method: "is_sleeping".to_string(),
|
||||
values: "[]".to_string(),
|
||||
})?;
|
||||
if results.iter().all(|&v| v == first) {
|
||||
Ok(first)
|
||||
} else {
|
||||
Err(Error::InconsistentUtilityResults {
|
||||
method: "is_sleeping".to_string(),
|
||||
values: format!("{results:?}"),
|
||||
})
|
||||
}
|
||||
self.call_utility_consensus("is_sleeping", ()).await
|
||||
}
|
||||
|
||||
/// Reset the multi-modal cache.
|
||||
@@ -643,22 +645,14 @@ impl EngineCoreClient {
|
||||
reset_running_requests: bool,
|
||||
reset_connector: bool,
|
||||
) -> Result<bool> {
|
||||
let results: Vec<bool> = self
|
||||
Ok(self
|
||||
.call_utility(
|
||||
"reset_prefix_cache",
|
||||
(reset_running_requests, reset_connector),
|
||||
)
|
||||
.await?;
|
||||
// `engine_count >= 1` is enforced during startup handshake, so `results`
|
||||
// is normally non-empty; fail loud rather than reporting a vacuous
|
||||
// success (`[].all() == true`) in case that invariant is ever bypassed.
|
||||
if results.is_empty() {
|
||||
return Err(Error::InconsistentUtilityResults {
|
||||
method: "reset_prefix_cache".to_string(),
|
||||
values: "[]".to_string(),
|
||||
});
|
||||
}
|
||||
Ok(results.into_iter().all(|ok| ok))
|
||||
.await?
|
||||
.into_iter()
|
||||
.all(|reset| reset))
|
||||
}
|
||||
|
||||
/// Load or refresh one LoRA adapter on every connected engine.
|
||||
@@ -680,7 +674,7 @@ impl EngineCoreClient {
|
||||
}
|
||||
|
||||
/// Put the engine to sleep.
|
||||
pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> {
|
||||
pub async fn sleep(&self, level: u32, mode: PauseMode) -> Result<()> {
|
||||
self.call_utility::<(), _>("sleep", (level, mode)).await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -692,6 +686,23 @@ impl EngineCoreClient {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Pause the scheduler so generation can be halted
|
||||
pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
|
||||
self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resume the scheduler after a pause
|
||||
pub async fn resume_scheduler(&self) -> Result<()> {
|
||||
self.call_utility::<(), _>("resume_scheduler", ()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Return whether the scheduler is currently in any pause state.
|
||||
pub async fn is_scheduler_paused(&self) -> Result<bool> {
|
||||
self.call_utility_consensus("is_scheduler_paused", ()).await
|
||||
}
|
||||
|
||||
/// Shut down local client tasks and close transport state.
|
||||
pub async fn shutdown(self) -> Result<()> {
|
||||
let Self {
|
||||
|
||||
@@ -162,6 +162,21 @@ pub enum RequestOutputKind {
|
||||
FinalOnly = 2,
|
||||
}
|
||||
|
||||
/// Structured-output backend selected for EngineCore grammar compilation.
|
||||
///
|
||||
/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
|
||||
/// validation. The Rust frontend currently always lowers structured-output
|
||||
/// requests to guidance, while ignoring any user-supplied `_backend` value.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum StructuredOutputBackend {
|
||||
Xgrammar,
|
||||
#[default]
|
||||
Guidance,
|
||||
Outlines,
|
||||
LmFormatEnforcer,
|
||||
}
|
||||
|
||||
/// The stop reason associated with a finished output.
|
||||
///
|
||||
/// Python models this as the union-typed `stop_reason: int | str | None`
|
||||
@@ -209,6 +224,17 @@ pub struct StructuredOutputsParams {
|
||||
pub whitespace_pattern: Option<String>,
|
||||
/// Structural tag configuration (JSON-encoded string).
|
||||
pub structural_tag: Option<String>,
|
||||
/// Structured-output backend, mirroring Python's internal `_backend`.
|
||||
///
|
||||
/// User-supplied values are ignored during deserialization. This matches
|
||||
/// Python's request boundary, where `_backend` is set by validation rather
|
||||
/// than accepted as a request-level backend selector.
|
||||
#[serde(
|
||||
default,
|
||||
rename = "_backend",
|
||||
deserialize_with = "serde_with::rust::deserialize_ignore_any"
|
||||
)]
|
||||
pub backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
/// Engine-core-facing sampling parameters for text generation.
|
||||
@@ -600,4 +626,18 @@ mod tests {
|
||||
|
||||
expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_backend_ignores_deserialized_value() {
|
||||
let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
|
||||
"json_object": true,
|
||||
"_backend": "xgrammar",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
assert_eq!(value["_backend"], "guidance");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,63 @@
|
||||
use std::any::type_name;
|
||||
use std::fmt;
|
||||
use std::{fmt, str::FromStr};
|
||||
|
||||
use rmpv::Value;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
use serde_with::{DeserializeFromStr, SerializeDisplay};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{OpaqueValue, default_opaque_value_nil};
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// How pause/sleep utility calls handle in-flight requests.
|
||||
///
|
||||
/// Use display/from-str serde so MessagePack utility args stay as Python
|
||||
/// literal strings instead of serde enum variant tuples.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, SerializeDisplay, DeserializeFromStr)]
|
||||
pub enum PauseMode {
|
||||
/// Abort all in-flight requests immediately.
|
||||
#[default]
|
||||
Abort,
|
||||
/// Wait for in-flight requests to complete.
|
||||
Wait,
|
||||
/// Freeze queued requests so they can resume later.
|
||||
Keep,
|
||||
}
|
||||
|
||||
impl PauseMode {
|
||||
/// Return the Python literal used on the utility-call wire.
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Abort => "abort",
|
||||
Self::Wait => "wait",
|
||||
Self::Keep => "keep",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for PauseMode {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for PauseMode {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(value: &str) -> std::result::Result<Self, Self::Err> {
|
||||
match value {
|
||||
"abort" => Ok(Self::Abort),
|
||||
"wait" => Ok(Self::Wait),
|
||||
"keep" => Ok(Self::Keep),
|
||||
other => Err(format!(
|
||||
"invalid pause mode `{other}`; expected one of: abort, wait, keep"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Utility call id as carried on the engine-core MessagePack wire.
|
||||
///
|
||||
/// Python emits utility ids as MessagePack integers, including values that may
|
||||
@@ -212,7 +260,7 @@ mod tests {
|
||||
use rmpv::Value;
|
||||
use serde::Serialize;
|
||||
|
||||
use super::{EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope};
|
||||
use super::{EngineCoreUtilityRequest, PauseMode, UtilityOutput, UtilityResultEnvelope};
|
||||
use crate::Error;
|
||||
use crate::protocol::{decode_msgpack, decode_value, encode_msgpack};
|
||||
|
||||
@@ -241,6 +289,26 @@ mod tests {
|
||||
assert_eq!(array[3], Value::Array(Vec::new()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pause_mode_serializes_as_python_literal() {
|
||||
let request =
|
||||
EngineCoreUtilityRequest::new(7, 42, "pause_scheduler", (PauseMode::Abort, true))
|
||||
.unwrap();
|
||||
|
||||
let encoded = encode_msgpack(&request).unwrap();
|
||||
let value = decode_value(&encoded).unwrap();
|
||||
let array = match value {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected utility request array, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(array[2], Value::from("pause_scheduler"));
|
||||
assert_eq!(
|
||||
array[3],
|
||||
Value::Array(vec![Value::from("abort"), Value::from(true)])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn utility_output_decodes_typed_result() {
|
||||
let output = UtilityOutput {
|
||||
|
||||
@@ -68,6 +68,11 @@ impl DelimitedReasoningParser {
|
||||
.unwrap_or(self.default_in_reasoning);
|
||||
}
|
||||
|
||||
/// Return whether the parser is currently inside a reasoning section.
|
||||
pub(crate) fn in_reasoning(&self) -> bool {
|
||||
self.current_in_reasoning
|
||||
}
|
||||
|
||||
/// Parse one decoded text delta and return its reasoning/content split.
|
||||
pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta {
|
||||
self.buffer.push_str(delta);
|
||||
|
||||
@@ -20,6 +20,8 @@ mod delimited;
|
||||
mod gemma4;
|
||||
mod kimi;
|
||||
mod qwen3;
|
||||
mod seed_oss;
|
||||
mod step3p5;
|
||||
|
||||
use thiserror::Error;
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
@@ -30,6 +32,8 @@ pub(crate) use self::delimited::DelimitedReasoningParser;
|
||||
pub use self::gemma4::Gemma4ReasoningParser;
|
||||
pub use self::kimi::KimiReasoningParser;
|
||||
pub use self::qwen3::Qwen3ReasoningParser;
|
||||
pub use self::seed_oss::SeedOssReasoningParser;
|
||||
pub use self::step3p5::Step3p5ReasoningParser;
|
||||
|
||||
/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
|
||||
pub type DeepSeekV3ReasoningParser = Qwen3ReasoningParser;
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
|
||||
|
||||
/// Reasoning parser for SeedOSS models using `<seed:think>`/`</seed:think>`
|
||||
/// delimiters.
|
||||
pub struct SeedOssReasoningParser {
|
||||
inner: DelimitedReasoningParser,
|
||||
}
|
||||
|
||||
impl SeedOssReasoningParser {
|
||||
/// Create a SeedOSS parser backed by the shared delimited state machine.
|
||||
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: DelimitedReasoningParser::new(
|
||||
tokenizer,
|
||||
"<seed:think>",
|
||||
"</seed:think>",
|
||||
false,
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ReasoningParser for SeedOssReasoningParser {
|
||||
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self::new(tokenizer)?))
|
||||
}
|
||||
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
|
||||
self.inner.initialize(prompt_token_ids);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
|
||||
Ok(self.inner.push(delta))
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ReasoningDelta> {
|
||||
Ok(self.inner.finish())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::SeedOssReasoningParser;
|
||||
use crate::{ReasoningParser, tests::FakeTokenizer};
|
||||
|
||||
#[test]
|
||||
fn without_prompt_markers_expects_start_token() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
assert_eq!(
|
||||
delta.content.as_deref(),
|
||||
Some("implicit reasoning</seed:think>answer")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn picks_up_prompt_start_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
|
||||
parser.initialize(&[10]).unwrap();
|
||||
|
||||
let delta = parser.push("reason</seed:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn respects_prompt_end_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt already closed reasoning with `</seed:think>` (id 11).
|
||||
parser.initialize(&[11]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_explicit_start_token() {
|
||||
// An explicit start delimiter must not leak into reasoning text.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streams_explicit_start_token_across_pushes() {
|
||||
// Start token, reasoning body, end token, and content arrive in separate
|
||||
// streaming deltas.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let mut reasoning = String::new();
|
||||
let mut content = String::new();
|
||||
for delta_str in [
|
||||
"<seed:think>",
|
||||
"Some ",
|
||||
"reasoning ",
|
||||
"content",
|
||||
"</seed:think>",
|
||||
"Final ",
|
||||
"answer",
|
||||
] {
|
||||
let delta = parser.push(delta_str).unwrap();
|
||||
if let Some(r) = delta.reasoning {
|
||||
reasoning.push_str(&r);
|
||||
}
|
||||
if let Some(c) = delta.content {
|
||||
content.push_str(&c);
|
||||
}
|
||||
}
|
||||
assert_eq!(reasoning, "Some reasoning content");
|
||||
assert_eq!(content, "Final answer");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_partial_delimiters_across_pushes() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[10]).unwrap();
|
||||
|
||||
// Closing delimiter `</seed:think>` arrives in two halves.
|
||||
let first = parser.push("reason</seed:").unwrap();
|
||||
assert_eq!(first.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(first.content, None);
|
||||
|
||||
let second = parser.push("think>answer").unwrap();
|
||||
assert_eq!(second.reasoning, None);
|
||||
assert_eq!(second.content.as_deref(), Some("answer"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
|
||||
|
||||
/// Reasoning parser for Step3p5 outputs.
|
||||
///
|
||||
/// Step3p5 uses standard `<think>`/`</think>` delimiters but emits a `\n`
|
||||
/// immediately before and/or after `</think>`. The parser drops these framing
|
||||
/// newlines on both sides of the boundary, holding a trailing `\n` from
|
||||
/// reasoning across pushes until either more reasoning text or `</think>`
|
||||
/// arrives, and dropping a leading `\n` from the first content delta after
|
||||
/// the boundary.
|
||||
pub struct Step3p5ReasoningParser {
|
||||
inner: DelimitedReasoningParser,
|
||||
/// `\n` at end of last reasoning delta, held in case `</think>` follows.
|
||||
pending_reasoning_newline: bool,
|
||||
/// Last push ended on `</think>` without emitting content; the next
|
||||
/// content delta's leading `\n` should be dropped.
|
||||
just_ended_reasoning: bool,
|
||||
}
|
||||
|
||||
impl Step3p5ReasoningParser {
|
||||
/// Create a Step3p5 parser backed by the shared delimited state machine.
|
||||
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false)?,
|
||||
pending_reasoning_newline: false,
|
||||
just_ended_reasoning: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Drop framing newlines around `</think>` and track held-newline state.
|
||||
fn process(
|
||||
&mut self,
|
||||
mut inner_delta: ReasoningDelta,
|
||||
was_in_reasoning: bool,
|
||||
now_in_reasoning: bool,
|
||||
) -> ReasoningDelta {
|
||||
// A `<think>...</think>` round-trip in one push still counts as a
|
||||
// transition: the inner emits reasoning while ending in content mode.
|
||||
let transitioned =
|
||||
!now_in_reasoning && (was_in_reasoning || inner_delta.reasoning.is_some());
|
||||
|
||||
// Replay or drop a previously-held trailing reasoning newline.
|
||||
if self.pending_reasoning_newline {
|
||||
if let Some(reasoning) = inner_delta.reasoning.as_mut() {
|
||||
reasoning.insert(0, '\n');
|
||||
self.pending_reasoning_newline = false;
|
||||
} else if transitioned {
|
||||
// The held `\n` was the one right before `</think>`: drop it.
|
||||
self.pending_reasoning_newline = false;
|
||||
}
|
||||
}
|
||||
|
||||
// Hold back a trailing reasoning `\n` until we know if `</think>` follows.
|
||||
if let Some(reasoning) = inner_delta.reasoning.as_mut()
|
||||
&& reasoning.ends_with('\n')
|
||||
{
|
||||
reasoning.pop();
|
||||
if !transitioned {
|
||||
self.pending_reasoning_newline = true;
|
||||
}
|
||||
}
|
||||
|
||||
// Drop a leading `\n` of content emitted right after `</think>`.
|
||||
if let Some(content) = inner_delta.content.as_mut()
|
||||
&& (transitioned || self.just_ended_reasoning)
|
||||
&& content.starts_with('\n')
|
||||
{
|
||||
content.remove(0);
|
||||
}
|
||||
|
||||
self.just_ended_reasoning = transitioned && inner_delta.content.is_none();
|
||||
|
||||
if inner_delta.reasoning.as_deref() == Some("") {
|
||||
inner_delta.reasoning = None;
|
||||
}
|
||||
if inner_delta.content.as_deref() == Some("") {
|
||||
inner_delta.content = None;
|
||||
}
|
||||
|
||||
inner_delta
|
||||
}
|
||||
}
|
||||
|
||||
impl ReasoningParser for Step3p5ReasoningParser {
|
||||
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self::new(tokenizer)?))
|
||||
}
|
||||
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
|
||||
self.inner.initialize(prompt_token_ids);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
|
||||
let was = self.inner.in_reasoning();
|
||||
let inner_delta = self.inner.push(delta);
|
||||
let now = self.inner.in_reasoning();
|
||||
Ok(self.process(inner_delta, was, now))
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ReasoningDelta> {
|
||||
let was = self.inner.in_reasoning();
|
||||
let inner_delta = self.inner.finish();
|
||||
let now = self.inner.in_reasoning();
|
||||
let mut delta = self.process(inner_delta, was, now);
|
||||
|
||||
// Emit a still-held newline rather than silently dropping it.
|
||||
if self.pending_reasoning_newline {
|
||||
match delta.reasoning.as_mut() {
|
||||
Some(existing) => existing.push('\n'),
|
||||
None => delta.reasoning = Some("\n".to_string()),
|
||||
}
|
||||
self.pending_reasoning_newline = false;
|
||||
}
|
||||
|
||||
Ok(delta)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::Step3p5ReasoningParser;
|
||||
use crate::{ReasoningParser, tests::FakeTokenizer};
|
||||
|
||||
#[test]
|
||||
fn picks_up_prompt_start_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt prefills `<think>` (id 1), opening reasoning before the stream.
|
||||
parser.initialize(&[1]).unwrap();
|
||||
|
||||
let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
|
||||
assert_eq!(
|
||||
delta.reasoning.as_deref(),
|
||||
Some("This is a reasoning section")
|
||||
);
|
||||
assert_eq!(delta.content.as_deref(), Some("This is the rest"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_unterminated_reasoning() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let pushed = parser.push("<think>reason without end").unwrap();
|
||||
assert_eq!(pushed.reasoning.as_deref(), Some("reason without end"));
|
||||
assert_eq!(pushed.content, None);
|
||||
|
||||
let flushed = parser.finish().unwrap();
|
||||
assert!(flushed.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_empty_input() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let pushed = parser.push("").unwrap();
|
||||
assert!(pushed.is_empty());
|
||||
let flushed = parser.finish().unwrap();
|
||||
assert!(flushed.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
|
||||
// Only the immediately-adjacent framing `\n` is dropped on each side of
|
||||
// `</think>`; surrounding newlines remain part of reasoning/content.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[1]).unwrap();
|
||||
|
||||
let delta = parser
|
||||
.push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
delta.reasoning.as_deref(),
|
||||
Some("\n This is a \n reasoning section\n\n")
|
||||
);
|
||||
assert_eq!(delta.content.as_deref(), Some("\nThis is the rest"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drops_framing_newlines_in_single_push() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drops_framing_newlines_across_pushes() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
// The trailing `\n` from the first push is held until we know whether
|
||||
// `</think>` follows.
|
||||
let first = parser.push("<think>reason\n").unwrap();
|
||||
assert_eq!(first.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(first.content, None);
|
||||
|
||||
// `</think>` arrives standalone; the held newline should be dropped.
|
||||
let second = parser.push("</think>").unwrap();
|
||||
assert!(second.is_empty());
|
||||
|
||||
// The leading newline of the first content delta is dropped.
|
||||
let third = parser.push("\nanswer").unwrap();
|
||||
assert_eq!(third.reasoning, None);
|
||||
assert_eq!(third.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replays_held_newline_when_more_reasoning_follows() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason\n").unwrap();
|
||||
assert_eq!(first.reasoning.as_deref(), Some("reason"));
|
||||
|
||||
let second = parser.push("more reason").unwrap();
|
||||
assert_eq!(second.reasoning.as_deref(), Some("\nmore reason"));
|
||||
assert_eq!(second.content, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finish_flushes_held_newline_in_unterminated_stream() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason\n").unwrap();
|
||||
assert_eq!(first.reasoning.as_deref(), Some("reason"));
|
||||
|
||||
let flushed = parser.finish().unwrap();
|
||||
assert_eq!(flushed.reasoning.as_deref(), Some("\n"));
|
||||
assert_eq!(flushed.content, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_inner_newlines_in_reasoning() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("line1\nline2"));
|
||||
assert_eq!(delta.content.as_deref(), Some("tail"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trims_only_one_trailing_reasoning_newline() {
|
||||
// Only the single framing newline immediately before `</think>` is
|
||||
// dropped; earlier newlines in the reasoning body are preserved.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason\n"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drops_only_first_content_newline_after_transition() {
|
||||
// The leading-`\n` drop applies only to the first content delta after
|
||||
// `</think>`; later deltas pass through untouched.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason</think>").unwrap();
|
||||
assert_eq!(first.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(first.content, None);
|
||||
|
||||
let second = parser.push("\nfirst").unwrap();
|
||||
assert_eq!(second.reasoning, None);
|
||||
assert_eq!(second.content.as_deref(), Some("first"));
|
||||
|
||||
// A `\n` arriving in a later content delta must NOT be dropped.
|
||||
let third = parser.push("\nsecond").unwrap();
|
||||
assert_eq!(third.reasoning, None);
|
||||
assert_eq!(third.content.as_deref(), Some("\nsecond"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_through_clean_boundary_without_framing_newlines() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason</think>tail").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("tail"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_empty_reasoning_section() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think></think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
}
|
||||
@@ -6,7 +6,7 @@ use super::{
|
||||
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
|
||||
};
|
||||
|
||||
struct FakeTokenizer;
|
||||
pub(crate) struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
@@ -32,6 +32,8 @@ impl Tokenizer for FakeTokenizer {
|
||||
"<|END_THINKING|>" => Some(4),
|
||||
"◁think▷" => Some(5),
|
||||
"◁/think▷" => Some(6),
|
||||
"<seed:think>" => Some(10),
|
||||
"</seed:think>" => Some(11),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ license.workspace = true
|
||||
anyhow.workspace = true
|
||||
asynk-strim-attr.workspace = true
|
||||
axum.workspace = true
|
||||
educe.workspace = true
|
||||
futures.workspace = true
|
||||
http-body.workspace = true
|
||||
itertools.workspace = true
|
||||
@@ -19,7 +20,9 @@ rmpv.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_with.workspace = true
|
||||
sha2.workspace = true
|
||||
socket2.workspace = true
|
||||
subtle.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
|
||||
@@ -70,6 +70,7 @@ async fn main() -> Result<()> {
|
||||
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
api_keys: Vec::new(),
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: Duration::ZERO,
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use educe::Educe;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
|
||||
@@ -33,7 +35,8 @@ pub enum CoordinatorMode {
|
||||
}
|
||||
|
||||
/// Normalized runtime configuration for the minimal OpenAI-compatible server.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
#[derive(Educe, Clone, PartialEq, Eq, Serialize)]
|
||||
#[educe(Debug)]
|
||||
pub struct Config {
|
||||
/// Frontend-to-engine transport setup.
|
||||
pub transport_mode: TransportMode,
|
||||
@@ -67,6 +70,10 @@ pub struct Config {
|
||||
pub enable_log_requests: bool,
|
||||
/// When `true`, set `X-Request-Id` on every HTTP response.
|
||||
pub enable_request_id_headers: bool,
|
||||
/// API keys accepted as bearer tokens for guarded routes.
|
||||
#[serde(skip_serializing)]
|
||||
#[educe(Debug(method(fmt_redacted_api_keys)))]
|
||||
pub api_keys: Vec<String>,
|
||||
/// When `true`, suppress periodic stats logging (throughput, queue depth,
|
||||
/// cache usage).
|
||||
pub disable_log_stats: bool,
|
||||
@@ -114,3 +121,19 @@ impl Config {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct RedactedApiKeys<'a>(&'a [String]);
|
||||
|
||||
impl fmt::Debug for RedactedApiKeys<'_> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
if self.0.is_empty() {
|
||||
f.debug_list().finish()
|
||||
} else {
|
||||
write!(f, "[<redacted>; {}]", self.0.len())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fmt_redacted_api_keys(api_keys: &[String], f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
fmt::Debug::fmt(&RedactedApiKeys(api_keys), f)
|
||||
}
|
||||
|
||||
@@ -14,7 +14,8 @@ mod utils;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use anyhow::{Context as _, Result};
|
||||
use axum::{Router, serve::ListenerExt as _};
|
||||
use axum::Router;
|
||||
use axum::serve::ListenerExt as _;
|
||||
pub use config::{Config, CoordinatorMode, HttpListenerMode};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::time::{Instant, sleep_until};
|
||||
@@ -92,7 +93,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
|
||||
AppState::new(served_model_names, chat)
|
||||
.with_log_requests(config.enable_log_requests)
|
||||
.with_request_id_headers(config.enable_request_id_headers)
|
||||
.with_server_info(ServerInfoSnapshot::from_config(config)),
|
||||
.with_server_info(ServerInfoSnapshot::from_config(config))
|
||||
.with_api_keys(config.api_keys.clone()),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::Json;
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::{HeaderValue, Method, StatusCode};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use serde_json::json;
|
||||
|
||||
use crate::state::{ApiKeyHash, AppState, hash_api_key};
|
||||
|
||||
const GUARDED_PREFIXES: &[&str] = &["/v1", "/v2", "/inference"];
|
||||
|
||||
/// Authenticate guarded HTTP routes with an OpenAI-compatible bearer token.
|
||||
///
|
||||
/// Mirrors Python `AuthenticationMiddleware`: OPTIONS requests and non-guarded
|
||||
/// helper endpoints such as `/health` are allowed through without a token.
|
||||
pub async fn authenticate_api_key(
|
||||
State(state): State<Arc<AppState>>,
|
||||
req: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
if req.method() == Method::OPTIONS || !requires_auth(req.uri().path()) {
|
||||
return next.run(req).await;
|
||||
}
|
||||
|
||||
if verify_token(req.headers().get(AUTHORIZATION), state.api_key_hashes()) {
|
||||
return next.run(req).await;
|
||||
}
|
||||
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({ "error": "Unauthorized" })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn requires_auth(path: &str) -> bool {
|
||||
GUARDED_PREFIXES.iter().any(|prefix| path.starts_with(prefix))
|
||||
}
|
||||
|
||||
fn verify_token(authorization: Option<&HeaderValue>, api_key_hashes: &[ApiKeyHash]) -> bool {
|
||||
let Some(authorization) = authorization else {
|
||||
return false;
|
||||
};
|
||||
let Ok(authorization) = authorization.to_str() else {
|
||||
return false;
|
||||
};
|
||||
let Some((scheme, token)) = authorization.split_once(' ') else {
|
||||
return false;
|
||||
};
|
||||
if !scheme.eq_ignore_ascii_case("bearer") {
|
||||
return false;
|
||||
}
|
||||
|
||||
let token_hash = hash_api_key(token);
|
||||
let mut token_match = false;
|
||||
for api_key_hash in api_key_hashes {
|
||||
token_match |= constant_time_eq(&token_hash, api_key_hash);
|
||||
}
|
||||
token_match
|
||||
}
|
||||
|
||||
fn constant_time_eq(left: &ApiKeyHash, right: &ApiKeyHash) -> bool {
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
bool::from(left.ct_eq(right))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::constant_time_eq;
|
||||
use crate::state::hash_api_key;
|
||||
|
||||
#[test]
|
||||
fn constant_time_eq_checks_sha256_digests() {
|
||||
assert!(constant_time_eq(
|
||||
&hash_api_key("secret"),
|
||||
&hash_api_key("secret")
|
||||
));
|
||||
assert!(!constant_time_eq(
|
||||
&hash_api_key("secret"),
|
||||
&hash_api_key("secrex")
|
||||
));
|
||||
assert!(!constant_time_eq(
|
||||
&hash_api_key("secret"),
|
||||
&hash_api_key("secret-more")
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,9 @@
|
||||
mod auth;
|
||||
mod load;
|
||||
mod metrics;
|
||||
mod request_id;
|
||||
|
||||
pub use auth::authenticate_api_key;
|
||||
pub use load::track_server_load;
|
||||
pub use metrics::track_http_metrics;
|
||||
pub use request_id::set_request_id_header;
|
||||
|
||||
@@ -6,8 +6,10 @@ mod load;
|
||||
mod lora;
|
||||
mod metrics;
|
||||
pub(crate) mod openai;
|
||||
mod pause;
|
||||
mod server_info;
|
||||
mod sleep;
|
||||
mod tokenize;
|
||||
mod version;
|
||||
|
||||
use std::sync::Arc;
|
||||
@@ -71,7 +73,9 @@ fn build_router_with_options(
|
||||
.route("/v1/models", get(openai::list_models))
|
||||
.route("/v1/completions", post(openai::completions))
|
||||
.route("/v1/chat/completions", post(openai::chat_completions))
|
||||
// vLLM specific inference endpoints
|
||||
// vLLM specific endpoints
|
||||
.route("/tokenize", post(tokenize::tokenize))
|
||||
.route("/detokenize", post(tokenize::detokenize))
|
||||
.route("/inference/v1/generate", post(inference::generate));
|
||||
|
||||
if runtime_lora_updating_enabled {
|
||||
@@ -90,15 +94,32 @@ fn build_router_with_options(
|
||||
.route("/sleep", post(sleep::sleep))
|
||||
.route("/wake_up", post(sleep::wake_up))
|
||||
.route("/is_sleeping", get(sleep::is_sleeping))
|
||||
.route("/pause", post(pause::pause))
|
||||
.route("/resume", post(pause::resume))
|
||||
.route("/is_paused", get(pause::is_paused))
|
||||
.route("/server_info", get(server_info::server_info))
|
||||
}
|
||||
|
||||
let enable_request_id_headers = state.enable_request_id_headers;
|
||||
let enable_api_key_auth = state.has_api_keys();
|
||||
let mut router = router
|
||||
.with_state(state.clone())
|
||||
.layer(from_fn_with_state(state, middleware::track_server_load))
|
||||
.layer(from_fn(middleware::track_http_metrics))
|
||||
.layer(TraceLayer::new_for_http());
|
||||
.layer(from_fn_with_state(
|
||||
state.clone(),
|
||||
middleware::track_server_load,
|
||||
))
|
||||
.layer(from_fn(middleware::track_http_metrics));
|
||||
|
||||
if enable_api_key_auth {
|
||||
router = router.layer(from_fn_with_state(
|
||||
state.clone(),
|
||||
middleware::authenticate_api_key,
|
||||
));
|
||||
}
|
||||
|
||||
// Later layers wrap earlier ones. Keep tracing outside auth so rejected
|
||||
// requests are visible, while metrics/load only see authenticated traffic.
|
||||
router = router.layer(TraceLayer::new_for_http());
|
||||
|
||||
if enable_request_id_headers {
|
||||
router = router.layer(from_fn(middleware::set_request_id_header));
|
||||
|
||||
@@ -22,7 +22,7 @@ use vllm_llm::{
|
||||
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
|
||||
};
|
||||
|
||||
use self::convert::prepare_generate_request;
|
||||
use self::convert::{ResponseOptions, prepare_generate_request};
|
||||
use self::types::{
|
||||
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
|
||||
GenerateResponseStreamChoice, GenerateStreamResponse,
|
||||
@@ -54,10 +54,7 @@ pub async fn generate(
|
||||
);
|
||||
|
||||
let log_request = state.enable_log_requests;
|
||||
let include_logprobs = prepared.include_logprobs;
|
||||
let include_prompt_logprobs = prepared.include_prompt_logprobs;
|
||||
let stream = prepared.stream;
|
||||
|
||||
let raw_stream = match state
|
||||
.chat
|
||||
.text()
|
||||
@@ -80,9 +77,7 @@ pub async fn generate(
|
||||
raw_stream,
|
||||
prepared.request_id,
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.include_continuous_usage,
|
||||
include_logprobs,
|
||||
prepared.options,
|
||||
);
|
||||
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
|
||||
|
||||
@@ -100,21 +95,11 @@ pub async fn generate(
|
||||
}
|
||||
};
|
||||
|
||||
if log_request {
|
||||
info!(
|
||||
parent: &request_span,
|
||||
prompt_tokens = collected.prompt_token_ids.len(),
|
||||
output_tokens = collected.token_ids.len(),
|
||||
finish_reason = collected.finish_reason.as_str(),
|
||||
"generate finished"
|
||||
);
|
||||
}
|
||||
|
||||
let response = match collect_generate(
|
||||
collected,
|
||||
prepared.request_id,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
log_request,
|
||||
prepared.options,
|
||||
) {
|
||||
Ok(response) => response,
|
||||
Err(error) => return error.into_response(),
|
||||
@@ -128,9 +113,13 @@ async fn generate_chunk_stream(
|
||||
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
|
||||
request_id: String,
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
include_continuous_usage: bool,
|
||||
include_logprobs: bool,
|
||||
ResponseOptions {
|
||||
include_usage,
|
||||
include_continuous_usage,
|
||||
include_logprobs,
|
||||
// Ignored: raw generate streaming has no prompt-logprobs wire shape.
|
||||
include_prompt_logprobs: _,
|
||||
}: ResponseOptions,
|
||||
mut y: TryYielder<GenerateStreamResponse, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
pin_mut!(stream);
|
||||
@@ -222,8 +211,15 @@ async fn generate_chunk_stream(
|
||||
fn collect_generate(
|
||||
collected: CollectedGenerateOutput,
|
||||
request_id: String,
|
||||
include_logprobs: bool,
|
||||
include_prompt_logprobs: bool,
|
||||
log_request: bool,
|
||||
ResponseOptions {
|
||||
// Ignored: non-streaming raw generate responses do not include usage.
|
||||
include_usage: _,
|
||||
// Ignored: continuous usage is a streaming-only option.
|
||||
include_continuous_usage: _,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
}: ResponseOptions,
|
||||
) -> Result<GenerateResponse, ApiError> {
|
||||
let logprobs = if include_logprobs {
|
||||
let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
|
||||
@@ -246,13 +242,23 @@ fn collect_generate(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let finish_reason = collected.finish_reason.as_str().to_string();
|
||||
|
||||
if log_request {
|
||||
info!(
|
||||
prompt_tokens = collected.prompt_token_ids.len(),
|
||||
output_tokens = collected.token_ids.len(),
|
||||
%finish_reason,
|
||||
"generate finished"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(GenerateResponse {
|
||||
request_id,
|
||||
choices: vec![GenerateResponseChoice {
|
||||
index: 0,
|
||||
logprobs,
|
||||
finish_reason: Some(collected.finish_reason.as_str().to_string()),
|
||||
finish_reason: Some(finish_reason),
|
||||
token_ids: collected.token_ids,
|
||||
}],
|
||||
prompt_logprobs,
|
||||
@@ -408,11 +414,19 @@ mod tests {
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks: Vec<_> =
|
||||
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("collect chunks");
|
||||
let chunks: Vec<_> = generate_chunk_stream(
|
||||
stream,
|
||||
"raw-stream".to_string(),
|
||||
false,
|
||||
ResponseOptions {
|
||||
include_usage: true,
|
||||
include_continuous_usage: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("collect chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(
|
||||
|
||||
@@ -8,19 +8,29 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
|
||||
|
||||
/// Lowered generate request plus the response request ID.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PreparedRequest {
|
||||
pub(super) struct PreparedRequest {
|
||||
pub request_id: String,
|
||||
pub text_request: TextRequest,
|
||||
pub stream: bool,
|
||||
/// Public response rendering options for route-layer helpers.
|
||||
pub options: ResponseOptions,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq)]
|
||||
pub(super) struct ResponseOptions {
|
||||
/// Whether the caller asked for the final streamed usage chunk.
|
||||
pub include_usage: bool,
|
||||
/// Whether the caller asked for usage on every streamed chunk.
|
||||
pub include_continuous_usage: bool,
|
||||
/// Whether the caller requested output logprobs on generate choices.
|
||||
pub include_logprobs: bool,
|
||||
/// Whether the caller requested top-level prompt logprobs.
|
||||
pub include_prompt_logprobs: bool,
|
||||
}
|
||||
|
||||
/// Validate and lower one raw generate request into the internal
|
||||
/// text-generation format.
|
||||
pub fn prepare_generate_request(
|
||||
pub(super) fn prepare_generate_request(
|
||||
request: GenerateRequest,
|
||||
lora_resolution: &LoraModelResolution,
|
||||
ctx: ResolvedRequestContext,
|
||||
@@ -65,10 +75,12 @@ pub fn prepare_generate_request(
|
||||
request_id: ctx.request_id,
|
||||
text_request,
|
||||
stream,
|
||||
include_usage,
|
||||
include_continuous_usage,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
options: ResponseOptions {
|
||||
include_usage,
|
||||
include_continuous_usage,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -158,7 +170,7 @@ mod tests {
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
assert!(!prepared.include_usage);
|
||||
assert!(!prepared.include_continuous_usage);
|
||||
assert!(!prepared.options.include_usage);
|
||||
assert!(!prepared.options.include_continuous_usage);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
pub mod convert;
|
||||
pub(crate) mod convert;
|
||||
mod types;
|
||||
mod validate;
|
||||
|
||||
@@ -23,8 +23,8 @@ use vllm_chat::{
|
||||
};
|
||||
use vllm_engine_core_client::protocol::StopReason;
|
||||
|
||||
use self::convert::{ResponseOptions, prepare_chat_request};
|
||||
use crate::error::{ApiError, bail_server_error, server_error};
|
||||
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
|
||||
use crate::routes::openai::chat_completions::types::{
|
||||
AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
|
||||
ChatCompletionResponse, ChatCompletionStreamChoice, ChatCompletionStreamResponse,
|
||||
@@ -83,12 +83,7 @@ pub async fn chat_completions(
|
||||
prepared.response_model,
|
||||
created,
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.requested_logprobs,
|
||||
prepared.include_reasoning,
|
||||
prepared.echo,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
prepared.options,
|
||||
);
|
||||
let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
|
||||
|
||||
@@ -99,12 +94,8 @@ pub async fn chat_completions(
|
||||
prepared.request_id,
|
||||
prepared.response_model,
|
||||
created,
|
||||
prepared.requested_logprobs,
|
||||
prepared.include_prompt_logprobs,
|
||||
prepared.include_reasoning,
|
||||
prepared.echo,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
log_request,
|
||||
prepared.options,
|
||||
)
|
||||
.instrument(request_span.clone())
|
||||
.await
|
||||
@@ -113,18 +104,6 @@ pub async fn chat_completions(
|
||||
Err(error) => return error.into_response(),
|
||||
};
|
||||
|
||||
if log_request {
|
||||
let usage = response.usage.as_ref();
|
||||
info!(
|
||||
parent: &request_span,
|
||||
model = %response.model,
|
||||
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
|
||||
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
|
||||
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
|
||||
"chat completion finished"
|
||||
);
|
||||
}
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
}
|
||||
@@ -134,12 +113,17 @@ async fn collect_chat_completion(
|
||||
request_id: String,
|
||||
response_model: String,
|
||||
created: u64,
|
||||
requested_logprobs: bool,
|
||||
include_prompt_logprobs: bool,
|
||||
include_reasoning: bool,
|
||||
echo: Option<String>,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
log_request: bool,
|
||||
ResponseOptions {
|
||||
// Ignored: non-streaming responses always include usage.
|
||||
include_usage: _,
|
||||
requested_logprobs,
|
||||
include_prompt_logprobs,
|
||||
include_reasoning,
|
||||
echo,
|
||||
return_token_ids,
|
||||
return_tokens_as_token_ids,
|
||||
}: ResponseOptions,
|
||||
) -> Result<ChatCompletionResponse, ApiError> {
|
||||
let collected = stream.collect_message().await.map_err(|error| {
|
||||
server_error!(
|
||||
@@ -201,6 +185,16 @@ async fn collect_chat_completion(
|
||||
};
|
||||
let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
|
||||
|
||||
if log_request {
|
||||
info!(
|
||||
model = %response_model,
|
||||
prompt_tokens = usage.prompt_tokens,
|
||||
output_tokens = usage.completion_tokens.unwrap_or(0),
|
||||
finish_reason = %finish_reason,
|
||||
"chat completion finished"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(ChatCompletionResponse {
|
||||
id: request_id,
|
||||
object: "chat.completion".to_string(),
|
||||
@@ -238,12 +232,16 @@ async fn chat_completion_chunk_stream(
|
||||
response_model: String,
|
||||
created: u64,
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
requested_logprobs: bool,
|
||||
include_reasoning: bool,
|
||||
echo: Option<String>,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
ResponseOptions {
|
||||
include_usage,
|
||||
requested_logprobs,
|
||||
// Ignored: chat streaming prompt logprobs are rejected for Python parity.
|
||||
include_prompt_logprobs: _,
|
||||
include_reasoning,
|
||||
echo,
|
||||
return_token_ids,
|
||||
return_tokens_as_token_ids,
|
||||
}: ResponseOptions,
|
||||
mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
let mut saw_tool_calls = false;
|
||||
@@ -806,7 +804,7 @@ mod tests {
|
||||
use vllm_engine_core_client::protocol::StopReason;
|
||||
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
|
||||
|
||||
use super::{block_delta_chunk, chat_completion_chunk_stream, final_chunk};
|
||||
use super::{ResponseOptions, block_delta_chunk, chat_completion_chunk_stream, final_chunk};
|
||||
|
||||
#[test]
|
||||
fn text_chunk_uses_content_only_delta() {
|
||||
@@ -932,12 +930,11 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
ResponseOptions {
|
||||
requested_logprobs: true,
|
||||
include_reasoning: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
@@ -996,12 +993,11 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
ResponseOptions {
|
||||
requested_logprobs: true,
|
||||
include_reasoning: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
@@ -1049,12 +1045,7 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
ResponseOptions::default(),
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
@@ -1132,12 +1123,11 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
ResponseOptions {
|
||||
requested_logprobs: true,
|
||||
return_token_ids: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
@@ -1263,12 +1253,11 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
ResponseOptions {
|
||||
requested_logprobs: true,
|
||||
return_token_ids: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
@@ -1342,12 +1331,10 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
ResponseOptions {
|
||||
include_reasoning: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
|
||||
@@ -18,11 +18,19 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
|
||||
/// Lowered chat request plus the public response metadata carried by every SSE
|
||||
/// chunk.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PreparedRequest {
|
||||
pub(super) struct PreparedRequest {
|
||||
/// Stable OpenAI-style request ID, reused as the external chat request ID.
|
||||
pub request_id: String,
|
||||
/// Public model ID echoed back to the client.
|
||||
pub response_model: String,
|
||||
/// Public response rendering options for route-layer helpers.
|
||||
pub options: ResponseOptions,
|
||||
/// Lowered chat request for `vllm-chat`.
|
||||
pub chat_request: ChatRequest,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub(super) struct ResponseOptions {
|
||||
/// Whether the caller asked for the final streamed usage chunk.
|
||||
pub include_usage: bool,
|
||||
/// Whether the caller requested output logprobs on chat choices.
|
||||
@@ -31,8 +39,6 @@ pub struct PreparedRequest {
|
||||
pub include_prompt_logprobs: bool,
|
||||
/// Whether to include reasoning content in OpenAI responses.
|
||||
pub include_reasoning: bool,
|
||||
/// Lowered chat request for `vllm-chat`.
|
||||
pub chat_request: ChatRequest,
|
||||
/// Last assistant-role message content to echo back when `echo=true`.
|
||||
pub echo: Option<String>,
|
||||
/// Whether to include token IDs alongside generated text.
|
||||
@@ -46,7 +52,7 @@ pub struct PreparedRequest {
|
||||
///
|
||||
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
|
||||
/// the base `model` field in responses when no LoRA adapter is selected.
|
||||
pub(crate) fn prepare_chat_request(
|
||||
pub(super) fn prepare_chat_request(
|
||||
request: ChatCompletionRequest,
|
||||
lora_resolution: &LoraModelResolution,
|
||||
ctx: ResolvedRequestContext,
|
||||
@@ -146,18 +152,19 @@ pub(crate) fn prepare_chat_request(
|
||||
Ok(PreparedRequest {
|
||||
request_id,
|
||||
response_model,
|
||||
include_usage,
|
||||
requested_logprobs,
|
||||
include_prompt_logprobs,
|
||||
include_reasoning,
|
||||
options: ResponseOptions {
|
||||
include_usage,
|
||||
requested_logprobs,
|
||||
include_prompt_logprobs,
|
||||
include_reasoning,
|
||||
echo,
|
||||
return_token_ids: request.return_token_ids.unwrap_or(false),
|
||||
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
|
||||
},
|
||||
chat_request,
|
||||
echo,
|
||||
return_token_ids: request.return_token_ids.unwrap_or(false),
|
||||
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_generation_prompt_mode(
|
||||
pub(crate) fn normalize_generation_prompt_mode(
|
||||
add_generation_prompt: Option<bool>,
|
||||
continue_final_message: bool,
|
||||
messages: &[VllmChatMessage],
|
||||
@@ -204,7 +211,7 @@ fn extract_last_assistant_content(messages: &[ChatMessage]) -> Option<String> {
|
||||
}
|
||||
|
||||
/// Lower one OpenAI chat message into the `vllm-chat` message shape.
|
||||
fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
|
||||
pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
|
||||
match message {
|
||||
ChatMessage::System { content, .. } => {
|
||||
Ok(VllmChatMessage::system(convert_content(content)?))
|
||||
@@ -316,7 +323,7 @@ fn convert_assistant_tool_calls(
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
|
||||
pub(crate) fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
|
||||
tools
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
@@ -498,7 +505,7 @@ mod tests {
|
||||
)
|
||||
.expect("request is valid");
|
||||
|
||||
assert!(!prepared.include_reasoning);
|
||||
assert!(!prepared.options.include_reasoning);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -868,8 +875,8 @@ mod tests {
|
||||
)
|
||||
.expect("request is valid");
|
||||
|
||||
assert!(prepared.requested_logprobs);
|
||||
assert!(prepared.include_prompt_logprobs);
|
||||
assert!(prepared.options.requested_logprobs);
|
||||
assert!(prepared.options.include_prompt_logprobs);
|
||||
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
|
||||
assert_eq!(
|
||||
prepared.chat_request.sampling_params.prompt_logprobs,
|
||||
@@ -895,7 +902,7 @@ mod tests {
|
||||
|
||||
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
|
||||
assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
|
||||
assert!(!prepared.include_prompt_logprobs);
|
||||
assert!(!prepared.options.include_prompt_logprobs);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -9,9 +9,9 @@ use vllm_chat::ReasoningEffort;
|
||||
|
||||
use crate::routes::openai::utils::structured_outputs::ResponseFormat;
|
||||
use crate::routes::openai::utils::types::{
|
||||
ChatLogProbs, ChatMessage, MessageContent, Normalizable, StreamOptions, StringOrArray, Tool,
|
||||
ToolCall, ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage,
|
||||
default_true, validate_stop, validate_top_p_value,
|
||||
ChatLogProbs, ChatMessage, Normalizable, StreamOptions, StringOrArray, Tool, ToolCall,
|
||||
ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage,
|
||||
default_true, validate_messages, validate_stop, validate_top_p_value,
|
||||
};
|
||||
|
||||
/// vLLM-compatible request type for the Chat Completions API.
|
||||
@@ -430,32 +430,6 @@ fn default_model() -> String {
|
||||
UNKNOWN_MODEL_ID.to_string()
|
||||
}
|
||||
|
||||
/// Validates messages array is not empty and has valid content
|
||||
fn validate_messages(messages: &[ChatMessage]) -> Result<(), validator::ValidationError> {
|
||||
if messages.is_empty() {
|
||||
return Err(validator::ValidationError::new("messages cannot be empty"));
|
||||
}
|
||||
|
||||
for msg in messages {
|
||||
if let ChatMessage::User { content, .. } = msg {
|
||||
match content {
|
||||
MessageContent::Text(text) if text.is_empty() => {
|
||||
return Err(validator::ValidationError::new(
|
||||
"message content cannot be empty",
|
||||
));
|
||||
}
|
||||
MessageContent::Parts(parts) if parts.is_empty() => {
|
||||
return Err(validator::ValidationError::new(
|
||||
"message content parts cannot be empty",
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Schema-level validation for cross-field dependencies
|
||||
fn validate_chat_cross_parameters(
|
||||
req: &ChatCompletionRequest,
|
||||
|
||||
@@ -18,13 +18,13 @@ use tracing::{debug, error, info, trace};
|
||||
use tracing_futures::Instrument as _;
|
||||
use vllm_text::{DecodedTextEvent, FinishReason, TextOutputStream, TextOutputStreamExt as _};
|
||||
|
||||
use self::convert::{ResponseOptions, prepare_completion_request};
|
||||
use super::utils::logprobs::{
|
||||
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
|
||||
text_len,
|
||||
};
|
||||
use super::utils::types::Usage;
|
||||
use crate::error::{ApiError, bail_server_error, server_error};
|
||||
use crate::routes::openai::completions::convert::prepare_completion_request;
|
||||
use crate::routes::openai::completions::types::{
|
||||
CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
|
||||
CompletionStreamChoice, CompletionStreamResponse,
|
||||
@@ -42,7 +42,6 @@ pub async fn completions(
|
||||
ValidatedJson(body): ValidatedJson<CompletionRequest>,
|
||||
) -> Response {
|
||||
let stream = body.stream;
|
||||
let logprobs = body.logprobs;
|
||||
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
|
||||
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
|
||||
|
||||
@@ -57,9 +56,7 @@ pub async fn completions(
|
||||
);
|
||||
|
||||
let created = unix_timestamp();
|
||||
let include_prompt_logprobs = prepared.text_request.sampling_params.prompt_logprobs.is_some();
|
||||
let log_request = state.enable_log_requests;
|
||||
|
||||
let text_stream = match state
|
||||
.chat
|
||||
.text()
|
||||
@@ -84,11 +81,7 @@ pub async fn completions(
|
||||
prepared.response_model,
|
||||
created,
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.echo,
|
||||
logprobs,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
prepared.options,
|
||||
);
|
||||
let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
|
||||
|
||||
@@ -99,11 +92,8 @@ pub async fn completions(
|
||||
prepared.request_id,
|
||||
prepared.response_model,
|
||||
created,
|
||||
prepared.echo,
|
||||
logprobs,
|
||||
include_prompt_logprobs,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
log_request,
|
||||
prepared.options,
|
||||
)
|
||||
.instrument(request_span.clone())
|
||||
.await
|
||||
@@ -112,18 +102,6 @@ pub async fn completions(
|
||||
Err(error) => return error.into_response(),
|
||||
};
|
||||
|
||||
if log_request {
|
||||
let usage = response.usage.as_ref();
|
||||
info!(
|
||||
parent: &request_span,
|
||||
model = %response.model,
|
||||
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
|
||||
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
|
||||
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
|
||||
"completion finished"
|
||||
);
|
||||
}
|
||||
|
||||
Json(response).into_response()
|
||||
}
|
||||
}
|
||||
@@ -133,11 +111,16 @@ async fn collect_completion(
|
||||
request_id: String,
|
||||
response_model: String,
|
||||
created: u64,
|
||||
echo: Option<String>,
|
||||
requested_logprobs: Option<u32>,
|
||||
include_prompt_logprobs: bool,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
log_request: bool,
|
||||
ResponseOptions {
|
||||
// Ignored: non-streaming responses always include usage.
|
||||
include_usage: _,
|
||||
echo,
|
||||
requested_logprobs,
|
||||
include_prompt_logprobs,
|
||||
return_token_ids,
|
||||
return_tokens_as_token_ids,
|
||||
}: ResponseOptions,
|
||||
) -> Result<CompletionResponse, ApiError> {
|
||||
let collected = stream
|
||||
.collect_output()
|
||||
@@ -175,6 +158,21 @@ async fn collect_completion(
|
||||
None => collected.text,
|
||||
Some(prompt) => format!("{prompt}{}", collected.text),
|
||||
};
|
||||
let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();
|
||||
let usage = Usage::from_counts(
|
||||
collected.prompt_token_ids.len() as u32,
|
||||
collected.token_ids.len() as u32,
|
||||
);
|
||||
|
||||
if log_request {
|
||||
info!(
|
||||
model = %response_model,
|
||||
prompt_tokens = usage.prompt_tokens,
|
||||
output_tokens = usage.completion_tokens.unwrap_or(0),
|
||||
%finish_reason,
|
||||
"completion finished"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(CompletionResponse {
|
||||
id: request_id,
|
||||
@@ -185,16 +183,13 @@ async fn collect_completion(
|
||||
index: 0,
|
||||
text,
|
||||
logprobs,
|
||||
finish_reason: Some(completion_finish_reason_to_openai(finish_reason)?.into()),
|
||||
finish_reason: Some(finish_reason),
|
||||
stop_reason,
|
||||
prompt_logprobs,
|
||||
token_ids: return_token_ids.then(|| collected.token_ids.clone()),
|
||||
prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()),
|
||||
}],
|
||||
usage: Some(Usage::from_counts(
|
||||
collected.prompt_token_ids.len() as u32,
|
||||
collected.token_ids.len() as u32,
|
||||
)),
|
||||
usage: Some(usage),
|
||||
system_fingerprint: None,
|
||||
kv_transfer_params: collected.kv_transfer_params,
|
||||
})
|
||||
@@ -208,11 +203,15 @@ async fn completion_chunk_stream(
|
||||
response_model: String,
|
||||
created: u64,
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
echo: Option<String>,
|
||||
requested_logprobs: Option<u32>,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
ResponseOptions {
|
||||
include_usage,
|
||||
echo,
|
||||
requested_logprobs,
|
||||
// Ignored: streaming prompt logprobs are rejected for Python parity.
|
||||
include_prompt_logprobs: _,
|
||||
return_token_ids,
|
||||
return_tokens_as_token_ids,
|
||||
}: ResponseOptions,
|
||||
mut y: TryYielder<CompletionSseChunk, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
pin_mut!(stream);
|
||||
@@ -432,7 +431,7 @@ mod tests {
|
||||
FinishReason, Finished,
|
||||
};
|
||||
|
||||
use super::{CompletionSseChunk, completion_chunk_stream, final_chunk};
|
||||
use super::{CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk};
|
||||
|
||||
#[test]
|
||||
fn final_chunk_maps_stop_finish_reason() {
|
||||
@@ -527,11 +526,10 @@ mod tests {
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(1),
|
||||
false,
|
||||
false,
|
||||
ResponseOptions {
|
||||
requested_logprobs: Some(1),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
@@ -10,18 +10,28 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
|
||||
/// Lowered completion request plus the public response metadata carried by
|
||||
/// every SSE chunk.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PreparedRequest {
|
||||
pub(super) struct PreparedRequest {
|
||||
/// Stable OpenAI-style request ID, reused as the external text request ID.
|
||||
pub request_id: String,
|
||||
/// Public model ID echoed back to the client.
|
||||
pub response_model: String,
|
||||
/// Whether the caller asked for the final streamed usage chunk.
|
||||
pub include_usage: bool,
|
||||
/// Public response rendering options for route-layer helpers.
|
||||
pub options: ResponseOptions,
|
||||
/// Lowered text request for the shared `vllm-text` facade.
|
||||
pub text_request: TextRequest,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub(super) struct ResponseOptions {
|
||||
/// Whether the caller asked for the final streamed usage chunk.
|
||||
pub include_usage: bool,
|
||||
/// Original text prompt that should be echoed back northbound when
|
||||
/// `echo=true`.
|
||||
pub echo: Option<String>,
|
||||
/// Whether the caller requested output logprobs on completion choices.
|
||||
pub requested_logprobs: Option<u32>,
|
||||
/// Whether the caller requested choice-level prompt logprobs.
|
||||
pub include_prompt_logprobs: bool,
|
||||
/// Whether to include token IDs alongside generated text.
|
||||
pub return_token_ids: bool,
|
||||
/// Whether to format logprob tokens as `token_id:{id}`.
|
||||
@@ -33,7 +43,7 @@ pub struct PreparedRequest {
|
||||
///
|
||||
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
|
||||
/// the base `model` field in responses when no LoRA adapter is selected.
|
||||
pub(crate) fn prepare_completion_request(
|
||||
pub(super) fn prepare_completion_request(
|
||||
request: CompletionRequest,
|
||||
lora_resolution: &LoraModelResolution,
|
||||
ctx: ResolvedRequestContext,
|
||||
@@ -64,6 +74,7 @@ pub(crate) fn prepare_completion_request(
|
||||
let include_usage = (request.stream_options.as_ref())
|
||||
.and_then(|options| options.include_usage)
|
||||
.unwrap_or(false);
|
||||
let include_prompt_logprobs = prompt_logprobs.is_some();
|
||||
let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();
|
||||
|
||||
let structured_outputs =
|
||||
@@ -116,11 +127,15 @@ pub(crate) fn prepare_completion_request(
|
||||
Ok(PreparedRequest {
|
||||
request_id,
|
||||
response_model,
|
||||
include_usage,
|
||||
options: ResponseOptions {
|
||||
include_usage,
|
||||
echo,
|
||||
requested_logprobs: request.logprobs,
|
||||
include_prompt_logprobs,
|
||||
return_token_ids: request.return_token_ids.unwrap_or(false),
|
||||
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
|
||||
},
|
||||
text_request,
|
||||
echo,
|
||||
return_token_ids: request.return_token_ids.unwrap_or(false),
|
||||
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -206,7 +221,7 @@ mod tests {
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
assert!(prepared.include_usage);
|
||||
assert!(prepared.options.include_usage);
|
||||
assert_eq!(
|
||||
prepared.text_request.prompt,
|
||||
Prompt::TokenIds(vec![11, 22, 33])
|
||||
@@ -250,7 +265,7 @@ mod tests {
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
assert_eq!(prepared.echo, Some("hello".to_string()));
|
||||
assert_eq!(prepared.options.echo, Some("hello".to_string()));
|
||||
assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
|
||||
}
|
||||
|
||||
|
||||
@@ -412,6 +412,41 @@ pub struct ListModelsResponse {
|
||||
pub data: Vec<ModelObject>,
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Shared validation helpers
|
||||
// ============================================================================
|
||||
|
||||
/// Validates a messages array is non-empty and has valid user-message content.
|
||||
///
|
||||
/// Used by both `POST /v1/chat/completions` and `POST /tokenize` (chat form)
|
||||
/// so validation behaviour stays in lockstep.
|
||||
pub(crate) fn validate_messages(
|
||||
messages: &[ChatMessage],
|
||||
) -> Result<(), validator::ValidationError> {
|
||||
if messages.is_empty() {
|
||||
return Err(validator::ValidationError::new("messages cannot be empty"));
|
||||
}
|
||||
|
||||
for msg in messages {
|
||||
if let ChatMessage::User { content, .. } = msg {
|
||||
match content {
|
||||
MessageContent::Text(text) if text.is_empty() => {
|
||||
return Err(validator::ValidationError::new(
|
||||
"message content cannot be empty",
|
||||
));
|
||||
}
|
||||
MessageContent::Parts(parts) if parts.is_empty() => {
|
||||
return Err(validator::ValidationError::new(
|
||||
"message content parts cannot be empty",
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Normalizable trait
|
||||
// ============================================================================
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::Json;
|
||||
use axum::extract::rejection::QueryRejection;
|
||||
use axum::extract::{Query, State};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use vllm_engine_core_client::protocol::utility::PauseMode;
|
||||
|
||||
use crate::error::ApiError;
|
||||
use crate::state::AppState;
|
||||
use crate::utils::utility_call_error;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct PauseParams {
|
||||
#[serde(default)]
|
||||
mode: PauseMode,
|
||||
#[serde(default = "default_clear_cache")]
|
||||
clear_cache: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(crate) struct StatusResponse {
|
||||
status: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(crate) struct IsPausedResponse {
|
||||
is_paused: bool,
|
||||
}
|
||||
|
||||
const fn default_clear_cache() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn invalid_query(error: QueryRejection) -> ApiError {
|
||||
ApiError::invalid_request(error.body_text(), Some("mode"))
|
||||
}
|
||||
|
||||
// TODO: the Python frontend also accepts the deprecated
|
||||
// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is
|
||||
// intentionally omitted here in favor of the `mode` parameter.
|
||||
|
||||
/// Pause the scheduler so generation can be halted (e.g. for weight updates).
|
||||
pub async fn pause(
|
||||
State(state): State<Arc<AppState>>,
|
||||
params: Result<Query<PauseParams>, QueryRejection>,
|
||||
) -> Result<Json<StatusResponse>, ApiError> {
|
||||
let Query(params) = params.map_err(invalid_query)?;
|
||||
|
||||
state
|
||||
.engine_core_client()
|
||||
.pause_scheduler(params.mode, params.clear_cache)
|
||||
.await
|
||||
.map_err(|error| utility_call_error("pause", error))?;
|
||||
|
||||
Ok(Json(StatusResponse { status: "paused" }))
|
||||
}
|
||||
|
||||
/// Resume the scheduler after a pause.
|
||||
pub async fn resume(State(state): State<Arc<AppState>>) -> Result<Json<StatusResponse>, ApiError> {
|
||||
state
|
||||
.engine_core_client()
|
||||
.resume_scheduler()
|
||||
.await
|
||||
.map_err(|error| utility_call_error("resume", error))?;
|
||||
|
||||
Ok(Json(StatusResponse { status: "resumed" }))
|
||||
}
|
||||
|
||||
/// Return whether the scheduler is currently paused.
|
||||
pub async fn is_paused(
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> Result<Json<IsPausedResponse>, ApiError> {
|
||||
let is_paused = state
|
||||
.engine_core_client()
|
||||
.is_scheduler_paused()
|
||||
.await
|
||||
.map_err(|error| utility_call_error("is_paused", error))?;
|
||||
|
||||
Ok(Json(IsPausedResponse { is_paused }))
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::Json;
|
||||
use axum::extract::rejection::QueryRejection;
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use vllm_engine_core_client::protocol::utility::PauseMode;
|
||||
|
||||
use crate::error::ApiError;
|
||||
use crate::state::AppState;
|
||||
@@ -18,8 +20,8 @@ pub(crate) struct IsSleepingResponse {
|
||||
pub(crate) struct SleepParams {
|
||||
#[serde(default = "default_sleep_level")]
|
||||
level: u32,
|
||||
#[serde(default = "default_sleep_mode")]
|
||||
mode: String,
|
||||
#[serde(default)]
|
||||
mode: PauseMode,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
@@ -32,18 +34,20 @@ const fn default_sleep_level() -> u32 {
|
||||
1
|
||||
}
|
||||
|
||||
fn default_sleep_mode() -> String {
|
||||
"abort".to_string()
|
||||
fn invalid_query(error: QueryRejection) -> ApiError {
|
||||
ApiError::invalid_request(error.body_text(), Some("mode"))
|
||||
}
|
||||
|
||||
/// Put the engine to sleep.
|
||||
pub async fn sleep(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Query(params): Query<SleepParams>,
|
||||
params: Result<Query<SleepParams>, QueryRejection>,
|
||||
) -> Result<StatusCode, ApiError> {
|
||||
let Query(params) = params.map_err(invalid_query)?;
|
||||
|
||||
state
|
||||
.engine_core_client()
|
||||
.sleep(params.level, ¶ms.mode)
|
||||
.sleep(params.level, params.mode)
|
||||
.await
|
||||
.map_err(|error| utility_call_error("sleep", error))?;
|
||||
|
||||
|
||||
@@ -14,15 +14,14 @@ use std::{fmt, fs};
|
||||
use axum::body::{Body, to_bytes};
|
||||
use axum::http::{Request, StatusCode};
|
||||
use bytes::Bytes;
|
||||
use futures::StreamExt as _;
|
||||
use rmpv::Value;
|
||||
use serde_json::json;
|
||||
use serial_test::serial;
|
||||
use tower::{Service as _, ServiceExt as _};
|
||||
use vllm_chat::{
|
||||
ChatBackend, ChatContent, ChatContentPart, ChatEvent, ChatLlm, ChatMessage, ChatRenderer,
|
||||
ChatRequest, ChatRole, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor,
|
||||
DynChatRenderer, NewChatOutputProcessorOptions, SamplingParams,
|
||||
ChatBackend, ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRenderer, ChatRequest,
|
||||
ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
|
||||
NewChatOutputProcessorOptions,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::logprobs::{
|
||||
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
|
||||
@@ -44,8 +43,6 @@ use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
|
||||
use crate::lora::LoraModelResolution;
|
||||
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
|
||||
use crate::state::AppState;
|
||||
|
||||
fn request_output(
|
||||
@@ -407,6 +404,9 @@ struct FakeChatBackend {
|
||||
multimodal_model_info: Option<vllm_chat::multimodal::MultimodalModelInfo>,
|
||||
}
|
||||
|
||||
/// Synthetic BOS id used when `add_special_tokens` is true in tests.
|
||||
const FAKE_BOS_TOKEN_ID: u32 = 1;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FakeChatTokenizer;
|
||||
|
||||
@@ -414,9 +414,12 @@ impl Tokenizer for FakeChatTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
let mut token_ids = Vec::new();
|
||||
if add_special_tokens {
|
||||
token_ids.push(FAKE_BOS_TOKEN_ID);
|
||||
}
|
||||
let mut rest = text;
|
||||
while !rest.is_empty() {
|
||||
if let Some(stripped) = rest.strip_prefix("<image>") {
|
||||
@@ -460,6 +463,7 @@ impl Tokenizer for FakeChatTokenizer {
|
||||
|
||||
fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
match id {
|
||||
FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
|
||||
999 => Some("<image>".to_string()),
|
||||
151655 => Some("<|image_pad|>".to_string()),
|
||||
0xF001 => Some("<think>".to_string()),
|
||||
@@ -468,6 +472,7 @@ impl Tokenizer for FakeChatTokenizer {
|
||||
0xF004 => Some("<|END_THINKING|>".to_string()),
|
||||
0xF005 => Some("◁think▷".to_string()),
|
||||
0xF006 => Some("◁/think▷".to_string()),
|
||||
id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -775,6 +780,19 @@ async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
|
||||
(app, engine_task)
|
||||
}
|
||||
|
||||
async fn test_app_with_api_keys(api_keys: Vec<String>) -> (axum::Router, MockEngineTask) {
|
||||
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
|
||||
b"engine-openai-api-key",
|
||||
default_stream_output_specs(),
|
||||
Arc::new(FakeChatBackend::new()),
|
||||
)
|
||||
.await;
|
||||
let app = build_router(Arc::new(
|
||||
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_keys(api_keys),
|
||||
));
|
||||
(app, engine_task)
|
||||
}
|
||||
|
||||
async fn test_health_app_with_engine_script<F>(
|
||||
script: F,
|
||||
) -> (axum::Router, Arc<AppState>, MockEngineTask)
|
||||
@@ -1073,6 +1091,101 @@ async fn request_id_header_echoes_incoming_header_when_enabled() {
|
||||
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn api_key_auth_rejects_missing_token_on_guarded_route() {
|
||||
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/v1/models")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
let json: serde_json::Value = serde_json::from_slice(&body).expect("json body");
|
||||
assert_eq!(json, json!({ "error": "Unauthorized" }));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn api_key_auth_rejects_wrong_token_on_guarded_route() {
|
||||
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/v1/models")
|
||||
.header("authorization", "Bearer wrong")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn api_key_auth_accepts_matching_bearer_token_on_guarded_route() {
|
||||
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/v1/models")
|
||||
.header("authorization", "Bearer secret")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn api_key_auth_allows_options_without_token() {
|
||||
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("OPTIONS")
|
||||
.uri("/v1/models")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_ne!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn api_key_auth_allows_unguarded_route_without_token() {
|
||||
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/health")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn version_returns_engine_vllm_version() {
|
||||
@@ -2334,7 +2447,8 @@ async fn non_stream_completions_echo_prepends_prompt_text() {
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hello",
|
||||
"echo": true,
|
||||
"stream": false
|
||||
"stream": false,
|
||||
"add_special_tokens": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -2537,7 +2651,8 @@ async fn non_stream_completions_include_prompt_logprobs() {
|
||||
"prompt": "hello",
|
||||
"stream": false,
|
||||
"echo": true,
|
||||
"logprobs": 1
|
||||
"logprobs": 1,
|
||||
"add_special_tokens": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -3318,7 +3433,8 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {
|
||||
"prompt": "hello",
|
||||
"echo": true,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true}
|
||||
"stream_options": {"include_usage": true},
|
||||
"add_special_tokens": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -3359,92 +3475,6 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {
|
||||
assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn chat_harness_streams_text_events() {
|
||||
let (chat, engine_task) = test_chat_with_engine_handle().await;
|
||||
let mut stream = chat
|
||||
.chat(ChatRequest {
|
||||
messages: vec![ChatMessage::text(ChatRole::User, "hello")],
|
||||
sampling_params: SamplingParams {
|
||||
max_tokens: Some(8),
|
||||
..Default::default()
|
||||
},
|
||||
request_id: "chat-harness".to_string(),
|
||||
..ChatRequest::for_test()
|
||||
})
|
||||
.await
|
||||
.expect("submit chat request");
|
||||
|
||||
let mut saw_text = false;
|
||||
let mut saw_done = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
match event.expect("chat event") {
|
||||
ChatEvent::BlockDelta { .. } => saw_text = true,
|
||||
ChatEvent::Done { .. } => {
|
||||
saw_done = true;
|
||||
break;
|
||||
}
|
||||
ChatEvent::Start { .. }
|
||||
| ChatEvent::LogprobsDelta { .. }
|
||||
| ChatEvent::BlockStart { .. }
|
||||
| ChatEvent::BlockEnd { .. }
|
||||
| ChatEvent::ToolCallStart { .. }
|
||||
| ChatEvent::ToolCallArgumentsDelta { .. }
|
||||
| ChatEvent::ToolCallEnd { .. } => {}
|
||||
}
|
||||
}
|
||||
engine_task.await.expect("mock engine task");
|
||||
|
||||
assert!(saw_text);
|
||||
assert!(saw_done);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn prepared_openai_request_streams_text_events() {
|
||||
let (chat, engine_task) = test_chat_with_engine_handle().await;
|
||||
let prepared = prepare_chat_request(
|
||||
serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"stream": true,
|
||||
"messages": [{"role": "user", "content": "hello"}]
|
||||
}))
|
||||
.expect("decode request"),
|
||||
&LoraModelResolution {
|
||||
model_names: vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
|
||||
lora_request: None,
|
||||
},
|
||||
crate::utils::ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("prepare request");
|
||||
|
||||
let mut stream = chat.chat(prepared.chat_request).await.expect("submit chat request");
|
||||
|
||||
let mut saw_text = false;
|
||||
let mut saw_done = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
match event.expect("chat event") {
|
||||
ChatEvent::BlockDelta { .. } => saw_text = true,
|
||||
ChatEvent::Done { .. } => {
|
||||
saw_done = true;
|
||||
break;
|
||||
}
|
||||
ChatEvent::Start { .. }
|
||||
| ChatEvent::LogprobsDelta { .. }
|
||||
| ChatEvent::BlockStart { .. }
|
||||
| ChatEvent::BlockEnd { .. }
|
||||
| ChatEvent::ToolCallStart { .. }
|
||||
| ChatEvent::ToolCallArgumentsDelta { .. }
|
||||
| ChatEvent::ToolCallEnd { .. } => {}
|
||||
}
|
||||
}
|
||||
engine_task.await.expect("mock engine task");
|
||||
|
||||
assert!(saw_text);
|
||||
assert!(saw_done);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {
|
||||
@@ -4228,6 +4258,161 @@ async fn is_sleeping_route_returns_json_payload() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn pause_route_uses_python_compatible_default_query_values() {
|
||||
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
|
||||
boxed_test_future(async move {
|
||||
let utility = recv_engine_message(dealer).await;
|
||||
assert_eq!(utility[0].as_ref(), &[0x03]);
|
||||
|
||||
let payload = decode_value(&utility[1]).expect("decode utility payload");
|
||||
let array = payload.as_array().expect("utility payload array");
|
||||
let call_id = array[1].as_u64().expect("call id");
|
||||
|
||||
assert_eq!(array[2], Value::from("pause_scheduler"));
|
||||
assert_eq!(
|
||||
array[3],
|
||||
Value::Array(vec![Value::from("abort"), Value::from(true)])
|
||||
);
|
||||
|
||||
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
|
||||
})
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/pause")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
|
||||
json!({ "status": "paused" })
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn pause_route_rejects_invalid_mode() {
|
||||
let (app, engine_task) =
|
||||
test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/pause?mode=banana")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(json["error"]["param"], "mode");
|
||||
engine_task.abort_and_join().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn resume_route_sends_no_args() {
|
||||
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
|
||||
boxed_test_future(async move {
|
||||
let utility = recv_engine_message(dealer).await;
|
||||
assert_eq!(utility[0].as_ref(), &[0x03]);
|
||||
|
||||
let payload = decode_value(&utility[1]).expect("decode utility payload");
|
||||
let array = payload.as_array().expect("utility payload array");
|
||||
let call_id = array[1].as_u64().expect("call id");
|
||||
|
||||
assert_eq!(array[2], Value::from("resume_scheduler"));
|
||||
assert_eq!(array[3], Value::Array(Vec::new()));
|
||||
|
||||
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
|
||||
})
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/resume")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
|
||||
json!({ "status": "resumed" })
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn is_paused_route_returns_json_payload() {
|
||||
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
|
||||
boxed_test_future(async move {
|
||||
let utility = recv_engine_message(dealer).await;
|
||||
assert_eq!(utility[0].as_ref(), &[0x03]);
|
||||
|
||||
let payload = decode_value(&utility[1]).expect("decode utility payload");
|
||||
let array = payload.as_array().expect("utility payload array");
|
||||
let call_id = array[1].as_u64().expect("call id");
|
||||
|
||||
assert_eq!(array[2], Value::from("is_scheduler_paused"));
|
||||
assert_eq!(array[3], Value::Array(Vec::new()));
|
||||
|
||||
send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
|
||||
})
|
||||
})
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/is_paused")
|
||||
.body(Body::empty())
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
|
||||
json!({ "is_paused": true })
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
|
||||
@@ -4244,6 +4429,9 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
|
||||
("GET", "/is_sleeping"),
|
||||
("POST", "/sleep"),
|
||||
("POST", "/wake_up"),
|
||||
("GET", "/is_paused"),
|
||||
("POST", "/pause"),
|
||||
("POST", "/resume"),
|
||||
("POST", "/collective_rpc"),
|
||||
("POST", "/reset_prefix_cache"),
|
||||
("POST", "/reset_mm_cache"),
|
||||
@@ -4598,3 +4786,392 @@ async fn completions_empty_stop_string_returns_validation_error() {
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
async fn post_json(
|
||||
app: &mut axum::Router,
|
||||
uri: &str,
|
||||
body: serde_json::Value,
|
||||
) -> (StatusCode, serde_json::Value) {
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri(uri)
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
let status = response.status();
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
let json: serde_json::Value = serde_json::from_slice(&bytes)
|
||||
.unwrap_or_else(|_| json!({ "raw": String::from_utf8_lossy(&bytes) }));
|
||||
(status, json)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_completion_round_trips_through_detokenize() {
|
||||
let mut app = test_app().await;
|
||||
let prompt = "Hello world";
|
||||
|
||||
let (_, tokenize_json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": prompt,
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let tokens = tokenize_json["tokens"]
|
||||
.as_array()
|
||||
.expect("tokens array")
|
||||
.iter()
|
||||
.map(|v| v.as_u64().expect("token id") as u32)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let (status, detokenize_json) = post_json(
|
||||
&mut app,
|
||||
"/detokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"tokens": tokens,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(detokenize_json["prompt"], prompt);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_completion_add_special_tokens_changes_ids() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (_, with_special) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hi",
|
||||
"add_special_tokens": true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let (_, without_special) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hi",
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let with_ids: Vec<u32> = with_special["tokens"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v.as_u64().unwrap() as u32)
|
||||
.collect();
|
||||
let without_ids: Vec<u32> = without_special["tokens"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v.as_u64().unwrap() as u32)
|
||||
.collect();
|
||||
|
||||
assert_ne!(with_ids, without_ids);
|
||||
assert_eq!(with_ids.first().copied(), Some(FAKE_BOS_TOKEN_ID));
|
||||
assert_eq!(without_ids.first().copied(), Some(b'h' as u32));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_completion_return_token_strs_matches_tokens() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hi",
|
||||
"add_special_tokens": false,
|
||||
"return_token_strs": true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
let tokens = json["tokens"].as_array().expect("tokens");
|
||||
let token_strs = json["token_strs"].as_array().expect("token_strs");
|
||||
assert_eq!(tokens.len(), token_strs.len());
|
||||
assert_eq!(token_strs.len(), json["count"].as_u64().unwrap() as usize);
|
||||
assert!(!token_strs[0].as_str().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_completion_count_and_max_model_len() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hello",
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(
|
||||
json["count"].as_u64().unwrap() as usize,
|
||||
json["tokens"].as_array().unwrap().len()
|
||||
);
|
||||
assert!(json["max_model_len"].as_u64().unwrap() > 0);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_includes_generation_prompt_in_token_count() {
|
||||
let mut app = test_app().await;
|
||||
let messages = json!([{"role": "user", "content": "hi"}]);
|
||||
|
||||
let (_, with_prompt) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": messages,
|
||||
"add_generation_prompt": true,
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let (_, without_prompt) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": messages,
|
||||
"add_generation_prompt": false,
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let with_len = with_prompt["tokens"].as_array().unwrap().len();
|
||||
let without_len = without_prompt["tokens"].as_array().unwrap().len();
|
||||
assert!(with_len > without_len);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_conflicting_generation_flags_returns_400() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"add_generation_prompt": true,
|
||||
"continue_final_message": true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_empty_messages_returns_400() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": [],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_empty_message_content_returns_400() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": [{"role": "user", "content": ""}],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_unknown_model_returns_404() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "does-not-exist",
|
||||
"prompt": "hello",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(json["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn detokenize_unknown_model_returns_404() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/detokenize",
|
||||
json!({
|
||||
"model": "does-not-exist",
|
||||
"tokens": [72, 101, 108, 108, 111],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(json["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn detokenize_empty_tokens_returns_empty_prompt() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/detokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"tokens": [],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(json["prompt"], "");
|
||||
}
|
||||
|
||||
/// Decode an explicit token sequence — pins `/detokenize` independently of
|
||||
/// `/tokenize` (the round-trip test alone would pass even if encode and decode
|
||||
/// were both wrong in mirrored ways).
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn detokenize_decodes_known_token_ids() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/detokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"tokens": [72, 101, 108, 108, 111],
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(json["prompt"], "Hello");
|
||||
}
|
||||
|
||||
/// `continue_final_message` without a trailing assistant message must 400.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_continue_without_assistant_returns_400() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let (status, json) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"add_generation_prompt": false,
|
||||
"continue_final_message": true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
}
|
||||
|
||||
/// `continue_final_message` must not append a new generation suffix vs `add_generation_prompt`.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn tokenize_chat_continue_final_vs_new_assistant_differs() {
|
||||
let mut app = test_app().await;
|
||||
let messages = json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "partial,"}
|
||||
]);
|
||||
|
||||
let (_, continue_final) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": messages,
|
||||
"add_generation_prompt": false,
|
||||
"continue_final_message": true,
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let (_, new_assistant) = post_json(
|
||||
&mut app,
|
||||
"/tokenize",
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"messages": messages,
|
||||
"add_generation_prompt": true,
|
||||
"continue_final_message": false,
|
||||
"add_special_tokens": false,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let continue_len = continue_final["tokens"].as_array().unwrap().len();
|
||||
let new_len = new_assistant["tokens"].as_array().unwrap().len();
|
||||
assert!(new_len > continue_len);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
//! `POST /tokenize` and `POST /detokenize` (root paths, matching Python).
|
||||
//!
|
||||
//! Encode/decode runs entirely in-process via [`DynTokenizer`]; the inference
|
||||
//! engine is not involved.
|
||||
|
||||
mod types;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::Json;
|
||||
use axum::extract::State;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use thiserror_ext::AsReport as _;
|
||||
|
||||
use crate::error::{ApiError, server_error};
|
||||
use crate::routes::openai::utils::validated_json::ValidatedJson;
|
||||
use crate::routes::tokenize::types::{
|
||||
DetokenizeRequest, DetokenizeResponse, TokenizeChatRequest, TokenizeCompletionRequest,
|
||||
TokenizeRequest, TokenizeResponse,
|
||||
};
|
||||
use crate::state::AppState;
|
||||
use crate::utils::resolve_base_request_id;
|
||||
|
||||
/// Match Python `tokenize-{base}` where base is `X-Request-Id` or a new UUID.
|
||||
fn tokenize_request_id(headers: &HeaderMap) -> String {
|
||||
let base = resolve_base_request_id(
|
||||
headers.get("X-Request-Id").and_then(|value| value.to_str().ok()),
|
||||
None,
|
||||
);
|
||||
format!("tokenize-{base}")
|
||||
}
|
||||
|
||||
/// Reject an unknown model name, matching the other handlers.
|
||||
fn check_model(state: &AppState, model: Option<&str>) -> Result<(), ApiError> {
|
||||
if let Some(model) = model
|
||||
&& !state.served_model_names().iter().any(|n| n == model)
|
||||
{
|
||||
return Err(ApiError::model_not_found(model.to_string()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build the `token_strs` vector when requested, via the tokenizer vocab.
|
||||
fn token_strs(tokenizer: &vllm_text::tokenizer::DynTokenizer, ids: &[u32]) -> Vec<String> {
|
||||
// Unknown IDs yield "" — intentional; matches Python's convert_ids_to_tokens behaviour.
|
||||
ids.iter().map(|&id| tokenizer.id_to_token(id).unwrap_or_default()).collect()
|
||||
}
|
||||
|
||||
pub async fn tokenize(
|
||||
State(state): State<Arc<AppState>>,
|
||||
headers: HeaderMap,
|
||||
ValidatedJson(body): ValidatedJson<TokenizeRequest>,
|
||||
) -> Response {
|
||||
let request_id = tokenize_request_id(&headers);
|
||||
let tokenizer = state.chat.text().tokenizer();
|
||||
let max_model_len = state.chat.engine_core_client().max_model_len();
|
||||
|
||||
let result = match body {
|
||||
// Completion form: encode the raw `prompt` string (no chat template).
|
||||
TokenizeRequest::Completion(req) => tokenize_completion(&state, &tokenizer, req),
|
||||
// Chat form: render `messages` through the template, then encode (see `tokenize_chat`).
|
||||
TokenizeRequest::Chat(req) => tokenize_chat(&state, &request_id, req).await,
|
||||
};
|
||||
|
||||
match result {
|
||||
Ok((tokens, want_strs)) => {
|
||||
let token_strs = want_strs.then(|| token_strs(&tokenizer, &tokens));
|
||||
Json(TokenizeResponse {
|
||||
count: tokens.len(),
|
||||
max_model_len,
|
||||
tokens,
|
||||
token_strs,
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
Err(error) => error.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
fn tokenize_completion(
|
||||
state: &AppState,
|
||||
tokenizer: &vllm_text::tokenizer::DynTokenizer,
|
||||
req: TokenizeCompletionRequest,
|
||||
) -> Result<(Vec<u32>, bool), ApiError> {
|
||||
check_model(state, req.model.as_deref())?;
|
||||
let tokens = tokenizer
|
||||
.encode(&req.prompt, req.add_special_tokens)
|
||||
.map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
|
||||
Ok((tokens, req.return_token_strs))
|
||||
}
|
||||
|
||||
/// HTTP adapter for the chat-shaped `/tokenize` body.
|
||||
///
|
||||
/// Not [`vllm_chat::ChatLlm::tokenize_chat`]: this checks the model name and maps
|
||||
/// errors to [`ApiError`]; the chat-crate method does render → finalize → encode.
|
||||
async fn tokenize_chat(
|
||||
state: &AppState,
|
||||
request_id: &str,
|
||||
req: TokenizeChatRequest,
|
||||
) -> Result<(Vec<u32>, bool), ApiError> {
|
||||
check_model(state, req.model.as_deref())?;
|
||||
let return_token_strs = req.return_token_strs;
|
||||
// `continue_final_message` / `add_generation_prompt` mutual exclusion is
|
||||
// enforced in `normalize_generation_prompt_mode` inside `into_chat_request`.
|
||||
let tokens = state
|
||||
.chat
|
||||
.tokenize_chat(req.into_chat_request(request_id.to_string())?)
|
||||
.await
|
||||
.map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
|
||||
Ok((tokens, return_token_strs))
|
||||
}
|
||||
|
||||
pub async fn detokenize(
|
||||
State(state): State<Arc<AppState>>,
|
||||
ValidatedJson(body): ValidatedJson<DetokenizeRequest>,
|
||||
) -> Response {
|
||||
if let Err(error) = check_model(&state, body.model.as_deref()) {
|
||||
return error.into_response();
|
||||
}
|
||||
let tokenizer = state.chat.text().tokenizer();
|
||||
match tokenizer.decode(&body.tokens, /* skip_special_tokens = */ false) {
|
||||
Ok(prompt) => Json(DetokenizeResponse { prompt }).into_response(),
|
||||
Err(e) => server_error!("detokenize failed: {}", e.to_report_string()).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::http::{HeaderMap, HeaderValue};
|
||||
|
||||
use super::tokenize_request_id;
|
||||
|
||||
#[test]
|
||||
fn tokenize_request_id_prefers_x_request_id_header() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("X-Request-Id", HeaderValue::from_static("client-req-1"));
|
||||
assert_eq!(tokenize_request_id(&headers), "tokenize-client-req-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokenize_request_id_generates_uuid_when_header_missing() {
|
||||
let headers = HeaderMap::new();
|
||||
let id = tokenize_request_id(&headers);
|
||||
assert!(id.starts_with("tokenize-"));
|
||||
assert_ne!(id, "tokenize-");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use itertools::Itertools as _;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use validator::{Validate, ValidationErrors};
|
||||
use vllm_chat::{ChatOptions, ChatRequest, ChatToolChoice, SamplingParams};
|
||||
use vllm_text::output::TextDecodeOptions;
|
||||
|
||||
use crate::error::ApiError;
|
||||
use crate::routes::openai::chat_completions::convert::{
|
||||
convert_message, convert_tools, normalize_generation_prompt_mode,
|
||||
};
|
||||
use crate::routes::openai::utils::types::{
|
||||
ChatMessage, Normalizable, Tool, default_true, validate_messages,
|
||||
};
|
||||
|
||||
/// `POST /tokenize` body. Untagged: a JSON object with `messages` parses as the
|
||||
/// chat variant; one with `prompt` parses as the completion variant.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum TokenizeRequest {
|
||||
Chat(TokenizeChatRequest),
|
||||
Completion(TokenizeCompletionRequest),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct TokenizeCompletionRequest {
|
||||
pub model: Option<String>,
|
||||
pub prompt: String,
|
||||
#[serde(default = "default_true")]
|
||||
pub add_special_tokens: bool,
|
||||
#[serde(default)]
|
||||
pub return_token_strs: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Validate)]
|
||||
pub struct TokenizeChatRequest {
|
||||
pub model: Option<String>,
|
||||
#[validate(custom(function = "validate_messages"))]
|
||||
pub messages: Vec<ChatMessage>,
|
||||
#[serde(default = "default_true")]
|
||||
pub add_generation_prompt: bool,
|
||||
#[serde(default)]
|
||||
pub continue_final_message: bool,
|
||||
#[serde(default)] // chat default is FALSE (template adds specials)
|
||||
pub add_special_tokens: bool,
|
||||
#[serde(default)]
|
||||
pub return_token_strs: bool,
|
||||
#[serde(default)]
|
||||
pub chat_template: Option<String>,
|
||||
#[serde(default)]
|
||||
pub chat_template_kwargs: Option<HashMap<String, Value>>,
|
||||
#[serde(default)]
|
||||
pub tools: Option<Vec<Tool>>,
|
||||
}
|
||||
|
||||
impl TokenizeChatRequest {
|
||||
/// Lower this tokenize body into a [`ChatRequest`] for template rendering.
|
||||
///
|
||||
/// Reuses [`convert_message`] and [`normalize_generation_prompt_mode`] from
|
||||
/// `chat_completions/convert` so message lowering and generation-prompt
|
||||
/// rules match chat completions. Only fields that affect rendering are set;
|
||||
/// `sampling_params`, `decode_options`, etc. stay at default because
|
||||
/// tokenize never generates.
|
||||
pub fn into_chat_request(self, request_id: String) -> Result<ChatRequest, ApiError> {
|
||||
let messages: Vec<_> = self.messages.into_iter().map(convert_message).try_collect()?;
|
||||
let generation_prompt_mode = normalize_generation_prompt_mode(
|
||||
Some(self.add_generation_prompt),
|
||||
self.continue_final_message,
|
||||
&messages,
|
||||
)?;
|
||||
|
||||
Ok(ChatRequest {
|
||||
request_id,
|
||||
messages,
|
||||
sampling_params: SamplingParams::default(),
|
||||
chat_options: ChatOptions {
|
||||
generation_prompt_mode,
|
||||
chat_template: self.chat_template,
|
||||
reasoning_effort: None,
|
||||
template_kwargs: self.chat_template_kwargs.unwrap_or_default(),
|
||||
},
|
||||
tools: convert_tools(self.tools)?,
|
||||
tool_choice: ChatToolChoice::Auto,
|
||||
decode_options: TextDecodeOptions::default(),
|
||||
intermediate: false,
|
||||
priority: 0,
|
||||
documents: None,
|
||||
cache_salt: None,
|
||||
add_special_tokens: self.add_special_tokens,
|
||||
data_parallel_rank: None,
|
||||
lora_request: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct DetokenizeRequest {
|
||||
pub model: Option<String>,
|
||||
pub tokens: Vec<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TokenizeResponse {
|
||||
pub count: usize,
|
||||
pub max_model_len: u32,
|
||||
pub tokens: Vec<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub token_strs: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct DetokenizeResponse {
|
||||
pub prompt: String,
|
||||
}
|
||||
|
||||
// ---- trait impls required by ValidatedJson ----
|
||||
impl Validate for TokenizeRequest {
|
||||
fn validate(&self) -> Result<(), ValidationErrors> {
|
||||
if let Self::Chat(req) = self {
|
||||
req.validate()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl Validate for DetokenizeRequest {
|
||||
fn validate(&self) -> Result<(), ValidationErrors> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
impl Normalizable for TokenizeRequest {} // default no-op normalize()
|
||||
impl Normalizable for DetokenizeRequest {}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
|
||||
use serde_json::json;
|
||||
use vllm_chat::ChatTool;
|
||||
|
||||
#[test]
|
||||
fn tokenize_request_converts_openai_tools() {
|
||||
// The untagged `TokenizeRequest` must resolve a messages+tools body to
|
||||
// the chat variant and accept standard OpenAI tool objects
|
||||
// (`{"type":"function",...}`), then convert them to `ChatTool`.
|
||||
let request: TokenizeRequest = serde_json::from_value(json!({
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}],
|
||||
}))
|
||||
.expect("OpenAI tool JSON deserializes to the chat variant");
|
||||
|
||||
let TokenizeRequest::Chat(req) = request else {
|
||||
panic!("messages+tools body should parse as the chat variant");
|
||||
};
|
||||
|
||||
let chat_request =
|
||||
req.into_chat_request("tokenize-test".to_string()).expect("request is valid");
|
||||
|
||||
assert_eq!(
|
||||
chat_request.tools,
|
||||
vec![ChatTool {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Get weather".to_string()),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
}),
|
||||
strict: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn into_chat_request_rejects_conflicting_generation_flags() {
|
||||
let req = TokenizeChatRequest {
|
||||
model: None,
|
||||
messages: vec![ChatMessage::User {
|
||||
content: MessageContent::Text("hi".to_string()),
|
||||
name: None,
|
||||
}],
|
||||
add_generation_prompt: true,
|
||||
continue_final_message: true,
|
||||
add_special_tokens: false,
|
||||
return_token_strs: false,
|
||||
chat_template: None,
|
||||
chat_template_kwargs: None,
|
||||
tools: None,
|
||||
};
|
||||
|
||||
let error = req
|
||||
.into_chat_request("tokenize-test".to_string())
|
||||
.expect_err("conflicting flags");
|
||||
assert_eq!(
|
||||
error.to_error_response().error.message,
|
||||
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::time::{Duration, Instant, sleep_until};
|
||||
use tracing::warn;
|
||||
use vllm_chat::ChatLlm;
|
||||
@@ -9,11 +10,16 @@ use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::lora::LoraRequest;
|
||||
|
||||
use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
|
||||
|
||||
use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
|
||||
|
||||
const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
|
||||
|
||||
pub(crate) type ApiKeyHash = [u8; 32];
|
||||
|
||||
pub(crate) fn hash_api_key(api_key: &str) -> ApiKeyHash {
|
||||
Sha256::digest(api_key.as_bytes()).into()
|
||||
}
|
||||
|
||||
/// Shared router state for the minimal single-model OpenAI server.
|
||||
pub struct AppState {
|
||||
/// All public model IDs served by this frontend. The first entry is the
|
||||
@@ -27,6 +33,8 @@ pub struct AppState {
|
||||
pub enable_request_id_headers: bool,
|
||||
/// Runtime server information returned by `/server_info`, when available.
|
||||
server_info: Option<ServerInfoSnapshot>,
|
||||
/// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
|
||||
api_key_hashes: Vec<ApiKeyHash>,
|
||||
/// Number of in-flight inference requests currently owned by this frontend.
|
||||
server_load: AtomicU64,
|
||||
/// Dynamic LoRA adapter registry.
|
||||
@@ -53,6 +61,7 @@ impl AppState {
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
server_info: None,
|
||||
api_key_hashes: Vec::new(),
|
||||
server_load: AtomicU64::new(0),
|
||||
lora_manager: LoraManager::new(),
|
||||
}
|
||||
@@ -84,6 +93,24 @@ impl AppState {
|
||||
self.server_info.as_ref().map(|server_info| server_info.response(config_format))
|
||||
}
|
||||
|
||||
/// Configure API keys accepted by guarded HTTP routes.
|
||||
pub fn with_api_keys(mut self, api_keys: Vec<String>) -> Self {
|
||||
self.api_key_hashes = api_keys
|
||||
.into_iter()
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(|key| hash_api_key(&key))
|
||||
.collect();
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn has_api_keys(&self) -> bool {
|
||||
!self.api_key_hashes.is_empty()
|
||||
}
|
||||
|
||||
pub(crate) fn api_key_hashes(&self) -> &[ApiKeyHash] {
|
||||
&self.api_key_hashes
|
||||
}
|
||||
|
||||
/// The primary model name echoed back in API responses (the first served
|
||||
/// name).
|
||||
pub fn primary_model_name(&self) -> &str {
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
[package]
|
||||
name = "vllm-tool-parser-py"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "_rust_tool_parser"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[dependencies]
|
||||
pyo3.workspace = true
|
||||
pythonize = { workspace = true, features = ["serde_json"] }
|
||||
serde_json.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
vllm-tool-parser.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,366 @@
|
||||
//! Thin PyO3 bindings for `vllm_tool_parser`.
|
||||
//!
|
||||
//! This crate exposes the Rust tool parser trait and data shapes to Python
|
||||
//! while keeping parser state, grammar, and schema-aware argument conversion in
|
||||
//! Rust. Python callers should use this module as a typed bridge and keep any
|
||||
//! vLLM protocol adaptation outside the binding.
|
||||
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyModule};
|
||||
use pythonize::{depythonize, pythonize};
|
||||
use serde_json::Value;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
macro_rules! tool_parser_factory {
|
||||
($($parser:ident),+ $(,)?) => {
|
||||
fn create_tool_parser(
|
||||
name: &str,
|
||||
tools: &[Tool],
|
||||
) -> PyResult<Box<dyn ToolParser>> {
|
||||
match name {
|
||||
$(
|
||||
stringify!($parser) => {
|
||||
<vllm_tool_parser::$parser as ToolParser>::create(tools)
|
||||
}
|
||||
)+
|
||||
_ => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"unsupported tool parser `{name}`"
|
||||
)));
|
||||
}
|
||||
}
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Export a tool parser to Python by registering it here.
|
||||
tool_parser_factory! {
|
||||
DeepSeekV4ToolParser, // for testing on Python side
|
||||
}
|
||||
|
||||
#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
struct PyTool(Tool);
|
||||
|
||||
#[pymethods]
|
||||
impl PyTool {
|
||||
#[new]
|
||||
#[pyo3(signature = (name, description, parameters, strict=None))]
|
||||
fn new(
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
parameters: &Bound<'_, PyAny>,
|
||||
strict: Option<bool>,
|
||||
) -> PyResult<Self> {
|
||||
let parameters = depythonize::<Value>(parameters).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert tool parameters from Python to JSON: {error}"
|
||||
))
|
||||
})?;
|
||||
Ok(Self(Tool {
|
||||
name,
|
||||
description,
|
||||
parameters,
|
||||
strict,
|
||||
}))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.0.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> Option<&str> {
|
||||
self.0.description.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert tool parameters from JSON to Python: {error}"
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn strict(&self) -> Option<bool> {
|
||||
self.0.strict
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(
|
||||
name = "ToolCallDelta",
|
||||
module = "vllm._rust_tool_parser",
|
||||
skip_from_py_object
|
||||
)]
|
||||
#[derive(Clone)]
|
||||
struct PyToolCallDelta(ToolCallDelta);
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolCallDelta {
|
||||
#[new]
|
||||
#[pyo3(signature = (tool_index, name, arguments))]
|
||||
fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
|
||||
Self(ToolCallDelta {
|
||||
tool_index,
|
||||
name,
|
||||
arguments,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn tool_index(&self) -> usize {
|
||||
self.0.tool_index
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> Option<&str> {
|
||||
self.0.name.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn arguments(&self) -> &str {
|
||||
&self.0.arguments
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(
|
||||
name = "ToolParserOutput",
|
||||
module = "vllm._rust_tool_parser",
|
||||
skip_from_py_object
|
||||
)]
|
||||
#[derive(Clone)]
|
||||
struct PyToolParserOutput(ToolParserOutput);
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolParserOutput {
|
||||
#[new]
|
||||
#[pyo3(signature = (normal_text="", calls=None))]
|
||||
fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
|
||||
let calls =
|
||||
calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
|
||||
Self(ToolParserOutput {
|
||||
normal_text: normal_text.to_owned(),
|
||||
calls,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn normal_text(&self) -> &str {
|
||||
&self.0.normal_text
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn calls(&self) -> Vec<PyToolCallDelta> {
|
||||
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
|
||||
}
|
||||
|
||||
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
|
||||
self.0.append(other.0.clone());
|
||||
}
|
||||
|
||||
fn coalesce_calls(&self) -> Self {
|
||||
Self(self.0.clone().coalesce_calls())
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
|
||||
struct PyToolParser(Box<dyn ToolParser>);
|
||||
|
||||
impl PyToolParser {
|
||||
fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
|
||||
self.0
|
||||
.parse_into(chunk, &mut output.0)
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolParser {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
|
||||
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
|
||||
create_tool_parser(parser_name, &tools).map(Self)
|
||||
}
|
||||
|
||||
fn parse_into(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
mut output: PyRefMut<'_, PyToolParserOutput>,
|
||||
) -> PyResult<()> {
|
||||
self.parse_into_output(chunk, &mut output)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> PyResult<PyToolParserOutput> {
|
||||
self.0
|
||||
.finish()
|
||||
.map(PyToolParserOutput)
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
|
||||
fn preserve_special_tokens(&self) -> bool {
|
||||
self.0.preserve_special_tokens()
|
||||
}
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyTool>()?;
|
||||
m.add_class::<PyToolCallDelta>()?;
|
||||
m.add_class::<PyToolParserOutput>()?;
|
||||
m.add_class::<PyToolParser>()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
|
||||
Python::initialize();
|
||||
Python::attach(f)
|
||||
}
|
||||
|
||||
fn tool_schema() -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "integer"},
|
||||
"shipping": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"zip": {"type": "integer"}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_call() -> String {
|
||||
r#"<|DSML|tool_calls>
|
||||
<|DSML|invoke name="create_order">
|
||||
<|DSML|parameter name="user_id" string="false">42</|DSML|parameter>
|
||||
<|DSML|parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</|DSML|parameter>
|
||||
</|DSML|invoke>
|
||||
</|DSML|tool_calls>"#
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
|
||||
let parameters = pythonize(py, &tool_schema()).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert test schema from JSON to Python: {error}"
|
||||
))
|
||||
})?;
|
||||
Py::new(
|
||||
py,
|
||||
PyTool::new(
|
||||
"create_order".to_owned(),
|
||||
Some("Create an order".to_owned()),
|
||||
¶meters,
|
||||
None,
|
||||
)?,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_round_trips_typed_fields() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let borrowed = tool.borrow(py);
|
||||
assert_eq!(borrowed.name(), "create_order");
|
||||
assert_eq!(borrowed.description(), Some("Create an order"));
|
||||
assert_eq!(borrowed.strict(), None);
|
||||
|
||||
let parameters = borrowed.parameters(py)?;
|
||||
let parameters = depythonize::<Value>(parameters.bind(py))?;
|
||||
assert_eq!(parameters, tool_schema());
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_append_and_coalesce_calls() {
|
||||
with_python(|py| {
|
||||
let first = Py::new(
|
||||
py,
|
||||
PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
|
||||
)?;
|
||||
let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
|
||||
let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
|
||||
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
|
||||
output.append(other.borrow(py));
|
||||
|
||||
let coalesced = output.coalesce_calls();
|
||||
assert_eq!(coalesced.normal_text(), "text");
|
||||
let calls = coalesced.calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].tool_index(), 0);
|
||||
assert_eq!(calls[0].name(), Some("create_order"));
|
||||
assert_eq!(calls[0].arguments(), "{\"a\":1}");
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_parse_finish_and_preserve_special_tokens() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
|
||||
assert!(parser.preserve_special_tokens());
|
||||
|
||||
let mut output = PyToolParserOutput::new(py, "", None);
|
||||
parser.parse_into_output(&build_call(), &mut output)?;
|
||||
let finish = Py::new(py, parser.finish()?)?;
|
||||
output.append(finish.borrow(py));
|
||||
let output = output.coalesce_calls();
|
||||
|
||||
assert_eq!(output.normal_text(), "");
|
||||
let calls = output.calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name(), Some("create_order"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
|
||||
json!({
|
||||
"user_id": 42,
|
||||
"shipping": {
|
||||
"city": "Singapore",
|
||||
"zip": 18956
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
assert_eq!(parser.reset(), "");
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_errors_for_unknown_name() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let error = match PyToolParser::new(py, "missing", vec![tool]) {
|
||||
Ok(_) => panic!("missing parser name unexpectedly succeeded"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let message = format!("{error}");
|
||||
assert!(message.contains("unsupported tool parser `missing`"));
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use winnow::ascii::{digit1, multispace0 as ws0};
|
||||
use winnow::combinator::{alt, eof, repeat, seq};
|
||||
use winnow::prelude::*;
|
||||
@@ -33,6 +35,7 @@ enum KimiK2Event {
|
||||
ToolCallsStart,
|
||||
ToolCallStart,
|
||||
ToolCallHeader {
|
||||
tool_call_id: String,
|
||||
function_name: String,
|
||||
function_index: usize,
|
||||
},
|
||||
@@ -60,6 +63,7 @@ pub struct KimiK2ToolParser {
|
||||
buffer: String,
|
||||
mode: KimiK2Mode,
|
||||
active_tool_index: Option<usize>,
|
||||
call_ids: BTreeMap<usize, String>,
|
||||
}
|
||||
|
||||
impl KimiK2ToolParser {
|
||||
@@ -69,6 +73,7 @@ impl KimiK2ToolParser {
|
||||
buffer: String::new(),
|
||||
mode: KimiK2Mode::Text,
|
||||
active_tool_index: None,
|
||||
call_ids: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -81,6 +86,7 @@ impl KimiK2ToolParser {
|
||||
KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,
|
||||
KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,
|
||||
KimiK2Event::ToolCallHeader {
|
||||
tool_call_id,
|
||||
function_name,
|
||||
function_index,
|
||||
} => {
|
||||
@@ -89,6 +95,7 @@ impl KimiK2ToolParser {
|
||||
self.mode = KimiK2Mode::Arguments {
|
||||
json_scan: JsonObjectScanState::default(),
|
||||
};
|
||||
self.call_ids.insert(tool_index, tool_call_id);
|
||||
output.calls.push(ToolCallDelta {
|
||||
tool_index,
|
||||
name: Some(function_name),
|
||||
@@ -123,6 +130,7 @@ impl KimiK2ToolParser {
|
||||
fn reset(&mut self) -> String {
|
||||
self.mode = KimiK2Mode::Text;
|
||||
self.active_tool_index = None;
|
||||
self.call_ids.clear();
|
||||
std::mem::take(&mut self.buffer)
|
||||
}
|
||||
}
|
||||
@@ -139,6 +147,10 @@ impl ToolParser for KimiK2ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
self.call_ids.get(&tool_index).map(String::as_str)
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
self.buffer.push_str(chunk);
|
||||
|
||||
@@ -232,16 +244,18 @@ fn tool_call_end_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event>
|
||||
|
||||
/// Parse a Kimi K2 tool-call header before the argument marker.
|
||||
fn tool_call_header_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
|
||||
let (header, _) = (
|
||||
let (raw_header, _) = (
|
||||
take_until(1.., TOOL_CALL_ARGUMENT_START),
|
||||
literal(TOOL_CALL_ARGUMENT_START),
|
||||
)
|
||||
.parse_next(input)?;
|
||||
|
||||
let mut header_input = header;
|
||||
let tool_call_id = raw_header.trim().to_string();
|
||||
let mut header_input = raw_header;
|
||||
let (header, _, _) = (tool_header, ws0, eof).parse_next(&mut header_input)?;
|
||||
|
||||
Ok(KimiK2Event::ToolCallHeader {
|
||||
tool_call_id,
|
||||
function_name: header.function_name,
|
||||
function_index: header.function_index,
|
||||
})
|
||||
@@ -502,6 +516,25 @@ mod tests {
|
||||
.assert_debug_eq(&output);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kimi_k2_preserves_model_generated_tool_call_ids() {
|
||||
let mut parser = KimiK2ToolParser::new(&test_tools());
|
||||
let input = build_tool_section(&[
|
||||
build_tool_call("get_weather", 0, r#"{"location":"Shanghai"}"#),
|
||||
build_tool_call("add", 1, r#"{"x":1,"y":2}"#),
|
||||
]);
|
||||
|
||||
for chunk in split_by_chars(&input, 7) {
|
||||
parser.parse_chunk(chunk).unwrap();
|
||||
}
|
||||
|
||||
// IDs are available after parsing but before finish(), which calls reset().
|
||||
assert_eq!(parser.tool_call_id(0), Some("functions.get_weather:0"));
|
||||
assert_eq!(parser.tool_call_id(1), Some("functions.add:1"));
|
||||
parser.finish().unwrap();
|
||||
assert_eq!(parser.tool_call_id(0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kimi_k2_accepts_non_functions_header_prefix() {
|
||||
let mut parser = KimiK2ToolParser::new(&test_tools());
|
||||
@@ -509,9 +542,10 @@ mod tests {
|
||||
"{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"
|
||||
);
|
||||
|
||||
let output = parser.parse_complete(&input).unwrap();
|
||||
let output = parser.parse_chunk(&input).unwrap().coalesce_calls();
|
||||
|
||||
assert_eq!(output.calls[0].tool_index, 42);
|
||||
assert_eq!(parser.tool_call_id(42), Some("api.tools.search:42"));
|
||||
assert_eq!(output.calls[0].name.as_deref(), Some("search"));
|
||||
assert_eq!(output.calls[0].arguments, "{}");
|
||||
}
|
||||
|
||||
@@ -121,6 +121,12 @@ pub trait ToolParser: Send {
|
||||
false
|
||||
}
|
||||
|
||||
/// Return the parser-provided ID for a tool call by index, if the model
|
||||
/// emitted one.
|
||||
fn tool_call_id(&self, _tool_index: usize) -> Option<&str> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Feed one decoded text delta into the parser, appending committed output
|
||||
/// into `output`.
|
||||
///
|
||||
|
||||
@@ -18,8 +18,6 @@ import torch
|
||||
from packaging.version import Version, parse
|
||||
from setuptools import Extension, setup
|
||||
from setuptools.command.build_ext import build_ext
|
||||
from setuptools_rust import Binding, RustExtension
|
||||
from setuptools_rust.build import build_rust
|
||||
from setuptools_scm import get_version
|
||||
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
|
||||
|
||||
@@ -35,11 +33,12 @@ def load_module_from_path(module_name, path):
|
||||
ROOT_DIR = Path(__file__).parent
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
|
||||
|
||||
# cannot import envs directly because it depends on vllm,
|
||||
# which is not installed yet
|
||||
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
|
||||
rust_build = load_module_from_path(
|
||||
"rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")
|
||||
)
|
||||
|
||||
VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
|
||||
USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED
|
||||
@@ -49,11 +48,6 @@ USE_PRECOMPILED_RUST_FRONTEND = (
|
||||
)
|
||||
|
||||
|
||||
def should_require_rust_frontend() -> bool:
|
||||
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
|
||||
return value.lower() not in ("", "0", "false", "no")
|
||||
|
||||
|
||||
if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
|
||||
logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
|
||||
VLLM_TARGET_DEVICE = "cpu"
|
||||
@@ -420,24 +414,6 @@ class precompiled_build_ext(build_ext):
|
||||
return
|
||||
|
||||
|
||||
class precompiled_build_rust(build_rust):
|
||||
"""Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
|
||||
|
||||
def run(self) -> None:
|
||||
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
|
||||
logger.info(
|
||||
"Skipping local Rust build: using precompiled %s",
|
||||
PRECOMPILED_RUST_FRONTEND_PATH,
|
||||
)
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"Precompiled wheel did not provide %s; falling back to local Rust build.",
|
||||
PRECOMPILED_RUST_FRONTEND_PATH,
|
||||
)
|
||||
super().run()
|
||||
|
||||
|
||||
class precompiled_wheel_utils:
|
||||
"""Extracts libraries and other files from an existing wheel."""
|
||||
|
||||
@@ -731,9 +707,6 @@ class precompiled_wheel_utils:
|
||||
"vllm/_rocm_C.abi3.so",
|
||||
}
|
||||
)
|
||||
if extract_rust_frontend:
|
||||
exact_members.add("vllm/vllm-rs")
|
||||
|
||||
flash_attn_regex = re.compile(
|
||||
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
|
||||
)
|
||||
@@ -756,6 +729,12 @@ class precompiled_wheel_utils:
|
||||
if member.filename in exact_members:
|
||||
file_members.append(member)
|
||||
continue
|
||||
if (
|
||||
extract_rust_frontend
|
||||
and rust_build.is_precompiled_artifact_member(member.filename)
|
||||
):
|
||||
file_members.append(member)
|
||||
continue
|
||||
|
||||
if not extract_extensions:
|
||||
continue
|
||||
@@ -1109,6 +1088,12 @@ package_data = {
|
||||
}
|
||||
|
||||
|
||||
def add_vllm_package_data(filename: str) -> None:
|
||||
vllm_files = package_data.setdefault("vllm", [])
|
||||
if filename not in vllm_files:
|
||||
vllm_files.append(filename)
|
||||
|
||||
|
||||
# If using precompiled artifacts, extract and patch package_data in advance.
|
||||
if USE_PRECOMPILED_RUST_FRONTEND:
|
||||
wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()
|
||||
@@ -1121,12 +1106,10 @@ if USE_PRECOMPILED_RUST_FRONTEND:
|
||||
for pkg, files in patch.items():
|
||||
package_data.setdefault(pkg, []).extend(files)
|
||||
|
||||
# If the rust frontend binary is already present in the source tree (e.g.,
|
||||
# pre-built in a separate Docker build stage), ship it as-is.
|
||||
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
|
||||
vllm_files = package_data.setdefault("vllm", [])
|
||||
if "vllm-rs" not in vllm_files:
|
||||
vllm_files.append("vllm-rs")
|
||||
# Rust artifacts already present in the source tree (e.g., pre-built in a
|
||||
# separate Docker build stage) are shipped as-is.
|
||||
for rust_artifact in rust_build.find_precompiled_artifacts():
|
||||
add_vllm_package_data(rust_artifact.name)
|
||||
|
||||
if _no_device():
|
||||
ext_modules = []
|
||||
@@ -1139,23 +1122,14 @@ else:
|
||||
if USE_PRECOMPILED_EXTENSIONS
|
||||
else cmake_build_ext,
|
||||
}
|
||||
if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():
|
||||
cmdclass["build_rust"] = precompiled_build_rust
|
||||
if USE_PRECOMPILED_RUST_FRONTEND or rust_build.find_precompiled_artifacts():
|
||||
cmdclass["build_rust"] = rust_build.precompiled_build_rust
|
||||
|
||||
# Rust frontend binary, built via setuptools-rust and installed into the
|
||||
# package directory alongside the Python modules.
|
||||
# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
|
||||
# requires extra work on using precompiled binaries.
|
||||
rust_extensions = [
|
||||
RustExtension(
|
||||
target="vllm.vllm-rs",
|
||||
path="rust/src/cmd/Cargo.toml",
|
||||
args=["--bin", "vllm-rs"],
|
||||
features=["native-tls-vendored"],
|
||||
binding=Binding.Exec,
|
||||
optional=not should_require_rust_frontend(),
|
||||
),
|
||||
]
|
||||
# Rust artifacts, built via setuptools-rust and installed into the package
|
||||
# directory alongside the Python modules.
|
||||
rust_extensions = rust_build.rust_extensions(
|
||||
optional=not rust_build.should_require_rust_frontend()
|
||||
)
|
||||
|
||||
setup(
|
||||
# static metadata should rather go in pyproject.toml
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer, PreTrainedTokenizerBase
|
||||
|
||||
from vllm.benchmarks.datasets import BFCLDataset, get_samples
|
||||
|
||||
|
||||
def _patch_hf_api(side_effect):
|
||||
"""Return a patch context that swaps `hf_api()` to a stub whose
|
||||
`.hf_hub_download` attribute uses `side_effect`."""
|
||||
fake_api = MagicMock()
|
||||
fake_api.hf_hub_download.side_effect = side_effect
|
||||
return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def hf_tokenizer() -> PreTrainedTokenizerBase:
|
||||
return AutoTokenizer.from_pretrained("gpt2")
|
||||
|
||||
|
||||
_FAKE_ROWS = {
|
||||
"simple": [
|
||||
{
|
||||
"id": "simple_0",
|
||||
"question": [
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is 2+2?",
|
||||
}
|
||||
]
|
||||
],
|
||||
"function": [
|
||||
{
|
||||
"name": "add",
|
||||
"description": "Add two numbers.",
|
||||
"parameters": {
|
||||
"type": "dict",
|
||||
"properties": {
|
||||
"a": {"type": "integer", "description": "first"},
|
||||
"b": {"type": "float", "description": "second"},
|
||||
},
|
||||
"required": ["a", "b"],
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
"live_simple": [
|
||||
{
|
||||
"id": "live_simple_0",
|
||||
"question": [[{"role": "user", "content": "Tell me the weather."}]],
|
||||
"function": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get weather.",
|
||||
"parameters": {
|
||||
"type": "dict",
|
||||
"properties": {
|
||||
"city": {"type": "any", "description": "city"},
|
||||
"coords": {"type": "tuple", "description": "coords"},
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _write_fake_files(tmp_path: Path) -> dict[str, Path]:
|
||||
"""Write fake BFCL JSONL files mimicking the HF repo layout."""
|
||||
paths = {}
|
||||
for category, rows in _FAKE_ROWS.items():
|
||||
p = tmp_path / f"BFCL_v3_{category}.json"
|
||||
with p.open("w") as f:
|
||||
for row in rows:
|
||||
f.write(json.dumps(row) + "\n")
|
||||
paths[category] = p
|
||||
return paths
|
||||
|
||||
|
||||
def _args_for_bfcl(categories: list[str] | None) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
dataset_name="hf",
|
||||
dataset_path="gorilla-llm/Berkeley-Function-Calling-Leaderboard",
|
||||
hf_name=None,
|
||||
hf_subset=None,
|
||||
hf_split=None,
|
||||
hf_output_len=64,
|
||||
disable_shuffle=True,
|
||||
num_prompts=2,
|
||||
no_oversample=False,
|
||||
no_stream=True,
|
||||
seed=0,
|
||||
request_id_prefix="",
|
||||
trust_remote_code=False,
|
||||
skip_chat_template=False,
|
||||
enable_multimodal_chat=False,
|
||||
backend="openai-chat",
|
||||
bfcl_categories=categories,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_dataset_translates_schema_and_attaches_tools(
|
||||
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
|
||||
) -> None:
|
||||
"""BFCLDataset should translate schemas to OpenAI tool format, set
|
||||
`messages` directly on SampleRequest, and attach tools/tool_choice via
|
||||
request_overrides."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
|
||||
return str(paths[category])
|
||||
|
||||
args = _args_for_bfcl(categories=["simple", "live_simple"])
|
||||
|
||||
with _patch_hf_api(fake_download):
|
||||
samples = get_samples(args, hf_tokenizer)
|
||||
|
||||
assert len(samples) == 2
|
||||
for s in samples:
|
||||
assert s.chat_messages is not None
|
||||
assert isinstance(s.chat_messages, list)
|
||||
assert s.chat_messages[0]["role"] == "user"
|
||||
assert s.request_overrides is not None
|
||||
assert "tools" in s.request_overrides
|
||||
assert s.request_overrides["tool_choice"] == "auto"
|
||||
# messages must NOT leak into request_overrides — it has its own
|
||||
# typed field on SampleRequest.
|
||||
assert "messages" not in s.request_overrides
|
||||
tools = s.request_overrides["tools"]
|
||||
assert len(tools) == 1
|
||||
tool = tools[0]
|
||||
assert tool["type"] == "function"
|
||||
# Translated schema: dict -> object, float -> number,
|
||||
# any -> string, tuple -> array.
|
||||
params = tool["function"]["parameters"]
|
||||
assert params["type"] == "object"
|
||||
for prop in params["properties"].values():
|
||||
assert prop["type"] in {"integer", "number", "string", "array"}
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_dataset_requires_openai_chat_backend(
|
||||
hf_tokenizer: PreTrainedTokenizerBase,
|
||||
) -> None:
|
||||
args = _args_for_bfcl(categories=["simple"])
|
||||
args.backend = "openai"
|
||||
|
||||
with pytest.raises(ValueError, match="openai-chat"):
|
||||
get_samples(args, hf_tokenizer)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_dataset_missing_category_raises_clear_error(
|
||||
hf_tokenizer: PreTrainedTokenizerBase,
|
||||
) -> None:
|
||||
"""A typo'd category should produce an actionable ValueError, not an
|
||||
opaque huggingface_hub exception."""
|
||||
from huggingface_hub.errors import EntryNotFoundError
|
||||
|
||||
args = _args_for_bfcl(categories=["simpl"]) # typo
|
||||
|
||||
def raise_missing(_repo, filename, **_kwargs):
|
||||
raise EntryNotFoundError(f"404 Not Found: {filename}")
|
||||
|
||||
with (
|
||||
_patch_hf_api(raise_missing),
|
||||
pytest.raises(ValueError, match=r"BFCL category 'simpl' not found"),
|
||||
):
|
||||
get_samples(args, hf_tokenizer)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_chat_backend_uses_messages_field_when_set() -> None:
|
||||
"""When RequestFuncInput.chat_messages is set, the chat backend must use
|
||||
it verbatim and skip default content construction from `prompt`."""
|
||||
import asyncio
|
||||
|
||||
from vllm.benchmarks.lib.endpoint_request_func import (
|
||||
RequestFuncInput,
|
||||
async_request_openai_chat_completions,
|
||||
)
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
class _FakeResp:
|
||||
status = 500
|
||||
reason = "stop-after-capture"
|
||||
content = None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, headers): # noqa: A002
|
||||
captured["url"] = url
|
||||
captured["payload"] = json
|
||||
return _FakeResp()
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "you are helpful"},
|
||||
{"role": "user", "content": "call add(3, 4)"},
|
||||
]
|
||||
req = RequestFuncInput(
|
||||
prompt="IGNORED",
|
||||
api_url="http://localhost:0/v1/chat/completions",
|
||||
prompt_len=10,
|
||||
output_len=16,
|
||||
model="test-model",
|
||||
chat_messages=messages,
|
||||
extra_body={"tools": [{"type": "function", "function": {"name": "add"}}]},
|
||||
)
|
||||
|
||||
asyncio.run(
|
||||
async_request_openai_chat_completions(
|
||||
request_func_input=req, session=_FakeSession()
|
||||
)
|
||||
)
|
||||
|
||||
payload = captured["payload"]
|
||||
assert payload["messages"] is messages, (
|
||||
"chat backend must forward RequestFuncInput.chat_messages verbatim "
|
||||
"instead of constructing a default user message from `prompt`"
|
||||
)
|
||||
# extra_body still merges in as before (shallow, per-request wins).
|
||||
assert payload["tools"][0]["function"]["name"] == "add"
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None:
|
||||
"""prompt_len must reflect tokens from both messages *and* tool schemas,
|
||||
so percentile buckets and input-distribution summaries aren't biased
|
||||
low for tool-heavy traffic."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
|
||||
return str(paths[category])
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
class _FakeTokenizer:
|
||||
def apply_chat_template(
|
||||
self, messages, tools=None, tokenize=False, add_generation_prompt=True
|
||||
):
|
||||
captured["tools"] = tools
|
||||
base = " ".join(m.get("content", "") for m in messages)
|
||||
tool_text = json.dumps(tools) if tools else ""
|
||||
return base + " " + tool_text
|
||||
|
||||
def __call__(self, text):
|
||||
# 1 "token" per whitespace-separated word.
|
||||
return type("Enc", (), {"input_ids": text.split()})()
|
||||
|
||||
fake = _FakeTokenizer()
|
||||
args = _args_for_bfcl(categories=["simple"])
|
||||
args.num_prompts = 1
|
||||
|
||||
with _patch_hf_api(fake_download):
|
||||
samples = get_samples(args, fake)
|
||||
|
||||
assert len(samples) == 1
|
||||
assert captured["tools"] is not None, (
|
||||
"apply_chat_template must be called with tools= so the schema "
|
||||
"contributes to the prompt-length estimate"
|
||||
)
|
||||
assert len(captured["tools"]) == 1
|
||||
assert captured["tools"][0]["function"]["name"] == "add"
|
||||
|
||||
# Sanity: prompt_len exceeds a messages-only estimate. The fake row's
|
||||
# user message is "What is 2+2?" (3 whitespace-separated tokens).
|
||||
assert samples[0].prompt_len > 3
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Older tokenizers don't accept tools=; fallback must still produce a
|
||||
non-zero prompt_len without crashing."""
|
||||
paths = _write_fake_files(tmp_path)
|
||||
|
||||
def fake_download(_repo, filename, **_kwargs):
|
||||
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
|
||||
return str(paths[category])
|
||||
|
||||
class _LegacyTokenizer:
|
||||
def apply_chat_template(self, messages, **kwargs):
|
||||
if "tools" in kwargs:
|
||||
raise TypeError("unexpected keyword argument 'tools'")
|
||||
return " ".join(m.get("content", "") for m in messages)
|
||||
|
||||
def __call__(self, text):
|
||||
return type("Enc", (), {"input_ids": text.split()})()
|
||||
|
||||
args = _args_for_bfcl(categories=["simple"])
|
||||
args.num_prompts = 1
|
||||
|
||||
with _patch_hf_api(fake_download):
|
||||
samples = get_samples(args, _LegacyTokenizer())
|
||||
|
||||
assert len(samples) == 1
|
||||
assert samples[0].prompt_len > 0
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_bfcl_schema_translation_is_recursive() -> None:
|
||||
"""_translate_schema must recurse into nested properties."""
|
||||
input_schema = {
|
||||
"type": "dict",
|
||||
"properties": {
|
||||
"nested": {
|
||||
"type": "dict",
|
||||
"properties": {
|
||||
"value": {"type": "float"},
|
||||
"tags": {"type": "tuple", "items": {"type": "any"}},
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
out = BFCLDataset._translate_schema(input_schema)
|
||||
assert out["type"] == "object"
|
||||
assert out["properties"]["nested"]["type"] == "object"
|
||||
assert out["properties"]["nested"]["properties"]["value"]["type"] == "number"
|
||||
assert out["properties"]["nested"]["properties"]["tags"]["type"] == "array"
|
||||
nested_props = out["properties"]["nested"]["properties"]
|
||||
assert nested_props["tags"]["items"]["type"] == "string"
|
||||
@@ -189,6 +189,118 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
|
||||
]
|
||||
|
||||
|
||||
class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
|
||||
"""Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
|
||||
|
||||
Four ``rms_norm`` sites that together hit every pattern registered by
|
||||
``RocmAiterAllReduceFusionPass`` for the per-group FP8 quant path:
|
||||
|
||||
* ``norm[0]``: ``all_reduce -> rms_norm -> group_fp8_quant`` (no residual)
|
||||
-> ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern``
|
||||
* ``norm[1]``: ``all_reduce -> fused_add_rms_norm -> group_fp8_quant``
|
||||
(single ``rms`` consumer)
|
||||
-> ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern``
|
||||
* ``norm[2..3]``: ``all_reduce -> fused_add_rms_norm
|
||||
-> (group_fp8_quant + rocm_unquantized_gemm)`` (two ``rms`` consumers,
|
||||
modeling the DSv3.2 indexer fan-out)
|
||||
-> ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern``
|
||||
|
||||
The chain feeds the next AllReduce by dequantizing the FP8 output (FP8
|
||||
cast back to bf16 multiplied by the per-group scale), which is enough to
|
||||
keep the matmul chain bf16 without depending on a real FP8 block-scaled
|
||||
GEMM kernel.
|
||||
"""
|
||||
|
||||
quant_group_size = 128
|
||||
indexer_out_dim = 8
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size=128,
|
||||
token_num=16,
|
||||
eps=1e-6,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
use_triton_quant: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.eps = eps
|
||||
self.use_triton_quant = use_triton_quant
|
||||
assert hidden_size % self.quant_group_size == 0, (
|
||||
f"hidden_size ({hidden_size}) must be a multiple of "
|
||||
f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant"
|
||||
)
|
||||
self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
|
||||
self.w = [torch.rand(hidden_size, hidden_size, dtype=dtype) for _ in range(3)]
|
||||
self.indexer_w = [
|
||||
torch.rand(self.indexer_out_dim, hidden_size, dtype=dtype) for _ in range(2)
|
||||
]
|
||||
|
||||
def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.use_triton_quant:
|
||||
return torch.ops.vllm.triton_per_token_group_quant_fp8(
|
||||
rms, self.quant_group_size
|
||||
)
|
||||
return torch.ops.vllm.rocm_aiter_group_fp8_quant.default(
|
||||
rms, self.quant_group_size
|
||||
)
|
||||
|
||||
def _dequantize_to_bf16(
|
||||
self, q: torch.Tensor, s: torch.Tensor, ref: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
# Broadcast the per-group scale across each group of `quant_group_size`
|
||||
# so we can chain the FP8 output back into a bf16 matmul. This avoids
|
||||
# depending on a real FP8 block-scaled GEMM kernel in the test.
|
||||
s_full = s.repeat_interleave(self.quant_group_size, dim=-1).to(ref.dtype)
|
||||
return q.to(ref.dtype) * s_full
|
||||
|
||||
def forward(self, hidden_states):
|
||||
z = torch.relu(hidden_states)
|
||||
x = resid = tensor_model_parallel_all_reduce(z)
|
||||
rms = self.norm[0](x)
|
||||
q0, s0 = self._group_quant(rms)
|
||||
y = self._dequantize_to_bf16(q0, s0, rms)
|
||||
|
||||
z2 = torch.mm(y, self.w[0])
|
||||
x2 = tensor_model_parallel_all_reduce(z2)
|
||||
rms2, resid = self.norm[1](x2, resid)
|
||||
q1, s1 = self._group_quant(rms2)
|
||||
y2 = self._dequantize_to_bf16(q1, s1, rms2)
|
||||
|
||||
z3 = torch.mm(y2, self.w[1])
|
||||
x3 = tensor_model_parallel_all_reduce(z3)
|
||||
rms3, resid = self.norm[2](x3, resid)
|
||||
q2, s2 = self._group_quant(rms3)
|
||||
# Second consumer of ``rms3``: forces the with-indexer pattern.
|
||||
idx2 = torch.ops.vllm.rocm_unquantized_gemm(rms3, self.indexer_w[0], None)
|
||||
y3 = self._dequantize_to_bf16(q2, s2, rms3)
|
||||
|
||||
z4 = torch.mm(y3, self.w[2])
|
||||
x4 = tensor_model_parallel_all_reduce(z4)
|
||||
rms4, resid = self.norm[3](x4, resid)
|
||||
q3, s3 = self._group_quant(rms4)
|
||||
# Second consumer of ``rms4``: forces the with-indexer pattern.
|
||||
idx3 = torch.ops.vllm.rocm_unquantized_gemm(rms4, self.indexer_w[1], None)
|
||||
y4 = self._dequantize_to_bf16(q3, s3, rms4)
|
||||
return y4, idx2, idx3
|
||||
|
||||
def ops_in_model_before(self):
|
||||
return [
|
||||
torch.ops.vllm.all_reduce.default,
|
||||
(
|
||||
torch.ops.vllm.triton_per_token_group_quant_fp8.default
|
||||
if self.use_triton_quant
|
||||
else torch.ops.vllm.rocm_aiter_group_fp8_quant.default
|
||||
),
|
||||
]
|
||||
|
||||
def ops_in_model_after(self):
|
||||
return [
|
||||
rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op(),
|
||||
rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op(), # noqa: E501
|
||||
]
|
||||
|
||||
|
||||
class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
|
||||
def __init__(
|
||||
self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
|
||||
@@ -463,3 +575,167 @@ def all_reduce_fusion_pass_on_test_model(
|
||||
assert fused_nodes
|
||||
assert all(n.kwargs.get("weight_bias") == 1.0 for n in fused_nodes)
|
||||
del all_reduce_fusion_pass
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("use_triton_quant", [True, False])
|
||||
@pytest.mark.parametrize("batch_size", [8])
|
||||
@pytest.mark.parametrize("seq_len", [8])
|
||||
@pytest.mark.parametrize("hidden_size", [128])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only",
|
||||
)
|
||||
@pytest.mark.skipif(not IS_AITER_FOUND, reason="aiter is not found")
|
||||
def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
enable_rms_norm_custom_op: bool,
|
||||
use_triton_quant: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Sibling of ``test_all_reduce_fusion_pass_replace`` for the new
|
||||
ROCm AITER AR+RMS+per-group-FP8-quant fusion patterns.
|
||||
|
||||
Validates the three new ``VllmPatternReplacement`` patterns added to
|
||||
``RocmAiterAllReduceFusionPass``:
|
||||
|
||||
* ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern`` (no-residual)
|
||||
* ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual,
|
||||
single ``rms`` consumer)
|
||||
* ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with-
|
||||
residual, DSv3.2 indexer fan-out; parametrized over both
|
||||
``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant``
|
||||
producers).
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group():
|
||||
pytest.skip(
|
||||
"aiter build is missing 'fused_ar_rms_per_group_quant' (needs "
|
||||
"ROCm/aiter PR #2823); the new patterns aren't registered."
|
||||
)
|
||||
|
||||
num_processes = 2
|
||||
|
||||
def run_torch_spawn(fn, nprocs):
|
||||
torch.multiprocessing.spawn(
|
||||
fn,
|
||||
args=(
|
||||
num_processes,
|
||||
TestAiterAllReduceRMSNormGroupQuantFP8Model,
|
||||
batch_size,
|
||||
seq_len,
|
||||
hidden_size,
|
||||
dtype,
|
||||
enable_rms_norm_custom_op,
|
||||
use_triton_quant,
|
||||
monkeypatch,
|
||||
),
|
||||
nprocs=nprocs,
|
||||
)
|
||||
|
||||
run_torch_spawn(rocm_aiter_group_quant_fusion_pass_on_test_model, num_processes)
|
||||
|
||||
|
||||
def rocm_aiter_group_quant_fusion_pass_on_test_model(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
test_model_cls: torch.nn.Module,
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
enable_rms_norm_custom_op: bool,
|
||||
use_triton_quant: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
set_random_seed(0)
|
||||
|
||||
device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
"VLLM_ROCM_USE_AITER": "1",
|
||||
}
|
||||
)
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
init_distributed_environment()
|
||||
|
||||
custom_ops = []
|
||||
if enable_rms_norm_custom_op:
|
||||
custom_ops.append("+rms_norm")
|
||||
# ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip``
|
||||
# only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at
|
||||
# the call site). The patterns in this PR are robust to both Triton and
|
||||
# rocm_aiter forms; we always enable +quant_fp8 so the matcher's example
|
||||
# trace finds the same form the test model uses.
|
||||
custom_ops.append("+quant_fp8")
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE, custom_ops=custom_ops
|
||||
)
|
||||
)
|
||||
vllm_config.compilation_config.pass_config = PassConfig(
|
||||
fuse_allreduce_rms=True, eliminate_noops=True
|
||||
)
|
||||
vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
|
||||
vllm_config.parallel_config.rank = local_rank
|
||||
|
||||
model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8"
|
||||
vllm_config.model_config = ModelConfig(
|
||||
model=model_name, trust_remote_code=True, dtype=dtype, seed=42
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
all_reduce_fusion_pass = RocmAiterAllReduceFusionPass(vllm_config)
|
||||
noop_pass = NoOpEliminationPass(vllm_config)
|
||||
func_pass = FixFunctionalizationPass(vllm_config)
|
||||
cleanup_pass = PostCleanupPass(vllm_config)
|
||||
|
||||
backend = TestBackend(
|
||||
noop_pass, all_reduce_fusion_pass, func_pass, cleanup_pass
|
||||
)
|
||||
|
||||
token_num = batch_size * seq_len
|
||||
model = test_model_cls(
|
||||
hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant
|
||||
)
|
||||
|
||||
hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
|
||||
|
||||
compiled_model = torch.compile(model, backend=backend)
|
||||
compiled_model(hidden_states)
|
||||
|
||||
results_unfused = model(hidden_states)
|
||||
results_fused = compiled_model(hidden_states)
|
||||
# The fused per-group AR+RMS+QUANT op is bit-equivalent to the unfused
|
||||
# chain modulo the small AllReduce + RMSNorm reordering inside aiter.
|
||||
# Per-group FP8 quant introduces step noise <=1 per group; use the
|
||||
# same tolerance as the sibling FP8 static test.
|
||||
torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
|
||||
|
||||
# Four pattern firings: norm[0] (no-add quant), norm[1] (add quant,
|
||||
# single ``rms`` consumer), norm[2..3] (add quant + indexer fan-out).
|
||||
assert all_reduce_fusion_pass.matched_count == 4, (
|
||||
f"{all_reduce_fusion_pass.matched_count=}"
|
||||
)
|
||||
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
|
||||
backend.check_after_ops(model.ops_in_model_after())
|
||||
del all_reduce_fusion_pass
|
||||
|
||||
@@ -382,11 +382,12 @@ def test_rope_kvcache_fusion(
|
||||
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
|
||||
# Cannot compare fp8_* directly here, cast to model dtype instead
|
||||
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
|
||||
torch.testing.assert_close(
|
||||
kv_cache_unfused.view(dtype),
|
||||
kv_cache_fused.view(dtype),
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
kv_cache_unfused.to(dtype),
|
||||
kv_cache_fused.to(dtype),
|
||||
atol=1e-1,
|
||||
rtol=1e-1,
|
||||
)
|
||||
|
||||
|
||||
@@ -569,17 +570,19 @@ def test_rope_static_qquant_kvcache_fusion(
|
||||
else:
|
||||
ATOL, RTOL = (1e-2, 1e-2)
|
||||
|
||||
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
|
||||
torch.testing.assert_close(
|
||||
q_unfused.to(torch.float32),
|
||||
q_fused.to(torch.float32),
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
atol=1e-1,
|
||||
rtol=1e-1,
|
||||
)
|
||||
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
|
||||
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
|
||||
torch.testing.assert_close(
|
||||
kv_cache_unfused.view(dtype),
|
||||
kv_cache_fused.view(dtype),
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
kv_cache_unfused.to(dtype),
|
||||
kv_cache_fused.to(dtype),
|
||||
atol=1e-1,
|
||||
rtol=1e-1,
|
||||
)
|
||||
|
||||
+18
-17
@@ -854,8 +854,13 @@ class HfRunner:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
from tests.utils import wait_for_rocm_memory_to_settle
|
||||
|
||||
del self.model
|
||||
cleanup_dist_env_and_memory()
|
||||
# ROCm frees VRAM lazily; wait so a runner started right after this HF
|
||||
# model exits does not OOM on its startup memory guard.
|
||||
wait_for_rocm_memory_to_settle()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@@ -1248,25 +1253,13 @@ class VllmRunner:
|
||||
return self
|
||||
|
||||
def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None:
|
||||
from tests.utils import wait_for_gpu_memory_to_clear
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
return
|
||||
|
||||
num_gpus = torch.accelerator.device_count()
|
||||
if num_gpus == 0:
|
||||
return
|
||||
from tests.utils import wait_for_rocm_memory_to_settle
|
||||
|
||||
# V1 startup requires free_memory >= total * gpu_memory_utilization.
|
||||
# Wait for the complementary used-memory ratio so the next runner does
|
||||
# not fail the startup guard immediately after this runner exits. Bound
|
||||
# the wait so cleanup failures fail this test instead of hanging.
|
||||
wait_for_gpu_memory_to_clear(
|
||||
devices=list(range(num_gpus)),
|
||||
threshold_ratio=1.0 - gpu_memory_utilization,
|
||||
timeout_s=120,
|
||||
)
|
||||
# not fail the startup guard immediately after this runner exits. The
|
||||
# wait is bounded so cleanup failures fail this test instead of hanging.
|
||||
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
# Explicitly shutdown the engine core to release GPU resources
|
||||
@@ -1276,8 +1269,16 @@ class VllmRunner:
|
||||
gpu_memory_utilization = (
|
||||
self.llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
try:
|
||||
self.llm.llm_engine.engine_core.shutdown()
|
||||
# Give the engine core time to run its own graceful shutdown
|
||||
# (model_executor teardown + empty_cache + process-group destroy)
|
||||
# before the process manager SIGKILLs it at the default 5s. On ROCm
|
||||
# a hard kill leaves the whole allocation for the driver's slow async
|
||||
# VRAM reclamation, which starves the next test's startup.
|
||||
shutdown_timeout = 60.0 if current_platform.is_rocm() else None
|
||||
self.llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)
|
||||
except Exception:
|
||||
# Ignore shutdown errors as cleanup will still proceed
|
||||
pass
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user