forked from Karylab-cklius/vllm
Compare commits
51
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5b4f6d5284 | ||
|
|
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}" \
|
||||
|
||||
+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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
+98
-13
@@ -1,6 +1,7 @@
|
||||
// A CUDAPluggableAllocator based on cumem* APIs.
|
||||
// Important: allocation size, CUdeviceptr and CUmemGenericAllocationHandle*
|
||||
// need to be unsigned long long
|
||||
#include <atomic>
|
||||
#include <iostream>
|
||||
|
||||
#include "cumem_allocator_compat.h"
|
||||
@@ -9,6 +10,7 @@
|
||||
static const char* PYARGS_PARSE = "KKKK";
|
||||
#else
|
||||
#include <cstdlib>
|
||||
#include <cstdint>
|
||||
#include <cerrno>
|
||||
#include <climits>
|
||||
|
||||
@@ -46,6 +48,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
|
||||
|
||||
@@ -92,6 +117,59 @@ void ensure_context(unsigned long long device) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Cached fabric handle probe (CUDA 12.4+, NVIDIA only):
|
||||
|
||||
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040
|
||||
// Per-device cache: 0 = not probed, 1 = supported, 2 = not supported
|
||||
static constexpr int MAX_DEVICES = 32;
|
||||
static std::atomic<int> fabric_support[MAX_DEVICES] = {};
|
||||
|
||||
static bool probe_fabric_support(unsigned long long device) {
|
||||
if (device >= MAX_DEVICES) return false;
|
||||
int cached = fabric_support[device].load(std::memory_order_acquire);
|
||||
if (cached != 0) return cached == 1;
|
||||
|
||||
int fab_flag = 0;
|
||||
CUresult r = cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
|
||||
if (r != CUDA_SUCCESS || !fab_flag) {
|
||||
fabric_support[device].store(2, std::memory_order_release);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Attribute says supported — verify with a real allocation.
|
||||
// cuDeviceGetAttribute can report supported even when IMEX is not
|
||||
// configured, so we need a real probe.
|
||||
CUmemAllocationProp probe_prop = {};
|
||||
probe_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
|
||||
probe_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
||||
probe_prop.location.id = device;
|
||||
probe_prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
|
||||
|
||||
size_t granularity;
|
||||
r = cuMemGetAllocationGranularity(&granularity, &probe_prop,
|
||||
CU_MEM_ALLOC_GRANULARITY_MINIMUM);
|
||||
if (r != CUDA_SUCCESS) {
|
||||
fabric_support[device].store(2, std::memory_order_release);
|
||||
return false;
|
||||
}
|
||||
|
||||
CUmemGenericAllocationHandle test_handle;
|
||||
r = cuMemCreate(&test_handle, granularity, &probe_prop, 0);
|
||||
if (r == CUDA_SUCCESS) {
|
||||
cuMemRelease(test_handle);
|
||||
fabric_support[device].store(1, std::memory_order_release);
|
||||
return true;
|
||||
}
|
||||
|
||||
fabric_support[device].store(2, std::memory_order_release);
|
||||
return false;
|
||||
}
|
||||
#endif
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
|
||||
#ifndef USE_ROCM
|
||||
CUmemGenericAllocationHandle* p_memHandle) {
|
||||
@@ -112,32 +190,40 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
|
||||
CUresult rdma_result = cuDeviceGetAttribute(
|
||||
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
|
||||
device);
|
||||
if (rdma_result == CUDA_SUCCESS &&
|
||||
flag) { // support GPUDirect RDMA if possible
|
||||
if (rdma_result == CUDA_SUCCESS && flag) {
|
||||
prop.allocFlags.gpuDirectRDMACapable = 1;
|
||||
}
|
||||
int fab_flag = 0;
|
||||
CUresult fab_result = cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
|
||||
if (fab_result == CUDA_SUCCESS &&
|
||||
fab_flag) { // support fabric handle if possible
|
||||
|
||||
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
|
||||
if (probe_fabric_support(device)) {
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
|
||||
} else {
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
|
||||
}
|
||||
#else
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Allocate memory using cuMemCreate
|
||||
CUresult ret = (CUresult)cuMemCreate(p_memHandle, size, &prop, 0);
|
||||
if (ret) {
|
||||
if (fab_flag &&
|
||||
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
|
||||
// Safety net: if fabric was probed as available but this allocation
|
||||
// still fails, fall back to POSIX FD and update the cache.
|
||||
if (device < MAX_DEVICES &&
|
||||
fabric_support[device].load(std::memory_order_acquire) == 1 &&
|
||||
(ret == CUDA_ERROR_NOT_PERMITTED || ret == CUDA_ERROR_NOT_SUPPORTED)) {
|
||||
// Fabric allocation may fail without multi-node nvlink,
|
||||
// fallback to POSIX file descriptor
|
||||
fabric_support[device].store(2, std::memory_order_release);
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
|
||||
CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));
|
||||
} else {
|
||||
CUDA_CHECK(ret);
|
||||
}
|
||||
#else
|
||||
CUDA_CHECK(ret);
|
||||
#endif
|
||||
}
|
||||
if (error_code != 0) {
|
||||
return;
|
||||
@@ -302,14 +388,13 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
|
||||
// first allocation, align the size, and reserve an address, and also allocate
|
||||
// a CUmemGenericAllocationHandle
|
||||
|
||||
// Define memory allocation properties
|
||||
// No handle type here; create_and_map sets fabric/POSIX as needed.
|
||||
CUmemAllocationProp prop = {};
|
||||
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
|
||||
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
||||
prop.location.id = device;
|
||||
prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
|
||||
|
||||
// Check if the allocation is supported
|
||||
size_t granularity;
|
||||
CUDA_CHECK(cuMemGetAllocationGranularity(&granularity, &prop,
|
||||
CU_MEM_ALLOC_GRANULARITY_MINIMUM));
|
||||
@@ -325,7 +410,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;
|
||||
}
|
||||
|
||||
@@ -118,9 +118,9 @@ 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 \
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
|
||||
--mount=type=cache,target=/workspace/rust/target,sharing=locked \
|
||||
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
|
||||
|
||||
######################### BUILD IMAGE #########################
|
||||
|
||||
+92
-79
@@ -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,19 +119,98 @@ 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
|
||||
@@ -147,75 +226,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"]
|
||||
|
||||
@@ -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:
|
||||
@@ -50,10 +47,6 @@ Now supports 9 types of connectors:
|
||||
--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.
|
||||
|
||||
@@ -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 ""
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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
+3
@@ -5799,6 +5799,7 @@ dependencies = [
|
||||
"axum",
|
||||
"bytes",
|
||||
"clap",
|
||||
"educe",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"http-body",
|
||||
@@ -5813,7 +5814,9 @@ dependencies = [
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"serial_test",
|
||||
"sha2",
|
||||
"socket2",
|
||||
"subtle",
|
||||
"thiserror-ext",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
|
||||
@@ -75,8 +75,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"
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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])
|
||||
}
|
||||
|
||||
@@ -170,8 +170,7 @@ roundtrip_tests! {
|
||||
glm47 => [reasoning_and_content, tool_call_mix],
|
||||
|
||||
// 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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -92,7 +92,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));
|
||||
|
||||
@@ -156,8 +156,7 @@ pub(crate) fn prepare_chat_request(
|
||||
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 +203,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 +315,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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))?;
|
||||
|
||||
|
||||
@@ -407,6 +407,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 +417,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 +466,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 +475,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 +783,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 +1094,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 +2450,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 +2654,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 +3436,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(),
|
||||
))
|
||||
@@ -4228,6 +4347,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 +4518,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 +4875,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;
|
||||
@@ -14,6 +15,12 @@ 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 +34,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 +62,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 +94,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 {
|
||||
|
||||
@@ -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`.
|
||||
///
|
||||
|
||||
@@ -248,6 +248,32 @@ def test_deep_sleep_async():
|
||||
asyncio.run(test())
|
||||
|
||||
|
||||
@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
|
||||
def test_cumem_without_sleep_mode():
|
||||
"""Verify cumem allocator works independently of sleep mode."""
|
||||
llm = LLM("hmellor/tiny-random-LlamaForCausalLM", enable_cumem_allocator=True)
|
||||
prompt = "How are you?"
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=10)
|
||||
output = llm.generate(prompt, sampling_params)
|
||||
assert output[0].outputs[0].text
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Sleep mode requires CUDA or ROCm",
|
||||
)
|
||||
def test_sleep_mode_auto_enables_cumem():
|
||||
"""Verify sleep mode automatically enables cumem allocator."""
|
||||
from vllm.config.model import ModelConfig
|
||||
|
||||
cfg = ModelConfig(
|
||||
"hmellor/tiny-random-LlamaForCausalLM",
|
||||
enable_sleep_mode=True,
|
||||
enable_cumem_allocator=False,
|
||||
)
|
||||
assert cfg.enable_cumem_allocator is True
|
||||
|
||||
|
||||
@requires_fp8
|
||||
def test_deep_sleep_fp8_kvcache():
|
||||
model = "Qwen/Qwen2-0.5B"
|
||||
|
||||
@@ -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
|
||||
|
||||
+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
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_oot_registration import (
|
||||
from tests.plugins_tests.test_oot_registration_online import (
|
||||
run_and_test_dummy_opt_api_server,
|
||||
)
|
||||
|
||||
|
||||
@@ -77,9 +77,9 @@ def make_fused_moe_layer(
|
||||
intermediate_size=test_config.intermediate_size,
|
||||
prefix=f"dummy_layer_{layer_idx}",
|
||||
activation="silu",
|
||||
is_act_and_mul=True,
|
||||
params_dtype=test_config.weight_dtype,
|
||||
)
|
||||
re = fml.routed_experts
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
@@ -92,12 +92,12 @@ def make_fused_moe_layer(
|
||||
tensor_device=device,
|
||||
)
|
||||
|
||||
assert isinstance(fml.w13_weight.data, torch.Tensor)
|
||||
assert isinstance(fml.w2_weight.data, torch.Tensor)
|
||||
fml.w13_weight.data = fml.w13_weight.data.to(device=device)
|
||||
fml.w2_weight.data = fml.w2_weight.data.to(device=device)
|
||||
w13_weight = fml.w13_weight.data
|
||||
w2_weight = fml.w2_weight.data
|
||||
assert isinstance(re.w13_weight.data, torch.Tensor)
|
||||
assert isinstance(re.w2_weight.data, torch.Tensor)
|
||||
re.w13_weight.data = re.w13_weight.data.to(device=device)
|
||||
re.w2_weight.data = re.w2_weight.data.to(device=device)
|
||||
w13_weight = re.w13_weight.data
|
||||
w2_weight = re.w2_weight.data
|
||||
assert w13_weight.size(0) == test_config.num_local_experts
|
||||
for i in range(test_config.num_local_experts):
|
||||
g_i = rank * test_config.num_local_experts + i
|
||||
@@ -172,10 +172,10 @@ def make_fused_moe_layer(
|
||||
assert not w2_weight_scale_inv.is_contiguous()
|
||||
|
||||
# Add scales to the parameter list
|
||||
fml.w13_weight_scale_inv = torch.nn.Parameter(
|
||||
re.w13_weight_scale_inv = torch.nn.Parameter(
|
||||
w13_weight_scale_inv, requires_grad=False
|
||||
)
|
||||
fml.w2_weight_scale_inv = torch.nn.Parameter(
|
||||
re.w2_weight_scale_inv = torch.nn.Parameter(
|
||||
w2_weight_scale_inv, requires_grad=False
|
||||
)
|
||||
|
||||
|
||||
@@ -61,7 +61,6 @@ def make_fused_moe_layer(
|
||||
intermediate_size=test_config.intermediate_size,
|
||||
prefix=f"dummy_layer_{layer_idx}",
|
||||
activation="silu",
|
||||
is_act_and_mul=True,
|
||||
params_dtype=torch.bfloat16,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
@@ -77,6 +76,7 @@ def make_fused_moe_layer(
|
||||
)
|
||||
|
||||
fml = fml.to(device)
|
||||
re = fml.routed_experts
|
||||
w1_q, w2_q, quant_config = make_test_quant_config(
|
||||
test_config.num_local_experts,
|
||||
test_config.intermediate_size,
|
||||
@@ -87,21 +87,21 @@ def make_fused_moe_layer(
|
||||
per_act_token_quant=False,
|
||||
)
|
||||
|
||||
fml.w13_weight.data = w1_q
|
||||
fml.w2_weight.data = w2_q
|
||||
re.w13_weight.data = w1_q
|
||||
re.w2_weight.data = w2_q
|
||||
|
||||
fml.w2_input_scale.data = torch.randn_like(fml.w2_input_scale.data) / 5
|
||||
fml.w13_input_scale.data = torch.randn_like(fml.w13_input_scale.data) / 5
|
||||
fml.w2_weight_scale_2.data = torch.randn_like(fml.w2_weight_scale_2.data) / 5
|
||||
fml.w13_weight_scale_2.data = torch.randn_like(fml.w13_weight_scale_2.data) / 5
|
||||
fml.w2_weight_scale.data = (
|
||||
torch.randn(fml.w2_weight_scale.data.shape, device=device) / 5
|
||||
).to(fml.w2_weight_scale.data.dtype)
|
||||
fml.w13_weight_scale.data = (
|
||||
torch.randn(fml.w13_weight_scale.data.shape, device=device) / 5
|
||||
).to(fml.w13_weight_scale.data.dtype)
|
||||
re.w2_input_scale.data = torch.randn_like(re.w2_input_scale.data) / 5
|
||||
re.w13_input_scale.data = torch.randn_like(re.w13_input_scale.data) / 5
|
||||
re.w2_weight_scale_2.data = torch.randn_like(re.w2_weight_scale_2.data) / 5
|
||||
re.w13_weight_scale_2.data = torch.randn_like(re.w13_weight_scale_2.data) / 5
|
||||
re.w2_weight_scale.data = (
|
||||
torch.randn(re.w2_weight_scale.data.shape, device=device) / 5
|
||||
).to(re.w2_weight_scale.data.dtype)
|
||||
re.w13_weight_scale.data = (
|
||||
torch.randn(re.w13_weight_scale.data.shape, device=device) / 5
|
||||
).to(re.w13_weight_scale.data.dtype)
|
||||
|
||||
nvfp4_fused_moe.process_weights_after_loading(fml)
|
||||
nvfp4_fused_moe.process_weights_after_loading(re)
|
||||
|
||||
fml.maybe_init_modular_kernel()
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from vllm.utils.flashinfer import (
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
has_flashinfer_nvlink_two_sided,
|
||||
)
|
||||
from vllm.utils.import_utils import has_deep_ep_v2
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
from ..utils import init_test_distributed_environment
|
||||
@@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif(
|
||||
not _has_sys_ptrace(),
|
||||
reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
|
||||
)
|
||||
requires_deep_ep_v2 = pytest.mark.skipif(
|
||||
not has_deep_ep_v2(),
|
||||
reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
|
||||
)
|
||||
|
||||
# NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
|
||||
# their own @requires_two_sided / @requires_one_sided decorators, and
|
||||
@@ -856,3 +861,76 @@ def _one_sided_data_worker(rank, world_size):
|
||||
def test_one_sided_dispatch_combine(world_size):
|
||||
"""Test FlashInfer one-sided dispatch/combine with actual data flow."""
|
||||
_spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
|
||||
# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
|
||||
#
|
||||
# Uses EP group because the DeepEP v2 manager is constructed with an
|
||||
# EP-scoped communicator in production. With tp=world_size the EP group
|
||||
# spans all ranks.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _deepep_v2_lifecycle_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
DeepEPV2All2AllManager,
|
||||
)
|
||||
|
||||
cpu_group = get_ep_group().cpu_group
|
||||
manager = DeepEPV2All2AllManager(cpu_group)
|
||||
|
||||
assert manager.rank == rank
|
||||
assert manager.world_size == world_size
|
||||
assert manager._num_sms is None
|
||||
|
||||
hidden_size = 7168
|
||||
num_experts = world_size * 32
|
||||
num_topk = 8
|
||||
max_tokens = 256
|
||||
|
||||
handle_kwargs = dict(
|
||||
num_max_tokens_per_rank=max_tokens,
|
||||
hidden=hidden_size,
|
||||
num_topk=num_topk,
|
||||
num_experts=num_experts,
|
||||
use_fp8_dispatch=False,
|
||||
)
|
||||
|
||||
handle = manager.get_handle(handle_kwargs)
|
||||
assert handle is not None
|
||||
assert manager._num_sms is not None
|
||||
assert manager._num_sms > 0
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# get_handle again with same args should return cached handle
|
||||
handle2 = manager.get_handle(dict(handle_kwargs))
|
||||
assert handle2 is handle
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Destroy clears the cache
|
||||
manager.destroy()
|
||||
assert len(manager.handle_cache._cache) == 0
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Re-create after destroy
|
||||
handle3 = manager.get_handle(dict(handle_kwargs))
|
||||
assert handle3 is not None
|
||||
|
||||
torch.distributed.barrier()
|
||||
manager.destroy()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_deep_ep_v2
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_deepep_v2_manager_lifecycle(world_size):
|
||||
"""Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
|
||||
_spawn_workers(_deepep_v2_lifecycle_worker, world_size)
|
||||
|
||||
@@ -4,7 +4,6 @@ import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.sampling_params import SamplingParams
|
||||
@@ -76,47 +75,6 @@ def test_multi_chat(text_llm):
|
||||
assert len(outputs) == 2
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def vision_llm():
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model="microsoft/Phi-3.5-vision-instruct",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=5,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
limit_mm_per_prompt={"image": 2},
|
||||
seed=0,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
|
||||
)
|
||||
def test_chat_multi_image(vision_llm, image_urls: list[str]):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*(
|
||||
{"type": "image_url", "image_url": {"url": image_url}}
|
||||
for image_url in image_urls
|
||||
),
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
],
|
||||
}
|
||||
]
|
||||
outputs = vision_llm.chat(messages)
|
||||
assert len(outputs) >= 0
|
||||
|
||||
|
||||
def test_llm_chat_tokenization_no_double_bos(text_llm):
|
||||
"""
|
||||
LLM.chat() should not add special tokens when using chat templates.
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
@@ -0,0 +1,50 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def vision_llm():
|
||||
# pytest caches the fixture so we use weakref.proxy to
|
||||
# enable garbage collection
|
||||
llm = LLM(
|
||||
model="microsoft/Phi-3.5-vision-instruct",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=5,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
limit_mm_per_prompt={"image": 2},
|
||||
seed=0,
|
||||
)
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
|
||||
)
|
||||
def test_chat_multi_image(vision_llm, image_urls: list[str]):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*(
|
||||
{"type": "image_url", "image_url": {"url": image_url}}
|
||||
for image_url in image_urls
|
||||
),
|
||||
{"type": "text", "text": "What's in this image?"},
|
||||
],
|
||||
}
|
||||
]
|
||||
outputs = vision_llm.chat(messages)
|
||||
assert len(outputs) >= 0
|
||||
+1
-1
@@ -15,7 +15,7 @@ import logging
|
||||
import pytest
|
||||
import regex as re
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
|
||||
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.renderers.params import ChatParams
|
||||
from vllm.v1.metrics import loggers as stat_loggers
|
||||
+1
-1
@@ -6,7 +6,7 @@ import logging
|
||||
import pytest
|
||||
import regex as re
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
|
||||
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM
|
||||
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
|
||||
from vllm.v1.metrics import loggers as stat_loggers
|
||||
+1
-1
@@ -52,7 +52,7 @@ async def client_with_image_embeds(server_with_image_embeds):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("dtype", [torch.half, torch.float16, torch.float32])
|
||||
async def test_completions_with_image_embeds(
|
||||
async def test_chat_completions_with_image_embeds(
|
||||
client_with_image_embeds: openai.AsyncOpenAI,
|
||||
model_name: str,
|
||||
image_assets: ImageTestAssets,
|
||||
+1
-8
@@ -8,6 +8,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
|
||||
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
from vllm.multimodal.media import MediaWithBytes
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
@@ -16,14 +17,6 @@ from vllm.platforms import current_platform
|
||||
MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
|
||||
MAXIMUM_IMAGES = 2
|
||||
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
|
||||
# Required terms for beam search validation
|
||||
# Each entry is a list of term groups - ALL groups must match
|
||||
# Each group is a list of alternatives - at least ONE term in the group must appear
|
||||
+1
-7
@@ -7,19 +7,13 @@ import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.utils import encode_image_url
|
||||
|
||||
# Use a small vision model for testing
|
||||
MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct"
|
||||
MAXIMUM_IMAGES = 2
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@@ -14,6 +14,7 @@ import torch
|
||||
from openai import BadRequestError
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
CHAT_TEMPLATE = VLLM_PATH / "examples/template_chatml.jinja"
|
||||
@@ -41,7 +42,11 @@ def server_args() -> list[str]:
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(server_args):
|
||||
def server(server_args, request):
|
||||
if current_platform.is_rocm():
|
||||
# Materialize HF embeddings before the server reserves ROCm VRAM.
|
||||
request.getfixturevalue("prompt_embeds_b64")
|
||||
request.getfixturevalue("aligned_content_and_embeds_b64")
|
||||
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@ import json
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL = "ibm-granite/granite-4.0-h-tiny"
|
||||
|
||||
+1
-2
@@ -9,12 +9,11 @@ import pytest_asyncio
|
||||
from huggingface_hub import snapshot_download
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.tool_parsers.abstract_tool_parser import ToolParser
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
|
||||
|
||||
TOOLS = [
|
||||
+1
-1
@@ -9,7 +9,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from rapidfuzz import fuzz
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "openai/gpt-oss-20b"
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user