forked from Karylab-cklius/vllm
Compare commits
58
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6e714a103c | ||
|
|
c9951fd5c7 | ||
|
|
a78c15616f | ||
|
|
5c4db60f01 | ||
|
|
4e5ca89cfe | ||
|
|
a22e0dfc69 | ||
|
|
cc56379e28 | ||
|
|
024b06b0dc | ||
|
|
e7d0fcbc09 | ||
|
|
aa8bb5562e | ||
|
|
fa4bec9056 | ||
|
|
dee5da1dec | ||
|
|
ed41aa270a | ||
|
|
77a9c5ae28 | ||
|
|
f651a8a9a4 | ||
|
|
8f82be5705 | ||
|
|
a461070d1c | ||
|
|
4470ae84de | ||
|
|
697c34b97b | ||
|
|
5b431b905c | ||
|
|
89e99202f2 | ||
|
|
b446792306 | ||
|
|
c3b1f9e827 | ||
|
|
df802a87b7 | ||
|
|
93d8f834dd | ||
|
|
aeb35b90f0 | ||
|
|
9a08a5118e | ||
|
|
c5200d3565 | ||
|
|
3c1396bab6 | ||
|
|
9969466a59 | ||
|
|
3406e8f83d | ||
|
|
a264e41975 | ||
|
|
f098ee70c7 | ||
|
|
9294dd27eb | ||
|
|
b1190d03cc | ||
|
|
92c7fac640 | ||
|
|
ac521f6237 | ||
|
|
28242824e0 | ||
|
|
68294739d1 | ||
|
|
c8d2f3cb14 | ||
|
|
345b28ff2f | ||
|
|
248d1fbb71 | ||
|
|
11b26c5528 | ||
|
|
20434c472e | ||
|
|
c8f9c156a5 | ||
|
|
953bba488d | ||
|
|
3a9784b82c | ||
|
|
3cecee40f3 | ||
|
|
a7732537f4 | ||
|
|
727971f1c1 | ||
|
|
25671cb520 | ||
|
|
27d5f78b63 | ||
|
|
7a341fa109 | ||
|
|
f41e8ddc97 | ||
|
|
245888ff77 | ||
|
|
e840f0d3f5 | ||
|
|
fcaa84efa7 | ||
|
|
9e84ec8648 |
@@ -1,25 +0,0 @@
|
||||
group: Zen CPU
|
||||
depends_on: []
|
||||
steps:
|
||||
- label: Zen-CPU-Kernel Tests
|
||||
depends_on: []
|
||||
soft_fail: false
|
||||
device: zen5
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
- setup.py
|
||||
- docker/Dockerfile.cpu
|
||||
- requirements/cpu.txt
|
||||
- requirements/build/cpu.txt
|
||||
- csrc/cpu/
|
||||
- cmake/cpu_extension.cmake
|
||||
- CMakeLists.txt
|
||||
- vllm/model_executor/layers/utils.py
|
||||
- vllm/platforms/cpu.py
|
||||
- vllm/platforms/zen_cpu.py
|
||||
- vllm/platforms/__init__.py
|
||||
- tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-zen-cpu-test.sh 20m "
|
||||
pytest -x -v -s tests/model_executor/test_cpu_unquantized_gemm_dispatch.py"
|
||||
@@ -103,6 +103,31 @@ steps:
|
||||
pytest -v -s v1/kv_offload &&
|
||||
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
|
||||
|
||||
- label: NixlConnector PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
num_devices: 2
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
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/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- vllm/platforms/xpu.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh'
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -22,7 +22,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
'pip install av &&
|
||||
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'
|
||||
@@ -47,8 +47,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
'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"
|
||||
@@ -71,8 +70,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
'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'
|
||||
|
||||
@@ -96,7 +94,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
'pip install av &&
|
||||
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'
|
||||
|
||||
@@ -121,7 +119,7 @@ steps:
|
||||
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 av matplotlib ftfy &&
|
||||
pip install open-clip-torch --no-deps &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
@@ -68,7 +68,6 @@ steps:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
|
||||
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
|
||||
pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
|
||||
pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
|
||||
|
||||
@@ -534,6 +534,20 @@ else
|
||||
echo "--- Single-node job"
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
|
||||
ulimit_core_hard=$(ulimit -H -c)
|
||||
if [[ "$ulimit_core_hard" == "unlimited" ]]; then
|
||||
# docker run can't pass "unlimited" to --ulimit
|
||||
ulimit_core_hard="-1"
|
||||
fi
|
||||
# Disable core dumps in the ROCm test container unless the ROCm debug agent is enabled
|
||||
coredump_flags="--ulimit core=0:$ulimit_core_hard"
|
||||
if [[ "$commands" == *"ROCm debug agent enabled"* ]]; then
|
||||
# Works around https://github.com/rocm/rocm-systems/issues/6206
|
||||
coredump_flags='-e HSA_COREDUMP_PATTERN="/tmp/gpucore.%p"'
|
||||
else
|
||||
echo "ROCm debug agent not enabled, coredumps are disabled in the test container."
|
||||
fi
|
||||
|
||||
docker run \
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
$RDMA_FLAGS \
|
||||
@@ -541,6 +555,7 @@ else
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
$coredump_flags \
|
||||
-e HF_TOKEN \
|
||||
-e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
|
||||
-e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script build the CPU docker image and run the offline inference inside the container.
|
||||
# It serves a sanity check for compilation and basic model usage.
|
||||
set -euox pipefail
|
||||
|
||||
# allow to bind to different cores
|
||||
CORE_RANGE=${CORE_RANGE:-48-95}
|
||||
NUMA_NODE=${NUMA_NODE:-1}
|
||||
IMAGE_NAME="zen-cpu-test-$NUMA_NODE"
|
||||
TIMEOUT_VAL=$1
|
||||
TEST_COMMAND=$2
|
||||
|
||||
# building the docker image
|
||||
echo "--- :docker: Building Docker image"
|
||||
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-zen-test -f docker/Dockerfile.cpu .
|
||||
|
||||
# Run the image, setting --shm-size=4g for tensor parallel.
|
||||
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
|
||||
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
+34
-50
@@ -114,26 +114,6 @@ steps:
|
||||
|
||||
#---------------------------------------------------------- mi250 · compile ----------------------------------------------------------#
|
||||
|
||||
- label: PyTorch Compilation Unit Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/worker/
|
||||
- vllm/v1/attention/
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- vllm/config/compilation.py
|
||||
- csrc/
|
||||
- tests/compile
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
@@ -259,7 +239,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
|
||||
@@ -438,7 +417,7 @@ steps:
|
||||
#----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------#
|
||||
|
||||
- label: Basic Correctness # TBD
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 95
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
@@ -456,7 +435,7 @@ steps:
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
- label: Distributed Model Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 110
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -498,6 +477,26 @@ steps:
|
||||
|
||||
#---------------------------------------------------------- mi300 · compile ----------------------------------------------------------#
|
||||
|
||||
- label: PyTorch Compilation Unit Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/worker/
|
||||
- vllm/v1/attention/
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- vllm/config/compilation.py
|
||||
- csrc/
|
||||
- tests/compile
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
- label: Fusion E2E Config Sweep (H100-MI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -678,7 +677,7 @@ steps:
|
||||
- pytest -v -s distributed/test_eplb_spec_decode.py
|
||||
|
||||
- label: Distributed Tests (2xH100-2xMI300) # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -1222,7 +1221,7 @@ steps:
|
||||
#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------#
|
||||
|
||||
- label: Examples # TBD
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1258,7 +1257,7 @@ steps:
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 55
|
||||
timeout_in_minutes: 100
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1292,7 +1291,7 @@ steps:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
|
||||
|
||||
- label: Kernels MoE Test %N # TBD
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 95
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1374,8 +1373,6 @@ steps:
|
||||
- tests/lora
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
- pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- pytest -v -s -x lora/test_llama_tp.py
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
@@ -1439,7 +1436,7 @@ steps:
|
||||
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
|
||||
|
||||
- label: Basic Models Tests (Other) # TBD
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1516,7 +1513,6 @@ steps:
|
||||
- tests/models/multimodal/generation
|
||||
- tests/models/multimodal/test_mapping.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
|
||||
@@ -1530,7 +1526,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
|
||||
|
||||
|
||||
@@ -1544,7 +1539,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2" # TBD
|
||||
@@ -1558,7 +1552,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
|
||||
@@ -1574,7 +1567,6 @@ steps:
|
||||
- tests/models/multimodal/generation
|
||||
- tests/models/multimodal/test_mapping.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
|
||||
@@ -1589,7 +1581,6 @@ steps:
|
||||
- tests/models/multimodal/generation
|
||||
- tests/models/multimodal/test_mapping.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
- pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
|
||||
@@ -1605,7 +1596,6 @@ steps:
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Processor (CPU) %N # TBD
|
||||
@@ -1621,7 +1611,6 @@ steps:
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
#----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------#
|
||||
@@ -1903,7 +1892,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: Spec Decode Eagle # TBD
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 90
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -2119,7 +2108,7 @@ steps:
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs) # TBD
|
||||
timeout_in_minutes: 20
|
||||
timeout_in_minutes: 65
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
optional: true
|
||||
@@ -2272,7 +2261,7 @@ steps:
|
||||
#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -2284,7 +2273,7 @@ steps:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
|
||||
|
||||
- label: Weight Loading Multiple GPU - Large Models # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -2469,7 +2458,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
|
||||
@@ -2831,7 +2819,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 100
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
@@ -2980,7 +2968,6 @@ steps:
|
||||
- tests/models/multimodal/generation
|
||||
- tests/models/multimodal/test_mapping.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
|
||||
@@ -2994,7 +2981,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Pooling) # TBD
|
||||
@@ -3020,7 +3006,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
|
||||
@@ -3035,7 +3020,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
- pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model
|
||||
@@ -3178,7 +3162,7 @@ steps:
|
||||
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
@@ -3190,7 +3174,7 @@ steps:
|
||||
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
|
||||
|
||||
- label: Weight Loading Multiple GPU - Large Models # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 75
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -233,7 +233,7 @@ steps:
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
|
||||
- pytest -v -s tests/distributed/test_nccl_symm_mem.py
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
|
||||
|
||||
|
||||
@@ -54,8 +54,8 @@ steps:
|
||||
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
|
||||
|
||||
- label: LM Eval Small Models (2xB200)
|
||||
key: lm-eval-small-models-2xb200
|
||||
- label: LM Eval Small Models (1xB200)
|
||||
key: lm-eval-small-models-1xb200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
@@ -65,9 +65,10 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Small Models (2xL4)
|
||||
key: lm-eval-small-models-tp
|
||||
timeout_in_minutes: 10
|
||||
- label: LM Eval Small Models Distributed (2xB200)
|
||||
key: lm-eval-small-models-distributed-2xb200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
num_devices: 2
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -10,7 +10,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
mirror:
|
||||
@@ -27,7 +26,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
mirror:
|
||||
@@ -44,7 +42,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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
|
||||
mirror:
|
||||
@@ -61,7 +58,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- 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/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
|
||||
@@ -83,7 +79,6 @@ steps:
|
||||
- tests/models/registry.py
|
||||
device: cpu-medium
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
@@ -95,7 +90,6 @@ steps:
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
|
||||
@@ -129,7 +123,6 @@ steps:
|
||||
- tests/models/multimodal/generation
|
||||
- tests/models/multimodal/test_mapping.py
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
mirror:
|
||||
@@ -146,7 +139,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 3)
|
||||
@@ -157,7 +149,6 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Pooling)
|
||||
|
||||
@@ -132,10 +132,8 @@ def benchmark_function(
|
||||
reset_memory_stats()
|
||||
|
||||
# Benchmark
|
||||
start_events = [
|
||||
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
|
||||
]
|
||||
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
|
||||
for i in range(benchmark_iters):
|
||||
logits_copy = logits.clone()
|
||||
|
||||
@@ -134,8 +134,8 @@ def benchmark_config(
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Benchmark
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(num_iters):
|
||||
with override_config(config):
|
||||
|
||||
@@ -170,8 +170,8 @@ def benchmark_config(
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
latencies: list[float] = []
|
||||
for _ in range(num_iters):
|
||||
start.record()
|
||||
|
||||
@@ -427,6 +427,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
"csrc/cpu/activation_lut_bf16.cpp"
|
||||
"csrc/cpu/cpu_tanhf_neon.hpp"
|
||||
"csrc/cpu/cpu_fused_moe.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
|
||||
GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
|
||||
activation_kernel<scalar_t, gelu_tanh_act, false>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_new(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_TANHF_NEON_HPP
|
||||
#define CPU_TANHF_NEON_HPP
|
||||
|
||||
#include <cstdint>
|
||||
#include <arm_neon.h>
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
namespace {
|
||||
|
||||
struct TanhfConstants {
|
||||
float32x4_t special_bound;
|
||||
float32x4_t two;
|
||||
float32x4_t c0;
|
||||
float32x4_t c2;
|
||||
int32x4_t exponent_bias;
|
||||
float c1;
|
||||
float c3;
|
||||
float two_over_ln2;
|
||||
float c4;
|
||||
float ln2_hi;
|
||||
float ln2_lo;
|
||||
};
|
||||
|
||||
const TanhfConstants kTanhfConstants = {
|
||||
// 9.01, above which tanhf rounds to 1 (or -1 for negative).
|
||||
.special_bound = vdupq_n_f32(0x1.205966p+3f),
|
||||
.two = vdupq_n_f32(0x1.0p+1f),
|
||||
.c0 = vdupq_n_f32(0x1.fffffep-2f),
|
||||
.c2 = vdupq_n_f32(0x1.555736p-5f),
|
||||
.exponent_bias = vdupq_n_s32(0x3f800000),
|
||||
.c1 = 0x1.5554aep-3f,
|
||||
.c3 = 0x1.12287cp-7f,
|
||||
.two_over_ln2 = 0x1.715476p+1f,
|
||||
.c4 = 0x1.6b55a2p-10f,
|
||||
.ln2_hi = 0x1.62e4p-1f,
|
||||
.ln2_lo = 0x1.7f7d1cp-20f,
|
||||
};
|
||||
|
||||
// Return the ptr but hide it's value from the compiler so accesses
|
||||
// through it can't be optimised based on contents.
|
||||
template <typename T>
|
||||
inline const T* ptr_barrier(const T* ptr) {
|
||||
const T* opaque_ptr = ptr;
|
||||
__asm__("" : "+r"(opaque_ptr));
|
||||
return opaque_ptr;
|
||||
}
|
||||
|
||||
// Check whether any lanes in the mask are set
|
||||
inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
|
||||
|
||||
// e^2x - 1 inline helper
|
||||
inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
|
||||
float32x2_t ln2 = vld1_f32(&d->ln2_hi);
|
||||
float32x4_t lane_consts = vld1q_f32(&d->c1);
|
||||
|
||||
// Reduce argument: f in [-ln2/2, ln2/2], i is exact.
|
||||
float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
|
||||
int32x4_t i = vcvtq_s32_f32(j);
|
||||
float32x4_t f = vaddq_f32(x, x);
|
||||
f = vfmsq_lane_f32(f, j, ln2, 0);
|
||||
f = vfmsq_lane_f32(f, j, ln2, 1);
|
||||
|
||||
// Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
|
||||
float32x4_t f2 = vmulq_f32(f, f);
|
||||
float32x4_t f4 = vmulq_f32(f2, f2);
|
||||
float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
|
||||
float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
|
||||
float32x4_t poly = vfmaq_f32(p01, f2, p23);
|
||||
poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
|
||||
poly = vfmaq_f32(f, f2, poly);
|
||||
|
||||
// scale = 2^i
|
||||
int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
|
||||
float32x4_t scale = vreinterpretq_f32_s32(u);
|
||||
return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
|
||||
}
|
||||
|
||||
// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
|
||||
inline float32x4_t special_case(float32x4_t x, float32x4_t q,
|
||||
uint32x4_t special) {
|
||||
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
|
||||
|
||||
float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
|
||||
uint32x4_t ix = vreinterpretq_u32_f32(x);
|
||||
uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
|
||||
uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
|
||||
uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
|
||||
float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
|
||||
return vbslq_f32(special, special_y, y);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
|
||||
// AdvSIMD)
|
||||
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
|
||||
//
|
||||
// Approximation for single-precision vector tanh(x), using a simplified
|
||||
// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
|
||||
// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
|
||||
inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
|
||||
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
|
||||
|
||||
// tanh(x) = (e^2x - 1) / (e^2x + 1)
|
||||
// q = e^2x -1
|
||||
float32x4_t q = e2xm1f_inline(x, d);
|
||||
|
||||
// Check for special cases
|
||||
uint32x4_t special = vcagtq_f32(x, d->special_bound);
|
||||
|
||||
// Fall back to vectorised special case for any lanes which would cause
|
||||
// expm1 to overflow
|
||||
if (any_u32(special)) {
|
||||
return special_case(x, q, special);
|
||||
}
|
||||
|
||||
// Complete fast path if no special lanes
|
||||
// tanh(x) = q / (q+2)
|
||||
return vdivq_f32(q, vaddq_f32(q, d->two));
|
||||
}
|
||||
|
||||
} // namespace vec_op
|
||||
|
||||
#endif // CPU_TANHF_NEON_HPP
|
||||
@@ -3,6 +3,8 @@
|
||||
|
||||
#include <arm_neon.h>
|
||||
|
||||
#include "cpu/cpu_tanhf_neon.hpp"
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
|
||||
explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
|
||||
|
||||
explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
|
||||
|
||||
FORCE_INLINE FP32Vec4 tanh() const {
|
||||
return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
|
||||
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
|
||||
reg.val[1] = Vectorized<float>(data.val[1]);
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 tanh() const {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
|
||||
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE float reduce_sum() const noexcept {
|
||||
float answer = 0;
|
||||
std::plus<VectorizedT> add;
|
||||
@@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
|
||||
reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
|
||||
};
|
||||
|
||||
FORCE_INLINE FP32Vec16 tanh() const {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
|
||||
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
|
||||
r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
|
||||
r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
|
||||
return r;
|
||||
}
|
||||
|
||||
static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
|
||||
FP32Vec16& odd) noexcept {
|
||||
const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
|
||||
|
||||
@@ -298,6 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
|
||||
|
||||
// GELU tanh implementation.
|
||||
ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
|
||||
|
||||
// GELU implementation used in GPT-2.
|
||||
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_new", torch::kCPU, &gelu_new);
|
||||
|
||||
@@ -804,35 +804,6 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
|
||||
} // namespace tensorrt_llm
|
||||
} // namespace vllm
|
||||
|
||||
torch::stable::Tensor minimax_allreduce_rms(
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
|
||||
int64_t const rank, int64_t const nranks, double const eps) {
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
|
||||
|
||||
allreduce_params.nranks = static_cast<int>(nranks);
|
||||
allreduce_params.rank = static_cast<int>(rank);
|
||||
allreduce_params.dtype = input.scalar_type();
|
||||
allreduce_params.size_q = static_cast<int>(input.numel());
|
||||
allreduce_params.hidden_dim = static_cast<int>(input.size(-1));
|
||||
allreduce_params.stride_q = allreduce_params.hidden_dim;
|
||||
allreduce_params.workspace =
|
||||
reinterpret_cast<void**>(workspace.mutable_data_ptr());
|
||||
allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
|
||||
allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
|
||||
allreduce_params.rms_eps = static_cast<float>(eps);
|
||||
allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
|
||||
allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();
|
||||
|
||||
vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
|
||||
|
||||
return rms_norm_out;
|
||||
}
|
||||
|
||||
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
|
||||
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
|
||||
torch::stable::Tensor const& norm_weight_q,
|
||||
|
||||
@@ -288,10 +288,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
int64_t cache_block_size);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::stable::Tensor minimax_allreduce_rms(
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
|
||||
int64_t const rank, int64_t const nranks, double const eps);
|
||||
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
|
||||
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
|
||||
torch::stable::Tensor const& norm_weight_q,
|
||||
|
||||
@@ -449,10 +449,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"int cache_block_size) -> ()");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def(
|
||||
"minimax_allreduce_rms("
|
||||
"Tensor input, Tensor norm_weight, Tensor workspace, "
|
||||
"int rank, int nranks, float eps) -> Tensor");
|
||||
ops.def(
|
||||
"minimax_allreduce_rms_qk("
|
||||
"Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
|
||||
@@ -705,7 +701,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
|
||||
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
|
||||
#endif
|
||||
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
|
||||
|
||||
@@ -35,6 +35,8 @@ void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void gelu_tanh(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void gelu_new(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void gelu_fast(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
@@ -318,30 +318,3 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install "vllm[zen]"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
######################### ZEN CPU TEST IMAGE #########################
|
||||
FROM vllm-openai-zen AS vllm-zen-test
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/test/cpu.txt
|
||||
|
||||
ADD ./tests/ ./tests/
|
||||
ADD ./examples/ ./examples/
|
||||
ADD ./benchmarks/ ./benchmarks/
|
||||
ADD ./vllm/collect_env.py .
|
||||
ADD ./docker/ ./docker/
|
||||
ADD ./.buildkite/ ./.buildkite/
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
ENV HF_XET_HIGH_PERFORMANCE=1
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT=60
|
||||
|
||||
ENTRYPOINT []
|
||||
|
||||
@@ -222,7 +222,7 @@ MLA decode backends are selected using the standard
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
|
||||
@@ -405,7 +405,6 @@ th {
|
||||
| `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ |
|
||||
| `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ |
|
||||
| `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
|
||||
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ |
|
||||
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
|
||||
| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
|
||||
| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ |
|
||||
@@ -477,7 +476,6 @@ th {
|
||||
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
|
||||
| `StableLMEpochForCausalLM` | StableLM Epoch | `stabilityai/stablelm-zephyr-3b`, etc. | | ✅︎ |
|
||||
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
|
||||
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
|
||||
| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ |
|
||||
| `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -490,7 +488,9 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
|
||||
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
|
||||
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
|
||||
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
|
||||
@@ -532,7 +532,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
|
||||
| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
|
||||
| `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
|
||||
| `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
|
||||
| `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
|
||||
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -579,7 +578,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
|
||||
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
|
||||
| `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
|
||||
| `LlavaNextForConditionalGeneration` | LLaVA-NeXT, Granite Vision | T + I<sup>E+</sup> | `llava-hf/llava-v1.6-mistral-7b-hf`, `llava-hf/llava-v1.6-vicuna-7b-hf`, `ibm-granite/granite-vision-3.3-2b`, etc. | | ✅︎ |
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
@@ -594,7 +593,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
|
||||
| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
|
||||
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
|
||||
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
|
||||
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
|
||||
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
|
||||
| `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
|
||||
@@ -678,9 +676,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
coordinate decoding and are not exposed by this vLLM implementation.
|
||||
See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes).
|
||||
|
||||
!!! note
|
||||
To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM.
|
||||
|
||||
!!! note
|
||||
The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now.
|
||||
For more details, please see: <https://github.com/vllm-project/vllm/pull/4087#issuecomment-2250397630>
|
||||
|
||||
@@ -58,7 +58,7 @@ class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
|
||||
|
||||
### High Level Weight Transfer API
|
||||
|
||||
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
|
||||
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle.
|
||||
|
||||
### Mid Level `reload_weights` API
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
|
||||
| ------- | --------- | -------- |
|
||||
| [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
|
||||
| [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
|
||||
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -41,7 +42,7 @@ vllm serve my-model \
|
||||
--weight-transfer-config '{"backend": "nccl"}'
|
||||
```
|
||||
|
||||
The `backend` field accepts `"nccl"` (default) or `"ipc"`.
|
||||
The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
|
||||
|
||||
## API Endpoints
|
||||
|
||||
@@ -69,7 +70,7 @@ Both backends provide static methods that the trainer calls to send weights. The
|
||||
EngineClass.trainer_init(init_info)
|
||||
|
||||
# 2. Start weight update on inference side
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
llm.start_weight_update()
|
||||
|
||||
# 3. Send weights to inference workers
|
||||
EngineClass.trainer_send_weights(
|
||||
|
||||
@@ -11,15 +11,23 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data
|
||||
|
||||
### Abstract Methods
|
||||
|
||||
Subclasses must implement these four methods:
|
||||
Subclasses must implement these methods:
|
||||
|
||||
| Method | Side | Description |
|
||||
| ------ | ---- | ----------- |
|
||||
| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
|
||||
| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
|
||||
| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
|
||||
| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
|
||||
| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
|
||||
| `shutdown()` | Inference | Clean up resources |
|
||||
| `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
|
||||
|
||||
The base class provides two methods:
|
||||
|
||||
1. `__init__` : Engines receive `config` (`WeightTransferConfig`), `vllm_config` (`VllmConfig`), `device` (`torch.device`) and `model` (`nn.Module`)
|
||||
2. `update_weights(update_info_dict)`: Thin wrapper for `receive_weights`: parses
|
||||
the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
|
||||
|
||||
### Request Classes
|
||||
|
||||
The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
|
||||
@@ -81,7 +89,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo):
|
||||
### 2. Implement the Engine
|
||||
|
||||
```python
|
||||
from collections.abc import Callable, Iterator
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
import torch
|
||||
|
||||
@@ -93,18 +101,25 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
|
||||
# Set up connection to trainer using init_info.endpoint, etc.
|
||||
...
|
||||
|
||||
def receive_weights(
|
||||
self,
|
||||
update_info: MyUpdateInfo,
|
||||
load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
|
||||
) -> None:
|
||||
# Receive each weight and call load_weights incrementally
|
||||
def start_weight_update(self) -> None:
|
||||
# Checkpoint-format engines: run initialize_layerwise_reload(self.model).
|
||||
# In-place engines: no-op
|
||||
...
|
||||
|
||||
def finish_weight_update(self) -> None:
|
||||
# Checkpoint-format engines: run finalize_layerwise_reload(...).
|
||||
# In-place engines: no-op
|
||||
...
|
||||
|
||||
def receive_weights(self, update_info: MyUpdateInfo) -> None:
|
||||
weights = []
|
||||
for name, dtype_name, shape in zip(
|
||||
update_info.names, update_info.dtype_names, update_info.shapes
|
||||
):
|
||||
dtype = getattr(torch, dtype_name)
|
||||
weight = self._fetch_weight(name, shape, dtype)
|
||||
load_weights([(name, weight)])
|
||||
weights.append((name, weight))
|
||||
self.model.load_weights(weights)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
# Clean up resources
|
||||
@@ -121,9 +136,6 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
|
||||
...
|
||||
```
|
||||
|
||||
!!! important
|
||||
The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
|
||||
|
||||
### 3. Register with the Factory
|
||||
|
||||
```python
|
||||
@@ -147,7 +159,7 @@ Once registered, users can select your backend via `WeightTransferConfig(backend
|
||||
|
||||
## WeightTransferEngineFactory
|
||||
|
||||
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
|
||||
The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
||||
@@ -155,7 +167,8 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
||||
# Create an engine from config
|
||||
engine = WeightTransferEngineFactory.create_engine(
|
||||
config=weight_transfer_config,
|
||||
parallel_config=parallel_config,
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
model=model,
|
||||
)
|
||||
```
|
||||
|
||||
@@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
|
||||
llm_handle=llm_actor_handle,
|
||||
)
|
||||
# start
|
||||
ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(llm_actor_handle.start_weight_update.remote())
|
||||
# send weights
|
||||
IPCWeightTransferEngine.trainer_send_weights(
|
||||
iterator=model.named_parameters(),
|
||||
@@ -80,7 +80,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
|
||||
# start
|
||||
base_url = "http://localhost:8000"
|
||||
url = f"{base_url}/start_weight_update"
|
||||
response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
response.raise_for_status()
|
||||
# send weights
|
||||
IPCWeightTransferEngine.trainer_send_weights(
|
||||
|
||||
@@ -11,7 +11,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
|
||||
## How It Works
|
||||
|
||||
1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
|
||||
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
|
||||
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
|
||||
3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
|
||||
|
||||
## Initialization
|
||||
@@ -93,7 +93,7 @@ remaining three steps are:
|
||||
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
|
||||
|
||||
# 1. Start the weight update
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
llm.start_weight_update()
|
||||
|
||||
# 2. Receive weights (can be called multiple times for chunked transfers)
|
||||
llm.update_weights(
|
||||
@@ -116,19 +116,22 @@ must match the order in which the trainer iterates over its parameters.
|
||||
|
||||
`start_weight_update` must be called before `update_weights`, and
|
||||
`finish_weight_update` must be called after all weight chunks have been
|
||||
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
|
||||
processing is applied (`True` for checkpoint-format weights, `False` for
|
||||
pre-processed kernel-format weights).
|
||||
transferred. The NCCL engine receives checkpoint-format weights and applies
|
||||
layerwise reload processing automatically inside `start_weight_update` /
|
||||
`finish_weight_update`.
|
||||
|
||||
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
|
||||
`update_info`, but they should be wrapped in
|
||||
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
|
||||
directly to runtime/kernel-format parameters. The current sparse MVP requires
|
||||
`TP=1` and `PP=1`.
|
||||
## Sparse NCCL
|
||||
|
||||
Sparse, flat-index weight patches use a separate backend,
|
||||
`WeightTransferConfig(backend="sparse_nccl")`, implemented by
|
||||
`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
|
||||
initialization with the dense engine; patches are applied directly in place to
|
||||
existing parameters (no layerwise reload). The current sparse MVP requires
|
||||
`TP=1` and `PP=1`. See the example below.
|
||||
|
||||
## Examples
|
||||
|
||||
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
|
||||
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
|
||||
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
|
||||
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
|
||||
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
|
||||
|
||||
@@ -326,6 +326,27 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the
|
||||
|
||||
**Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).
|
||||
|
||||
## gRPC Interface
|
||||
|
||||
vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server.
|
||||
|
||||
**Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network.
|
||||
|
||||
### Security Implications
|
||||
|
||||
An attacker who can reach the gRPC port can:
|
||||
|
||||
1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials
|
||||
2. **Consume GPU and compute resources** by submitting unbounded generation requests
|
||||
3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.
|
||||
|
||||
### Recommendations
|
||||
|
||||
- Only enable `--grpc-port` when you have a specific need for gRPC-based inference
|
||||
- Ensure the gRPC port is only accessible from trusted hosts or services
|
||||
- Use firewall rules to block external access to the gRPC port
|
||||
- Consider deploying the gRPC interface on a dedicated internal network interface
|
||||
|
||||
## Cache Directory Security
|
||||
|
||||
vLLM assumes that its cache directories are **private and trusted**. Cache contents are loaded without cryptographic integrity verification, including formats that support arbitrary code execution. If an untrusted user or process can write to vLLM's cache directories, they may be able to crash vLLM or cause it to execute arbitrary code.
|
||||
|
||||
@@ -91,44 +91,6 @@ def run_cohere_asr(question: str, audio_count: int) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# MusicFlamingo
|
||||
def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData:
|
||||
model_name = "nvidia/music-flamingo-2601-hf"
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
limit_mm_per_prompt={"audio": audio_count},
|
||||
enforce_eager=True,
|
||||
)
|
||||
|
||||
# MusicFlamingo prompt placeholders use <sound>; vLLM's MusicFlamingo
|
||||
# multimodal processor expands each one into <|sound_bos|> + audio tokens +
|
||||
# <|sound_eos|> based on extracted audio feature lengths.
|
||||
audio_placeholder = "<sound>" * audio_count
|
||||
system_prompt = (
|
||||
"You are Music Flamingo, a multimodal assistant for language and music. "
|
||||
"On each turn you receive an audio clip which contains music and optional "
|
||||
"text, you will receive at least one or both; use your world knowledge and "
|
||||
"reasoning to help the user with any task. Interpret the entirety of the "
|
||||
"content any input music--regardlenss of whether the user calls it audio, "
|
||||
"music, or sound."
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\n"
|
||||
f"{system_prompt}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{audio_placeholder}{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
)
|
||||
|
||||
|
||||
# Gemma3N
|
||||
def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:
|
||||
model_name = "google/gemma-3n-E2B-it"
|
||||
@@ -565,7 +527,6 @@ model_example_map = {
|
||||
"kimi_audio": run_kimi_audio,
|
||||
"midashenglm": run_midashenglm,
|
||||
"minicpmo": run_minicpmo,
|
||||
"musicflamingo": run_musicflamingo,
|
||||
"phi4_mm": run_phi4mm,
|
||||
"qwen2_audio": run_qwen2_audio,
|
||||
"qwen2_5_omni": run_qwen2_5_omni,
|
||||
|
||||
@@ -74,39 +74,6 @@ def load_aria(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
def load_aya_vision(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "CohereLabs/aya-vision-8b"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_num_seqs=2,
|
||||
limit_mm_per_prompt={"image": len(image_urls)},
|
||||
)
|
||||
|
||||
placeholders = [{"type": "image", "image": url} for url in image_urls]
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
*placeholders,
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
processor = AutoProcessor.from_pretrained(model_name)
|
||||
|
||||
prompt = processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
image_data=[fetch_image(url) for url in image_urls],
|
||||
)
|
||||
|
||||
|
||||
def load_bee(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "Open-Bee/Bee-8B-RL"
|
||||
|
||||
@@ -1420,7 +1387,6 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
|
||||
model_example_map = {
|
||||
"aria": load_aria,
|
||||
"aya_vision": load_aya_vision,
|
||||
"bee": load_bee,
|
||||
"command_a_vision": load_command_a_vision,
|
||||
"deepseek_vl_v2": load_deepseek_vl2,
|
||||
|
||||
@@ -68,28 +68,6 @@ def run_aria(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# Aya Vision
|
||||
def run_aya_vision(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
model_name = "CohereLabs/aya-vision-8b"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=2048,
|
||||
max_num_seqs=2,
|
||||
mm_processor_kwargs={"crop_to_patches": True},
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
prompts = [
|
||||
f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|><image>{question}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
|
||||
for question in questions
|
||||
]
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# Bee-8B
|
||||
def run_bee(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
@@ -1377,28 +1355,6 @@ def run_llava_onevision(questions: list[str], modality: str) -> ModelRequestData
|
||||
)
|
||||
|
||||
|
||||
# Mantis
|
||||
def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501
|
||||
prompts = [llama3_template.format(f"{question}\n<image>") for question in questions]
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model="TIGER-Lab/Mantis-8B-siglip-llama3",
|
||||
max_model_len=4096,
|
||||
hf_overrides={"architectures": ["MantisForConditionalGeneration"]},
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
stop_token_ids = [128009]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
stop_token_ids=stop_token_ids,
|
||||
)
|
||||
|
||||
|
||||
# MiniCPM-V
|
||||
def run_minicpmv_base(questions: list[str], modality: str, model_name):
|
||||
assert modality in ["image", "video", "image+video"]
|
||||
@@ -2349,7 +2305,6 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
|
||||
model_example_map = {
|
||||
"aria": run_aria,
|
||||
"aya_vision": run_aya_vision,
|
||||
"bagel": run_bagel,
|
||||
"cheers": run_cheers,
|
||||
"bee": run_bee,
|
||||
@@ -2390,7 +2345,6 @@ model_example_map = {
|
||||
"llava-next": run_llava_next,
|
||||
"llava-next-video": run_llava_next_video,
|
||||
"llava-onevision": run_llava_onevision,
|
||||
"mantis": run_mantis,
|
||||
"minicpmo": run_minicpmo,
|
||||
"minicpmv": run_minicpmv,
|
||||
"mistral3": run_mistral3,
|
||||
|
||||
@@ -306,7 +306,7 @@ gen_futures = [
|
||||
|
||||
ray.get(llm.pause_after_n_tokens.remote())
|
||||
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(llm.start_weight_update.remote())
|
||||
|
||||
inference_handle = llm.update_weights.remote(
|
||||
WeightTransferUpdateRequest(
|
||||
|
||||
@@ -80,14 +80,10 @@ def init_weight_transfer_engine(base_url: str) -> None:
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def start_weight_update(
|
||||
base_url: str,
|
||||
is_checkpoint_format: bool = True,
|
||||
) -> None:
|
||||
def start_weight_update(base_url: str) -> None:
|
||||
"""Start a weight update via HTTP endpoint."""
|
||||
url = f"{base_url}/start_weight_update"
|
||||
payload = {"is_checkpoint_format": is_checkpoint_format}
|
||||
response = requests.post(url, json=payload, timeout=60)
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
@@ -170,7 +166,7 @@ def main():
|
||||
pause_generation(BASE_URL)
|
||||
|
||||
# Start weight update, broadcast via IPC, then finish
|
||||
start_weight_update(BASE_URL, is_checkpoint_format=False)
|
||||
start_weight_update(BASE_URL)
|
||||
|
||||
print("Broadcasting weights via CUDA IPC (HTTP)...")
|
||||
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
|
||||
|
||||
@@ -83,14 +83,10 @@ def init_weight_transfer_engine(
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def start_weight_update(
|
||||
base_url: str,
|
||||
is_checkpoint_format: bool = True,
|
||||
) -> None:
|
||||
def start_weight_update(base_url: str) -> None:
|
||||
"""Start a weight update via HTTP endpoint."""
|
||||
url = f"{base_url}/start_weight_update"
|
||||
payload = {"is_checkpoint_format": is_checkpoint_format}
|
||||
response = requests.post(url, json=payload, timeout=60)
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
@@ -223,7 +219,7 @@ def main():
|
||||
shapes.append(list(p.shape))
|
||||
|
||||
# Start weight update
|
||||
start_weight_update(BASE_URL, is_checkpoint_format=True)
|
||||
start_weight_update(BASE_URL)
|
||||
|
||||
# Start the update_weights call in a separate thread since it will block
|
||||
# waiting for NCCL broadcasts
|
||||
|
||||
@@ -139,7 +139,7 @@ ray.get(llm.sleep.remote(level=0))
|
||||
|
||||
ray.get(train_model.init_weight_transfer.remote())
|
||||
# Start weight update, sync weights, then finish
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(llm.start_weight_update.remote())
|
||||
ray.get(train_model.broadcast_weights.remote(llm))
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
|
||||
|
||||
@@ -277,15 +277,8 @@ class DataParallelInferenceEngine:
|
||||
]
|
||||
)
|
||||
|
||||
def start_weight_update(self, is_checkpoint_format: bool = True):
|
||||
ray.get(
|
||||
[
|
||||
actor.start_weight_update.remote(
|
||||
is_checkpoint_format=is_checkpoint_format
|
||||
)
|
||||
for actor in self.llm_actors
|
||||
]
|
||||
)
|
||||
def start_weight_update(self):
|
||||
ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
|
||||
|
||||
def finish_weight_update(self):
|
||||
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
|
||||
@@ -392,7 +385,7 @@ def main():
|
||||
ray.get(inference_engine.wake_up.remote(tags=["weights"]))
|
||||
|
||||
print("[sync] Starting weight update...")
|
||||
ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(inference_engine.start_weight_update.remote())
|
||||
|
||||
print("[sync] Packed IPC transfer FSDP → vLLM...")
|
||||
ray.get(
|
||||
|
||||
@@ -202,7 +202,7 @@ ray.get([train_handle, inference_handle])
|
||||
names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
|
||||
|
||||
# Start weight update
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(llm.start_weight_update.remote())
|
||||
|
||||
# Issue update_weights call with NCCL-specific update info
|
||||
# packed=True enables efficient batched tensor broadcasting
|
||||
|
||||
@@ -299,7 +299,7 @@ async def main():
|
||||
print(f"[sync] Got metadata for {len(names)} parameters.")
|
||||
|
||||
print("[sync] Starting weight update...")
|
||||
await engine.start_weight_update(is_checkpoint_format=True)
|
||||
await engine.start_weight_update()
|
||||
|
||||
print("[sync] Broadcasting weights from FSDP → vLLM...")
|
||||
broadcast_handles = [
|
||||
|
||||
@@ -44,11 +44,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer.base import SparseWeightPatch
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLTrainerSendWeightsArgs,
|
||||
NCCLWeightTransferEngine,
|
||||
)
|
||||
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
|
||||
SparseNCCLWeightTransferEngine,
|
||||
SparseWeightPatch,
|
||||
)
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
@@ -244,7 +247,6 @@ class TrainModel:
|
||||
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
|
||||
shapes=[list(self.patched_param.shape)],
|
||||
num_updates_list=[flat_indices.numel()],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
|
||||
|
||||
@@ -271,7 +273,7 @@ class TrainModel:
|
||||
raise RuntimeError("Sparse patch has not been prepared")
|
||||
|
||||
start = time.perf_counter()
|
||||
NCCLWeightTransferEngine.trainer_send_sparse_weights(
|
||||
SparseNCCLWeightTransferEngine.trainer_send_weights(
|
||||
iter(self.pending_sparse_patches),
|
||||
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
|
||||
)
|
||||
@@ -282,6 +284,7 @@ class TrainModel:
|
||||
|
||||
def launch_llm(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
backend: str = "nccl",
|
||||
):
|
||||
return ray.remote(
|
||||
num_cpus=0,
|
||||
@@ -293,7 +296,7 @@ def launch_llm(
|
||||
tensor_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
gpu_memory_utilization=0.7,
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
weight_transfer_config=WeightTransferConfig(backend=backend),
|
||||
)
|
||||
|
||||
|
||||
@@ -332,7 +335,7 @@ def run_dense_phase(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
llm = launch_llm(scheduling_inference, backend="nccl")
|
||||
try:
|
||||
dense_before = collect_vllm_generations(llm)
|
||||
|
||||
@@ -351,7 +354,7 @@ def run_dense_phase(
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(llm.start_weight_update.remote())
|
||||
|
||||
dense_update_info, dense_payload_bytes = ray.get(
|
||||
train_model.get_dense_update_info.remote()
|
||||
@@ -391,7 +394,7 @@ def run_sparse_phase(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
llm = launch_llm(scheduling_inference, backend="sparse_nccl")
|
||||
try:
|
||||
sparse_before = collect_vllm_generations(llm)
|
||||
|
||||
@@ -410,7 +413,7 @@ def run_sparse_phase(
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
|
||||
ray.get(llm.start_weight_update.remote())
|
||||
|
||||
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
|
||||
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
|
||||
|
||||
Generated
+16
-1
@@ -5099,6 +5099,7 @@ dependencies = [
|
||||
"uuid",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-model-files",
|
||||
"vllm-parser",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
@@ -5236,6 +5237,20 @@ dependencies = [
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-model-files"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"hf-hub",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-parser"
|
||||
version = "0.1.0"
|
||||
@@ -5324,7 +5339,6 @@ dependencies = [
|
||||
"enum-as-inner",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"hf-hub",
|
||||
"itertools 0.14.0",
|
||||
"reqwest",
|
||||
"serde",
|
||||
@@ -5339,6 +5353,7 @@ dependencies = [
|
||||
"trait-set",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-model-files",
|
||||
"vllm-tokenizer",
|
||||
]
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ members = [
|
||||
"src/llm",
|
||||
"src/managed-engine",
|
||||
"src/metrics",
|
||||
"src/model-files",
|
||||
"src/mock-engine",
|
||||
"src/parser",
|
||||
"src/parser/python",
|
||||
@@ -126,6 +127,7 @@ vllm-chat = { path = "src/chat" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
vllm-managed-engine = { path = "src/managed-engine" }
|
||||
vllm-model-files = { path = "src/model-files" }
|
||||
vllm-metrics = { path = "src/metrics" }
|
||||
vllm-parser = { path = "src/parser" }
|
||||
vllm-server = { path = "src/server" }
|
||||
|
||||
@@ -32,6 +32,7 @@ trait-set.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-model-files.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
vllm-text.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use thiserror::Error;
|
||||
use thiserror_ext::Macro;
|
||||
use vllm_model_files::Error as ModelFilesError;
|
||||
|
||||
type BoxedError = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
@@ -69,6 +70,8 @@ pub enum Error {
|
||||
#[error(transparent)]
|
||||
Text(#[from] vllm_text::Error),
|
||||
#[error(transparent)]
|
||||
ModelFiles(#[from] ModelFilesError),
|
||||
#[error(transparent)]
|
||||
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
|
||||
}
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ use thiserror_ext::AsReport as _;
|
||||
use tracing::{info, trace, warn};
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::{
|
||||
HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
|
||||
ResolvedModelFiles, SpecialTokens, TokenizerConfig, load_tokenizer_config,
|
||||
};
|
||||
|
||||
use self::format::{
|
||||
@@ -42,7 +42,7 @@ pub struct HfChatRenderer {
|
||||
default_template: Option<CompiledChatTemplate>,
|
||||
default_template_kwargs: HashMap<String, JsonValue>,
|
||||
content_format: ContentFormatOption,
|
||||
special_tokens: Option<HfSpecialTokens>,
|
||||
special_tokens: Option<SpecialTokens>,
|
||||
multimodal: Option<MultimodalRenderInfo>,
|
||||
}
|
||||
|
||||
@@ -67,7 +67,7 @@ impl HfChatRenderer {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
|
||||
pub fn with_special_tokens(mut self, special_tokens: Option<SpecialTokens>) -> Self {
|
||||
self.special_tokens = special_tokens;
|
||||
self
|
||||
}
|
||||
@@ -83,7 +83,7 @@ impl HfChatRenderer {
|
||||
options: LoadModelBackendsOptions,
|
||||
multimodal: Option<MultimodalRenderInfo>,
|
||||
) -> Result<Self> {
|
||||
let HfTokenizerConfig {
|
||||
let TokenizerConfig {
|
||||
special_tokens,
|
||||
chat_template,
|
||||
..
|
||||
@@ -451,7 +451,7 @@ mod tests {
|
||||
use expect_test::expect;
|
||||
use serde_json::Value;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
|
||||
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
|
||||
|
||||
use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
|
||||
use crate::request::{
|
||||
@@ -675,7 +675,7 @@ mod tests {
|
||||
#[test]
|
||||
fn chat_template_injects_special_tokens_into_context() {
|
||||
let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<bos>".to_string())),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
//!
|
||||
//! This module is inlined from SMG's tokenizer crate with local adaptations:
|
||||
//! - thinking-related detection/state is removed
|
||||
//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens`
|
||||
//! - special tokens are wired to `vllm_text::backends::hf::SpecialTokens`
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
@@ -11,7 +11,7 @@ use std::path::Path;
|
||||
use minijinja::Environment;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{self};
|
||||
use vllm_text::backend::hf::HfSpecialTokens;
|
||||
use vllm_text::backend::hf::SpecialTokens;
|
||||
|
||||
use super::error::TemplateError;
|
||||
use super::format::{
|
||||
@@ -46,7 +46,7 @@ pub(super) struct TemplateContext<'a> {
|
||||
pub(super) tools: Option<&'a [TemplateTool]>,
|
||||
pub(super) documents: Option<&'a [serde_json::Value]>,
|
||||
#[serde(flatten)]
|
||||
pub(super) special_tokens: Option<&'a HfSpecialTokens>,
|
||||
pub(super) special_tokens: Option<&'a SpecialTokens>,
|
||||
#[serde(flatten)]
|
||||
pub(super) template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
@@ -133,7 +133,7 @@ mod tests {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
|
||||
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -170,7 +170,7 @@ mod tests {
|
||||
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
|
||||
.unwrap();
|
||||
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
|
||||
eos_token: Some(NamedSpecialToken::Text("</s>".to_string())),
|
||||
..Default::default()
|
||||
@@ -205,7 +205,7 @@ mod tests {
|
||||
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
|
||||
.unwrap();
|
||||
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
|
||||
eos_token: None,
|
||||
..Default::default()
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "vllm-model-files"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
hf-hub.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,21 @@
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::json::read_json_file;
|
||||
|
||||
/// Minimal subset of `tokenizer_config.json` needed by tokenizer selection.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub(crate) struct TokenizerConfig {
|
||||
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
|
||||
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
|
||||
/// containing "Tiktoken" which can be used as a hint for backend
|
||||
/// selection.
|
||||
pub tokenizer_class: Option<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
/// Error returned while resolving or reading model files.
|
||||
#[derive(Debug, ThisError)]
|
||||
#[error("model file error: {0}")]
|
||||
pub struct Error(String);
|
||||
|
||||
impl Error {
|
||||
pub(crate) fn new(message: impl Into<String>) -> Self {
|
||||
Self(message.into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Result type used by model-file discovery helpers.
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
@@ -0,0 +1,31 @@
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
use thiserror_ext::AsReport as _;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Read an optional JSON file into `T`, returning `T::default()` when absent.
|
||||
pub fn read_json_file<T>(path: Option<&Path>) -> Result<T>
|
||||
where
|
||||
T: for<'de> Deserialize<'de> + Default,
|
||||
{
|
||||
let Some(path) = path else {
|
||||
return Ok(T::default());
|
||||
};
|
||||
let content = fs::read_to_string(path).map_err(|error| {
|
||||
Error::new(format!(
|
||||
"failed to read {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})?;
|
||||
serde_json::from_str(&content).map_err(|error| {
|
||||
Error::new(format!(
|
||||
"failed to parse {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Hugging Face model file discovery shared by Rust frontend crates.
|
||||
|
||||
mod config;
|
||||
mod error;
|
||||
mod json;
|
||||
mod model_files;
|
||||
|
||||
pub use error::{Error, Result};
|
||||
pub use json::read_json_file;
|
||||
pub use model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
+70
-125
@@ -4,7 +4,7 @@ use hf_hub::Cache;
|
||||
use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
|
||||
use thiserror_ext::AsReport as _;
|
||||
|
||||
use super::config::{HfTokenizerConfig, load_tokenizer_config};
|
||||
use crate::config::{TokenizerConfig, load_tokenizer_config};
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
const HF_TOKEN_ENV: &str = "HF_TOKEN";
|
||||
@@ -26,6 +26,28 @@ pub enum TokenizerSource {
|
||||
}
|
||||
|
||||
impl TokenizerSource {
|
||||
/// Select a tokenizer source from a tokenizer file path.
|
||||
pub fn from_path(path: impl Into<PathBuf>) -> Result<Self> {
|
||||
let path = path.into();
|
||||
let file_name = path.file_name().and_then(|name| name.to_str()).ok_or_else(|| {
|
||||
Error::new(format!(
|
||||
"tokenizer path has no file name: {}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
match file_name {
|
||||
"tekken.json" => Ok(Self::Tekken(path)),
|
||||
"tokenizer.json" => Ok(Self::HuggingFace(path)),
|
||||
_ if is_tiktoken_file(&path) => Ok(Self::Tiktoken(path)),
|
||||
_ => Err(Error::new(format!(
|
||||
"unsupported tokenizer file '{}'",
|
||||
path.display()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the local filesystem path for this tokenizer source.
|
||||
pub fn path(&self) -> &Path {
|
||||
match self {
|
||||
Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path,
|
||||
@@ -38,10 +60,15 @@ impl TokenizerSource {
|
||||
pub struct ResolvedModelFiles {
|
||||
/// The selected tokenizer source for this model.
|
||||
pub tokenizer: TokenizerSource,
|
||||
/// Path to `tokenizer_config.json` when present.
|
||||
pub tokenizer_config_path: Option<PathBuf>,
|
||||
/// Path to `generation_config.json` when present.
|
||||
pub generation_config_path: Option<PathBuf>,
|
||||
/// Path to `preprocessor_config.json` when present.
|
||||
pub preprocessor_config_path: Option<PathBuf>,
|
||||
/// Path to a discovered chat template file when present.
|
||||
pub chat_template_path: Option<PathBuf>,
|
||||
/// Path to `config.json` when present.
|
||||
pub config_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
@@ -76,10 +103,10 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
|
||||
}
|
||||
|
||||
async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
|
||||
let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
|
||||
let api = build_api().map_err(|error| Error::new(format!("{}", error.as_report())))?;
|
||||
let repo = api.model(model_id.to_string());
|
||||
let info = repo.info().await.map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"failed to fetch model '{model_id}': {}",
|
||||
error.as_report()
|
||||
))
|
||||
@@ -138,9 +165,10 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let model_dir = tokenizer.path().parent().ok_or_else(|| {
|
||||
Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
|
||||
})?;
|
||||
let model_dir = tokenizer
|
||||
.path()
|
||||
.parent()
|
||||
.ok_or_else(|| Error::new("resolved tokenizer file has no parent directory"))?;
|
||||
let generation_config_path = cache_repo.get("generation_config.json");
|
||||
let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
|
||||
let chat_template_path = discover_chat_template_in_dir(model_dir);
|
||||
@@ -162,113 +190,88 @@ async fn resolve_remote_tokenizer_source(
|
||||
siblings: &std::collections::BTreeSet<&str>,
|
||||
tokenizer_class: Option<&str>,
|
||||
) -> Result<TokenizerSource> {
|
||||
if let Some(tekken_path) = download_if_present(repo, model_id, siblings, "tekken.json").await? {
|
||||
return Ok(TokenizerSource::Tekken(tekken_path));
|
||||
}
|
||||
|
||||
let tokenizer_path = if siblings.contains("tokenizer.json") {
|
||||
let tokenizer_path = if siblings.contains("tekken.json") {
|
||||
download_known_file(repo, model_id, "tekken.json").await?
|
||||
} else if siblings.contains("tokenizer.json") {
|
||||
download_known_file(repo, model_id, "tokenizer.json").await?
|
||||
} else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) {
|
||||
download_known_file(repo, model_id, tiktoken_name).await?
|
||||
} else {
|
||||
return Err(Error::Tokenizer(format!(
|
||||
return Err(Error::new(format!(
|
||||
"model '{model_id}' does not expose a supported tokenizer file \
|
||||
(tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face"
|
||||
)));
|
||||
};
|
||||
|
||||
Ok(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_class,
|
||||
None,
|
||||
))
|
||||
resolve_tokenizer_source(tokenizer_path, tokenizer_class)
|
||||
}
|
||||
|
||||
fn resolve_cached_tokenizer_source(
|
||||
cache_repo: &hf_hub::CacheRepo,
|
||||
tokenizer_config: &HfTokenizerConfig,
|
||||
tokenizer_config: &TokenizerConfig,
|
||||
) -> Result<Option<TokenizerSource>> {
|
||||
let tekken_path = cache_repo.get("tekken.json");
|
||||
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return Ok(Some(TokenizerSource::Tekken(tekken_path)));
|
||||
}
|
||||
|
||||
let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
|
||||
// tiktoken.model is the most common name, try it first.
|
||||
cache_repo.get("tiktoken.model").or_else(|| {
|
||||
// Scan for any *.tiktoken file in the cache snapshot directory.
|
||||
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
|
||||
discover_tiktoken_in_dir(&snapshot_dir)
|
||||
let Some(tokenizer_path) = cache_repo
|
||||
.get("tekken.json")
|
||||
.or_else(|| cache_repo.get("tokenizer.json"))
|
||||
.or_else(|| {
|
||||
// tiktoken.model is the most common name, try it first.
|
||||
cache_repo.get("tiktoken.model").or_else(|| {
|
||||
// Scan for any *.tiktoken file in the cache snapshot directory.
|
||||
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
|
||||
discover_tiktoken_in_dir(&snapshot_dir)
|
||||
})
|
||||
})
|
||||
}) else {
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_config.tokenizer_class.as_deref(),
|
||||
None,
|
||||
)))
|
||||
)?))
|
||||
}
|
||||
|
||||
fn resolve_local_tokenizer_source(
|
||||
model_dir: &Path,
|
||||
tokenizer_config: &HfTokenizerConfig,
|
||||
tokenizer_config: &TokenizerConfig,
|
||||
) -> Result<TokenizerSource> {
|
||||
let tekken_path = local_file_if_exists(model_dir, "tekken.json");
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return Ok(TokenizerSource::Tekken(tekken_path));
|
||||
}
|
||||
|
||||
let tokenizer_path = local_file_if_exists(model_dir, "tokenizer.json")
|
||||
let tokenizer_path = local_file_if_exists(model_dir, "tekken.json")
|
||||
.or_else(|| local_file_if_exists(model_dir, "tokenizer.json"))
|
||||
.or_else(|| local_file_if_exists(model_dir, "tiktoken.model"))
|
||||
.or_else(|| discover_tiktoken_in_dir(model_dir))
|
||||
.ok_or_else(|| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"local model directory '{}' does not contain a supported tokenizer file \
|
||||
(tokenizer.json, tiktoken.model, or *.tiktoken)",
|
||||
model_dir.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_config.tokenizer_class.as_deref(),
|
||||
None,
|
||||
))
|
||||
resolve_tokenizer_source(tokenizer_path, tokenizer_config.tokenizer_class.as_deref())
|
||||
}
|
||||
|
||||
/// Choose the tokenizer.
|
||||
///
|
||||
/// Selection order:
|
||||
/// 1. `tekken.json` — Mistral native tokenizer (preferred over HF `tokenizer.json` because the HF
|
||||
/// version has a known regex bug for Mistral models).
|
||||
/// 2. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
|
||||
/// 3. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
|
||||
/// 1. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
|
||||
/// 2. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
|
||||
/// insensitive) trigger tiktoken loading from a sibling BPE file.
|
||||
/// 4. Default — `tokenizer.json` in HuggingFace format.
|
||||
/// 3. Default — `tokenizer.json` in HuggingFace format.
|
||||
fn resolve_tokenizer_source(
|
||||
tokenizer_path: PathBuf,
|
||||
tokenizer_class: Option<&str>,
|
||||
tekken_path: Option<PathBuf>,
|
||||
) -> TokenizerSource {
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return TokenizerSource::Tekken(tekken_path);
|
||||
}
|
||||
) -> Result<TokenizerSource> {
|
||||
let tokenizer = TokenizerSource::from_path(tokenizer_path)?;
|
||||
|
||||
if is_tiktoken_file(&tokenizer_path) {
|
||||
return TokenizerSource::Tiktoken(tokenizer_path);
|
||||
}
|
||||
|
||||
if tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
|
||||
&& let Some(dir) = tokenizer_path.parent()
|
||||
if let TokenizerSource::HuggingFace(path) = &tokenizer
|
||||
&& tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
|
||||
&& let Some(dir) = path.parent()
|
||||
&& let Some(tiktoken_path) = discover_tiktoken_in_dir(dir)
|
||||
{
|
||||
return TokenizerSource::Tiktoken(tiktoken_path);
|
||||
return Ok(TokenizerSource::Tiktoken(tiktoken_path));
|
||||
}
|
||||
|
||||
TokenizerSource::HuggingFace(tokenizer_path)
|
||||
Ok(tokenizer)
|
||||
}
|
||||
|
||||
/// Download `filename` only if it exists in `siblings`.
|
||||
@@ -286,7 +289,7 @@ async fn download_if_present(
|
||||
|
||||
async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result<PathBuf> {
|
||||
repo.get(filename).await.map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"failed to download '{filename}' for model '{model_id}': {}",
|
||||
error.as_report()
|
||||
))
|
||||
@@ -317,7 +320,7 @@ fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) ->
|
||||
}
|
||||
|
||||
/// Discover a tiktoken model file in a local directory.
|
||||
pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
|
||||
fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
|
||||
let tiktoken_model = dir.join("tiktoken.model");
|
||||
if tiktoken_model.exists() {
|
||||
return Some(tiktoken_model);
|
||||
@@ -337,7 +340,7 @@ pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf>
|
||||
}
|
||||
|
||||
/// Returns `true` if `path` points to a tiktoken-format file (by name).
|
||||
pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
|
||||
fn is_tiktoken_file(path: &std::path::Path) -> bool {
|
||||
path.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
|
||||
@@ -368,7 +371,6 @@ mod tests {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::tempdir;
|
||||
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
|
||||
|
||||
use super::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
@@ -399,61 +401,4 @@ mod tests {
|
||||
Some(dir.path().join("tokenizer_config.json"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "too slow for CI and requires network access to Hugging Face"]
|
||||
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
|
||||
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
|
||||
.await
|
||||
.expect("resolve real Kimi K2.5 model files");
|
||||
|
||||
let tokenizer_path = match &files.tokenizer {
|
||||
TokenizerSource::Tiktoken(path) => path.clone(),
|
||||
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
|
||||
};
|
||||
|
||||
for backend in [
|
||||
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
|
||||
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
|
||||
] {
|
||||
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
|
||||
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
|
||||
let tool_section_id = backend
|
||||
.token_to_id("<|tool_calls_section_begin|>")
|
||||
.expect("resolve tool call section marker");
|
||||
let contraction_heavy_text =
|
||||
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
|
||||
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(think_id, end_think_id, tool_section_id),
|
||||
(163606, 163607, 163595)
|
||||
);
|
||||
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
|
||||
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
|
||||
assert_eq!(
|
||||
backend.decode(&[tool_section_id], true).unwrap(),
|
||||
"<|tool_calls_section_begin|>"
|
||||
);
|
||||
|
||||
// This demonstrates that we're using Kimi's custom BPE pattern.
|
||||
// With CL100K this will be 23 tokens instead.
|
||||
assert_eq!(
|
||||
contraction_heavy_ids,
|
||||
vec![
|
||||
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
|
||||
1573, 56229, 13922, 13,
|
||||
]
|
||||
);
|
||||
assert_eq!(contraction_heavy_ids.len(), 17);
|
||||
assert_eq!(
|
||||
backend.decode(&contraction_heavy_ids, false).unwrap(),
|
||||
contraction_heavy_text
|
||||
);
|
||||
|
||||
// Special-looking text that is not actually registered should fail gracefully.
|
||||
assert_eq!(backend.token_to_id("◁think▷"), None);
|
||||
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -238,6 +238,9 @@ mod tests {
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[
|
||||
r#"tool parser parsing failed: near "tool<|tool▁sep|>get_weather\n```json\n{}": "#
|
||||
]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -242,6 +242,7 @@ mod tests {
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "<|tool▁sep|>{}": "#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -540,8 +540,10 @@ mod tests {
|
||||
.parse_chunk(r#"<tool_call>{"name":"f","arguments":42}</tool_call>"#)
|
||||
.unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: invalid Granite4 arguments"]
|
||||
.assert_eq(&error.to_report_string());
|
||||
expect![[
|
||||
r#"tool parser parsing failed: near "42}</tool_call>": invalid Granite4 arguments"#
|
||||
]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -335,7 +335,7 @@ mod tests {
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid InternLM2
|
||||
tool parser parsing failed: near "{\"name\":\"get_weather\",\"params\":{\"location\":\"Tokyo\"}}<|action_end|>": invalid InternLM2
|
||||
expected `parameters`, `arguments`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
@@ -336,7 +336,7 @@ mod tests {
|
||||
.unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid Llama JSON
|
||||
tool parser parsing failed: near "{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Tokyo\"}}": invalid Llama JSON
|
||||
expected `parameters`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
@@ -474,7 +474,7 @@ mod tests {
|
||||
let error = parser.parse_chunk(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid Llama JSON
|
||||
tool parser parsing failed: near "{\"parameters\":{},\"name\":\"get_weather\"}": invalid Llama JSON
|
||||
expected `name`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
@@ -489,7 +489,7 @@ mod tests {
|
||||
))
|
||||
.unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: invalid Llama JSON"]
|
||||
expect![[r#"tool parser parsing failed: near " trailing": invalid Llama JSON"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -240,7 +240,7 @@ mod tests {
|
||||
.unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid Mistral
|
||||
tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}]": invalid Mistral
|
||||
expected `name`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
@@ -280,7 +280,7 @@ mod tests {
|
||||
.unwrap_err();
|
||||
|
||||
expect![[r#"
|
||||
tool parser parsing failed: invalid Qwen XML
|
||||
tool parser parsing failed: near "{\"arguments\":{},\"name\":\"get_weather\"}\n</tool_call>": invalid Qwen XML
|
||||
expected `name`"#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
@@ -594,6 +594,9 @@ mod tests {
|
||||
|
||||
let error = parser.parse_chunk(&input).unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[
|
||||
r#"tool parser parsing failed: near "get_weather<|tool_call_argument_begin|>{}": "#
|
||||
]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -593,6 +593,7 @@ mod tests {
|
||||
let mut parser = MinimaxM2ToolParser::new(&test_tools());
|
||||
let error = parser.parse_chunk("<minimax:tool_call><bad></minimax:tool_call>").unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "<bad></minimax:tool_call>": "#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -878,7 +878,10 @@ mod tests {
|
||||
))
|
||||
.unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[
|
||||
r#"tool parser parsing failed: near "]<]minimax[>[<bad>]<]minimax[>[</tool_call>": "#
|
||||
]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -685,7 +685,8 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let error = parser.parse_chunk("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "\n<bad>\n</tool_call>": "#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -697,7 +698,7 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "\n<function=get_weather>\n<parameter=location>SF</function>\n</tool_call>": "#]].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -395,9 +395,9 @@ pub fn parse_buffered_event<E>(
|
||||
Ok(event) => event,
|
||||
Err(ErrMode::Incomplete(_)) => return Ok(None),
|
||||
Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => {
|
||||
// TODO: enrich context for error reporting
|
||||
let snippet = buffer.char_indices().nth(80).map_or(buffer, |(i, _)| &buffer[..i]);
|
||||
return Err(ToolParserError::ParsingFailed {
|
||||
message: e.to_string(),
|
||||
message: format!("near {snippet:?}: {e}"),
|
||||
});
|
||||
}
|
||||
};
|
||||
@@ -423,8 +423,9 @@ mod tests {
|
||||
use winnow::stream::{Offset, Partial, Stream};
|
||||
|
||||
use super::{
|
||||
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len,
|
||||
safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker,
|
||||
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, parse_buffered_event,
|
||||
partial_prefix_len, safe_text_len, safe_text_len_mul, take_json_object, take_json_string,
|
||||
take_until_marker,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -832,4 +833,27 @@ mod tests {
|
||||
|
||||
assert!(matches!(error, ErrMode::Incomplete(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_buffered_event_error_includes_input_snippet() {
|
||||
let result = parse_buffered_event(" {\"x\":1}", |input| {
|
||||
take_json_object(input, &mut JsonObjectScanState::default())
|
||||
});
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("near \""), "error must include snippet");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_buffered_event_error_truncates_long_input() {
|
||||
let long_input = format!(" {}", "x".repeat(100));
|
||||
let result = parse_buffered_event(&long_input, |input| {
|
||||
take_json_object(input, &mut JsonObjectScanState::default())
|
||||
});
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("near \""), "error must include snippet");
|
||||
assert!(
|
||||
!err.contains(&long_input),
|
||||
"snippet must be truncated for long input"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -199,6 +199,7 @@ mod tests {
|
||||
use super::prepare_completion_request;
|
||||
use crate::lora::LoraModelResolution;
|
||||
use crate::routes::openai::completions::types::CompletionRequest;
|
||||
use crate::routes::openai::utils::types::Normalizable;
|
||||
use crate::utils::{ResolvedRequestContext, resolve_request_context};
|
||||
|
||||
fn request_context(headers: &HeaderMap, request_id: Option<&str>) -> ResolvedRequestContext {
|
||||
@@ -249,6 +250,28 @@ mod tests {
|
||||
assert!(request.ignore_eos);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_coerces_null_max_tokens_to_default() {
|
||||
// An absent `max_tokens` already gets the serde default.
|
||||
let absent: CompletionRequest =
|
||||
serde_json::from_value(base_request_json()).expect("parse request");
|
||||
assert_eq!(absent.max_tokens, Some(16));
|
||||
|
||||
// An explicit `null` deserializes to `None`, bypassing the default;
|
||||
// `normalize` must coerce it back to match Python vLLM.
|
||||
let mut request: CompletionRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hello",
|
||||
"stream": true,
|
||||
"max_tokens": null
|
||||
}))
|
||||
.expect("parse request");
|
||||
assert_eq!(request.max_tokens, None);
|
||||
|
||||
request.normalize();
|
||||
assert_eq!(request.max_tokens, Some(16));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_completion_request_maps_sampling_fields() {
|
||||
let request: CompletionRequest = serde_json::from_value(json!({
|
||||
|
||||
@@ -179,7 +179,17 @@ pub struct CompletionRequest {
|
||||
pub other: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl Normalizable for CompletionRequest {}
|
||||
impl Normalizable for CompletionRequest {
|
||||
/// Normalize the request by applying defaults.
|
||||
fn normalize(&mut self) {
|
||||
// An explicit `"max_tokens": null` deserializes to `None`, bypassing the
|
||||
// serde field default. Coerce it back to the default so it behaves like
|
||||
// an absent field, matching Python vLLM's `normalize_null_max_tokens`.
|
||||
if self.max_tokens.is_none() {
|
||||
self.max_tokens = default_completion_max_tokens();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `CompletionResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
|
||||
@@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true
|
||||
easy-ext.workspace = true
|
||||
enum-as-inner.workspace = true
|
||||
futures.workspace = true
|
||||
hf-hub.workspace = true
|
||||
itertools.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
@@ -22,6 +21,7 @@ tracing.workspace = true
|
||||
trait-set.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-model-files.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
@@ -1,24 +1,18 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_model_files::read_json_file;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct HfTokenizerConfig {
|
||||
pub struct TokenizerConfig {
|
||||
#[serde(flatten)]
|
||||
pub special_tokens: HfSpecialTokens,
|
||||
pub special_tokens: SpecialTokens,
|
||||
pub chat_template: Option<String>,
|
||||
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
|
||||
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
|
||||
/// containing "Tiktoken" which can be used as a hint for backend
|
||||
/// selection.
|
||||
pub tokenizer_class: Option<String>,
|
||||
}
|
||||
|
||||
/// Hugging Face named special tokens may be serialized as a string or an
|
||||
@@ -61,14 +55,14 @@ impl NamedSpecialToken {
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct HfSpecialTokens {
|
||||
pub struct SpecialTokens {
|
||||
pub bos_token: Option<NamedSpecialToken>,
|
||||
pub eos_token: Option<NamedSpecialToken>,
|
||||
pub unk_token: Option<NamedSpecialToken>,
|
||||
pub pad_token: Option<NamedSpecialToken>,
|
||||
}
|
||||
|
||||
impl HfSpecialTokens {
|
||||
impl SpecialTokens {
|
||||
/// Returns true if we don't discover any special tokens in the config.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.bos_token.is_none()
|
||||
@@ -234,42 +228,19 @@ impl ModelConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the tokenizer-side EOS metadata if a config file is present.
|
||||
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<HfTokenizerConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
|
||||
/// Load the generation-side EOS metadata if a config file is present.
|
||||
pub(super) fn load_generation_config(path: Option<&Path>) -> Result<GenerationConfig> {
|
||||
read_json_file(path)
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
/// Load the tokenizer-side EOS metadata if a config file is present.
|
||||
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
/// Load the model-side config (`config.json`) if present.
|
||||
pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
|
||||
fn read_json_file<T>(path: Option<&Path>) -> Result<T>
|
||||
where
|
||||
T: for<'de> Deserialize<'de> + Default,
|
||||
{
|
||||
let Some(path) = path else {
|
||||
return Ok(T::default());
|
||||
};
|
||||
let content = fs::read_to_string(path).map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
"failed to read {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})?;
|
||||
serde_json::from_str(&content).map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
"failed to parse {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
mod config;
|
||||
mod model_files;
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::Arc;
|
||||
@@ -9,12 +8,12 @@ use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, Tiktok
|
||||
|
||||
use self::config::{GenerationConfig, load_generation_config};
|
||||
pub use self::config::{
|
||||
HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
|
||||
ModelConfig, NamedSpecialToken, SpecialTokens, TokenizerConfig, load_model_config,
|
||||
load_tokenizer_config,
|
||||
};
|
||||
pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
use crate::backend::{SamplingHints, TextBackend};
|
||||
use crate::error::Result;
|
||||
pub use vllm_model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
fn load_tokenizer(tokenizer: &TokenizerSource) -> Result<DynTokenizer> {
|
||||
match tokenizer {
|
||||
@@ -125,3 +124,67 @@ impl TextBackend for HfTextBackend {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
|
||||
|
||||
use super::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "too slow for CI and requires network access to Hugging Face"]
|
||||
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
|
||||
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
|
||||
.await
|
||||
.expect("resolve real Kimi K2.5 model files");
|
||||
|
||||
let tokenizer_path = match &files.tokenizer {
|
||||
TokenizerSource::Tiktoken(path) => path.clone(),
|
||||
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
|
||||
};
|
||||
|
||||
for backend in [
|
||||
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
|
||||
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
|
||||
] {
|
||||
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
|
||||
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
|
||||
let tool_section_id = backend
|
||||
.token_to_id("<|tool_calls_section_begin|>")
|
||||
.expect("resolve tool call section marker");
|
||||
let contraction_heavy_text =
|
||||
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
|
||||
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(think_id, end_think_id, tool_section_id),
|
||||
(163606, 163607, 163595)
|
||||
);
|
||||
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
|
||||
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
|
||||
assert_eq!(
|
||||
backend.decode(&[tool_section_id], true).unwrap(),
|
||||
"<|tool_calls_section_begin|>"
|
||||
);
|
||||
|
||||
// This demonstrates that we're using Kimi's custom BPE pattern.
|
||||
// With CL100K this will be 23 tokens instead.
|
||||
assert_eq!(
|
||||
contraction_heavy_ids,
|
||||
vec![
|
||||
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
|
||||
1573, 56229, 13922, 13,
|
||||
]
|
||||
);
|
||||
assert_eq!(contraction_heavy_ids.len(), 17);
|
||||
assert_eq!(
|
||||
backend.decode(&contraction_heavy_ids, false).unwrap(),
|
||||
contraction_heavy_text
|
||||
);
|
||||
|
||||
// Special-looking text that is not actually registered should fail gracefully.
|
||||
assert_eq!(backend.token_to_id("◁think▷"), None);
|
||||
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use thiserror::Error;
|
||||
use vllm_engine_core_client::Error as EngineCoreError;
|
||||
use vllm_llm::Error as LlmError;
|
||||
use vllm_model_files::Error as ModelFilesError;
|
||||
|
||||
pub use crate::lower::logprobs::LogprobsError;
|
||||
pub use crate::lower::token_ids::TokenIdsError;
|
||||
@@ -20,6 +21,8 @@ pub enum Error {
|
||||
Logprobs(#[from] LogprobsError),
|
||||
#[error(transparent)]
|
||||
TokenIds(#[from] TokenIdsError),
|
||||
#[error(transparent)]
|
||||
ModelFiles(#[from] ModelFilesError),
|
||||
#[error(
|
||||
"`min_tokens` must be less than or equal to `max_tokens`, \
|
||||
got min_tokens={min_tokens}, max_tokens={max_tokens}"
|
||||
|
||||
@@ -122,7 +122,7 @@ class QKNormRoPETestModel(torch.nn.Module):
|
||||
@pytest.mark.parametrize("enable_rope_custom_op", [True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="Only test on cuda and rocm platform",
|
||||
)
|
||||
def test_qk_norm_rope_fusion(
|
||||
@@ -136,7 +136,7 @@ def test_qk_norm_rope_fusion(
|
||||
if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
|
||||
pytest.skip("fused_qk_norm_rope custom op not available")
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_device(current_platform.device_type)
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
import typing
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import ensure_current_vllm_config
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
|
||||
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
|
||||
from vllm.distributed.device_communicators.pynccl_allocator import (
|
||||
get_nccl_mem_pool,
|
||||
is_symmetric_memory_enabled,
|
||||
)
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
torch.manual_seed(42)
|
||||
random.seed(44)
|
||||
|
||||
test_size_elements = 4 * 1024 * 1024
|
||||
|
||||
|
||||
def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{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",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
pynccl_comm = cuda_communicator.pynccl_comm
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory allreduce is disabled.")
|
||||
|
||||
register_nccl_symmetric_ops(pynccl_comm)
|
||||
input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
|
||||
input_clone = input.clone()
|
||||
output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
|
||||
assert output is not None
|
||||
|
||||
group = get_tp_group().device_group
|
||||
dist.all_reduce(input_clone, group=group)
|
||||
torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
# Enable SymmMemCommunicator
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{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": "12346",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory is disabled.")
|
||||
|
||||
per_rank_size = test_size_elements // world_size
|
||||
input_tensor = torch.randint(
|
||||
1, 23, (per_rank_size,), dtype=dtype, device=device
|
||||
)
|
||||
output = cuda_communicator.all_gatherv(input_tensor, dim=0)
|
||||
|
||||
group = get_tp_group().device_group
|
||||
expected = torch.empty(test_size_elements, dtype=dtype, device=device)
|
||||
dist.all_gather_into_tensor(expected, input_tensor, group=group)
|
||||
torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCL symmetric memory is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_allgather(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{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": "12347",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory is disabled.")
|
||||
|
||||
per_rank_size = test_size_elements // world_size
|
||||
input_tensor = torch.randint(
|
||||
1, 23, (test_size_elements,), dtype=dtype, device=device
|
||||
)
|
||||
input_clone = input_tensor.clone()
|
||||
output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
|
||||
|
||||
group = get_tp_group().device_group
|
||||
expected = torch.empty(per_rank_size, dtype=dtype, device=device)
|
||||
dist.reduce_scatter_tensor(expected, input_clone, group=group)
|
||||
torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCL symmetric memory is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_reduce_scatter_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
@@ -1,96 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
import typing
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import ensure_current_vllm_config
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
|
||||
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
|
||||
from vllm.distributed.device_communicators.pynccl_allocator import (
|
||||
get_nccl_mem_pool,
|
||||
is_symmetric_memory_enabled,
|
||||
)
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
torch.manual_seed(42)
|
||||
random.seed(44)
|
||||
|
||||
test_size_elements = 4 * 1024 * 1024
|
||||
|
||||
|
||||
def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{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",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
pynccl_comm = cuda_communicator.pynccl_comm
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory allreduce is disabled.")
|
||||
|
||||
register_nccl_symmetric_ops(pynccl_comm)
|
||||
input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
|
||||
input_clone = input.clone()
|
||||
output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
|
||||
assert output is not None
|
||||
|
||||
group = get_tp_group().device_group
|
||||
dist.all_reduce(input_clone, group=group)
|
||||
torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
# Enable SymmMemCommunicator
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
@@ -115,7 +115,6 @@ TEXT_GENERATION_MODELS = {
|
||||
"google/gemma-1.1-2b-it": PPTestSettings.fast(),
|
||||
"google/gemma-2-9b": PPTestSettings.fast(),
|
||||
"gpt2": PPTestSettings.fast(),
|
||||
"bigcode/starcoder": PPTestSettings.fast(),
|
||||
"EleutherAI/gpt-j-6b": PPTestSettings.fast(),
|
||||
"EleutherAI/pythia-1.4b": PPTestSettings.fast(),
|
||||
"ibm/PowerLM-3b": PPTestSettings.fast(),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -251,7 +251,13 @@ def _validate_field_consistency(events: list) -> None:
|
||||
"response.reasoning_part.added",
|
||||
):
|
||||
_assert_item_fields(event, etype, active_item_id, active_output_index)
|
||||
active_content_index = getattr(event, "content_index", None)
|
||||
content_index = getattr(event, "content_index", None)
|
||||
if active_content_index is None:
|
||||
assert content_index == 0, (
|
||||
f"{etype} for a new item must start at content_index 0, "
|
||||
f"got {content_index}"
|
||||
)
|
||||
active_content_index = content_index
|
||||
continue
|
||||
|
||||
# --- all other item-level events --------------------------
|
||||
|
||||
@@ -454,6 +454,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool):
|
||||
if event.type == "response.output_item.added":
|
||||
assert event.item.id != current_item_id
|
||||
current_item_id = event.item.id
|
||||
current_content_index = -1
|
||||
elif event.type in [
|
||||
"response.output_text.delta",
|
||||
"response.reasoning_text.delta",
|
||||
@@ -465,7 +466,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool):
|
||||
"response.content_part.added",
|
||||
"response.reasoning_part.added",
|
||||
]:
|
||||
assert event.content_index != current_content_index
|
||||
assert event.content_index == current_content_index + 1
|
||||
current_content_index = event.content_index
|
||||
elif event.type in [
|
||||
"response.output_text.delta",
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for vllm.entrypoints.openai.responses.harmony."""
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import (
|
||||
ResponseFunctionToolCall,
|
||||
ResponseFunctionWebSearch,
|
||||
ResponseOutputMessage,
|
||||
ResponseReasoningItem,
|
||||
)
|
||||
@@ -12,7 +14,6 @@ from openai_harmony import Author, Message, Role, TextContent
|
||||
|
||||
from vllm.entrypoints.openai.responses.harmony import (
|
||||
harmony_to_response_output,
|
||||
parser_state_to_response_output,
|
||||
response_previous_input_to_harmony,
|
||||
)
|
||||
|
||||
@@ -95,7 +96,8 @@ class TestResponsePreviousInputToHarmony:
|
||||
class TestHarmonyToResponseOutput:
|
||||
"""Tests for harmony_to_response_output function."""
|
||||
|
||||
def test_commentary_with_no_recipient_creates_message(self):
|
||||
@pytest.mark.parametrize("incomplete", [False, True])
|
||||
def test_commentary_with_no_recipient_creates_message(self, incomplete):
|
||||
"""Test that commentary with recipient=None (preambles) creates message items.
|
||||
|
||||
Per Harmony format, preambles are intended to be shown to end-users,
|
||||
@@ -108,13 +110,15 @@ class TestHarmonyToResponseOutput:
|
||||
message = message.with_channel("commentary")
|
||||
# recipient is None by default, representing a preamble
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(
|
||||
message, frozenset(), incomplete=incomplete
|
||||
)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseOutputMessage)
|
||||
assert output_items[0].type == "message"
|
||||
assert output_items[0].role == "assistant"
|
||||
assert output_items[0].status == "completed"
|
||||
assert output_items[0].status == ("incomplete" if incomplete else "completed")
|
||||
assert len(output_items[0].content) == 1
|
||||
assert output_items[0].content[0].type == "output_text"
|
||||
assert (
|
||||
@@ -122,82 +126,148 @@ class TestHarmonyToResponseOutput:
|
||||
== "I will now search for the weather information."
|
||||
)
|
||||
|
||||
def test_commentary_with_function_recipient_creates_function_call(self):
|
||||
"""Test commentary with recipient='functions.X' creates function calls."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
|
||||
@pytest.mark.parametrize(
|
||||
("recipient", "fn_names", "expected_name"),
|
||||
[
|
||||
("functions.get_weather", frozenset(), "get_weather"),
|
||||
("get_weather", frozenset({"get_weather"}), "get_weather"),
|
||||
("math.sum", frozenset({"math.sum"}), "math.sum"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("incomplete", [False, True])
|
||||
def test_function_recipient_creates_function_call(
|
||||
self, channel, recipient, fn_names, expected_name, incomplete
|
||||
):
|
||||
"""Function recipients create function calls across channels."""
|
||||
content = '{"location": "San Francisco"}'
|
||||
if recipient == "math.sum":
|
||||
content = '{"a": 1, "b": 2}'
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, content)
|
||||
message = message.with_channel(channel)
|
||||
message = message.with_recipient(recipient)
|
||||
|
||||
output_items = harmony_to_response_output(
|
||||
message, fn_names, incomplete=incomplete
|
||||
)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].type == "function_call"
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert (
|
||||
output_items[0].arguments
|
||||
== '{"location": "San Francisco", "units": "celsius"}'
|
||||
)
|
||||
assert output_items[0].name == expected_name
|
||||
assert output_items[0].arguments == content
|
||||
assert output_items[0].call_id.startswith("call_")
|
||||
assert output_items[0].id.startswith("fc_")
|
||||
assert output_items[0].status == ("incomplete" if incomplete else "completed")
|
||||
|
||||
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
|
||||
@pytest.mark.parametrize(
|
||||
("recipient", "content"),
|
||||
[
|
||||
("python", "import numpy as np\nprint(np.array([1, 2, 3]))"),
|
||||
("browser", "Navigating to the specified URL"),
|
||||
("container", "Running command in container"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("incomplete", [False, True])
|
||||
def test_builtin_recipient_creates_reasoning(
|
||||
self, channel, recipient, content, incomplete
|
||||
):
|
||||
"""Built-in recipients create reasoning items."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, content)
|
||||
message = message.with_channel(channel)
|
||||
message = message.with_recipient(recipient)
|
||||
|
||||
output_items = harmony_to_response_output(
|
||||
message, frozenset(), incomplete=incomplete
|
||||
)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == content
|
||||
assert output_items[0].status is None
|
||||
|
||||
@pytest.mark.parametrize("channel", ["commentary", "comment", "analysis", "final"])
|
||||
@pytest.mark.parametrize(
|
||||
("recipient", "fn_names", "content", "expected_name", "expected_server_label"),
|
||||
[
|
||||
(
|
||||
"get_weather",
|
||||
frozenset(),
|
||||
'{"arg": "value"}',
|
||||
"get_weather",
|
||||
"get_weather",
|
||||
),
|
||||
(
|
||||
"not_get_weather",
|
||||
frozenset({"get_weather"}),
|
||||
'{"arg": "value"}',
|
||||
"not_get_weather",
|
||||
"not_get_weather",
|
||||
),
|
||||
("repo_browser.list", frozenset(), '{"cmd": "ls"}', "list", "repo_browser"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("incomplete", [False, True])
|
||||
def test_non_function_non_builtin_recipient_creates_mcp_call(
|
||||
self,
|
||||
channel,
|
||||
recipient,
|
||||
fn_names,
|
||||
content,
|
||||
expected_name,
|
||||
expected_server_label,
|
||||
incomplete,
|
||||
):
|
||||
"""Non-function, non-built-in recipients create MCP calls."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, content)
|
||||
message = message.with_channel(channel)
|
||||
message = message.with_recipient(recipient)
|
||||
|
||||
output_items = harmony_to_response_output(
|
||||
message, fn_names, incomplete=incomplete
|
||||
)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == expected_name
|
||||
assert output_items[0].server_label == expected_server_label
|
||||
assert output_items[0].arguments == content
|
||||
assert output_items[0].status == ("incomplete" if incomplete else "completed")
|
||||
|
||||
@pytest.mark.parametrize("incomplete", [False, True])
|
||||
def test_browser_search_recipient_respects_incomplete(self, incomplete):
|
||||
"""browser.search emits a web search call unless the item is incomplete."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, '{"query": "weather in San Francisco"}'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("browser.search")
|
||||
|
||||
output_items = harmony_to_response_output(
|
||||
message, frozenset(), incomplete=incomplete
|
||||
)
|
||||
|
||||
if incomplete:
|
||||
assert output_items == []
|
||||
return
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionWebSearch)
|
||||
assert output_items[0].type == "web_search_call"
|
||||
assert output_items[0].status == "completed"
|
||||
|
||||
def test_commentary_with_python_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='python' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("python")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert (
|
||||
output_items[0].content[0].text
|
||||
== "import numpy as np\nprint(np.array([1, 2, 3]))"
|
||||
)
|
||||
|
||||
def test_commentary_with_browser_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='browser' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Navigating to the specified URL"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("browser")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Navigating to the specified URL"
|
||||
|
||||
def test_commentary_with_container_recipient_creates_reasoning(self):
|
||||
"""Test that commentary with recipient='container' creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, "Running command in container"
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("container")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
assert output_items[0].type == "reasoning"
|
||||
assert output_items[0].content[0].text == "Running command in container"
|
||||
assert output_items[0].action.type == "search"
|
||||
assert output_items[0].action.query == "cursor:weather in San Francisco"
|
||||
|
||||
def test_commentary_with_empty_content_and_no_recipient(self):
|
||||
"""Test edge case: empty commentary with recipient=None."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, "")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseOutputMessage)
|
||||
@@ -212,7 +282,7 @@ class TestHarmonyToResponseOutput:
|
||||
message = Message.from_role_and_contents(Role.ASSISTANT, contents)
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
# _parse_final_message returns single ResponseOutputMessage with
|
||||
# multiple contents
|
||||
@@ -232,7 +302,7 @@ class TestHarmonyToResponseOutput:
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 2
|
||||
assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items)
|
||||
@@ -241,21 +311,6 @@ class TestHarmonyToResponseOutput:
|
||||
assert output_items[0].arguments == '{"location": "San Francisco"}'
|
||||
assert output_items[1].arguments == '{"location": "New York"}'
|
||||
|
||||
def test_commentary_with_unknown_recipient_creates_mcp_call(self):
|
||||
"""Test that commentary with unknown recipient creates MCP call."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("custom_tool")
|
||||
|
||||
fn_names = frozenset({"other_tool"})
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "custom_tool"
|
||||
assert output_items[0].server_label == "custom_tool"
|
||||
|
||||
def test_analysis_channel_creates_reasoning(self):
|
||||
"""Test that analysis channel creates reasoning items."""
|
||||
message = Message.from_role_and_content(
|
||||
@@ -263,7 +318,7 @@ class TestHarmonyToResponseOutput:
|
||||
)
|
||||
message = message.with_channel("analysis")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseReasoningItem)
|
||||
@@ -283,352 +338,6 @@ class TestHarmonyToResponseOutput:
|
||||
"The weather is sunny, 72°F",
|
||||
)
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 0
|
||||
|
||||
|
||||
class TestHarmonyToResponseOutputWithFunctionToolNames:
|
||||
"""Tests for bare function name handling with function_tool_names."""
|
||||
|
||||
def test_bare_name_creates_function_call_when_in_tool_names(self):
|
||||
"""Bare function name matching a known tool creates function call."""
|
||||
message = Message.from_role_and_content(
|
||||
Role.ASSISTANT, '{"location": "San Francisco"}'
|
||||
)
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("get_weather")
|
||||
|
||||
fn_names = frozenset({"get_weather"})
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].type == "function_call"
|
||||
assert output_items[0].name == "get_weather"
|
||||
assert output_items[0].arguments == '{"location": "San Francisco"}'
|
||||
|
||||
def test_bare_name_creates_mcp_call_when_not_in_tool_names(self):
|
||||
"""Bare name not matching any known tool creates MCP call."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("custom_tool")
|
||||
|
||||
fn_names = frozenset({"get_weather"})
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
|
||||
def test_dotted_function_name_creates_function_call(self):
|
||||
"""Dotted function name in tool names creates function call."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"a": 1, "b": 2}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("math.sum")
|
||||
|
||||
fn_names = frozenset({"math.sum"})
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].name == "math.sum"
|
||||
|
||||
def test_empty_tool_names_defaults_to_mcp(self):
|
||||
"""With empty function_tool_names, bare names become MCP calls."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
|
||||
def test_prefixed_name_always_function_call(self):
|
||||
"""functions. prefix always creates function call even with empty tool names."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
message = message.with_channel("commentary")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message, frozenset())
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].name == "get_weather"
|
||||
|
||||
|
||||
class TestParserStateWithFunctionToolNames:
|
||||
"""Tests for parser_state_to_response_output with function_tool_names."""
|
||||
|
||||
def test_bare_name_creates_function_call(self):
|
||||
from unittest.mock import Mock
|
||||
|
||||
parser = Mock()
|
||||
parser.current_content = '{"arg": "value"}'
|
||||
parser.current_role = Role.ASSISTANT
|
||||
parser.current_channel = "commentary"
|
||||
parser.current_recipient = "get_weather"
|
||||
|
||||
fn_names = frozenset({"get_weather"})
|
||||
items = parser_state_to_response_output(parser, fn_names)
|
||||
|
||||
assert len(items) == 1
|
||||
assert isinstance(items[0], ResponseFunctionToolCall)
|
||||
assert items[0].name == "get_weather"
|
||||
assert items[0].status == "in_progress"
|
||||
|
||||
def test_bare_name_creates_mcp_when_not_in_tool_names(self):
|
||||
from unittest.mock import Mock
|
||||
|
||||
parser = Mock()
|
||||
parser.current_content = '{"arg": "value"}'
|
||||
parser.current_role = Role.ASSISTANT
|
||||
parser.current_channel = "commentary"
|
||||
parser.current_recipient = "unknown_tool"
|
||||
|
||||
fn_names = frozenset({"get_weather"})
|
||||
items = parser_state_to_response_output(parser, fn_names)
|
||||
|
||||
assert len(items) == 1
|
||||
assert isinstance(items[0], McpCall)
|
||||
assert items[0].name == "unknown_tool"
|
||||
|
||||
|
||||
class TestToolCallsOnNonStandardChannels:
|
||||
"""Tests verifying tool calls are detected regardless of channel."""
|
||||
|
||||
def test_function_call_on_comment_channel(self):
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"query": "weather"}')
|
||||
message = message.with_channel("comment")
|
||||
message = message.with_recipient("functions.get_weather")
|
||||
|
||||
output_items = harmony_to_response_output(message)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].type == "function_call"
|
||||
assert output_items[0].name == "get_weather"
|
||||
|
||||
def test_bare_function_on_comment_channel(self):
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"query": "weather"}')
|
||||
message = message.with_channel("comment")
|
||||
message = message.with_recipient("get_weather")
|
||||
|
||||
fn_names = frozenset({"get_weather"})
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], ResponseFunctionToolCall)
|
||||
assert output_items[0].name == "get_weather"
|
||||
|
||||
def test_parser_state_comment_channel_function(self):
|
||||
from unittest.mock import Mock
|
||||
|
||||
parser = Mock()
|
||||
parser.current_content = '{"arg": "value"}'
|
||||
parser.current_role = Role.ASSISTANT
|
||||
parser.current_channel = "comment"
|
||||
parser.current_recipient = "functions.get_weather"
|
||||
|
||||
items = parser_state_to_response_output(parser)
|
||||
|
||||
assert len(items) == 1
|
||||
assert isinstance(items[0], ResponseFunctionToolCall)
|
||||
assert items[0].name == "get_weather"
|
||||
|
||||
def test_parser_state_comment_channel_mcp(self):
|
||||
from unittest.mock import Mock
|
||||
|
||||
parser = Mock()
|
||||
parser.current_content = '{"arg": "value"}'
|
||||
parser.current_role = Role.ASSISTANT
|
||||
parser.current_channel = "comment"
|
||||
parser.current_recipient = "mcp.server.tool"
|
||||
|
||||
fn_names: frozenset[str] = frozenset()
|
||||
items = parser_state_to_response_output(parser, fn_names)
|
||||
|
||||
assert len(items) == 1
|
||||
assert isinstance(items[0], McpCall)
|
||||
|
||||
|
||||
def test_parse_mcp_call_basic() -> None:
|
||||
"""Test that MCP calls are parsed with correct type and server_label."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}')
|
||||
message = message.with_recipient("filesystem")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
fn_names: frozenset[str] = frozenset()
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].type == "mcp_call"
|
||||
assert output_items[0].name == "filesystem"
|
||||
assert output_items[0].server_label == "filesystem"
|
||||
assert output_items[0].arguments == '{"path": "/tmp"}'
|
||||
assert output_items[0].status == "completed"
|
||||
|
||||
|
||||
def test_parse_mcp_call_dotted_recipient() -> None:
|
||||
"""Test that dotted recipients extract the tool name correctly."""
|
||||
message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}')
|
||||
message = message.with_recipient("repo_browser.list")
|
||||
message = message.with_channel("commentary")
|
||||
|
||||
fn_names: frozenset[str] = frozenset()
|
||||
output_items = harmony_to_response_output(message, fn_names)
|
||||
|
||||
assert len(output_items) == 1
|
||||
assert isinstance(output_items[0], McpCall)
|
||||
assert output_items[0].name == "list"
|
||||
assert output_items[0].server_label == "repo_browser"
|
||||
|
||||
|
||||
def test_mcp_vs_function_call() -> None:
|
||||
"""Test that function calls are not parsed as MCP calls."""
|
||||
func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}')
|
||||
func_message = func_message.with_recipient("functions.my_tool")
|
||||
func_message = func_message.with_channel("commentary")
|
||||
|
||||
func_items = harmony_to_response_output(func_message)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
|
||||
|
||||
def test_mcp_vs_builtin_tools() -> None:
|
||||
"""Test that built-in tools (python, container) are not parsed as MCP calls."""
|
||||
# Test python (built-in tool) - should be reasoning, not MCP
|
||||
python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')")
|
||||
python_message = python_message.with_recipient("python")
|
||||
python_message = python_message.with_channel("commentary")
|
||||
|
||||
python_items = harmony_to_response_output(python_message)
|
||||
|
||||
assert len(python_items) == 1
|
||||
assert not isinstance(python_items[0], McpCall)
|
||||
assert python_items[0].type == "reasoning"
|
||||
|
||||
|
||||
def test_parser_state_to_response_output_commentary_channel() -> None:
|
||||
"""Test parser_state_to_response_output with commentary
|
||||
channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
# Test 1: functions.* recipient -> should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "commentary"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parser_state_to_response_output(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) -> should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"path": "/tmp"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "commentary"
|
||||
parser_mcp.current_recipient = "filesystem"
|
||||
|
||||
fn_names: frozenset[str] = frozenset()
|
||||
mcp_items = parser_state_to_response_output(parser_mcp, fn_names)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "filesystem"
|
||||
assert mcp_items[0].server_label == "filesystem"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (python)
|
||||
# should NOT return MCP call, returns reasoning (internal tool interaction)
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "print('hello')"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "commentary"
|
||||
parser_builtin.current_recipient = "python"
|
||||
|
||||
builtin_items = parser_state_to_response_output(parser_builtin)
|
||||
|
||||
# Built-in tools explicitly return reasoning
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
|
||||
# Test 4: No recipient (preamble) → should return message, not reasoning
|
||||
parser_preamble = Mock()
|
||||
parser_preamble.current_content = "I'll search for that information now."
|
||||
parser_preamble.current_role = Role.ASSISTANT
|
||||
parser_preamble.current_channel = "commentary"
|
||||
parser_preamble.current_recipient = None
|
||||
|
||||
preamble_items = parser_state_to_response_output(parser_preamble)
|
||||
|
||||
assert len(preamble_items) == 1
|
||||
assert isinstance(preamble_items[0], ResponseOutputMessage)
|
||||
assert preamble_items[0].type == "message"
|
||||
assert preamble_items[0].content[0].text == "I'll search for that information now."
|
||||
assert preamble_items[0].status == "incomplete" # streaming
|
||||
|
||||
|
||||
def test_parser_state_to_response_output_analysis_channel() -> None:
|
||||
"""Test parser_state_to_response_output with analysis
|
||||
channel and various recipients."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
# Test 1: functions.* recipient -> should return function tool call
|
||||
parser_func = Mock()
|
||||
parser_func.current_content = '{"arg": "value"}'
|
||||
parser_func.current_role = Role.ASSISTANT
|
||||
parser_func.current_channel = "analysis"
|
||||
parser_func.current_recipient = "functions.my_tool"
|
||||
|
||||
func_items = parser_state_to_response_output(parser_func)
|
||||
|
||||
assert len(func_items) == 1
|
||||
assert not isinstance(func_items[0], McpCall)
|
||||
assert func_items[0].type == "function_call"
|
||||
assert func_items[0].name == "my_tool"
|
||||
assert func_items[0].status == "in_progress"
|
||||
|
||||
# Test 2: MCP tool (not builtin) -> should return MCP call
|
||||
parser_mcp = Mock()
|
||||
parser_mcp.current_content = '{"query": "test"}'
|
||||
parser_mcp.current_role = Role.ASSISTANT
|
||||
parser_mcp.current_channel = "analysis"
|
||||
parser_mcp.current_recipient = "database"
|
||||
|
||||
fn_names: frozenset[str] = frozenset()
|
||||
mcp_items = parser_state_to_response_output(parser_mcp, fn_names)
|
||||
|
||||
assert len(mcp_items) == 1
|
||||
assert isinstance(mcp_items[0], McpCall)
|
||||
assert mcp_items[0].type == "mcp_call"
|
||||
assert mcp_items[0].name == "database"
|
||||
assert mcp_items[0].server_label == "database"
|
||||
assert mcp_items[0].status == "in_progress"
|
||||
|
||||
# Test 3: Built-in tool (container)
|
||||
# should NOT return MCP call, falls through to reasoning
|
||||
parser_builtin = Mock()
|
||||
parser_builtin.current_content = "docker run"
|
||||
parser_builtin.current_role = Role.ASSISTANT
|
||||
parser_builtin.current_channel = "analysis"
|
||||
parser_builtin.current_recipient = "container"
|
||||
|
||||
builtin_items = parser_state_to_response_output(parser_builtin)
|
||||
|
||||
# Should fall through to reasoning logic
|
||||
assert len(builtin_items) == 1
|
||||
assert not isinstance(builtin_items[0], McpCall)
|
||||
assert builtin_items[0].type == "reasoning"
|
||||
|
||||
@@ -51,6 +51,7 @@ from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
)
|
||||
from vllm.inputs import tokens_input
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.parser.harmony import Segment
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
|
||||
@@ -534,13 +535,9 @@ class TestHarmonyPreambleStreaming:
|
||||
"""Tests for preamble (commentary with no recipient) streaming events."""
|
||||
|
||||
@staticmethod
|
||||
def _make_ctx(*, channel, recipient, delta="hello"):
|
||||
"""Build a lightweight mock StreamingHarmonyContext."""
|
||||
ctx = MagicMock()
|
||||
ctx.last_content_delta = delta
|
||||
ctx.parser.current_channel = channel
|
||||
ctx.parser.current_recipient = recipient
|
||||
return ctx
|
||||
def _make_segment(*, channel, recipient, delta="hello"):
|
||||
"""Build a lightweight segment for Harmony streaming tests."""
|
||||
return Segment(channel=channel, recipient=recipient, delta=delta)
|
||||
|
||||
@staticmethod
|
||||
def _make_previous_item(*, channel, recipient, text="preamble text"):
|
||||
@@ -559,10 +556,10 @@ class TestHarmonyPreambleStreaming:
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(channel="commentary", recipient=None)
|
||||
segment = self._make_segment(channel="commentary", recipient=None)
|
||||
state = StreamingState()
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
events = emit_content_delta_events(segment, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" in type_names
|
||||
@@ -574,13 +571,13 @@ class TestHarmonyPreambleStreaming:
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(channel="commentary", recipient=None, delta="w")
|
||||
segment = self._make_segment(channel="commentary", recipient=None, delta="w")
|
||||
state = StreamingState()
|
||||
state.sent_output_item_added = True
|
||||
state.current_item_id = "msg_test"
|
||||
state.current_content_index = 0
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
events = emit_content_delta_events(segment, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" in type_names
|
||||
@@ -592,13 +589,13 @@ class TestHarmonyPreambleStreaming:
|
||||
emit_content_delta_events,
|
||||
)
|
||||
|
||||
ctx = self._make_ctx(
|
||||
segment = self._make_segment(
|
||||
channel="commentary",
|
||||
recipient="functions.get_weather",
|
||||
)
|
||||
state = StreamingState()
|
||||
|
||||
events = emit_content_delta_events(ctx, state)
|
||||
events = emit_content_delta_events(segment, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.delta" not in type_names
|
||||
@@ -612,6 +609,7 @@ class TestHarmonyPreambleStreaming:
|
||||
|
||||
previous = self._make_previous_item(channel="commentary", recipient=None)
|
||||
state = StreamingState()
|
||||
state.sent_output_item_added = True
|
||||
state.current_item_id = "msg_test"
|
||||
state.current_output_index = 0
|
||||
state.current_content_index = 0
|
||||
@@ -634,13 +632,53 @@ class TestHarmonyPreambleStreaming:
|
||||
channel="commentary", recipient="functions.get_weather"
|
||||
)
|
||||
state = StreamingState()
|
||||
state.is_first_function_call_delta = True
|
||||
state.current_item_id = "fc_test"
|
||||
state.current_call_id = "call_test"
|
||||
|
||||
events = emit_previous_item_done_events(previous, state)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_text.done" not in type_names
|
||||
|
||||
@pytest.mark.xfail(
|
||||
reason=(
|
||||
"TODO: Ensure added/in-progress events are emitted for zero-delta items."
|
||||
"So we can safely emit done events for zero-delta items."
|
||||
),
|
||||
strict=True,
|
||||
)
|
||||
def test_zero_delta_items_should_preserve_streaming_lifecycle(
|
||||
self,
|
||||
) -> None:
|
||||
"""Zero-delta Harmony items should still produce a coherent lifecycle."""
|
||||
from vllm.entrypoints.openai.responses.streaming_events import (
|
||||
emit_previous_item_done_events,
|
||||
)
|
||||
|
||||
cases: list[tuple[str, str | None, str]] = [
|
||||
("commentary", None, "msg_stale"),
|
||||
("analysis", None, "msg_stale"),
|
||||
("commentary", "functions.get_weather", "fc_stale"),
|
||||
("commentary", "python", "tool_stale"),
|
||||
("commentary", "repo_browser.list", "mcp_stale"),
|
||||
]
|
||||
|
||||
for channel, recipient, current_item_id in cases:
|
||||
previous = self._make_previous_item(channel=channel, recipient=recipient)
|
||||
state = StreamingState()
|
||||
state.current_item_id = current_item_id
|
||||
state.current_call_id = "call_stale"
|
||||
state.current_content_index = 0
|
||||
|
||||
events = emit_previous_item_done_events(
|
||||
previous, state, function_tool_names=None
|
||||
)
|
||||
|
||||
type_names = [e.type for e in events]
|
||||
assert "response.output_item.added" in type_names
|
||||
assert "response.output_item.done" in type_names
|
||||
|
||||
|
||||
def _make_simple_context_with_output(text, token_ids, response_parser=None):
|
||||
"""Create a SimpleContext with a RequestOutput containing the given text."""
|
||||
|
||||
@@ -176,6 +176,49 @@ def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
|
||||
assert child_args.api_server_count == 1
|
||||
|
||||
|
||||
def test_run_vllm_dp_server_uses_python_server_by_default(monkeypatch):
|
||||
calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
|
||||
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
|
||||
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
|
||||
monkeypatch.setattr(dp_sup.envs, "VLLM_RUST_FRONTEND_PATH", None, raising=False)
|
||||
monkeypatch.setattr(
|
||||
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
|
||||
)
|
||||
|
||||
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
|
||||
|
||||
assert calls == ["python"]
|
||||
|
||||
|
||||
def test_run_vllm_dp_server_uses_rust_frontend_when_enabled(monkeypatch):
|
||||
calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(dp_sup.os, "setpgrp", lambda: None)
|
||||
monkeypatch.setattr(dp_sup, "set_process_title", lambda *_args: None)
|
||||
monkeypatch.setattr(dp_sup, "decorate_logs", lambda *_args: None)
|
||||
monkeypatch.setattr(
|
||||
dp_sup.envs,
|
||||
"VLLM_RUST_FRONTEND_PATH",
|
||||
"/tmp/vllm-rs",
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dp_sup, "_run_python_vllm_dp_server", lambda _args: calls.append("python")
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dp_sup, "_run_rust_vllm_dp_server", lambda _args: calls.append("rust")
|
||||
)
|
||||
|
||||
dp_sup._run_vllm_dp_server(_make_unit_args(data_parallel_rank=4))
|
||||
|
||||
assert calls == ["rust"]
|
||||
|
||||
|
||||
def test_validate_multi_port_external_lb_args_allows_ssl():
|
||||
args = _make_unit_args(
|
||||
ssl_keyfile="/tmp/server.key",
|
||||
|
||||
@@ -12,15 +12,21 @@ Tests real customer usage scenarios:
|
||||
- Priority: env vars > decorators > customer script files > framework
|
||||
defaults
|
||||
|
||||
Note: These tests focus on validating server responses rather than directly calling
|
||||
get_ping_handler() and get_invoke_handler() to ensure full integration testing.
|
||||
The handler-override scenarios exercise the real vLLM SageMaker router and
|
||||
bootstrap path via an in-process FastAPI ``TestClient`` instead of launching a
|
||||
model server. These scenarios fully replace the ``/ping`` and ``/invocations``
|
||||
endpoints with customer handlers, so no inference engine is required to
|
||||
validate override behavior. Avoiding the model server also keeps the tests
|
||||
fast and deterministic rather than depending on the FastAPI version resolved
|
||||
into the test environment at runtime.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
@@ -29,6 +35,23 @@ from .conftest import (
|
||||
)
|
||||
|
||||
|
||||
def _build_sagemaker_test_client() -> TestClient:
|
||||
"""Build a TestClient over the real SageMaker router and bootstrap path.
|
||||
|
||||
``attach_router`` is called with empty supported tasks because the override
|
||||
tests replace the endpoints with customer handlers, so no framework
|
||||
invocation handler (and therefore no engine) is exercised.
|
||||
"""
|
||||
from vllm.entrypoints.serve.sagemaker.api_router import (
|
||||
attach_router,
|
||||
sagemaker_standards_bootstrap,
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
attach_router(app, ())
|
||||
return TestClient(sagemaker_standards_bootstrap(app))
|
||||
|
||||
|
||||
class TestHandlerOverrideIntegration:
|
||||
"""Integration tests simulating real customer usage scenarios.
|
||||
|
||||
@@ -89,8 +112,7 @@ class TestHandlerOverrideIntegration:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_customer_script_functions_auto_loaded(self):
|
||||
def test_customer_script_functions_auto_loaded(self, monkeypatch, tmp_path):
|
||||
"""Test customer scenario: script functions automatically override
|
||||
framework defaults."""
|
||||
try:
|
||||
@@ -101,15 +123,15 @@ class TestHandlerOverrideIntegration:
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Customer writes a script file with ping() and invoke() functions
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
async def custom_sagemaker_ping_handler():
|
||||
return {
|
||||
"status": "healthy",
|
||||
"source": "customer_override",
|
||||
"source": "customer_override",
|
||||
"message": "Custom ping from customer script"
|
||||
}
|
||||
|
||||
@@ -119,62 +141,39 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
"source": "customer_override"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
# Customer sets SageMaker environment variables to point to their script
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
# Customer tests their server and sees their overrides work
|
||||
# automatically
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
|
||||
# Customer sets SageMaker environment variables to point to their script
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
}
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
# Customer sees their functions are used
|
||||
assert ping_data["source"] == "customer_override"
|
||||
assert ping_data["message"] == "Custom ping from customer script"
|
||||
assert invoke_data["source"] == "customer_override"
|
||||
assert invoke_data["predictions"] == [
|
||||
"Custom response from customer script"
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
# Customer tests their server and sees their overrides work
|
||||
# automatically
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# Customer sees their functions are used
|
||||
assert ping_data["source"] == "customer_override"
|
||||
assert ping_data["message"] == "Custom ping from customer script"
|
||||
assert invoke_data["source"] == "customer_override"
|
||||
assert invoke_data["predictions"] == [
|
||||
"Custom response from customer script"
|
||||
]
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_customer_decorator_usage(self):
|
||||
def test_customer_decorator_usage(self, monkeypatch, tmp_path):
|
||||
"""Test customer scenario: using @custom_ping_handler and
|
||||
@custom_invocation_handler decorators."""
|
||||
try:
|
||||
@@ -185,9 +184,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Customer writes a script file with decorators
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
||||
from fastapi import Request
|
||||
|
||||
@@ -198,62 +197,39 @@ async def my_ping():
|
||||
"source": "customer_decorator"
|
||||
}
|
||||
|
||||
@sagemaker_standards.custom_invocation_handler
|
||||
@sagemaker_standards.custom_invocation_handler
|
||||
async def my_invoke(request: Request):
|
||||
return {
|
||||
"type": "invoke",
|
||||
"type": "invoke",
|
||||
"source": "customer_decorator"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
# Customer sees their handlers are used by the server
|
||||
assert ping_data["source"] == "customer_decorator"
|
||||
assert invoke_data["source"] == "customer_decorator"
|
||||
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
}
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# Customer sees their handlers are used by the server
|
||||
assert ping_data["source"] == "customer_decorator"
|
||||
assert invoke_data["source"] == "customer_decorator"
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_priority_order(self):
|
||||
def test_handler_priority_order(self, monkeypatch, tmp_path):
|
||||
"""Test priority: @custom_ping_handler/@custom_invocation_handler
|
||||
decorators vs script functions."""
|
||||
try:
|
||||
@@ -264,9 +240,9 @@ async def my_invoke(request: Request):
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Customer writes a script with both decorator and regular functions
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
||||
from fastapi import Request
|
||||
|
||||
@@ -275,7 +251,7 @@ from fastapi import Request
|
||||
async def decorated_ping():
|
||||
return {
|
||||
"status": "healthy",
|
||||
"source": "ping_decorator_in_script",
|
||||
"source": "ping_decorator_in_script",
|
||||
"priority": "decorator"
|
||||
}
|
||||
|
||||
@@ -296,60 +272,37 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
"priority": "function"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
# @custom_ping_handler decorator has higher priority than
|
||||
# script function
|
||||
assert ping_data["source"] == "ping_decorator_in_script"
|
||||
assert ping_data["priority"] == "decorator"
|
||||
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
}
|
||||
# Script function is used for invoke
|
||||
assert invoke_data["source"] == "script_invoke_function"
|
||||
assert invoke_data["priority"] == "function"
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# @custom_ping_handler decorator has higher priority than
|
||||
# script function
|
||||
assert ping_data["source"] == "ping_decorator_in_script"
|
||||
assert ping_data["priority"] == "decorator"
|
||||
|
||||
# Script function is used for invoke
|
||||
assert invoke_data["source"] == "script_invoke_function"
|
||||
assert invoke_data["priority"] == "function"
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_environment_variable_script_loading(self):
|
||||
def test_environment_variable_script_loading(self, monkeypatch, tmp_path):
|
||||
"""Test that environment variables correctly specify script location
|
||||
and loading."""
|
||||
try:
|
||||
@@ -360,9 +313,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Customer writes a script in a specific directory
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
async def custom_sagemaker_ping_handler():
|
||||
@@ -379,60 +332,43 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
"method": "environment_variable_loading"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
# Test environment variable script loading
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
|
||||
# Test environment variable script loading
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
}
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# Verify that the script was loaded via environment variables
|
||||
assert ping_data["source"] == "env_loaded_script"
|
||||
assert ping_data["method"] == "environment_variable_loading"
|
||||
assert invoke_data["source"] == "env_loaded_script"
|
||||
assert invoke_data["method"] == "environment_variable_loading"
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
# Verify that the script was loaded via environment variables
|
||||
assert ping_data["source"] == "env_loaded_script"
|
||||
assert ping_data["method"] == "environment_variable_loading"
|
||||
assert invoke_data["source"] == "env_loaded_script"
|
||||
assert invoke_data["method"] == "environment_variable_loading"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_framework_default_handlers(self):
|
||||
"""Test that framework default handlers work when no customer
|
||||
overrides exist."""
|
||||
overrides exist.
|
||||
|
||||
This scenario exercises the real inference path (default
|
||||
``/invocations``), so it keeps using a live model server rather than
|
||||
the in-process TestClient.
|
||||
"""
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
@@ -478,8 +414,7 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_env_var_override(self):
|
||||
def test_handler_env_var_override(self, monkeypatch, tmp_path):
|
||||
"""Test CUSTOM_FASTAPI_PING_HANDLER and CUSTOM_FASTAPI_INVOCATION_HANDLER
|
||||
environment variable overrides."""
|
||||
try:
|
||||
@@ -493,9 +428,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Create a script with both env var handlers and script functions
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
from fastapi import Request, Response
|
||||
import json
|
||||
|
||||
@@ -533,68 +468,47 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
"method": "script_function"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
# Set environment variables to override both handlers
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
monkeypatch.setenv(
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
|
||||
f"{script_path.name}:env_var_ping_handler",
|
||||
)
|
||||
monkeypatch.setenv(
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
|
||||
f"{script_path.name}:env_var_invoke_handler",
|
||||
)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
# Test ping handler override
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
# Environment variable should override script function
|
||||
assert ping_data["method"] == "environment_variable"
|
||||
assert ping_data["source"] == "env_var_ping"
|
||||
|
||||
# Test invocation handler override
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
# Environment variable should override script function
|
||||
assert invoke_data["method"] == "environment_variable"
|
||||
assert invoke_data["source"] == "env_var_invoke"
|
||||
|
||||
# Set environment variables to override both handlers
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
|
||||
f"{script_name}:env_var_ping_handler"
|
||||
),
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
|
||||
f"{script_name}:env_var_invoke_handler"
|
||||
),
|
||||
}
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
# Test ping handler override
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
# Environment variable should override script function
|
||||
assert ping_data["method"] == "environment_variable"
|
||||
assert ping_data["source"] == "env_var_ping"
|
||||
|
||||
# Test invocation handler override
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# Environment variable should override script function
|
||||
assert invoke_data["method"] == "environment_variable"
|
||||
assert invoke_data["source"] == "env_var_invoke"
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_env_var_priority_over_decorator_and_script(self):
|
||||
def test_env_var_priority_over_decorator_and_script(self, monkeypatch, tmp_path):
|
||||
"""Test that environment variables have highest priority over decorators
|
||||
and script functions for both ping and invocation handlers."""
|
||||
try:
|
||||
@@ -608,9 +522,9 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
pytest.skip("model-hosting-container-standards not available")
|
||||
|
||||
# Create a script with all three handler types for both ping and invocation
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
|
||||
f.write(
|
||||
"""
|
||||
script_path = tmp_path / "model.py"
|
||||
script_path.write_text(
|
||||
"""
|
||||
import model_hosting_container_standards.sagemaker as sagemaker_standards
|
||||
from fastapi import Request, Response
|
||||
import json
|
||||
@@ -674,62 +588,42 @@ async def custom_sagemaker_invocation_handler(request: Request):
|
||||
"priority": "script_function"
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
# Set environment variables to specify highest priority handlers
|
||||
monkeypatch.setenv(SageMakerEnvVars.SAGEMAKER_MODEL_PATH, str(tmp_path))
|
||||
monkeypatch.setenv(SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME, script_path.name)
|
||||
monkeypatch.setenv(
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER,
|
||||
f"{script_path.name}:env_priority_ping",
|
||||
)
|
||||
monkeypatch.setenv(
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER,
|
||||
f"{script_path.name}:env_priority_invoke",
|
||||
)
|
||||
|
||||
with _build_sagemaker_test_client() as client:
|
||||
# Test ping handler priority
|
||||
ping_response = client.get("/ping")
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
# Environment variable has highest priority and should be used
|
||||
assert ping_data["priority"] == "environment_variable"
|
||||
assert ping_data["source"] == "env_var"
|
||||
|
||||
# Test invocation handler priority
|
||||
invoke_response = client.post(
|
||||
"/invocations",
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
script_path = f.name
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
try:
|
||||
script_dir = os.path.dirname(script_path)
|
||||
script_name = os.path.basename(script_path)
|
||||
|
||||
# Set environment variables to specify highest priority handlers
|
||||
env_vars = {
|
||||
SageMakerEnvVars.SAGEMAKER_MODEL_PATH: script_dir,
|
||||
SageMakerEnvVars.CUSTOM_SCRIPT_FILENAME: script_name,
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_PING_HANDLER: (
|
||||
f"{script_name}:env_priority_ping"
|
||||
),
|
||||
FastAPIEnvVars.CUSTOM_FASTAPI_INVOCATION_HANDLER: (
|
||||
f"{script_name}:env_priority_invoke"
|
||||
),
|
||||
}
|
||||
|
||||
args = [
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME_SMOLLM, args, env_dict=env_vars
|
||||
) as server:
|
||||
# Test ping handler priority
|
||||
ping_response = requests.get(server.url_for("ping"))
|
||||
assert ping_response.status_code == 200
|
||||
ping_data = ping_response.json()
|
||||
|
||||
# Environment variable has highest priority and should be used
|
||||
assert ping_data["priority"] == "environment_variable"
|
||||
assert ping_data["source"] == "env_var"
|
||||
|
||||
# Test invocation handler priority
|
||||
invoke_response = requests.post(
|
||||
server.url_for("invocations"),
|
||||
json={
|
||||
"model": MODEL_NAME_SMOLLM,
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 5,
|
||||
},
|
||||
)
|
||||
assert invoke_response.status_code == 200
|
||||
invoke_data = invoke_response.json()
|
||||
|
||||
# Environment variable has highest priority and should be used
|
||||
assert invoke_data["priority"] == "environment_variable"
|
||||
assert invoke_data["source"] == "env_var"
|
||||
|
||||
finally:
|
||||
os.unlink(script_path)
|
||||
# Environment variable has highest priority and should be used
|
||||
assert invoke_data["priority"] == "environment_variable"
|
||||
assert invoke_data["source"] == "env_var"
|
||||
|
||||
@@ -1,18 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from openai_harmony import Author, Message, Role, StreamState, TextContent
|
||||
from openai_harmony import Author, Message, Role, TextContent
|
||||
|
||||
from vllm.entrypoints.openai.responses.context import (
|
||||
HarmonyContext,
|
||||
SimpleContext,
|
||||
StreamingHarmonyContext,
|
||||
TurnMetrics,
|
||||
)
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.parser.harmony import ChunkResult, HarmonyParser, Segment
|
||||
|
||||
|
||||
def create_mock_request_output(
|
||||
@@ -68,25 +68,59 @@ async def generate_mock_outputs(
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_parser():
|
||||
"""Set up a mock parser for tests."""
|
||||
with patch(
|
||||
"vllm.entrypoints.openai.responses.context.get_streamable_parser_for_assistant"
|
||||
) as mock_parser_factory:
|
||||
# Create a mock parser object
|
||||
parser = MagicMock()
|
||||
parser.messages = []
|
||||
parser.current_channel = None
|
||||
parser.state = StreamState.EXPECT_START
|
||||
mock_parser_factory.return_value = parser
|
||||
yield parser
|
||||
class FakeHarmonyParser(HarmonyParser):
|
||||
def __init__(self):
|
||||
# Skip HarmonyParser initialization and script outputs directly.
|
||||
self.reasoning_parser = None
|
||||
self.tool_parser = None
|
||||
self._chunk_results: list[ChunkResult] = []
|
||||
self._flush_results: list[Segment | None] = []
|
||||
self.processed_chunks: list[list[int]] = []
|
||||
|
||||
def enqueue_chunk_result(
|
||||
self,
|
||||
segments: list[Segment] | None = None,
|
||||
reasoning_token_count: int = 0,
|
||||
) -> None:
|
||||
self._chunk_results.append(
|
||||
ChunkResult(
|
||||
segments=[] if segments is None else segments,
|
||||
reasoning_token_count=reasoning_token_count,
|
||||
)
|
||||
)
|
||||
|
||||
def enqueue_flush_result(self, segment: Segment | None) -> None:
|
||||
self._flush_results.append(segment)
|
||||
|
||||
def process_chunk(self, token_ids) -> ChunkResult:
|
||||
self.processed_chunks.append(list(token_ids))
|
||||
if self._chunk_results:
|
||||
return self._chunk_results.pop(0)
|
||||
return ChunkResult(segments=[], reasoning_token_count=0)
|
||||
|
||||
def flush(self) -> Segment | None:
|
||||
if self._flush_results:
|
||||
return self._flush_results.pop(0)
|
||||
return None
|
||||
|
||||
|
||||
def make_harmony_context(
|
||||
messages=None, available_tools=None, function_tool_names=None
|
||||
) -> tuple[HarmonyContext, FakeHarmonyParser]:
|
||||
fake_parser = FakeHarmonyParser()
|
||||
context = HarmonyContext(
|
||||
messages=[] if messages is None else messages,
|
||||
available_tools=[] if available_tools is None else available_tools,
|
||||
function_tool_names=function_tool_names,
|
||||
response_parser=fake_parser,
|
||||
)
|
||||
return context, fake_parser
|
||||
|
||||
|
||||
def test_single_turn_token_counting():
|
||||
"""Test token counting behavior for a single turn."""
|
||||
# Create a context
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
context, _ = make_harmony_context()
|
||||
|
||||
# Create a mock RequestOutput with specific token counts
|
||||
mock_output = create_mock_request_output(
|
||||
@@ -118,7 +152,7 @@ def test_single_turn_token_counting():
|
||||
async def test_multi_turn_token_counting():
|
||||
"""Test token counting behavior across multiple turns with tool output."""
|
||||
# Create a context
|
||||
context = HarmonyContext(messages=[], available_tools=["browser"])
|
||||
context, _ = make_harmony_context(available_tools=["browser"])
|
||||
|
||||
# Simulate a conversation with 3 turns
|
||||
# Turn 1: prefill 5, decode 3, tool 7
|
||||
@@ -177,7 +211,7 @@ async def test_multi_turn_token_counting():
|
||||
|
||||
def test_empty_output_tokens():
|
||||
"""Test behavior when RequestOutput has empty output tokens."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
context, _ = make_harmony_context()
|
||||
|
||||
# Create a RequestOutput with empty output tokens
|
||||
mock_output = create_mock_request_output(
|
||||
@@ -197,7 +231,7 @@ def test_empty_output_tokens():
|
||||
|
||||
def test_missing_prompt_token_ids():
|
||||
"""Test behavior when RequestOutput has None prompt_token_ids."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
context, _ = make_harmony_context()
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=None, # No prompt token IDs
|
||||
@@ -216,12 +250,10 @@ def test_missing_prompt_token_ids():
|
||||
assert context.num_tool_output_tokens == 0
|
||||
|
||||
|
||||
def test_reasoning_tokens_counting(mock_parser):
|
||||
def test_reasoning_tokens_counting():
|
||||
"""Test that reasoning tokens are counted correctly."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
# Mock parser to simulate reasoning channel
|
||||
mock_parser.current_channel = "analysis" # Reasoning channel
|
||||
context, parser = make_harmony_context()
|
||||
parser.enqueue_chunk_result(reasoning_token_count=4)
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
@@ -236,13 +268,11 @@ def test_reasoning_tokens_counting(mock_parser):
|
||||
assert context.num_output_tokens == 4
|
||||
|
||||
|
||||
def test_preamble_tokens_not_counted_as_reasoning(mock_parser):
|
||||
def test_preamble_tokens_not_counted_as_reasoning():
|
||||
"""Preambles (commentary with no recipient) are visible user text,
|
||||
not hidden reasoning. They must NOT inflate num_reasoning_tokens."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
mock_parser.current_channel = "commentary"
|
||||
mock_parser.current_recipient = None # preamble
|
||||
context, parser = make_harmony_context()
|
||||
parser.enqueue_chunk_result(reasoning_token_count=0)
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
@@ -255,13 +285,11 @@ def test_preamble_tokens_not_counted_as_reasoning(mock_parser):
|
||||
assert context.num_output_tokens == 3
|
||||
|
||||
|
||||
def test_commentary_with_recipient_counted_as_reasoning(mock_parser):
|
||||
def test_commentary_with_recipient_counted_as_reasoning():
|
||||
"""Commentary directed at a tool (recipient != None) is hidden from
|
||||
the user, so it should still count as reasoning tokens."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
|
||||
mock_parser.current_channel = "commentary"
|
||||
mock_parser.current_recipient = "python"
|
||||
context, parser = make_harmony_context()
|
||||
parser.enqueue_chunk_result(reasoning_token_count=3)
|
||||
|
||||
mock_output = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
@@ -276,7 +304,7 @@ def test_commentary_with_recipient_counted_as_reasoning(mock_parser):
|
||||
|
||||
def test_zero_tokens_edge_case():
|
||||
"""Test behavior with all zero token counts."""
|
||||
context = HarmonyContext(messages=[], available_tools=[])
|
||||
context, _ = make_harmony_context()
|
||||
|
||||
# Create a request with empty lists (not None) for both prompt and
|
||||
# output tokens
|
||||
@@ -299,10 +327,7 @@ def test_zero_tokens_edge_case():
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_turn_no_tool_output():
|
||||
"""Test that first turn never generates tool output tokens."""
|
||||
context = HarmonyContext(
|
||||
messages=[],
|
||||
available_tools=["browser"], # Tools available
|
||||
)
|
||||
context, _ = make_harmony_context(available_tools=["browser"])
|
||||
|
||||
# Even with large prompt in first turn, no tool tokens should be counted
|
||||
mock_output = create_mock_request_output(
|
||||
@@ -324,7 +349,7 @@ async def test_negative_tool_tokens_edge_case():
|
||||
tokens. We should log an error and clamp the value to 0."""
|
||||
# Use patch to check if logger.error was called
|
||||
with patch("vllm.entrypoints.openai.responses.context.logger.error") as mock_log:
|
||||
context = HarmonyContext(messages=[], available_tools=["browser"])
|
||||
context, _ = make_harmony_context(available_tools=["browser"])
|
||||
|
||||
# First turn
|
||||
mock_output1 = create_mock_request_output(
|
||||
@@ -360,15 +385,15 @@ async def test_negative_tool_tokens_edge_case():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
async def test_streaming_multi_turn_token_counting():
|
||||
"""Test token counting for streaming multi-turn conversations.
|
||||
|
||||
This test focuses on how StreamingHarmonyContext counts tokens in a
|
||||
This test focuses on how HarmonyContext counts tokens in a
|
||||
multi-turn conversation with streaming (token-by-token) outputs and
|
||||
message boundaries.
|
||||
"""
|
||||
# Create a streaming context
|
||||
context = StreamingHarmonyContext(messages=[], available_tools=["browser"])
|
||||
context, parser = make_harmony_context(available_tools=["browser"])
|
||||
|
||||
num_prompt_tokens = [3, 8, 13]
|
||||
num_output_tokens = [3, 3, 2]
|
||||
@@ -413,10 +438,8 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
assert context.num_tool_output_tokens == 0 # No tool output in first turn
|
||||
assert context.first_tok_of_message is True # Ready for next message
|
||||
|
||||
# Second turn: reasoning tokens in analysis channel
|
||||
mock_parser.current_channel = "analysis" # Set to reasoning channel
|
||||
|
||||
# First token of second turn
|
||||
parser.enqueue_chunk_result(reasoning_token_count=1)
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
prompt_token_ids=[
|
||||
@@ -436,6 +459,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
)
|
||||
|
||||
# More tokens in reasoning channel
|
||||
parser.enqueue_chunk_result(reasoning_token_count=1)
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
output_token_ids=[202],
|
||||
@@ -443,6 +467,7 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
)
|
||||
)
|
||||
|
||||
parser.enqueue_chunk_result(reasoning_token_count=1)
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
output_token_ids=[203],
|
||||
@@ -460,9 +485,6 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
expected_tool_tokens = 8 - 3 - 3 # = 2
|
||||
assert context.num_tool_output_tokens == expected_tool_tokens
|
||||
|
||||
# Third turn: regular output channel
|
||||
mock_parser.current_channel = "final" # Switch back to regular channel
|
||||
|
||||
# Third turn (with more cached tokens)
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
@@ -520,13 +542,8 @@ async def test_streaming_multi_turn_token_counting(mock_parser):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_message_synchronization(mock_parser):
|
||||
"""Test message synchronization logic from lines 413-417 in context.py.
|
||||
|
||||
This test verifies that when parser.messages contains more messages than
|
||||
the context's _messages (minus initial messages), the context properly
|
||||
extends its message list with the new parser messages.
|
||||
"""
|
||||
async def test_streaming_message_synchronization():
|
||||
"""Completed messages from append-local and flush segments sync into context."""
|
||||
|
||||
# Create a streaming context with some initial messages
|
||||
initial_messages = [
|
||||
@@ -536,23 +553,30 @@ async def test_streaming_message_synchronization(mock_parser):
|
||||
recipient=Role.ASSISTANT,
|
||||
)
|
||||
]
|
||||
context = StreamingHarmonyContext(messages=initial_messages, available_tools=[])
|
||||
context, parser = make_harmony_context(messages=initial_messages)
|
||||
|
||||
# Verify initial state
|
||||
assert len(context._messages) == 1
|
||||
assert context.num_init_messages == 1
|
||||
|
||||
# Mock parser to have more messages than context
|
||||
# Simulate parser having processed 3 new messages
|
||||
mock_parser.messages = [
|
||||
Message(
|
||||
author=Author(role=Role.ASSISTANT, name="assistant"),
|
||||
content=[TextContent(text="Response 1")],
|
||||
recipient=Role.USER,
|
||||
),
|
||||
]
|
||||
response_text = "First response"
|
||||
message = Message(
|
||||
author=Author(role=Role.ASSISTANT, name="assistant"),
|
||||
content=[TextContent(text=response_text)],
|
||||
recipient=Role.USER,
|
||||
)
|
||||
parser.enqueue_chunk_result(
|
||||
segments=[
|
||||
Segment(
|
||||
channel="commentary",
|
||||
recipient=None,
|
||||
delta="",
|
||||
completed_message=message,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# This should trigger the message synchronization logic
|
||||
# This should sync the completed message from the latest append
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3], output_token_ids=[101], finished=False
|
||||
@@ -563,36 +587,39 @@ async def test_streaming_message_synchronization(mock_parser):
|
||||
assert len(context._messages) == 2
|
||||
|
||||
# Verify the new messages were added correctly
|
||||
assert context._messages[1].content[0].text == "Response 1"
|
||||
assert context._messages[1].content[0].text == response_text
|
||||
|
||||
# Test the specific condition from line 413-414:
|
||||
# len(self._messages) - self.num_init_messages < len(self.parser.messages)
|
||||
messages_minus_init = len(context._messages) - context.num_init_messages
|
||||
parser_messages_count = len(mock_parser.messages)
|
||||
assert messages_minus_init == 1
|
||||
|
||||
# After synchronization, they should be equal (no longer less than)
|
||||
assert messages_minus_init == parser_messages_count
|
||||
response_text = "Second response"
|
||||
message = Message(
|
||||
author=Author(role=Role.ASSISTANT, name="assistant"),
|
||||
content=[TextContent(text=response_text)],
|
||||
recipient=Role.USER,
|
||||
)
|
||||
flush_segment = Segment(
|
||||
channel="commentary",
|
||||
recipient=None,
|
||||
delta="",
|
||||
completed_message=message,
|
||||
)
|
||||
parser.enqueue_flush_result(flush_segment)
|
||||
|
||||
# Test edge case: add one more parser message
|
||||
mock_parser.messages.append(
|
||||
Message(
|
||||
author=Author(role=Role.ASSISTANT, name="assistant"),
|
||||
content=[TextContent(text="Response 4")],
|
||||
recipient=Role.USER,
|
||||
# Create another output to trigger synchronization via flush()
|
||||
context.append_output(
|
||||
create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3], output_token_ids=[102], finished=True
|
||||
)
|
||||
)
|
||||
|
||||
# Create another output to trigger synchronization again
|
||||
mock_output2 = create_mock_request_output(
|
||||
prompt_token_ids=[1, 2, 3], output_token_ids=[102], finished=True
|
||||
)
|
||||
|
||||
context.append_output(mock_output2)
|
||||
|
||||
# Verify the fourth message was added, num_init_messages is still 1
|
||||
# Verify the flushed response was added, num_init_messages is still 1
|
||||
assert len(context._messages) == 3
|
||||
assert context.num_init_messages == 1
|
||||
assert context._messages[2].content[0].text == "Response 4"
|
||||
assert context._messages[2].content[0].text == response_text
|
||||
assert context.last_append_flush_status is True
|
||||
assert len(context.last_append_segments) == 1
|
||||
assert context.last_append_segments[0].completed_message is message
|
||||
|
||||
|
||||
def test_turn_metrics_copy_and_reset():
|
||||
|
||||
@@ -8,7 +8,6 @@ actual NCCL communication.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -48,7 +47,6 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
|
||||
names: list[str] | None = None
|
||||
dtype_names: list[str] | None = None
|
||||
shapes: list[list[int]] | None = None
|
||||
num_updates_list: list[int] | None = None
|
||||
|
||||
|
||||
class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
|
||||
@@ -59,16 +57,20 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
|
||||
|
||||
# Class-level tracking for verification across processes
|
||||
init_transfer_engine_called: bool = False
|
||||
start_called: bool = False
|
||||
receive_weights_called: bool = False
|
||||
finish_called: bool = False
|
||||
shutdown_called: bool = False
|
||||
last_init_info: MockInitInfo | None = None
|
||||
last_update_info: MockUpdateInfo | None = None
|
||||
|
||||
def __init__(self, config, parallel_config, model):
|
||||
super().__init__(config, parallel_config, model)
|
||||
def __init__(self, config, vllm_config, device, model):
|
||||
super().__init__(config, vllm_config, device, model)
|
||||
# Reset tracking on init
|
||||
MockWeightTransferEngine.init_transfer_engine_called = False
|
||||
MockWeightTransferEngine.start_called = False
|
||||
MockWeightTransferEngine.receive_weights_called = False
|
||||
MockWeightTransferEngine.finish_called = False
|
||||
MockWeightTransferEngine.shutdown_called = False
|
||||
MockWeightTransferEngine.last_init_info = None
|
||||
MockWeightTransferEngine.last_update_info = None
|
||||
@@ -77,37 +79,28 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
|
||||
MockWeightTransferEngine.init_transfer_engine_called = True
|
||||
MockWeightTransferEngine.last_init_info = init_info
|
||||
|
||||
def receive_weights(
|
||||
self,
|
||||
update_info: MockUpdateInfo,
|
||||
load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
|
||||
) -> None:
|
||||
MockWeightTransferEngine.receive_weights_called = True
|
||||
MockWeightTransferEngine.last_update_info = update_info
|
||||
# Simulate loading weights by calling load_weights with empty list
|
||||
# (In real implementation, this would receive and load actual weights)
|
||||
load_weights([])
|
||||
def start_weight_update(self) -> None:
|
||||
MockWeightTransferEngine.start_called = True
|
||||
|
||||
def receive_sparse_weights(
|
||||
self,
|
||||
update_info: MockUpdateInfo,
|
||||
apply_patches: Callable[[list], None],
|
||||
) -> None:
|
||||
def finish_weight_update(self) -> None:
|
||||
MockWeightTransferEngine.finish_called = True
|
||||
|
||||
def receive_weights(self, update_info: MockUpdateInfo) -> None:
|
||||
MockWeightTransferEngine.receive_weights_called = True
|
||||
MockWeightTransferEngine.last_update_info = update_info
|
||||
apply_patches([])
|
||||
|
||||
def shutdown(self) -> None:
|
||||
MockWeightTransferEngine.shutdown_called = True
|
||||
|
||||
def trainer_send_weights(self, *args, **kwargs):
|
||||
@staticmethod
|
||||
def trainer_send_weights(*args, **kwargs):
|
||||
"""Mock method to simulate trainer sending weights."""
|
||||
pass
|
||||
|
||||
|
||||
def mock_create_engine(config, parallel_config, model):
|
||||
def mock_create_engine(config, vllm_config, device, model):
|
||||
"""Mock factory function that returns our mock engine."""
|
||||
return MockWeightTransferEngine(config, parallel_config, model)
|
||||
return MockWeightTransferEngine(config, vllm_config, device, model)
|
||||
|
||||
|
||||
# --- Tests ---
|
||||
@@ -208,7 +201,7 @@ def test_update_weights_calls_engine():
|
||||
llm.init_weight_transfer_engine(
|
||||
WeightTransferInitRequest(init_info={"test_param": "init"})
|
||||
)
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
llm.start_weight_update()
|
||||
|
||||
# Call update_weights
|
||||
test_names = ["layer.weight", "layer.bias"]
|
||||
@@ -243,61 +236,6 @@ def test_update_weights_calls_engine():
|
||||
llm.finish_weight_update()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_update_weights_passes_sparse_metadata():
|
||||
"""Test sparse update metadata is forwarded unchanged to the engine."""
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("Need at least 1 GPU for this test")
|
||||
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
|
||||
with patch(
|
||||
"vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
|
||||
mock_create_engine,
|
||||
):
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
load_format="dummy",
|
||||
tensor_parallel_size=1,
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
)
|
||||
|
||||
llm.init_weight_transfer_engine(
|
||||
WeightTransferInitRequest(init_info={"test_param": "init"})
|
||||
)
|
||||
llm.start_weight_update(is_checkpoint_format=False)
|
||||
|
||||
llm.update_weights(
|
||||
WeightTransferUpdateRequest(
|
||||
update_info={
|
||||
"names": ["layer.weight"],
|
||||
"dtype_names": ["bfloat16"],
|
||||
"shapes": [[100]],
|
||||
"num_updates_list": [3],
|
||||
"update_kind": "sparse_flat",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
def check_sparse_update_called(self):
|
||||
engine = self.weight_transfer_engine
|
||||
if not engine.receive_weights_called:
|
||||
return None
|
||||
info = engine.last_update_info
|
||||
return (
|
||||
info.update_kind,
|
||||
info.num_updates_list,
|
||||
)
|
||||
|
||||
results = llm.collective_rpc(check_sparse_update_called)
|
||||
for result in results:
|
||||
assert result == ("sparse_flat", [3])
|
||||
|
||||
llm.finish_weight_update()
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_full_weight_transfer_flow():
|
||||
"""Test the complete weight transfer flow: init -> start -> update -> finish."""
|
||||
@@ -327,7 +265,7 @@ def test_full_weight_transfer_flow():
|
||||
)
|
||||
|
||||
# Step 2: Start weight update
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
llm.start_weight_update()
|
||||
|
||||
# Step 3: Update weights
|
||||
llm.update_weights(
|
||||
|
||||
@@ -3,6 +3,7 @@ accuracy_threshold: 0.84
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
startup_max_wait_seconds: 1200
|
||||
use_chat_completions: true
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
|
||||
@@ -110,6 +110,39 @@ async def call_vllm_api(
|
||||
return "", 0
|
||||
|
||||
|
||||
async def call_vllm_chat_api(
|
||||
session: aiohttp.ClientSession,
|
||||
model: str,
|
||||
prompt: str,
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
stop: list[str] | None = None,
|
||||
url: str | None = None,
|
||||
seed: int | None = None,
|
||||
) -> tuple[str, int]:
|
||||
"""Call vLLM's OpenAI-compatible chat completions endpoint."""
|
||||
data = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"stop": stop,
|
||||
}
|
||||
if seed is not None:
|
||||
data["seed"] = seed
|
||||
|
||||
try:
|
||||
async with session.post(f"{url}/v1/chat/completions", json=data) as response:
|
||||
response.raise_for_status()
|
||||
result = await response.json()
|
||||
text = result["choices"][0]["message"]["content"] or ""
|
||||
completion_tokens = result.get("usage", {}).get("completion_tokens", 0)
|
||||
return text, completion_tokens
|
||||
except Exception as e:
|
||||
print(f"Error calling vLLM chat API ({type(e).__name__}): {e}")
|
||||
return "", 0
|
||||
|
||||
|
||||
def _build_gsm8k_prompts(
|
||||
num_questions: int = 1319,
|
||||
num_shots: int = 5,
|
||||
@@ -173,6 +206,8 @@ def evaluate_gsm8k(
|
||||
num_questions: int = 1319,
|
||||
num_shots: int = 5,
|
||||
max_tokens: int = 256,
|
||||
model: str | None = None,
|
||||
use_chat_completions: bool = False,
|
||||
host: str = "http://127.0.0.1",
|
||||
port: int = 8000,
|
||||
temperature: float = 0.0,
|
||||
@@ -193,15 +228,30 @@ def evaluate_gsm8k(
|
||||
output_tokens: list[int] = [0] * num_questions
|
||||
|
||||
async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]:
|
||||
answer, tokens = await call_vllm_api(
|
||||
session=session,
|
||||
prompt=prompts[i],
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop=["Question", "Assistant:", "<|separator|>"],
|
||||
url=base_url,
|
||||
seed=seed,
|
||||
)
|
||||
stop = ["Question", "Assistant:", "<|separator|>"]
|
||||
if use_chat_completions:
|
||||
if model is None:
|
||||
raise ValueError("model is required for chat completions")
|
||||
answer, tokens = await call_vllm_chat_api(
|
||||
session=session,
|
||||
model=model,
|
||||
prompt=prompts[i],
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop=stop,
|
||||
url=base_url,
|
||||
seed=seed,
|
||||
)
|
||||
else:
|
||||
answer, tokens = await call_vllm_api(
|
||||
session=session,
|
||||
prompt=prompts[i],
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop=stop,
|
||||
url=base_url,
|
||||
seed=seed,
|
||||
)
|
||||
states[i] = answer
|
||||
output_tokens[i] = tokens
|
||||
return answer, tokens
|
||||
|
||||
@@ -65,6 +65,8 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
|
||||
num_questions=eval_config["num_questions"],
|
||||
num_shots=eval_config["num_fewshot"],
|
||||
max_tokens=eval_config.get("max_tokens", 256),
|
||||
model=eval_config["model_name"],
|
||||
use_chat_completions=eval_config.get("use_chat_completions", False),
|
||||
host=host,
|
||||
port=port,
|
||||
request_timeout_seconds=request_timeout_seconds,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -109,3 +110,103 @@ def test_cpu_unary_activation(
|
||||
if not (activation_cls is GELU and dtype != torch.bfloat16):
|
||||
raw_out = torch.empty_like(x)
|
||||
opcheck(fn, (raw_out, x, *op_args))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@torch.inference_mode()
|
||||
def test_cpu_gelu_tanh_and_mul(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
gate = torch.tensor(
|
||||
[
|
||||
[
|
||||
-12.0,
|
||||
-10.0,
|
||||
-9.01,
|
||||
-5.0,
|
||||
-2.0,
|
||||
-1.0,
|
||||
-0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
2.0,
|
||||
5.0,
|
||||
9.01,
|
||||
10.0,
|
||||
12.0,
|
||||
11.0,
|
||||
],
|
||||
[
|
||||
-7.5,
|
||||
-4.5,
|
||||
-3.0,
|
||||
-1.5,
|
||||
-0.75,
|
||||
-0.25,
|
||||
0.25,
|
||||
0.75,
|
||||
1.5,
|
||||
3.0,
|
||||
4.5,
|
||||
7.5,
|
||||
-11.0,
|
||||
11.0,
|
||||
8.75,
|
||||
-8.75,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
val = torch.tensor(
|
||||
[
|
||||
[
|
||||
0.25,
|
||||
-0.5,
|
||||
0.75,
|
||||
-1.0,
|
||||
1.25,
|
||||
-1.5,
|
||||
1.75,
|
||||
-2.0,
|
||||
2.25,
|
||||
-2.5,
|
||||
2.75,
|
||||
-3.0,
|
||||
3.25,
|
||||
-3.5,
|
||||
3.75,
|
||||
-4.0,
|
||||
],
|
||||
[
|
||||
-0.4,
|
||||
0.6,
|
||||
-0.8,
|
||||
1.0,
|
||||
-1.2,
|
||||
1.4,
|
||||
-1.6,
|
||||
1.8,
|
||||
-2.0,
|
||||
2.2,
|
||||
-2.4,
|
||||
2.6,
|
||||
-2.8,
|
||||
3.0,
|
||||
-3.2,
|
||||
3.4,
|
||||
],
|
||||
],
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
x = torch.cat((val, gate), dim=-1).contiguous()
|
||||
kernel_out = torch.empty_like(val)
|
||||
torch.ops._C.gelu_tanh_and_mul(kernel_out, x)
|
||||
|
||||
torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate
|
||||
|
||||
atol = get_default_atol(kernel_out)
|
||||
rtol = get_default_rtol(kernel_out)
|
||||
torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol)
|
||||
|
||||
@@ -292,6 +292,76 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):
|
||||
assert (topk == 7).all(), "topk buffer should be untouched on shared layer"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])
|
||||
def test_fused_norm_rope_ds_mla(num_tokens: int):
|
||||
"""fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).
|
||||
|
||||
Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, dynamic float32
|
||||
scale) | 4 float32 scales | 64 bf16 (unquantized) RoPE.
|
||||
"""
|
||||
torch.manual_seed(5)
|
||||
dev = "cuda"
|
||||
max_pos = 8192
|
||||
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
|
||||
|
||||
q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)
|
||||
kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)
|
||||
k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)
|
||||
qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)
|
||||
kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)
|
||||
mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
|
||||
|
||||
bs = max_pos
|
||||
mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8)
|
||||
slot = torch.arange(num_tokens, device=dev, dtype=torch.int64)
|
||||
topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)
|
||||
|
||||
q_out = K.fused_norm_rope(
|
||||
pos,
|
||||
q_c,
|
||||
qw,
|
||||
EPS,
|
||||
kv_c,
|
||||
kvw,
|
||||
EPS,
|
||||
k_pe,
|
||||
mla_cos_sin,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
EPS,
|
||||
None,
|
||||
topk,
|
||||
slot_mapping=slot,
|
||||
indexer_k_cache=None,
|
||||
mla_kv_cache=mla_cache,
|
||||
mla_kv_cache_dtype="fp8_ds_mla",
|
||||
mla_k_scale=None,
|
||||
has_indexer=False,
|
||||
index_rope_interleave=False,
|
||||
)
|
||||
|
||||
assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (ds_mla)")
|
||||
|
||||
kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32
|
||||
kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64]
|
||||
tiles = kv_ref.view(num_tokens, 4, 128)
|
||||
ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1.1754944e-38)
|
||||
ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, KV_LORA).to(FP8)
|
||||
|
||||
cache = mla_cache[0, :num_tokens] # [N, 656] uint8
|
||||
nope = cache[:, :KV_LORA].view(FP8)
|
||||
scales = cache.view(torch.float32)[:, KV_LORA // 4 : KV_LORA // 4 + 4]
|
||||
rope_off = KV_LORA // 2 + 8
|
||||
rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + ROPE_DIM]
|
||||
|
||||
torch.testing.assert_close(scales, ref_scale, rtol=1e-2, atol=1e-6)
|
||||
assert_fp8(nope, ref_nope, "ds_mla NoPE fp8")
|
||||
assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16")
|
||||
# No indexer on this call: top-k buffer must be untouched.
|
||||
assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
|
||||
|
||||
|
||||
# ── fused_q ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -400,6 +470,75 @@ def test_fused_q_no_indexer(num_tokens: int):
|
||||
assert_fp8(mqa[:, :, KV_LORA:], (qpe_ref / s).to(FP8), "mqa q_pe")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
|
||||
@pytest.mark.parametrize("has_indexer", [True, False])
|
||||
def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
|
||||
"""bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
|
||||
produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
|
||||
torch.manual_seed(6)
|
||||
dev = "cuda"
|
||||
max_pos = 8192
|
||||
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
|
||||
|
||||
q_pe = torch.randn(
|
||||
num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16
|
||||
)
|
||||
ql_nope = torch.randn(
|
||||
num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16
|
||||
)
|
||||
q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
|
||||
q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
|
||||
|
||||
index_q = index_w = idx_cos_sin = None
|
||||
if has_indexer:
|
||||
index_q = torch.randn(
|
||||
num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16
|
||||
)
|
||||
index_w = torch.randn(num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32)
|
||||
idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)
|
||||
|
||||
iq_fp8, iw_out, q_pe_out = K.fused_q(
|
||||
pos,
|
||||
q_pe,
|
||||
q_cos_sin,
|
||||
index_q,
|
||||
idx_cos_sin,
|
||||
ql_nope,
|
||||
q_scale,
|
||||
index_w,
|
||||
INDEX_HEAD_DIM**-0.5,
|
||||
INDEX_HEADS**-0.5,
|
||||
has_indexer=has_indexer,
|
||||
index_rope_interleave=False,
|
||||
quantize_mqa=False,
|
||||
)
|
||||
|
||||
# MQA query: only the RoPE'd q_pe, bf16, unquantized.
|
||||
assert q_pe_out.dtype == torch.bfloat16
|
||||
assert q_pe_out.shape == (num_tokens, NUM_HEADS, ROPE_DIM)
|
||||
qpe_ref = rope(
|
||||
q_pe.float(),
|
||||
pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),
|
||||
q_cos_sin,
|
||||
interleave=True,
|
||||
)
|
||||
assert_bf16(q_pe_out, qpe_ref, "bf16 q_pe RoPE")
|
||||
|
||||
# Indexer-Q is unchanged on this path (still UE8M0 fp8 + folded weights).
|
||||
if has_indexer:
|
||||
assert index_q is not None
|
||||
iq_ref = rope(
|
||||
index_q.float(),
|
||||
pos.unsqueeze(-1).expand(num_tokens, INDEX_HEADS),
|
||||
idx_cos_sin,
|
||||
interleave=False,
|
||||
)
|
||||
q_ref, scale_ref = ue8m0_quant(iq_ref)
|
||||
assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
|
||||
iw_ref = index_w * scale_ref * (INDEX_HEAD_DIM**-0.5) * (INDEX_HEADS**-0.5)
|
||||
torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -42,7 +42,11 @@ VLLM_RUNNER_BASE_KWARGS = {
|
||||
}
|
||||
|
||||
|
||||
def run_test(vllm_runner, audio_assets, lora_request, expected_suffix, **kwargs):
|
||||
def run_test(
|
||||
vllm_runner, audio_assets, monkeypatch, lora_request, expected_suffix, **kwargs
|
||||
):
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
inputs = [([AUDIO_PROMPT], [audio_assets[0].audio_and_sample_rate[0]])]
|
||||
|
||||
# Apply any additional kwargs as overrides to the base kwargs
|
||||
@@ -66,11 +70,13 @@ def run_test(vllm_runner, audio_assets, lora_request, expected_suffix, **kwargs)
|
||||
def test_active_default_mm_lora(
|
||||
vllm_runner: type[VllmRunner],
|
||||
audio_assets: AudioTestAssets,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Ensure that we can use the default audio lora."""
|
||||
run_test(
|
||||
vllm_runner,
|
||||
audio_assets,
|
||||
monkeypatch,
|
||||
lora_request=None,
|
||||
default_mm_loras={"audio": AUDIO_LORA_PATH},
|
||||
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
|
||||
@@ -84,12 +90,14 @@ def test_active_default_mm_lora(
|
||||
def test_inactive_default_mm_lora(
|
||||
vllm_runner: type[VllmRunner],
|
||||
audio_assets: AudioTestAssets,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Ensure that modalities are filtered properly."""
|
||||
# Default image lora won't be active since we only pass audio
|
||||
run_test(
|
||||
vllm_runner,
|
||||
audio_assets,
|
||||
monkeypatch,
|
||||
lora_request=None,
|
||||
default_mm_loras={"image": IMAGE_LORA_PATH},
|
||||
expected_suffix=RESPONSE_SUFFIX_WITHOUT_LORA,
|
||||
@@ -103,11 +111,13 @@ def test_inactive_default_mm_lora(
|
||||
def test_default_mm_lora_succeeds_with_redundant_lora_request(
|
||||
vllm_runner: type[VllmRunner],
|
||||
audio_assets: AudioTestAssets,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Ensure that redundantly providing the lora works."""
|
||||
run_test(
|
||||
vllm_runner,
|
||||
audio_assets,
|
||||
monkeypatch,
|
||||
lora_request=LoRARequest("audio", 1, AUDIO_LORA_PATH),
|
||||
default_mm_loras={"audio": AUDIO_LORA_PATH},
|
||||
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
|
||||
@@ -121,12 +131,14 @@ def test_default_mm_lora_succeeds_with_redundant_lora_request(
|
||||
def test_default_mm_lora_fails_with_overridden_lora_request(
|
||||
vllm_runner: type[VllmRunner],
|
||||
audio_assets: AudioTestAssets,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Ensure that if the lora_request conflicts with default_mm_loras,
|
||||
we use the lora_request."""
|
||||
run_test(
|
||||
vllm_runner,
|
||||
audio_assets,
|
||||
monkeypatch,
|
||||
lora_request=LoRARequest("speech", 2, AUDIO_LORA_PATH),
|
||||
default_mm_loras={"audio": IMAGE_LORA_PATH},
|
||||
expected_suffix=RESPONSE_SUFFIX_WITH_LORA,
|
||||
@@ -134,7 +146,10 @@ def test_default_mm_lora_fails_with_overridden_lora_request(
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner):
|
||||
def test_default_mm_lora_does_not_expand_string_reqs(vllm_runner, monkeypatch):
|
||||
# See run_test: force spawn to avoid the forked-child CUDA re-init crash.
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
class MockEngineException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@@ -1,7 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import importlib.metadata
|
||||
from importlib.util import find_spec
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
import vllm
|
||||
from vllm.lora.request import LoRARequest
|
||||
@@ -9,6 +14,22 @@ from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import multi_gpu_test
|
||||
|
||||
# Require amd-quark >= 0.12 on torch >= 2.11.
|
||||
# Earlier torch releases work with older quark versions. See
|
||||
# https://github.com/amd/Quark/issues/34
|
||||
# TODO: Remove once amd-quark>=0.12.0
|
||||
QUARK_TORCH_COMPATIBLE = find_spec("quark") is not None and (
|
||||
version.parse(importlib.metadata.version("amd-quark")) >= version.parse("0.12.0")
|
||||
if version.parse(torch.__version__.split("+")[0]) >= version.parse("2.11")
|
||||
else True
|
||||
)
|
||||
|
||||
if current_platform.is_rocm() and not QUARK_TORCH_COMPATIBLE:
|
||||
pytest.skip(
|
||||
"This test requires amd-quark >= 0.12 on torch >= 2.11.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
MODEL_PATH = "openai/gpt-oss-20b"
|
||||
|
||||
PROMPT_TEMPLATE = """<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.
|
||||
|
||||
@@ -326,39 +326,6 @@ VLM_TEST_SETTINGS = {
|
||||
large_gpu_mark(min_gb=64),
|
||||
],
|
||||
),
|
||||
"aya_vision": VLMTestInfo(
|
||||
models=["CohereLabs/aya-vision-8b"],
|
||||
test_type=(VLMTestType.IMAGE),
|
||||
prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
|
||||
single_image_prompts=IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "<image>What's the content in the center of the image?",
|
||||
"cherry_blossom": "<image>What is the season?",
|
||||
}
|
||||
),
|
||||
multi_image_prompt="<image><image>Describe the two images in detail.",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
|
||||
),
|
||||
"aya_vision-multi_image": VLMTestInfo(
|
||||
models=["CohereLabs/aya-vision-8b"],
|
||||
test_type=(VLMTestType.MULTI_IMAGE),
|
||||
prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
|
||||
single_image_prompts=IMAGE_ASSETS.prompts(
|
||||
{
|
||||
"stop_sign": "<image>What's the content in the center of the image?",
|
||||
"cherry_blossom": "<image>What is the season?",
|
||||
}
|
||||
),
|
||||
multi_image_prompt="<image><image>Describe the two images in detail.",
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
"blip2": VLMTestInfo(
|
||||
models=["Salesforce/blip2-opt-2.7b"],
|
||||
test_type=VLMTestType.IMAGE,
|
||||
@@ -763,16 +730,6 @@ VLM_TEST_SETTINGS = {
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_output_post_proc=model_utils.llava_video_vllm_to_hf_output,
|
||||
),
|
||||
"mantis": VLMTestInfo(
|
||||
models=["TIGER-Lab/Mantis-8B-siglip-llama3"],
|
||||
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
|
||||
prompt_formatter=lambda img_prompt: f"<|start_header_id|>user<|end_header_id|>\n\n{img_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", # noqa: E501
|
||||
max_model_len=4096,
|
||||
get_stop_token_ids=lambda tok: [128009],
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_output_post_proc=model_utils.mantis_vllm_to_hf_output,
|
||||
patch_hf_runner=model_utils.mantis_patch_hf_runner,
|
||||
),
|
||||
"minicpmv_25": VLMTestInfo(
|
||||
models=["openbmb/MiniCPM-Llama3-V-2_5"],
|
||||
test_type=VLMTestType.IMAGE,
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
MODEL_NAME = "nvidia/music-flamingo-2601-hf"
|
||||
SINGLE_CONVERSATION = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this track in full detail - tell me the "
|
||||
"genre, tempo, and key, then dive into the instruments, "
|
||||
"production style, and overall mood it creates.",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
|
||||
"resolve/main/assets/song_1.mp3",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
BATCHED_CONVERSATIONS = [
|
||||
SINGLE_CONVERSATION,
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Generate a structured lyric sheet from the input music.",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/"
|
||||
"AudioSkills/resolve/main/assets/song_2.mp3",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
def get_fixture_path(filename):
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "../../fixtures/musicflamingo", filename
|
||||
)
|
||||
|
||||
|
||||
def load_expected_fixture(filename):
|
||||
fixture_path = get_fixture_path(filename)
|
||||
with open(fixture_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def assert_output_matches(output, expected_text, expected_token_ids):
|
||||
generated = output.outputs[0]
|
||||
assert generated.text == expected_text
|
||||
actual_token_ids = list(generated.token_ids)
|
||||
assert (
|
||||
actual_token_ids == expected_token_ids
|
||||
or actual_token_ids == expected_token_ids[:-1]
|
||||
or actual_token_ids[:-1] == expected_token_ids
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
model=MODEL_NAME,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=True,
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
except Exception as e:
|
||||
pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
|
||||
|
||||
# ROCm may compile decoder kernels on the first inference pass; warm up
|
||||
# once so exact fixture assertions cover the steady-state path.
|
||||
llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=SamplingParams(temperature=0.0, max_tokens=1),
|
||||
use_tqdm=False,
|
||||
)
|
||||
|
||||
return llm
|
||||
|
||||
|
||||
def test_single_generation(llm):
|
||||
expected = load_expected_fixture("expected_results_single.json")
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
|
||||
)
|
||||
|
||||
assert_output_matches(
|
||||
outputs[0],
|
||||
expected["transcriptions"][0],
|
||||
expected["token_ids"][0],
|
||||
)
|
||||
|
||||
|
||||
def test_batched_generation(llm):
|
||||
expected = load_expected_fixture("expected_results_batched.json")
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
|
||||
)
|
||||
|
||||
for i, output in enumerate(outputs):
|
||||
assert_output_matches(
|
||||
output,
|
||||
expected["transcriptions"][i],
|
||||
expected["token_ids"][i],
|
||||
)
|
||||
|
||||
|
||||
def test_single_and_batched_generation_match(llm):
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=50)
|
||||
|
||||
single_output = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
batched_output = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
|
||||
assert single_output.outputs[0].text == batched_output.outputs[0].text
|
||||
assert list(single_output.outputs[0].token_ids) == list(
|
||||
batched_output.outputs[0].token_ids
|
||||
)
|
||||
@@ -15,11 +15,14 @@ from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
|
||||
from vllm import LLM, EngineArgs, SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.kv_cache_interface import SlidingWindowSpec
|
||||
|
||||
from ....utils import ROCM_ENGINE_KWARGS
|
||||
|
||||
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
|
||||
AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
|
||||
ENGINE_CONFIG = {
|
||||
"model": MODEL_NAME,
|
||||
"max_model_len": 8192,
|
||||
@@ -60,6 +63,31 @@ def _normalize(texts: list[str]) -> list[str]:
|
||||
return texts
|
||||
|
||||
|
||||
def assert_encoder_kv_cache_spec(engine: LLM) -> None:
|
||||
vllm_config = engine.llm_engine.vllm_config
|
||||
audio_config = vllm_config.model_config.hf_config.audio_config
|
||||
kv_cache_specs_per_rank = engine.llm_engine.model_executor.get_kv_cache_specs()
|
||||
|
||||
assert len(kv_cache_specs_per_rank) == 1
|
||||
kv_cache_specs = kv_cache_specs_per_rank[0]
|
||||
assert AUDIO_LAYER_NAME in kv_cache_specs, kv_cache_specs.keys()
|
||||
spec = kv_cache_specs[AUDIO_LAYER_NAME]
|
||||
|
||||
assert audio_config.sliding_window == 750
|
||||
assert audio_config.block_pool_size == 4
|
||||
assert isinstance(spec, SlidingWindowSpec)
|
||||
assert spec.block_size == 16
|
||||
assert spec.num_kv_heads == 128
|
||||
assert spec.sliding_window == cdiv(750, 4) == 188
|
||||
assert (
|
||||
spec.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=1,
|
||||
max_model_len=vllm_config.model_config.max_model_len,
|
||||
)
|
||||
== 13
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def audio_assets() -> list[AudioAsset]:
|
||||
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
@@ -71,7 +99,9 @@ def tokenizer() -> MistralTokenizer:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def engine():
|
||||
def engine(monkeypatch: pytest.MonkeyPatch):
|
||||
# Disable multiprocessing allows us to access model executor from LLM engine
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
engine_args = EngineArgs(**ENGINE_CONFIG)
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
try:
|
||||
@@ -95,6 +125,7 @@ async def async_engine():
|
||||
|
||||
|
||||
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
|
||||
assert_encoder_kv_cache_spec(engine)
|
||||
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
|
||||
|
||||
def from_file(file_path: str):
|
||||
|
||||
@@ -167,15 +167,6 @@ def llava_onevision_vllm_to_hf_output(
|
||||
return hf_output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
def mantis_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutput:
|
||||
"""Sanitize vllm output [mantis] to compare with hf output."""
|
||||
output_ids, output_str, out_logprobs = vllm_output
|
||||
|
||||
hf_output_str = output_str + "<|eot_id|>"
|
||||
|
||||
return output_ids, hf_output_str, out_logprobs
|
||||
|
||||
|
||||
def phi3v_vllm_to_hf_output(vllm_output: RunnerOutput, model: str) -> RunnerOutput:
|
||||
"""Sanitize vllm output [phi3v] to be comparable with hf output."""
|
||||
_, output_str, out_logprobs = vllm_output
|
||||
@@ -940,29 +931,6 @@ def _internvl_generate(
|
||||
return outputs
|
||||
|
||||
|
||||
def mantis_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
from mantis.models.mllava import MLlavaProcessor
|
||||
|
||||
hf_model.processor = MLlavaProcessor.from_pretrained(hf_model.model_name)
|
||||
|
||||
orig_generate = hf_model.model.generate
|
||||
tokenizer = hf_model.processor.tokenizer
|
||||
|
||||
def _generate(self, *args, **kwargs):
|
||||
return orig_generate(
|
||||
*args,
|
||||
**kwargs,
|
||||
eos_token_id=[
|
||||
tokenizer.eos_token_id,
|
||||
tokenizer.convert_tokens_to_ids("<|eot_id|>"),
|
||||
],
|
||||
)
|
||||
|
||||
hf_model.model.generate = types.MethodType(_generate, hf_model.model)
|
||||
|
||||
return hf_model
|
||||
|
||||
|
||||
def minicpmv_25_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
orig_generate = hf_model.model.generate
|
||||
|
||||
|
||||
@@ -1,228 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 2026 The vLLM team.
|
||||
# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
|
||||
# reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from importlib.metadata import version
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
|
||||
|
||||
class MockMusicFlamingoConfig(PretrainedConfig):
|
||||
model_type = "musicflamingo"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.audio_config = PretrainedConfig()
|
||||
self.text_config = PretrainedConfig()
|
||||
|
||||
|
||||
class MockMusicFlamingoProcessor:
|
||||
def __init__(self):
|
||||
self.audio_token = "<sound>"
|
||||
self.audio_token_id = 12345
|
||||
self.audio_bos_token = "<|sound_bos|>"
|
||||
self.audio_bos_token_id = 12346
|
||||
self.audio_eos_token = "<|sound_eos|>"
|
||||
self.audio_eos_token_id = 12347
|
||||
self.max_audio_len = 1200
|
||||
self.feature_extractor = MockFeatureExtractor()
|
||||
|
||||
def __call__(self, text=None, audio=None, **kwargs):
|
||||
return {
|
||||
"input_ids": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"input_features": torch.zeros((3, 80, 3000)),
|
||||
"input_features_mask": torch.ones((3, 3000), dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
class MockFeatureExtractor:
|
||||
def __init__(self):
|
||||
self.sampling_rate = 16000
|
||||
self.chunk_length = 30
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ctx():
|
||||
config = MockMusicFlamingoConfig()
|
||||
|
||||
ctx = MagicMock()
|
||||
ctx.get_hf_config.return_value = config
|
||||
ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
|
||||
ctx.call_hf_processor.side_effect = lambda processor, data, kwargs: processor(
|
||||
**data, **kwargs
|
||||
)
|
||||
ctx.model_config.hf_config = config
|
||||
return ctx
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_transformers_version():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
|
||||
def test_musicflamingo_chunk_counting_without_rote_timestamps(mock_ctx):
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoDummyInputsBuilder,
|
||||
MusicFlamingoMultiModalProcessor,
|
||||
MusicFlamingoProcessingInfo,
|
||||
)
|
||||
|
||||
info = MusicFlamingoProcessingInfo(mock_ctx)
|
||||
processor = MusicFlamingoMultiModalProcessor(
|
||||
info, MusicFlamingoDummyInputsBuilder(info)
|
||||
)
|
||||
|
||||
sr = 16000
|
||||
audio_1 = np.zeros(30 * sr)
|
||||
audio_2 = np.zeros(45 * sr)
|
||||
|
||||
mm_data = {"audio": [audio_1, audio_2]}
|
||||
prompt = "<|user|>Listen.<|end|>"
|
||||
|
||||
processed = processor._call_hf_processor(prompt, mm_data, {}, {})
|
||||
|
||||
chunk_counts = processed["chunk_counts"]
|
||||
|
||||
assert chunk_counts.tolist() == [1, 2]
|
||||
assert "rote_timestamps" not in processed
|
||||
assert processed["feature_attention_mask"].shape == (3, 3000)
|
||||
|
||||
|
||||
def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoDummyInputsBuilder,
|
||||
MusicFlamingoProcessingInfo,
|
||||
)
|
||||
|
||||
info = MusicFlamingoProcessingInfo(mock_ctx)
|
||||
builder = MusicFlamingoDummyInputsBuilder(info)
|
||||
|
||||
assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
Version(version("transformers")) >= Version("5.5"),
|
||||
reason="transformers v5.5 added native MusicFlamingoForConditionalGeneration "
|
||||
"with a different get_audio_features signature (requires input_ids)",
|
||||
)
|
||||
def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
|
||||
from transformers.models.musicflamingo import (
|
||||
modeling_musicflamingo as hf_musicflamingo_modeling,
|
||||
)
|
||||
from transformers.models.musicflamingo.configuration_musicflamingo import (
|
||||
MusicFlamingoConfig,
|
||||
)
|
||||
|
||||
from vllm.model_executor.models.audioflamingo3 import (
|
||||
_build_audio_encoder_attention_mask,
|
||||
_flatten_valid_audio_embeddings,
|
||||
)
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoEncoder,
|
||||
MusicFlamingoMultiModalProjector,
|
||||
MusicFlamingoRotaryEmbedding,
|
||||
apply_rotary_time_emb,
|
||||
)
|
||||
|
||||
text_config = {
|
||||
"model_type": "qwen2",
|
||||
"intermediate_size": 64,
|
||||
"initializer_range": 0.02,
|
||||
"hidden_size": 32,
|
||||
"max_position_embeddings": 1024,
|
||||
"num_hidden_layers": 2,
|
||||
"num_attention_heads": 4,
|
||||
"num_key_value_heads": 2,
|
||||
"vocab_size": 128,
|
||||
"pad_token_id": 1,
|
||||
"use_mrope": False,
|
||||
}
|
||||
audio_config = {
|
||||
"hidden_size": 16,
|
||||
"num_attention_heads": 4,
|
||||
"intermediate_size": 32,
|
||||
"num_hidden_layers": 2,
|
||||
"num_mel_bins": 80,
|
||||
"max_source_positions": 1500,
|
||||
"dropout": 0.0,
|
||||
"attention_dropout": 0.0,
|
||||
"activation_dropout": 0.0,
|
||||
"encoder_layerdrop": 0.0,
|
||||
}
|
||||
|
||||
torch.manual_seed(0)
|
||||
config = MusicFlamingoConfig(
|
||||
text_config=text_config,
|
||||
audio_config=audio_config,
|
||||
audio_token_id=0,
|
||||
head_dim=8,
|
||||
rope_parameters={"rope_type": "default", "rope_theta": 2048},
|
||||
)
|
||||
hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
|
||||
config
|
||||
).eval()
|
||||
|
||||
vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
|
||||
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
|
||||
|
||||
vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
|
||||
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
|
||||
|
||||
vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
|
||||
vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)
|
||||
|
||||
input_features = torch.randn(3, 80, 3000)
|
||||
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
|
||||
feature_attention_mask[0, :3000] = True
|
||||
feature_attention_mask[1, :2500] = True
|
||||
feature_attention_mask[2, :1500] = True
|
||||
rote_timestamps = (
|
||||
torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
|
||||
)
|
||||
|
||||
hf_output = hf_model.get_audio_features(
|
||||
input_features,
|
||||
feature_attention_mask,
|
||||
rote_timestamps=rote_timestamps,
|
||||
return_dict=True,
|
||||
).pooler_output
|
||||
vllm_attention_mask = _build_audio_encoder_attention_mask(
|
||||
feature_attention_mask,
|
||||
dtype=vllm_encoder.conv1.weight.dtype,
|
||||
device=vllm_encoder.conv1.weight.device,
|
||||
)
|
||||
vllm_hidden_states = vllm_encoder(
|
||||
input_features,
|
||||
attention_mask=vllm_attention_mask,
|
||||
)
|
||||
cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
|
||||
vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
|
||||
vllm_output, _ = _flatten_valid_audio_embeddings(
|
||||
vllm_projector(vllm_hidden_states),
|
||||
feature_attention_mask,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(vllm_output, hf_output)
|
||||
@@ -780,14 +780,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
|
||||
},
|
||||
),
|
||||
"MusicFlamingoForConditionalGeneration": _HfExamplesInfo(
|
||||
"nvidia/music-flamingo-2601-hf",
|
||||
min_transformers_version="5.5.0",
|
||||
transformers_version_reason={
|
||||
"vllm": "Needs https://github.com/huggingface/transformers/pull/43538"
|
||||
},
|
||||
),
|
||||
"AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
|
||||
"BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
|
||||
"BeeForConditionalGeneration": _HfExamplesInfo(
|
||||
"Open-Bee/Bee-8B-RL",
|
||||
@@ -1059,12 +1051,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"LlavaOnevisionForConditionalGeneration": _HfExamplesInfo(
|
||||
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
|
||||
),
|
||||
"MantisForConditionalGeneration": _HfExamplesInfo(
|
||||
"TIGER-Lab/Mantis-8B-siglip-llama3",
|
||||
max_transformers_version="4.48",
|
||||
transformers_version_reason={"hf": "HF model is not compatible."},
|
||||
hf_overrides={"architectures": ["MantisForConditionalGeneration"]},
|
||||
),
|
||||
"MiDashengLMModel": _HfExamplesInfo(
|
||||
"mispeech/midashenglm-7b", trust_remote_code=True
|
||||
),
|
||||
@@ -1408,7 +1394,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
# [DFlash]
|
||||
"DFlashDraftModel": _HfExamplesInfo(
|
||||
"Qwen/Qwen3.5-4B",
|
||||
speculative_model="z-lab/Qwen3.5-4B-DFlash",
|
||||
speculative_model="z-lab/Qwen3-4B-DFlash-b16",
|
||||
use_original_num_layers=True, # Need all layers since DFlash has >1 layer,
|
||||
max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
|
||||
max_num_seqs=32,
|
||||
|
||||
@@ -10,7 +10,11 @@ import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm.assets.base import get_vllm_public_assets
|
||||
from vllm.assets.video import video_to_ndarrays, video_to_pil_images_list
|
||||
from vllm.assets.video import (
|
||||
video_get_metadata,
|
||||
video_to_ndarrays,
|
||||
video_to_pil_images_list,
|
||||
)
|
||||
from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
|
||||
from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader
|
||||
|
||||
@@ -112,6 +116,20 @@ def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext:
|
||||
assert np.nanmean(sim) > 0.99
|
||||
|
||||
|
||||
def test_opencv_video_metadata_matches_sampled_frame_timeline(tmp_path):
|
||||
image_path = f"{tmp_path}/test_metadata_image.png"
|
||||
Image.new("RGB", (8, 8), color=(255, 0, 0)).save(image_path)
|
||||
video_path = f"{tmp_path}/test_metadata_video.mp4"
|
||||
create_video_from_image(image_path, video_path, num_frames=10, fps=5.0)
|
||||
|
||||
metadata = video_get_metadata(video_path, num_frames=4)
|
||||
|
||||
assert metadata["fps"] == pytest.approx(5.0)
|
||||
assert metadata["duration"] == pytest.approx(2.0)
|
||||
assert metadata["frames_indices"] == [0, 3, 6, 9]
|
||||
assert metadata["total_num_frames"] == 4
|
||||
|
||||
|
||||
NUM_FRAMES = 10
|
||||
FAKE_OUTPUT_1 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
|
||||
FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user