forked from Karylab-cklius/vllm
Compare commits
86
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0116e1cedc | ||
|
|
80f66afd81 | ||
|
|
1fa4c3adb7 | ||
|
|
4787f2dd1b | ||
|
|
8cfeb84dba | ||
|
|
5fd442187c | ||
|
|
00eb7cefa3 | ||
|
|
c8bdcc0116 | ||
|
|
f5a8d73377 | ||
|
|
63fcce4de1 | ||
|
|
c638f9216a | ||
|
|
13c49f9845 | ||
|
|
f1cf6b0086 | ||
|
|
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 | ||
|
|
d8f483dc30 | ||
|
|
dc148dc4d7 | ||
|
|
7cf7cbcd95 | ||
|
|
c231d1f290 | ||
|
|
db808b3961 | ||
|
|
00ebf19cca | ||
|
|
ded6676458 | ||
|
|
7a327f0b4f | ||
|
|
1ab9522935 | ||
|
|
0fc2512094 | ||
|
|
62c7d8009f | ||
|
|
ab80b3dff4 | ||
|
|
91055efd36 | ||
|
|
3675bcff67 | ||
|
|
bdbd7278b6 | ||
|
|
5dc36a4fa5 | ||
|
|
aab7af0bcb |
@@ -81,7 +81,7 @@ steps:
|
||||
'cd tests &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
set -o pipefail &&
|
||||
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
|
||||
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]" --deselect="tests/lora/test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-32000-32-4-4]"'
|
||||
|
||||
- label: LoRA Punica FP8/XPU Ops
|
||||
timeout_in_minutes: 45
|
||||
|
||||
@@ -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}" \
|
||||
|
||||
+35
-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
|
||||
|
||||
@@ -2828,9 +2816,10 @@ steps:
|
||||
- rocm-smi
|
||||
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
- 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
|
||||
@@ -2979,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
|
||||
|
||||
@@ -2993,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
|
||||
@@ -3019,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
|
||||
|
||||
@@ -3034,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
|
||||
@@ -3177,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
|
||||
@@ -3189,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)
|
||||
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
runs-on: [self-hosted, linux, x64, vllm-runners]
|
||||
steps:
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
|
||||
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
# Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ |
|
||||
@@ -626,8 +624,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
|
||||
| `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
|
||||
| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
|
||||
| `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
|
||||
| `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
|
||||
| `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |
|
||||
| `UnlimitedOCRForCausalLM` | Unlimited-OCR | T + I<sup>+</sup> | `baidu/Unlimited-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -680,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.
|
||||
|
||||
@@ -327,6 +327,9 @@ async def handle_request(api: str, request: Request):
|
||||
session, decode_response = await decode_request_task
|
||||
stream_generator = stream_decode_response(session, decode_response, request_id)
|
||||
response = await make_response(stream_generator)
|
||||
response.headers["Content-Type"] = decode_response.headers.get(
|
||||
"Content-Type", "application/json"
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
logger.exception("An error occurred while handling the request: %s", e)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -1275,55 +1242,6 @@ def load_step_vl(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "omni-research/Tarsier-7b"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
limit_mm_per_prompt={"image": len(image_urls)},
|
||||
)
|
||||
|
||||
prompt = f"USER: {'<image>' * len(image_urls)}\n{question}\n ASSISTANT:"
|
||||
image_data = [fetch_image(url) for url in image_urls]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
image_data=image_data,
|
||||
)
|
||||
|
||||
|
||||
def load_tarsier2(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "omni-research/Tarsier2-Recap-7b"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
trust_remote_code=True,
|
||||
max_model_len=32768,
|
||||
limit_mm_per_prompt={"image": len(image_urls)},
|
||||
hf_overrides={
|
||||
"architectures": ["Tarsier2ForConditionalGeneration"],
|
||||
"model_type": "tarsier2",
|
||||
},
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
f"<|im_start|>user\n<|vision_start|>{'<|image_pad|>' * len(image_urls)}"
|
||||
f"<|vision_end|>{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
image_data = [fetch_image(url) for url in image_urls]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
image_data=image_data,
|
||||
)
|
||||
|
||||
|
||||
# GLM-4.1V
|
||||
def load_glm4_1v(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "zai-org/GLM-4.1V-9B-Thinking"
|
||||
@@ -1469,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,
|
||||
@@ -1507,8 +1424,6 @@ model_example_map = {
|
||||
"smolvlm": load_smolvlm,
|
||||
"step3": load_step3,
|
||||
"stepvl": load_step_vl,
|
||||
"tarsier": load_tarsier,
|
||||
"tarsier2": load_tarsier2,
|
||||
"glm4_1v": load_glm4_1v,
|
||||
"glm4_5v": load_glm4_5v,
|
||||
"glm4_5v_fp8": load_glm4_5v_fp8,
|
||||
|
||||
@@ -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"]
|
||||
@@ -2347,68 +2303,8 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# omni-research/Tarsier-7b
|
||||
def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
model_name = "omni-research/Tarsier-7b"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
)
|
||||
prompts = [(f"USER: <image>\n{question} ASSISTANT:") for question in questions]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
|
||||
model_name = "omni-research/Tarsier2-Recap-7b"
|
||||
|
||||
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
hf_overrides={
|
||||
"architectures": ["Tarsier2ForConditionalGeneration"],
|
||||
"model_type": "tarsier2",
|
||||
},
|
||||
limit_mm_per_prompt=mm_limit,
|
||||
)
|
||||
|
||||
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
|
||||
|
||||
if modality == "image":
|
||||
placeholder = image_placeholder
|
||||
elif modality == "video":
|
||||
placeholder = video_placeholder
|
||||
elif modality == "image+video":
|
||||
placeholder = image_placeholder + video_placeholder
|
||||
|
||||
prompts = [
|
||||
(
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
f"<|im_start|>user\n{placeholder}"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
for question in questions
|
||||
]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
model_example_map = {
|
||||
"aria": run_aria,
|
||||
"aya_vision": run_aya_vision,
|
||||
"bagel": run_bagel,
|
||||
"cheers": run_cheers,
|
||||
"bee": run_bee,
|
||||
@@ -2449,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,
|
||||
@@ -2479,8 +2374,6 @@ model_example_map = {
|
||||
"smolvlm": run_smolvlm,
|
||||
"step3": run_step3,
|
||||
"stepvl": run_step_vl,
|
||||
"tarsier": run_tarsier,
|
||||
"tarsier2": run_tarsier2,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -8,7 +8,7 @@ torch==2.11.0
|
||||
torchaudio==2.11.0
|
||||
# These must be updated alongside torch
|
||||
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
PyNvVideoCodec==2.0.4
|
||||
PyNvVideoCodec==2.1.0
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
flashinfer-python==0.6.13
|
||||
flashinfer-cubin==0.6.13
|
||||
|
||||
Generated
+1
@@ -5310,6 +5310,7 @@ dependencies = [
|
||||
"vllm-llm",
|
||||
"vllm-metrics",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
|
||||
@@ -50,6 +50,7 @@ tokio.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-engine-core-client = { workspace = true, features = ["test-util"] }
|
||||
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
|
||||
zeromq.workspace = true
|
||||
|
||||
[lints]
|
||||
|
||||
@@ -154,7 +154,8 @@ mod tests {
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::TokenizerSource;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::HfChatBackend;
|
||||
use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
|
||||
@@ -196,32 +197,8 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct TestTokenizer;
|
||||
|
||||
impl Tokenizer for TestTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
_text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
_token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn test_tokenizer() -> DynTokenizer {
|
||||
Arc::new(TestTokenizer)
|
||||
Arc::new(TestTokenizer::new())
|
||||
}
|
||||
|
||||
fn backend_for_selection(
|
||||
|
||||
@@ -50,7 +50,8 @@ mod request;
|
||||
mod stream;
|
||||
|
||||
use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::{ModelDtype, ReasoningParserKwargs};
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ use llm_multimodal::{
|
||||
TrackedMedia,
|
||||
};
|
||||
use tracing::warn;
|
||||
use vllm_engine_core_client::protocol::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::{
|
||||
MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem,
|
||||
MmSharedField, MmSlice, PlaceholderRange, SliceSpec,
|
||||
@@ -563,7 +563,7 @@ mod tests {
|
||||
|
||||
use llm_multimodal::TokenId;
|
||||
use vllm_engine_core_client::protocol::tensor::WireArrayData;
|
||||
use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -574,60 +574,14 @@ mod tests {
|
||||
const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093;
|
||||
const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094;
|
||||
|
||||
struct TestTokenizer;
|
||||
|
||||
impl Tokenizer for TestTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> std::result::Result<Vec<u32>, TokenizerError> {
|
||||
Ok(match text {
|
||||
"<|image|>" => vec![LLAMA4_IMAGE_ID],
|
||||
text => text.bytes().map(u32::from).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
_token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> std::result::Result<String, TokenizerError> {
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<|image_start|>" => Some(LLAMA4_IMAGE_START_ID),
|
||||
"<|image_end|>" => Some(LLAMA4_IMAGE_END_ID),
|
||||
"<|image|>" => Some(LLAMA4_IMAGE_ID),
|
||||
"<|patch|>" => Some(LLAMA4_PATCH_ID),
|
||||
"<|tile_x_separator|>" => Some(LLAMA4_TILE_X_SEPARATOR_ID),
|
||||
"<|tile_y_separator|>" => Some(LLAMA4_TILE_Y_SEPARATOR_ID),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
match id {
|
||||
LLAMA4_IMAGE_START_ID => Some("<|image_start|>".to_string()),
|
||||
LLAMA4_IMAGE_END_ID => Some("<|image_end|>".to_string()),
|
||||
LLAMA4_IMAGE_ID => Some("<|image|>".to_string()),
|
||||
LLAMA4_PATCH_ID => Some("<|patch|>".to_string()),
|
||||
LLAMA4_TILE_X_SEPARATOR_ID => Some("<|tile_x_separator|>".to_string()),
|
||||
LLAMA4_TILE_Y_SEPARATOR_ID => Some("<|tile_y_separator|>".to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_decode_stream(
|
||||
&self,
|
||||
_prompt_token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
_min_bytes_to_buffer: usize,
|
||||
) -> Box<dyn IncrementalDecoder + '_> {
|
||||
unreachable!("not used")
|
||||
}
|
||||
fn llama4_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<|image_start|>", LLAMA4_IMAGE_START_ID)
|
||||
.with_regular_token("<|image_end|>", LLAMA4_IMAGE_END_ID)
|
||||
.with_regular_token("<|image|>", LLAMA4_IMAGE_ID)
|
||||
.with_regular_token("<|patch|>", LLAMA4_PATCH_ID)
|
||||
.with_regular_token("<|tile_x_separator|>", LLAMA4_TILE_X_SEPARATOR_ID)
|
||||
.with_regular_token("<|tile_y_separator|>", LLAMA4_TILE_Y_SEPARATOR_ID)
|
||||
}
|
||||
|
||||
fn test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo {
|
||||
@@ -635,7 +589,7 @@ mod tests {
|
||||
model_id: format!("{model_type}-test"),
|
||||
model_type: Some(model_type.to_string()),
|
||||
config,
|
||||
tokenizer: TokenizerResolver(Arc::new(TestTokenizer)),
|
||||
tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())),
|
||||
};
|
||||
let spec = context
|
||||
.resolve_model_spec()
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::HashMap;
|
||||
|
||||
use half::{bf16, f16};
|
||||
use llm_multimodal::{ModelSpecificValue, PreprocessedImages};
|
||||
use vllm_engine_core_client::protocol::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
|
||||
use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
|
||||
|
||||
|
||||
@@ -189,46 +189,19 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::DefaultChatOutputProcessor;
|
||||
use crate::Error;
|
||||
use crate::parser::ParserSelection;
|
||||
use crate::request::ChatRequest;
|
||||
|
||||
struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.chars().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(token_ids
|
||||
.iter()
|
||||
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<|channel>" => Some(1),
|
||||
"<channel|>" => Some(2),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tokenizer() -> Arc<FakeTokenizer> {
|
||||
Arc::new(FakeTokenizer)
|
||||
fn tokenizer() -> Arc<TestTokenizer> {
|
||||
Arc::new(
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<|channel>", 256)
|
||||
.with_regular_token("<channel|>", 257),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
//! Applies xgrammar structural-tag constraints for strict tool calling.
|
||||
|
||||
use thiserror_ext::AsReport;
|
||||
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::StructuralTagModel;
|
||||
use xgrammar_structural_tag::{
|
||||
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
|
||||
@@ -76,7 +78,9 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagTool
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -1,14 +1,8 @@
|
||||
//! Harmony output tests share the upstream `openai-harmony` tiktoken cache.
|
||||
//!
|
||||
//! Use a file lock for tests that load the encoding so `cargo nextest` cannot
|
||||
//! start multiple processes that concurrently populate the same cache file.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::executor::block_on;
|
||||
use futures::{TryStreamExt as _, stream};
|
||||
use openai_harmony::chat::{Message, Role};
|
||||
use serial_test::file_serial;
|
||||
use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
|
||||
|
||||
use super::*;
|
||||
@@ -91,7 +85,6 @@ fn request_with_tools() -> ChatRequest {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn interrupted_final_message_is_preserved() {
|
||||
let tokens = completion_tokens(&[text_message("final", "hello")]);
|
||||
let events = block_on(collect_events(
|
||||
@@ -127,7 +120,6 @@ fn interrupted_final_message_is_preserved() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn eos_flush_preserves_trailing_replacement_text() {
|
||||
let mut tokens = completion_tokens(&[text_message("final", "Hi")]);
|
||||
tokens.pop();
|
||||
@@ -153,7 +145,6 @@ fn eos_flush_preserves_trailing_replacement_text() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn interrupted_analysis_message_is_preserved() {
|
||||
let tokens = completion_tokens(&[text_message("analysis", "think")]);
|
||||
let events = block_on(collect_events(
|
||||
@@ -189,7 +180,6 @@ fn interrupted_analysis_message_is_preserved() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
|
||||
let tokens = completion_tokens(&[
|
||||
text_message("commentary", "Let me check."),
|
||||
@@ -217,7 +207,6 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn multiple_messages_get_newline_separators() {
|
||||
let tokens = completion_tokens(&[
|
||||
text_message("analysis", "first think"),
|
||||
@@ -249,7 +238,6 @@ fn multiple_messages_get_newline_separators() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
|
||||
let tokens = completion_tokens(&[tool_message(
|
||||
"get_weather",
|
||||
@@ -302,7 +290,6 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn semantic_events_precede_same_update_logprobs() {
|
||||
let tokens = completion_tokens(&[text_message("final", "hello")]);
|
||||
let events = block_on(collect_events(
|
||||
@@ -353,7 +340,6 @@ fn rejects_generic_parser_overrides() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[file_serial(harmony_tiktoken_cache)]
|
||||
fn allows_auto_auto_only() {
|
||||
validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap();
|
||||
let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap();
|
||||
|
||||
@@ -1,32 +1,9 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::{ReasoningParserFactory, names};
|
||||
|
||||
struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.chars().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(token_ids
|
||||
.iter()
|
||||
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_contains_and_lists_registered_parsers() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
@@ -107,7 +84,7 @@ fn factory_resolves_minimax_m3_before_generic_minimax() {
|
||||
|
||||
#[test]
|
||||
fn factory_rejects_unknown_parser_names() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(TestTokenizer::new());
|
||||
let factory = ReasoningParserFactory::new();
|
||||
let error = match factory.create("missing", tokenizer) {
|
||||
Ok(_) => panic!("expected parser lookup to fail"),
|
||||
|
||||
@@ -75,39 +75,14 @@ impl UnifiedParserFactory {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::{UnifiedParserFactory, names};
|
||||
|
||||
struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.chars().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(token_ids
|
||||
.iter()
|
||||
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<|channel>" => Some(1),
|
||||
"<channel|>" => Some(2),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
fn tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<|channel>", 256)
|
||||
.with_regular_token("<channel|>", 257)
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -119,6 +94,6 @@ mod tests {
|
||||
factory.resolve_name_for_model("google/gemma-4-27b-it"),
|
||||
Some(names::GEMMA4)
|
||||
);
|
||||
factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap();
|
||||
factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
+37
-105
@@ -15,21 +15,24 @@ use vllm_chat::{
|
||||
use vllm_engine_core_client::protocol::logprobs::{
|
||||
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
|
||||
use vllm_engine_core_client::protocol::output::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, StopReason,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::{
|
||||
DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTokenLogprob, Prompt,
|
||||
TextBackend,
|
||||
};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
const SPECIAL_STOP_TOKEN_ID: u32 = 256;
|
||||
const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
|
||||
|
||||
fn request_output(
|
||||
request_id: &str,
|
||||
@@ -158,45 +161,18 @@ async fn connect_chat_llm_with_ipc(
|
||||
struct FakeChatBackend {
|
||||
has_template: bool,
|
||||
model_id: String,
|
||||
tokenizer: DynTokenizer,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FakeChatTokenizer;
|
||||
|
||||
impl Tokenizer for FakeChatTokenizer {
|
||||
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.bytes().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
let bytes = token_ids
|
||||
.iter()
|
||||
.filter_map(|id| {
|
||||
if skip_special_tokens && *id == SPECIAL_STOP_TOKEN_ID {
|
||||
None
|
||||
} else {
|
||||
Some(*id as u8)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(String::from_utf8_lossy(&bytes).into_owned())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<think>" => Some(0xF001),
|
||||
"</think>" => Some(0xF002),
|
||||
"<|START_THINKING|>" => Some(0xF003),
|
||||
"<|END_THINKING|>" => Some(0xF004),
|
||||
"◁think▷" => Some(0xF005),
|
||||
"◁/think▷" => Some(0xF006),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
fn fake_chat_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_special_token("<stop>", SPECIAL_STOP_TOKEN_ID)
|
||||
.with_regular_token("<think>", 0xF001)
|
||||
.with_regular_token("</think>", 0xF002)
|
||||
.with_regular_token("<|START_THINKING|>", 0xF003)
|
||||
.with_regular_token("<|END_THINKING|>", 0xF004)
|
||||
.with_regular_token("◁think▷", 0xF005)
|
||||
.with_regular_token("◁/think▷", 0xF006)
|
||||
}
|
||||
|
||||
impl fmt::Debug for FakeChatBackend {
|
||||
@@ -210,6 +186,7 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: true,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -217,6 +194,7 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: false,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,13 +202,19 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: true,
|
||||
model_id: model_id.into(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
fn with_tokenizer(mut self, tokenizer: DynTokenizer) -> Self {
|
||||
self.tokenizer = tokenizer;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl TextBackend for FakeChatBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FakeChatTokenizer)
|
||||
Arc::clone(&self.tokenizer)
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
@@ -282,65 +266,6 @@ impl ChatRenderer for FakeChatBackend {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct FailingDecodeBackend {
|
||||
inner: FakeChatBackend,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FailingDecodeTokenizer;
|
||||
|
||||
impl Tokenizer for FailingDecodeTokenizer {
|
||||
fn encode(&self, text: &str, add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
FakeChatTokenizer.encode(text, add_special_tokens)
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
if token_ids.contains(&(b'i' as u32)) {
|
||||
return Err(vllm_tokenizer::TokenizerError("decode failed".to_string()));
|
||||
}
|
||||
FakeChatTokenizer.decode(token_ids, skip_special_tokens)
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
FakeChatTokenizer.token_to_id(token)
|
||||
}
|
||||
}
|
||||
|
||||
impl TextBackend for FailingDecodeBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FailingDecodeTokenizer)
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
self.inner.model_id()
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatBackend for FailingDecodeBackend {
|
||||
fn chat_renderer(&self) -> DynChatRenderer {
|
||||
Arc::new(self.clone())
|
||||
}
|
||||
|
||||
fn new_chat_output_processor(
|
||||
&self,
|
||||
_request: &mut ChatRequest,
|
||||
_options: NewChatOutputProcessorOptions<'_>,
|
||||
) -> vllm_chat::Result<DynChatOutputProcessor> {
|
||||
Ok(Box::new(DefaultChatOutputProcessor::plain_text_only()))
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatRenderer for FailingDecodeBackend {
|
||||
fn render(&self, request: &ChatRequest) -> vllm_chat::Result<RenderedPrompt> {
|
||||
self.inner.render(request)
|
||||
}
|
||||
}
|
||||
|
||||
/// Skip `LogprobsDelta` events that carry only token_ids (no logprobs),
|
||||
/// returning the next semantically interesting event.
|
||||
async fn next_semantic<S>(stream: &mut S) -> Option<Result<ChatEvent, vllm_chat::Error>>
|
||||
@@ -738,7 +663,12 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output("chat-4", vec![b'i' as u32], None, None)],
|
||||
outputs: vec![request_output(
|
||||
"chat-4",
|
||||
vec![UNKNOWN_DECODE_TOKEN_ID],
|
||||
None,
|
||||
None,
|
||||
)],
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
@@ -747,9 +677,8 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
|
||||
},
|
||||
);
|
||||
|
||||
let backend: Arc<dyn ChatTextBackend> = Arc::new(FailingDecodeBackend {
|
||||
inner: FakeChatBackend::new(),
|
||||
});
|
||||
let backend: Arc<dyn ChatTextBackend> =
|
||||
Arc::new(FakeChatBackend::new().with_tokenizer(Arc::new(TestTokenizer::new())));
|
||||
let chat = connect_chat_llm_with_ipc(
|
||||
EngineCoreClientConfig::new_single(handshake_address),
|
||||
&ipc,
|
||||
@@ -769,7 +698,10 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
|
||||
|
||||
match timeout(Duration::from_secs(2), stream.next()).await.unwrap() {
|
||||
Some(Err(vllm_chat::Error::Text(vllm_text::Error::Tokenizer(message)))) => {
|
||||
assert_eq!(message, "decode failed");
|
||||
assert_eq!(
|
||||
message,
|
||||
format!("test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}")
|
||||
);
|
||||
}
|
||||
other => panic!("unexpected event after close: {other:?}"),
|
||||
}
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
|
||||
//! Roundtrip tests for the real chat-template and output-processor pairing.
|
||||
//!
|
||||
//! The invariant under test is that a structured assistant message rendered as history can be
|
||||
//! parsed from the generated assistant completion and then rendered back to the exact same
|
||||
//! assistant-completion text.
|
||||
//! assistant completion.
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
@@ -18,6 +18,10 @@ use vllm_chat::{
|
||||
RendererSelection, load_model_backends,
|
||||
};
|
||||
use vllm_text::{DecodedTextEvent, Finished, Prompt};
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
const TEXT_COMPLETION_CHUNK_CHARS: usize = 7;
|
||||
const TOKEN_COMPLETION_CHUNK_TOKENS: usize = 1;
|
||||
|
||||
/// One model/parser configuration used to run the fixed roundtrip fixtures.
|
||||
#[derive(Clone)]
|
||||
@@ -191,14 +195,28 @@ impl RoundtripCase {
|
||||
sort_json_keys: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// GPT-OSS Harmony token-id renderer and native Harmony output processor.
|
||||
fn gpt_oss() -> Self {
|
||||
Self {
|
||||
model_id: "openai/gpt-oss-20b",
|
||||
assistant_stop_suffix: "", // not applicable for token-id cases
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Always { value: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
sort_json_keys: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! roundtrip_tests {
|
||||
($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
|
||||
($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
|
||||
paste::paste! {
|
||||
$(
|
||||
$(
|
||||
#[tokio::test]
|
||||
$(#[$fixture_attr])*
|
||||
#[file_serial([<hf_ $case>])]
|
||||
async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
|
||||
[<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
|
||||
@@ -217,9 +235,9 @@ roundtrip_tests! {
|
||||
glm47 => [reasoning_and_content, tool_call_mix],
|
||||
seed_oss => [reasoning_and_content],
|
||||
step3p5 => [reasoning_and_content],
|
||||
|
||||
gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
|
||||
kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
|
||||
gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
|
||||
}
|
||||
|
||||
/// Run the fixed reasoning+content fixture for one model/parser case.
|
||||
@@ -421,10 +439,10 @@ struct RoundtripResult {
|
||||
parsed_message: AssistantMessage,
|
||||
/// Assistant-completion suffix cut from rendering the expected assistant as
|
||||
/// history.
|
||||
closed_completion: String,
|
||||
closed_completion: Prompt,
|
||||
/// Assistant-completion suffix cut after rendering the parsed assistant
|
||||
/// back as history.
|
||||
rerendered_closed_completion: String,
|
||||
rerendered_closed_completion: Prompt,
|
||||
}
|
||||
|
||||
/// Render, parse, and rerender one assistant turn through the production
|
||||
@@ -436,60 +454,59 @@ async fn run_roundtrip(
|
||||
assistant: AssistantMessage,
|
||||
) -> Result<RoundtripResult> {
|
||||
let renderer = backends.chat_backend.chat_renderer();
|
||||
let (prompt, closed_completion_text) =
|
||||
render_closed_completion(renderer.as_ref(), request, &assistant)?;
|
||||
let completion_body = closed_completion_text
|
||||
.strip_suffix(case.assistant_stop_suffix)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"closed assistant completion did not end with {:?}: {:?}",
|
||||
case.assistant_stop_suffix, closed_completion_text
|
||||
)
|
||||
})?;
|
||||
let rendered = render_closed_completion(renderer.as_ref(), request, &assistant)?;
|
||||
|
||||
let parsed_message =
|
||||
parse_completion(case, backends, request, &prompt, completion_body).await?;
|
||||
let (_, rerendered_closed_completion) =
|
||||
render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
|
||||
let parsed_message = parse_completion(case, backends, request, &rendered).await?;
|
||||
let rerendered = render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
|
||||
|
||||
Ok(RoundtripResult {
|
||||
parsed_message,
|
||||
closed_completion: closed_completion_text,
|
||||
rerendered_closed_completion,
|
||||
closed_completion: rendered.completion,
|
||||
rerendered_closed_completion: rerendered.completion,
|
||||
})
|
||||
}
|
||||
|
||||
/// Rendered prompt/completion artifacts at the renderer boundary.
|
||||
struct RenderedTurn {
|
||||
prompt: Prompt,
|
||||
completion: Prompt,
|
||||
}
|
||||
|
||||
/// Render `history` as a production prompt and `history + assistant` as closed
|
||||
/// history, then return the production prompt and assistant-completion suffix.
|
||||
fn render_closed_completion(
|
||||
renderer: &dyn vllm_chat::ChatRenderer,
|
||||
base_request: &ChatRequest,
|
||||
assistant: &AssistantMessage,
|
||||
) -> Result<(String, String)> {
|
||||
) -> Result<RenderedTurn> {
|
||||
let mut prompt_request = base_request.clone();
|
||||
prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
|
||||
let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
|
||||
let prompt = renderer.render(&prompt_request).context("failed to render prompt")?.prompt;
|
||||
|
||||
let mut full_request = base_request.clone();
|
||||
full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
|
||||
full_request.messages.push(ChatMessage::from(assistant.clone()));
|
||||
let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
|
||||
let full = renderer.render(&full_request).context("failed to render full prompt")?.prompt;
|
||||
|
||||
ensure!(
|
||||
full.starts_with(&prompt),
|
||||
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
|
||||
);
|
||||
let completion = full[prompt.len()..].to_string();
|
||||
let completion = match (&prompt, full) {
|
||||
(Prompt::Text(prompt), Prompt::Text(full)) => {
|
||||
ensure!(
|
||||
full.starts_with(prompt),
|
||||
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
|
||||
);
|
||||
Prompt::Text(full[prompt.len()..].to_string())
|
||||
}
|
||||
(Prompt::TokenIds(prompt), Prompt::TokenIds(full)) => {
|
||||
ensure!(
|
||||
full.starts_with(prompt),
|
||||
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
|
||||
);
|
||||
Prompt::TokenIds(full[prompt.len()..].to_vec())
|
||||
}
|
||||
(prompt, full) => bail!("prompt kind changed between renders: {prompt:?} vs {full:?}"),
|
||||
};
|
||||
|
||||
Ok((prompt, completion))
|
||||
}
|
||||
|
||||
/// Render one chat request and require a text prompt.
|
||||
fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
|
||||
match renderer.render(request)?.prompt {
|
||||
Prompt::Text(text) => Ok(text),
|
||||
other => bail!("roundtrip tests expect text prompts, got {other:?}"),
|
||||
}
|
||||
Ok(RenderedTurn { prompt, completion })
|
||||
}
|
||||
|
||||
/// Feed one rendered assistant completion body into the real output processor
|
||||
@@ -498,13 +515,15 @@ async fn parse_completion(
|
||||
case: &RoundtripCase,
|
||||
backends: &vllm_chat::LoadedModelBackends,
|
||||
base_request: &ChatRequest,
|
||||
prompt: &str,
|
||||
completion_body: &str,
|
||||
rendered: &RenderedTurn,
|
||||
) -> Result<AssistantMessage> {
|
||||
let tokenizer = backends.text_backend.tokenizer();
|
||||
let prompt_token_ids = tokenizer
|
||||
.encode(prompt, base_request.add_special_tokens)
|
||||
.context("failed to encode rendered prompt")?;
|
||||
let prompt_token_ids = match &rendered.prompt {
|
||||
Prompt::Text(prompt) => tokenizer
|
||||
.encode(prompt, base_request.add_special_tokens)
|
||||
.context("failed to encode rendered prompt")?,
|
||||
Prompt::TokenIds(token_ids) => token_ids.clone(),
|
||||
};
|
||||
|
||||
let mut request = base_request.clone();
|
||||
let processor = backends.chat_backend.new_chat_output_processor(
|
||||
@@ -515,7 +534,12 @@ async fn parse_completion(
|
||||
},
|
||||
)?;
|
||||
|
||||
let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
|
||||
let decoded = decoded_completion_stream(
|
||||
tokenizer.as_ref(),
|
||||
prompt_token_ids,
|
||||
&rendered.completion,
|
||||
case.assistant_stop_suffix,
|
||||
)?;
|
||||
let mut events = processor.process(decoded)?;
|
||||
|
||||
while let Some(event) = events.next().await {
|
||||
@@ -538,16 +562,46 @@ async fn parse_completion(
|
||||
/// split into small chunks to exercise streaming parser state across marker
|
||||
/// and JSON boundaries.
|
||||
fn decoded_completion_stream(
|
||||
tokenizer: &dyn Tokenizer,
|
||||
prompt_token_ids: Vec<u32>,
|
||||
completion_body: &str,
|
||||
) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
|
||||
let prompt_token_count = prompt_token_ids.len();
|
||||
completion: &Prompt,
|
||||
assistant_stop_suffix: &str,
|
||||
) -> Result<Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>>> {
|
||||
let mut events = vec![DecodedTextEvent::Start {
|
||||
prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
|
||||
prompt_token_ids: Arc::from(prompt_token_ids.clone().into_boxed_slice()),
|
||||
prompt_logprobs: None,
|
||||
}];
|
||||
|
||||
let chunks = split_by_chars(completion_body, 7);
|
||||
let chunks = match completion {
|
||||
Prompt::Text(text) => {
|
||||
let body = text.strip_suffix(assistant_stop_suffix).with_context(|| {
|
||||
format!(
|
||||
"closed assistant completion did not end with {:?}: {:?}",
|
||||
assistant_stop_suffix, text
|
||||
)
|
||||
})?;
|
||||
split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
|
||||
.into_iter()
|
||||
.map(|delta| DecodedCompletionChunk {
|
||||
delta,
|
||||
token_ids: Vec::new(), // unused for text-level roundtrip cases
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
Prompt::TokenIds(token_ids) => {
|
||||
ensure!(
|
||||
assistant_stop_suffix.is_empty(),
|
||||
"token-id roundtrip cases do not support text stop suffixes"
|
||||
);
|
||||
incremental_decode_chunks(
|
||||
tokenizer,
|
||||
&prompt_token_ids,
|
||||
token_ids,
|
||||
TOKEN_COMPLETION_CHUNK_TOKENS,
|
||||
)?
|
||||
}
|
||||
};
|
||||
|
||||
if chunks.is_empty() {
|
||||
events.push({
|
||||
DecodedTextEvent::TextDelta {
|
||||
@@ -555,11 +609,7 @@ fn decoded_completion_stream(
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished: Some(Finished {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 0,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
usage: Default::default(),
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
@@ -569,24 +619,26 @@ fn decoded_completion_stream(
|
||||
let last_index = chunks.len() - 1;
|
||||
for (index, chunk) in chunks.into_iter().enumerate() {
|
||||
let finished = (index == last_index).then(|| Finished {
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count,
|
||||
output_token_count: completion_body.chars().count(),
|
||||
cached_token_count: 0,
|
||||
},
|
||||
usage: Default::default(),
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
});
|
||||
events.push(DecodedTextEvent::TextDelta {
|
||||
delta: chunk,
|
||||
token_ids: Vec::new(),
|
||||
delta: chunk.delta,
|
||||
token_ids: chunk.token_ids,
|
||||
logprobs: None,
|
||||
finished,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
stream::iter(events).map(Ok).boxed()
|
||||
Ok(stream::iter(events).map(Ok).boxed())
|
||||
}
|
||||
|
||||
/// One decoded completion chunk fed into the output processor.
|
||||
struct DecodedCompletionChunk {
|
||||
delta: String,
|
||||
token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Split text into chunks containing at most `chunk_chars` Unicode scalar
|
||||
@@ -612,6 +664,49 @@ fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Split token ids into chunks containing at most `chunk_size` ids.
|
||||
fn split_by_count(token_ids: &[u32], chunk_size: usize) -> Vec<Vec<u32>> {
|
||||
token_ids.chunks(chunk_size).map(<[u32]>::to_vec).collect()
|
||||
}
|
||||
|
||||
/// Decode token ids incrementally using the production tokenizer stream.
|
||||
fn incremental_decode_chunks(
|
||||
tokenizer: &dyn Tokenizer,
|
||||
prompt_token_ids: &[u32],
|
||||
token_ids: &[u32],
|
||||
chunk_size: usize,
|
||||
) -> Result<Vec<DecodedCompletionChunk>> {
|
||||
let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
|
||||
let mut chunks = Vec::new();
|
||||
for chunk_token_ids in split_by_count(token_ids, chunk_size) {
|
||||
let mut delta = String::new();
|
||||
for token_id in chunk_token_ids.iter().copied() {
|
||||
decoder.push_token(token_id)?;
|
||||
while let Some(chunk) = decoder.next_chunk() {
|
||||
delta.push_str(&chunk);
|
||||
}
|
||||
}
|
||||
chunks.push(DecodedCompletionChunk {
|
||||
delta,
|
||||
token_ids: chunk_token_ids,
|
||||
});
|
||||
}
|
||||
|
||||
let (last_chunk, _) = decoder.flush(None)?;
|
||||
if let Some(last_chunk) = last_chunk {
|
||||
if let Some(delta) = chunks.last_mut() {
|
||||
delta.delta.push_str(&last_chunk);
|
||||
} else {
|
||||
chunks.push(DecodedCompletionChunk {
|
||||
delta: last_chunk,
|
||||
token_ids: Vec::new(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Build a chat request fixture with parser-enabling tool-choice semantics.
|
||||
fn roundtrip_request(
|
||||
request_id: impl Into<String>,
|
||||
|
||||
@@ -5,9 +5,9 @@ use clap::Parser;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::{
|
||||
EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
|
||||
};
|
||||
|
||||
@@ -11,10 +11,11 @@ use tracing::{debug, info, trace};
|
||||
use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
|
||||
use crate::coordinator::CoordinatorHandle;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
use crate::protocol::handshake::EngineCoreReadyResponse;
|
||||
use crate::protocol::lora::LoraRequest;
|
||||
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
|
||||
use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
|
||||
use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
|
||||
use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
|
||||
use crate::transport::{self, ConnectedEngine};
|
||||
|
||||
|
||||
@@ -16,12 +16,11 @@ use crate::client::stream::EngineCoreStreamOutput;
|
||||
use crate::client::{AbortCause, AbortRequest};
|
||||
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
|
||||
use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
|
||||
use crate::protocol::encode_msgpack;
|
||||
use crate::protocol::output::{ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs};
|
||||
use crate::protocol::request::EngineCoreRequestType;
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
use crate::protocol::{
|
||||
ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequestType,
|
||||
encode_msgpack,
|
||||
};
|
||||
use crate::transport::{ConnectedEngine, EngineId};
|
||||
use crate::{Error, Result, transport};
|
||||
|
||||
|
||||
@@ -7,9 +7,9 @@ use tracing::trace;
|
||||
use crate::EngineId;
|
||||
use crate::client::stream::EngineCoreStreamOutput;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
use crate::protocol::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::transport::ConnectedEngine;
|
||||
|
||||
pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
|
||||
@@ -452,7 +452,7 @@ mod tests {
|
||||
EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry,
|
||||
};
|
||||
use crate::mock_engine::default_ready_response;
|
||||
use crate::protocol::{
|
||||
use crate::protocol::output::{
|
||||
EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
|
||||
};
|
||||
use crate::transport::ConnectedEngine;
|
||||
|
||||
@@ -10,7 +10,7 @@ use tracing::{debug, error, warn};
|
||||
|
||||
use crate::client::AbortRequest;
|
||||
use crate::client::state::OutputReceiver;
|
||||
use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::{AbortCause, Error, Result};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
|
||||
@@ -10,10 +10,9 @@ use zeromq::{XPubSocket, ZmqMessage};
|
||||
use crate::client::imp::ClientInner;
|
||||
use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
|
||||
use crate::error::{Error, Result, bail_unexpected_coordinator_output};
|
||||
use crate::protocol::{
|
||||
ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType,
|
||||
encode_msgpack,
|
||||
};
|
||||
use crate::protocol::encode_msgpack;
|
||||
use crate::protocol::output::{ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs};
|
||||
use crate::protocol::request::EngineCoreRequestType;
|
||||
|
||||
/// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the
|
||||
/// engine-facing coordinator socket.
|
||||
|
||||
@@ -8,8 +8,9 @@ use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
|
||||
|
||||
use crate::EngineId;
|
||||
use crate::error::{Error, Result, bail_unexpected_handshake_message};
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
|
||||
use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
|
||||
/// Default model length advertised by reusable mock engine helpers.
|
||||
pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
|
||||
|
||||
@@ -2,7 +2,8 @@ use std::collections::BTreeMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::protocol::{ModelDtype, OpaqueValue};
|
||||
use crate::protocol::OpaqueValue;
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
|
||||
/// Decoded engine startup-handshake payload sent on the handshake socket.
|
||||
///
|
||||
|
||||
@@ -9,8 +9,7 @@ use enum_as_inner::EnumAsInner;
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
|
||||
use self::wire::*;
|
||||
use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack};
|
||||
use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode};
|
||||
use crate::error::{Error, Result, bail_ext_value_decode};
|
||||
use crate::protocol::tensor::{WireArrayData, WireNdArray};
|
||||
|
||||
/// One token candidate and its logprob metadata for a single sequence position.
|
||||
@@ -160,7 +159,7 @@ impl Serialize for MaybeWireLogprobs {
|
||||
impl MaybeWireLogprobs {
|
||||
/// Resolve the wire representation into decoded logprobs by looking up aux
|
||||
/// frames and decoding raw views as needed.
|
||||
fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
|
||||
pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
@@ -171,37 +170,6 @@ impl MaybeWireLogprobs {
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineCoreOutputs {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
for output in &mut self.outputs {
|
||||
output.resolve_in_place(frames)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
self.new_logprobs = (self.new_logprobs.take())
|
||||
.map(|value| value.resolve(frames, "new_logprobs"))
|
||||
.transpose()?;
|
||||
self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
|
||||
.map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl WireLogprobs {
|
||||
/// Convert semantic per-position logprobs into the Python wire tuple shape.
|
||||
///
|
||||
@@ -315,16 +283,3 @@ impl WireLogprobs {
|
||||
Ok(Logprobs { positions })
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one ordinary or multipart engine-core output message into the strong
|
||||
/// typed public protocol shape.
|
||||
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
|
||||
|
||||
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
|
||||
outputs.resolve_in_place(frames)?;
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
@@ -3,8 +3,8 @@ use std::collections::BTreeSet;
|
||||
use bytes::Bytes;
|
||||
use rmpv::Value;
|
||||
|
||||
use super::{Logprobs, PositionLogprobs, TokenLogprob, decode_engine_core_outputs};
|
||||
use crate::protocol::EngineCoreFinishReason;
|
||||
use super::{Logprobs, PositionLogprobs, TokenLogprob};
|
||||
use crate::protocol::output::{EngineCoreFinishReason, decode_engine_core_outputs};
|
||||
|
||||
fn encode_value(value: &Value) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
|
||||
@@ -1,28 +1,11 @@
|
||||
use std::any::type_name;
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::Bytes;
|
||||
use rmpv::Value;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_repr::{Deserialize_repr, Serialize_repr};
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::logprobs::MaybeWireLogprobs;
|
||||
use crate::protocol::multimodal::MmFeatures;
|
||||
use crate::protocol::stats::{PrefillStats, SchedulerStats};
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
|
||||
// TODO: This module currently mixes reusable frontend-facing semantic types
|
||||
// (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future
|
||||
// cleaned-up frontend sampling types) with engine-core-specific wire DTOs and
|
||||
// handshake/control messages. While the Rust frontend is still evolving
|
||||
// quickly, keep them co-located here for iteration speed. Once the higher-level
|
||||
// API boundary stabilizes, move the truly reusable semantic types into a
|
||||
// lower-level common crate and keep the engine transport/wire messages here.
|
||||
|
||||
/// Dynamic msgpack value used for schema positions that are preserved but not
|
||||
/// yet strongly typed in the early-stage Rust client.
|
||||
@@ -36,499 +19,18 @@ fn is_false(v: &bool) -> bool {
|
||||
!v
|
||||
}
|
||||
|
||||
fn default_top_p() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repetition_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_temperature() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
16
|
||||
}
|
||||
|
||||
mod classified_outputs;
|
||||
pub mod dtype;
|
||||
pub mod handshake;
|
||||
pub mod logprobs;
|
||||
pub mod lora;
|
||||
pub mod multimodal;
|
||||
pub mod output;
|
||||
pub mod request;
|
||||
pub mod sampling;
|
||||
pub mod stats;
|
||||
pub mod structured_outputs;
|
||||
pub mod tensor;
|
||||
pub mod utility;
|
||||
pub use classified_outputs::{
|
||||
ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput,
|
||||
};
|
||||
pub use dtype::ModelDtype;
|
||||
pub use logprobs::decode_engine_core_outputs;
|
||||
|
||||
/// Request types are encoded as single-byte protocol constants so they can be
|
||||
/// sent over the ZMQ socket without an extra encoding step.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreRequestType {
|
||||
Add = 0,
|
||||
Abort = 1,
|
||||
StartDpWave = 2,
|
||||
Utility = 3,
|
||||
}
|
||||
|
||||
impl EngineCoreRequestType {
|
||||
/// Decode the single-byte request type frame used on the engine input
|
||||
/// socket. Returns `None` for unrecognized values.
|
||||
pub fn from_frame(frame: &[u8]) -> Option<Self> {
|
||||
let [value] = frame else {
|
||||
return None;
|
||||
};
|
||||
|
||||
match value {
|
||||
0 => Some(Self::Add),
|
||||
1 => Some(Self::Abort),
|
||||
2 => Some(Self::StartDpWave),
|
||||
3 => Some(Self::Utility),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode the request type as the single-byte frame used on the engine
|
||||
/// input socket.
|
||||
pub fn to_frame(self) -> Bytes {
|
||||
Bytes::from_static(match self {
|
||||
Self::Add => b"\x00",
|
||||
Self::Abort => b"\x01",
|
||||
Self::StartDpWave => b"\x02",
|
||||
Self::Utility => b"\x03",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Reason a request finished: stop, length, abort, error, or repetition.
|
||||
///
|
||||
/// This mirrors the Python enum and uses integer encoding for compact wire
|
||||
/// representation.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreFinishReason {
|
||||
/// A stop string was emitted.
|
||||
Stop = 0,
|
||||
/// `max_tokens` or `max_model_len` was reached.
|
||||
Length = 1,
|
||||
/// The request was aborted by the client.
|
||||
Abort = 2,
|
||||
/// A retryable request-level internal error occurred.
|
||||
Error = 3,
|
||||
/// A repetitive token pattern was detected.
|
||||
Repetition = 4,
|
||||
}
|
||||
|
||||
/// Event types emitted by engine-core for one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreEventType {
|
||||
Queued = 1,
|
||||
Scheduled = 2,
|
||||
Preempted = 3,
|
||||
}
|
||||
|
||||
/// A timestamped engine-core event associated with one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EngineCoreEvent {
|
||||
pub r#type: EngineCoreEventType,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
/// Controls how intermediate outputs are returned to the frontend.
|
||||
///
|
||||
/// `Cumulative = 0` is intentionally not supported in Rust frontend.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L146-L152>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum RequestOutputKind {
|
||||
/// Return only token deltas in each update.
|
||||
#[default]
|
||||
Delta = 1,
|
||||
/// Suppress intermediate updates and return only the final output.
|
||||
FinalOnly = 2,
|
||||
}
|
||||
|
||||
/// Structured-output backend selected for EngineCore grammar compilation.
|
||||
///
|
||||
/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
|
||||
/// validation. The Rust frontend currently always lowers structured-output
|
||||
/// requests to guidance, while ignoring any user-supplied `_backend` value.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum StructuredOutputBackend {
|
||||
Xgrammar,
|
||||
#[default]
|
||||
Guidance,
|
||||
Outlines,
|
||||
LmFormatEnforcer,
|
||||
}
|
||||
|
||||
/// The stop reason associated with a finished output.
|
||||
///
|
||||
/// Python models this as the union-typed `stop_reason: int | str | None`
|
||||
/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
|
||||
///
|
||||
/// Original Python field:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum StopReason {
|
||||
TokenId(u32),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
///
|
||||
/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
|
||||
/// `json_object`, or `structural_tag`) should be set. The engine-core
|
||||
/// backend selects the appropriate grammar compiler based on which field
|
||||
/// is present.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct StructuredOutputsParams {
|
||||
/// JSON schema (as a dict/object or JSON string) constraining the output.
|
||||
pub json: Option<serde_json::Value>,
|
||||
/// Regular expression the output must match.
|
||||
pub regex: Option<String>,
|
||||
/// List of allowed output strings (the model must produce one of these).
|
||||
pub choice: Option<Vec<String>>,
|
||||
/// Context-free grammar (in EBNF-like notation) the output must conform to.
|
||||
pub grammar: Option<String>,
|
||||
/// When `true`, output must be valid JSON (free-form, no schema).
|
||||
pub json_object: Option<bool>,
|
||||
/// Disable any additional whitespace in guided JSON output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_any_whitespace: bool,
|
||||
/// Disable `additionalProperties` in JSON schema output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_additional_properties: bool,
|
||||
/// Custom whitespace pattern for guided JSON output.
|
||||
pub whitespace_pattern: Option<String>,
|
||||
/// Structural tag configuration (JSON-encoded string).
|
||||
pub structural_tag: Option<String>,
|
||||
/// Structured-output backend, mirroring Python's internal `_backend`.
|
||||
///
|
||||
/// User-supplied values are ignored during deserialization. This matches
|
||||
/// Python's request boundary, where `_backend` is set by validation rather
|
||||
/// than accepted as a request-level backend selector.
|
||||
#[serde(
|
||||
default,
|
||||
rename = "_backend",
|
||||
deserialize_with = "serde_with::rust::deserialize_ignore_any"
|
||||
)]
|
||||
pub backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
/// Engine-core-facing sampling parameters for text generation.
|
||||
///
|
||||
/// This is the normalized southbound subset used by the Rust frontend when it
|
||||
/// talks to Python engine-core over the wire. User-facing request semantics
|
||||
/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
|
||||
/// intentionally handled by higher layers before values reach this DTO.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
|
||||
// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
|
||||
// default-valued keys; default the whole struct. Per-field fns cover the
|
||||
// non-zero defaults.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
|
||||
#[serde(default)]
|
||||
pub struct EngineCoreSamplingParams {
|
||||
/// Controls randomness. Lower values are more deterministic; zero means
|
||||
/// greedy sampling.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: f32,
|
||||
/// Cumulative probability threshold for nucleus sampling.
|
||||
#[serde(default = "default_top_p")]
|
||||
pub top_p: f32,
|
||||
/// Maximum number of top tokens to consider. `0` means all tokens.
|
||||
pub top_k: u32,
|
||||
/// Random seed used by the sampler when present.
|
||||
pub seed: Option<i64>,
|
||||
/// Maximum number of tokens to generate per output sequence.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: u32,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` means unlimited; the
|
||||
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
|
||||
/// reaching this DTO, so only non-negative values are sent. Enforced
|
||||
/// engine-side (and only when a reasoning parser is configured).
|
||||
pub thinking_token_budget: Option<u64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
pub logprobs: Option<i32>,
|
||||
/// Number of log probabilities to return per prompt token.
|
||||
///
|
||||
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
|
||||
pub prompt_logprobs: Option<i32>,
|
||||
/// Minimum probability threshold for token sampling.
|
||||
pub min_p: f32,
|
||||
/// Frequency penalty applied by the sampler.
|
||||
pub frequency_penalty: f32,
|
||||
/// Presence penalty applied by the sampler.
|
||||
pub presence_penalty: f32,
|
||||
/// Repetition penalty applied by the sampler.
|
||||
#[serde(default = "default_repetition_penalty")]
|
||||
pub repetition_penalty: f32,
|
||||
/// Token IDs that stop generation.
|
||||
pub stop_token_ids: Vec<u32>,
|
||||
/// Primary EOS token ID used by engine-core's dedicated EOS stop path.
|
||||
///
|
||||
/// This mirrors Python's internal `_eos_token_id` field and is derived by
|
||||
/// the frontend from tokenizer/model metadata rather than supplied directly
|
||||
/// by end users.
|
||||
#[serde(rename = "_eos_token_id")]
|
||||
pub eos_token_id: Option<u32>,
|
||||
/// Complete stop-token set used by engine-core for `min_tokens` masking.
|
||||
///
|
||||
/// This mirrors Python's internal `_all_stop_token_ids` field and should
|
||||
/// contain explicit `stop_token_ids` plus any frontend-derived EOS token
|
||||
/// IDs.
|
||||
#[serde(rename = "_all_stop_token_ids")]
|
||||
pub all_stop_token_ids: BTreeSet<u32>,
|
||||
/// Logit biases to apply during sampling.
|
||||
/// Keys are token IDs
|
||||
pub logit_bias: Option<HashMap<u32, f32>>,
|
||||
/// Restrict output to these token IDs only.
|
||||
pub allowed_token_ids: Option<Vec<u32>>,
|
||||
/// Tokenized bad words to avoid during generation.
|
||||
#[serde(rename = "_bad_words_token_ids")]
|
||||
pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
pub structured_outputs: Option<StructuredOutputsParams>,
|
||||
/// Specific token IDs for which log probabilities should be returned at
|
||||
/// each position.
|
||||
///
|
||||
/// When set, the engine returns logprobs for exactly these tokens in
|
||||
/// addition to the sampled/scored token. Mutually exclusive with the
|
||||
/// `logprobs` count field in practice.
|
||||
pub logprob_token_ids: Option<Vec<u32>>,
|
||||
/// If `Some(true)`, the request will not attempt to read from the prefix
|
||||
/// cache; newly computed blocks may still populate the cache. `None`
|
||||
/// defers to engine-core defaults.
|
||||
pub skip_reading_prefix_cache: Option<bool>,
|
||||
/// Additional request parameters for custom extensions (from `vllm_xargs`).
|
||||
pub extra_args: Option<HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl EngineCoreSamplingParams {
|
||||
/// Constructs a default sampling params for testing purposes only.
|
||||
pub fn for_test() -> Self {
|
||||
Self {
|
||||
temperature: 1.0,
|
||||
top_p: 1.0,
|
||||
top_k: 0,
|
||||
seed: None,
|
||||
max_tokens: 65536,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
repetition_penalty: 1.0,
|
||||
stop_token_ids: Vec::new(),
|
||||
eos_token_id: None,
|
||||
all_stop_token_ids: BTreeSet::new(),
|
||||
logit_bias: None,
|
||||
allowed_token_ids: None,
|
||||
bad_words_token_ids: None,
|
||||
structured_outputs: None,
|
||||
logprob_token_ids: None,
|
||||
skip_reading_prefix_cache: None,
|
||||
extra_args: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extra kwargs consumed by engine-side reasoning parsers.
|
||||
///
|
||||
/// Original Python construction point:
|
||||
/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ReasoningParserKwargs {
|
||||
/// Effective kwargs visible to the chat template for this request.
|
||||
pub chat_template_kwargs: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Engine-core add-request payload sent from frontend to engine.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreRequest {
|
||||
pub request_id: String,
|
||||
pub prompt_token_ids: Option<Vec<u32>>,
|
||||
/// Multimodal features attached to the request.
|
||||
pub mm_features: Option<MmFeatures>,
|
||||
pub sampling_params: Option<EngineCoreSamplingParams>,
|
||||
/// Pooling parameters are preserved in the schema but not yet strongly
|
||||
/// typed.
|
||||
pub pooling_params: Option<OpaqueValue>,
|
||||
pub arrival_time: f64,
|
||||
#[serde(default)]
|
||||
pub lora_request: Option<lora::LoraRequest>,
|
||||
#[serde(default)]
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub data_parallel_rank: Option<u32>,
|
||||
/// Unsupported in the first-stage Rust client because Python uses a custom
|
||||
/// tensor/aux-frame encoding path for this field.
|
||||
#[serde(default)]
|
||||
pub prompt_embeds: Option<OpaqueValue>,
|
||||
/// Per-position mask for mixed-mode inputs (e.g. chat completion with
|
||||
/// `prompt_embeds` content parts). `Some(true)` means real token id;
|
||||
/// `Some(false)` means the position uses a pre-computed entry from
|
||||
/// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
|
||||
#[serde(default)]
|
||||
pub prompt_is_token_ids: Option<Vec<bool>>,
|
||||
/// Index of the client, used to ensure outputs are sent back to the same
|
||||
/// client when scaling out the frontend.
|
||||
#[serde(default)]
|
||||
pub client_index: u32,
|
||||
/// In DP mode, indicates which wave this request is expected to belong to.
|
||||
#[serde(default)]
|
||||
pub current_wave: u32,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<BTreeMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub resumable: bool,
|
||||
/// Original user-provided request ID, used for output reporting and aborts.
|
||||
#[serde(default)]
|
||||
pub external_req_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub reasoning_ended: Option<bool>,
|
||||
/// Reasoning-parser kwargs forwarded from the frontend to the
|
||||
/// structured-output backend.
|
||||
#[serde(default)]
|
||||
pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
|
||||
/// If `true`, the request should be added to the scheduler's waiting queue
|
||||
/// and immediately aborted, so connector-side cleanup runs via the
|
||||
/// standard `request_finished` hook.
|
||||
#[serde(default)]
|
||||
pub abort_immediately: bool,
|
||||
}
|
||||
|
||||
impl EngineCoreRequest {
|
||||
/// Validate fields intentionally not supported in the first-stage client.
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.prompt_embeds.is_some() {
|
||||
return Err(Error::UnsupportedField {
|
||||
context: "EngineCoreRequest",
|
||||
field: "prompt_embeds",
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Engine-core output for a single request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutput {
|
||||
pub request_id: String,
|
||||
pub new_token_ids: Vec<u32>,
|
||||
/// Decoded sample logprobs for the newly generated positions in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_logprobs: Option<MaybeWireLogprobs>,
|
||||
/// Decoded prompt logprobs for the scored prompt positions emitted in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
|
||||
#[serde(default)]
|
||||
pub pooling_output: Option<OpaqueValue>,
|
||||
#[serde(default)]
|
||||
pub finish_reason: Option<EngineCoreFinishReason>,
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<StopReason>,
|
||||
#[serde(default)]
|
||||
pub events: Option<Vec<EngineCoreEvent>>,
|
||||
#[serde(default)]
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<OpaqueValue>,
|
||||
/// Breakdown of the scheduled prefill computation, set on the first output
|
||||
/// of a newly scheduled prefill and elided for subsequent decode outputs.
|
||||
#[serde(default)]
|
||||
pub prefill_stats: Option<PrefillStats>,
|
||||
#[serde(default)]
|
||||
pub routed_experts: Option<OpaqueValue>,
|
||||
/// Number of NaNs seen in logits. Values above zero indicate corruption.
|
||||
#[serde(default)]
|
||||
pub num_nans_in_logits: u32,
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Returns whether this output is terminal for the request.
|
||||
pub fn finished(&self) -> bool {
|
||||
self.finish_reason.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch of engine-core outputs returned to a frontend client.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutputs {
|
||||
#[serde(default)]
|
||||
pub engine_index: u32,
|
||||
/// Outputs grouped for this client in the current engine tick.
|
||||
#[serde(default)]
|
||||
pub outputs: Vec<EngineCoreOutput>,
|
||||
#[serde(default)]
|
||||
pub scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
#[serde(default)]
|
||||
pub timestamp: f64,
|
||||
#[serde(default)]
|
||||
pub utility_output: Option<UtilityOutput>,
|
||||
#[serde(default)]
|
||||
pub finished_requests: Option<BTreeSet<String>>,
|
||||
/// In DP mode, signals that the current wave finished and engines are
|
||||
/// paused.
|
||||
#[serde(default)]
|
||||
pub wave_complete: Option<u32>,
|
||||
/// In DP mode, signals that a request arrived for an old wave and the next
|
||||
/// wave needs to start in other engines.
|
||||
#[serde(default)]
|
||||
pub start_wave: Option<u32>,
|
||||
}
|
||||
|
||||
/// Encode a Rust value into msgpack using the protocol crate's serde model.
|
||||
pub fn encode_msgpack<T>(value: &T) -> Result<Vec<u8>>
|
||||
@@ -564,81 +66,17 @@ where
|
||||
})
|
||||
}
|
||||
|
||||
/// Decode a msgpack payload into a dynamic value for diagnostics and tests.
|
||||
pub fn decode_value(bytes: &[u8]) -> Result<Value> {
|
||||
Ok(rmpv::decode::read_value(&mut Cursor::new(bytes))?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn engine_core_request_serializes_as_full_array() {
|
||||
let request = EngineCoreRequest {
|
||||
request_id: "req-1".to_string(),
|
||||
prompt_token_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: Some(EngineCoreSamplingParams {
|
||||
max_tokens: 8,
|
||||
..EngineCoreSamplingParams::for_test()
|
||||
}),
|
||||
arrival_time: 1234.5,
|
||||
client_index: 7,
|
||||
..EngineCoreRequest::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&request).unwrap();
|
||||
let value = decode_value(&encoded).unwrap();
|
||||
let array = match value {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected array, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(array.len(), 20);
|
||||
assert_eq!(array[0], Value::from("req-1"));
|
||||
assert_eq!(array[2], Value::Nil);
|
||||
assert_eq!(array[4], Value::Nil);
|
||||
assert_eq!(array[10], Value::Nil);
|
||||
assert_eq!(array[11], Value::from(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_roundtrip_finished_fields() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![42],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(EngineCoreFinishReason::Length),
|
||||
stop_reason: Some(StopReason::Text("stop".to_string())),
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&outputs).unwrap();
|
||||
let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
|
||||
|
||||
assert_eq!(decoded.outputs.len(), 1);
|
||||
assert_eq!(
|
||||
decoded.outputs[0].finish_reason,
|
||||
Some(EngineCoreFinishReason::Length)
|
||||
);
|
||||
assert_eq!(
|
||||
decoded.finished_requests,
|
||||
Some(BTreeSet::from(["req-1".to_string()]))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_msgpack_includes_type_name_and_value_fallback() {
|
||||
let error = decode_msgpack::<u64>(
|
||||
@@ -648,72 +86,4 @@ mod tests {
|
||||
|
||||
expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_backend_ignores_deserialized_value() {
|
||||
let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
|
||||
"json_object": true,
|
||||
"_backend": "xgrammar",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
assert_eq!(value["_backend"], "guidance");
|
||||
}
|
||||
|
||||
/// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
|
||||
/// must fall back to defaults. `python_compat` can't catch this since Rust
|
||||
/// encodes full maps (see `engine_core_request_serializes_as_full_array`).
|
||||
#[test]
|
||||
fn decodes_sampling_params_with_omitted_defaults() {
|
||||
let sampling_params = Value::Map(vec![
|
||||
(
|
||||
Value::from("stop_token_ids"),
|
||||
Value::Array(vec![Value::from(151643u32)]),
|
||||
),
|
||||
(Value::from("skip_reading_prefix_cache"), Value::from(false)),
|
||||
]);
|
||||
let request = Value::Array(vec![
|
||||
Value::from("req-omit-defaults"),
|
||||
Value::Array(vec![
|
||||
Value::from(1u32),
|
||||
Value::from(2u32),
|
||||
Value::from(3u32),
|
||||
]),
|
||||
Value::Nil,
|
||||
sampling_params,
|
||||
Value::Nil,
|
||||
Value::from(1.0f64),
|
||||
]);
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
rmpv::encode::write_value(&mut bytes, &request).unwrap();
|
||||
|
||||
let decoded: EngineCoreRequest = decode_msgpack(&bytes)
|
||||
.expect("a real omit_defaults request must decode (regression: missing field)");
|
||||
|
||||
assert_eq!(decoded.request_id, "req-omit-defaults");
|
||||
let sampling = decoded.sampling_params.expect("sampling params present");
|
||||
|
||||
assert_eq!(sampling.stop_token_ids, vec![151643]);
|
||||
assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
|
||||
|
||||
// Omitted fields -> Python defaults.
|
||||
assert_eq!(sampling.temperature, 1.0);
|
||||
assert_eq!(sampling.top_p, 1.0);
|
||||
assert_eq!(sampling.top_k, 0);
|
||||
assert_eq!(sampling.seed, None);
|
||||
assert_eq!(sampling.max_tokens, 16);
|
||||
assert_eq!(sampling.min_tokens, 0);
|
||||
assert_eq!(sampling.min_p, 0.0);
|
||||
assert_eq!(sampling.frequency_penalty, 0.0);
|
||||
assert_eq!(sampling.presence_penalty, 0.0);
|
||||
assert_eq!(sampling.repetition_penalty, 1.0);
|
||||
assert_eq!(sampling.logprobs, None);
|
||||
assert_eq!(sampling.prompt_logprobs, None);
|
||||
assert_eq!(sampling.eos_token_id, None);
|
||||
assert!(sampling.all_stop_token_ids.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
+219
-3
@@ -1,10 +1,164 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use enum_as_inner::EnumAsInner;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_repr::{Deserialize_repr, Serialize_repr};
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
|
||||
use super::utility::UtilityOutput;
|
||||
use super::{EngineCoreOutput, EngineCoreOutputs};
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::error::{Error, Result, ext_value_decode};
|
||||
use crate::protocol::logprobs::MaybeWireLogprobs;
|
||||
use crate::protocol::stats::{PrefillStats, SchedulerStats};
|
||||
use crate::protocol::{OpaqueValue, decode_msgpack};
|
||||
|
||||
/// The stop reason associated with a finished output.
|
||||
///
|
||||
/// Python models this as the union-typed `stop_reason: int | str | None`
|
||||
/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
|
||||
///
|
||||
/// Original Python field:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum StopReason {
|
||||
TokenId(u32),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Reason a request finished: stop, length, abort, error, or repetition.
|
||||
///
|
||||
/// This mirrors the Python enum and uses integer encoding for compact wire
|
||||
/// representation.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreFinishReason {
|
||||
/// A stop string was emitted.
|
||||
Stop = 0,
|
||||
/// `max_tokens` or `max_model_len` was reached.
|
||||
Length = 1,
|
||||
/// The request was aborted by the client.
|
||||
Abort = 2,
|
||||
/// A retryable request-level internal error occurred.
|
||||
Error = 3,
|
||||
/// A repetitive token pattern was detected.
|
||||
Repetition = 4,
|
||||
}
|
||||
|
||||
/// Event types emitted by engine-core for one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreEventType {
|
||||
Queued = 1,
|
||||
Scheduled = 2,
|
||||
Preempted = 3,
|
||||
}
|
||||
|
||||
/// A timestamped engine-core event associated with one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EngineCoreEvent {
|
||||
pub r#type: EngineCoreEventType,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
/// Engine-core output for a single request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutput {
|
||||
pub request_id: String,
|
||||
pub new_token_ids: Vec<u32>,
|
||||
/// Decoded sample logprobs for the newly generated positions in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_logprobs: Option<MaybeWireLogprobs>,
|
||||
/// Decoded prompt logprobs for the scored prompt positions emitted in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
|
||||
#[serde(default)]
|
||||
pub pooling_output: Option<OpaqueValue>,
|
||||
#[serde(default)]
|
||||
pub finish_reason: Option<EngineCoreFinishReason>,
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<StopReason>,
|
||||
#[serde(default)]
|
||||
pub events: Option<Vec<EngineCoreEvent>>,
|
||||
#[serde(default)]
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<OpaqueValue>,
|
||||
/// Breakdown of the scheduled prefill computation, set on the first output
|
||||
/// of a newly scheduled prefill and elided for subsequent decode outputs.
|
||||
#[serde(default)]
|
||||
pub prefill_stats: Option<PrefillStats>,
|
||||
#[serde(default)]
|
||||
pub routed_experts: Option<OpaqueValue>,
|
||||
/// Number of NaNs seen in logits. Values above zero indicate corruption.
|
||||
#[serde(default)]
|
||||
pub num_nans_in_logits: u32,
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Returns whether this output is terminal for the request.
|
||||
pub fn finished(&self) -> bool {
|
||||
self.finish_reason.is_some()
|
||||
}
|
||||
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
self.new_logprobs = (self.new_logprobs.take())
|
||||
.map(|value| value.resolve(frames, "new_logprobs"))
|
||||
.transpose()?;
|
||||
self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
|
||||
.map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch of engine-core outputs returned to a frontend client.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutputs {
|
||||
#[serde(default)]
|
||||
pub engine_index: u32,
|
||||
/// Outputs grouped for this client in the current engine tick.
|
||||
#[serde(default)]
|
||||
pub outputs: Vec<EngineCoreOutput>,
|
||||
#[serde(default)]
|
||||
pub scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
#[serde(default)]
|
||||
pub timestamp: f64,
|
||||
#[serde(default)]
|
||||
pub utility_output: Option<UtilityOutput>,
|
||||
#[serde(default)]
|
||||
pub finished_requests: Option<BTreeSet<String>>,
|
||||
/// In DP mode, signals that the current wave finished and engines are
|
||||
/// paused.
|
||||
#[serde(default)]
|
||||
pub wave_complete: Option<u32>,
|
||||
/// In DP mode, signals that a request arrived for an old wave and the next
|
||||
/// wave needs to start in other engines.
|
||||
#[serde(default)]
|
||||
pub start_wave: Option<u32>,
|
||||
}
|
||||
|
||||
/// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -49,6 +203,18 @@ pub enum ClassifiedEngineCoreOutputs {
|
||||
}
|
||||
|
||||
impl EngineCoreOutputs {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
for output in &mut self.outputs {
|
||||
output.resolve_in_place(frames)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Classify the raw wire message into a more semantic Rust enum.
|
||||
pub fn classify(self) -> ClassifiedEngineCoreOutputs {
|
||||
let has_request_payload = !self.outputs.is_empty()
|
||||
@@ -92,12 +258,62 @@ impl EngineCoreOutputs {
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one ordinary or multipart engine-core output message into the strong
|
||||
/// typed public protocol shape.
|
||||
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
|
||||
|
||||
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
|
||||
outputs.resolve_in_place(frames)?;
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::EngineCoreOutput;
|
||||
use crate::protocol::output::EngineCoreOutput;
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_roundtrip_finished_fields() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![42],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(EngineCoreFinishReason::Length),
|
||||
stop_reason: Some(StopReason::Text("stop".to_string())),
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&outputs).unwrap();
|
||||
let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
|
||||
|
||||
assert_eq!(decoded.outputs.len(), 1);
|
||||
assert_eq!(
|
||||
decoded.outputs[0].finish_reason,
|
||||
Some(EngineCoreFinishReason::Length)
|
||||
);
|
||||
assert_eq!(
|
||||
decoded.finished_requests,
|
||||
Some(BTreeSet::from(["req-1".to_string()]))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_request_batch() {
|
||||
@@ -0,0 +1,175 @@
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
use bytes::Bytes;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
|
||||
use crate::protocol::multimodal::MmFeatures;
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::{OpaqueValue, lora};
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Request types are encoded as single-byte protocol constants so they can be
|
||||
/// sent over the ZMQ socket without an extra encoding step.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreRequestType {
|
||||
Add = 0,
|
||||
Abort = 1,
|
||||
StartDpWave = 2,
|
||||
Utility = 3,
|
||||
}
|
||||
|
||||
impl EngineCoreRequestType {
|
||||
/// Decode the single-byte request type frame used on the engine input
|
||||
/// socket. Returns `None` for unrecognized values.
|
||||
pub fn from_frame(frame: &[u8]) -> Option<Self> {
|
||||
let [value] = frame else {
|
||||
return None;
|
||||
};
|
||||
|
||||
match value {
|
||||
0 => Some(Self::Add),
|
||||
1 => Some(Self::Abort),
|
||||
2 => Some(Self::StartDpWave),
|
||||
3 => Some(Self::Utility),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode the request type as the single-byte frame used on the engine
|
||||
/// input socket.
|
||||
pub fn to_frame(self) -> Bytes {
|
||||
Bytes::from_static(match self {
|
||||
Self::Add => b"\x00",
|
||||
Self::Abort => b"\x01",
|
||||
Self::StartDpWave => b"\x02",
|
||||
Self::Utility => b"\x03",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Extra kwargs consumed by engine-side reasoning parsers.
|
||||
///
|
||||
/// Original Python construction point:
|
||||
/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ReasoningParserKwargs {
|
||||
/// Effective kwargs visible to the chat template for this request.
|
||||
pub chat_template_kwargs: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Engine-core add-request payload sent from frontend to engine.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreRequest {
|
||||
pub request_id: String,
|
||||
pub prompt_token_ids: Option<Vec<u32>>,
|
||||
/// Multimodal features attached to the request.
|
||||
pub mm_features: Option<MmFeatures>,
|
||||
pub sampling_params: Option<EngineCoreSamplingParams>,
|
||||
/// Pooling parameters are preserved in the schema but not yet strongly
|
||||
/// typed.
|
||||
pub pooling_params: Option<OpaqueValue>,
|
||||
pub arrival_time: f64,
|
||||
#[serde(default)]
|
||||
pub lora_request: Option<lora::LoraRequest>,
|
||||
#[serde(default)]
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub data_parallel_rank: Option<u32>,
|
||||
/// Unsupported in the first-stage Rust client because Python uses a custom
|
||||
/// tensor/aux-frame encoding path for this field.
|
||||
#[serde(default)]
|
||||
pub prompt_embeds: Option<OpaqueValue>,
|
||||
/// Per-position mask for mixed-mode inputs (e.g. chat completion with
|
||||
/// `prompt_embeds` content parts). `Some(true)` means real token id;
|
||||
/// `Some(false)` means the position uses a pre-computed entry from
|
||||
/// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
|
||||
#[serde(default)]
|
||||
pub prompt_is_token_ids: Option<Vec<bool>>,
|
||||
/// Index of the client, used to ensure outputs are sent back to the same
|
||||
/// client when scaling out the frontend.
|
||||
#[serde(default)]
|
||||
pub client_index: u32,
|
||||
/// In DP mode, indicates which wave this request is expected to belong to.
|
||||
#[serde(default)]
|
||||
pub current_wave: u32,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<BTreeMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub resumable: bool,
|
||||
/// Original user-provided request ID, used for output reporting and aborts.
|
||||
#[serde(default)]
|
||||
pub external_req_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub reasoning_ended: Option<bool>,
|
||||
/// Reasoning-parser kwargs forwarded from the frontend to the
|
||||
/// structured-output backend.
|
||||
#[serde(default)]
|
||||
pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
|
||||
/// If `true`, the request should be added to the scheduler's waiting queue
|
||||
/// and immediately aborted, so connector-side cleanup runs via the
|
||||
/// standard `request_finished` hook.
|
||||
#[serde(default)]
|
||||
pub abort_immediately: bool,
|
||||
}
|
||||
|
||||
impl EngineCoreRequest {
|
||||
/// Validate fields intentionally not supported in the first-stage client.
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.prompt_embeds.is_some() {
|
||||
return Err(Error::UnsupportedField {
|
||||
context: "EngineCoreRequest",
|
||||
field: "prompt_embeds",
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rmpv::Value;
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::{decode_value, encode_msgpack};
|
||||
|
||||
#[test]
|
||||
fn engine_core_request_serializes_as_full_array() {
|
||||
let request = EngineCoreRequest {
|
||||
request_id: "req-1".to_string(),
|
||||
prompt_token_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: Some(EngineCoreSamplingParams {
|
||||
max_tokens: 8,
|
||||
..EngineCoreSamplingParams::for_test()
|
||||
}),
|
||||
arrival_time: 1234.5,
|
||||
client_index: 7,
|
||||
..EngineCoreRequest::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&request).unwrap();
|
||||
let value = decode_value(&encoded).unwrap();
|
||||
let array = match value {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected array, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(array.len(), 20);
|
||||
assert_eq!(array[0], Value::from("req-1"));
|
||||
assert_eq!(array[2], Value::Nil);
|
||||
assert_eq!(array[4], Value::Nil);
|
||||
assert_eq!(array[10], Value::Nil);
|
||||
assert_eq!(array[11], Value::from(7));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
|
||||
use crate::protocol::structured_outputs::StructuredOutputsParams;
|
||||
|
||||
fn default_top_p() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repetition_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_temperature() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
16
|
||||
}
|
||||
|
||||
/// Engine-core-facing sampling parameters for text generation.
|
||||
///
|
||||
/// This is the normalized southbound subset used by the Rust frontend when it
|
||||
/// talks to Python engine-core over the wire. User-facing request semantics
|
||||
/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
|
||||
/// intentionally handled by higher layers before values reach this DTO.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
|
||||
// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
|
||||
// default-valued keys; default the whole struct. Per-field fns cover the
|
||||
// non-zero defaults.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
|
||||
#[serde(default)]
|
||||
pub struct EngineCoreSamplingParams {
|
||||
/// Controls randomness. Lower values are more deterministic; zero means
|
||||
/// greedy sampling.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: f32,
|
||||
/// Cumulative probability threshold for nucleus sampling.
|
||||
#[serde(default = "default_top_p")]
|
||||
pub top_p: f32,
|
||||
/// Maximum number of top tokens to consider. `0` means all tokens.
|
||||
pub top_k: u32,
|
||||
/// Random seed used by the sampler when present.
|
||||
pub seed: Option<i64>,
|
||||
/// Maximum number of tokens to generate per output sequence.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: u32,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` means unlimited; the
|
||||
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
|
||||
/// reaching this DTO, so only non-negative values are sent. Enforced
|
||||
/// engine-side (and only when a reasoning parser is configured).
|
||||
pub thinking_token_budget: Option<u64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
pub logprobs: Option<i32>,
|
||||
/// Number of log probabilities to return per prompt token.
|
||||
///
|
||||
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
|
||||
pub prompt_logprobs: Option<i32>,
|
||||
/// Minimum probability threshold for token sampling.
|
||||
pub min_p: f32,
|
||||
/// Frequency penalty applied by the sampler.
|
||||
pub frequency_penalty: f32,
|
||||
/// Presence penalty applied by the sampler.
|
||||
pub presence_penalty: f32,
|
||||
/// Repetition penalty applied by the sampler.
|
||||
#[serde(default = "default_repetition_penalty")]
|
||||
pub repetition_penalty: f32,
|
||||
/// Token IDs that stop generation.
|
||||
pub stop_token_ids: Vec<u32>,
|
||||
/// Primary EOS token ID used by engine-core's dedicated EOS stop path.
|
||||
///
|
||||
/// This mirrors Python's internal `_eos_token_id` field and is derived by
|
||||
/// the frontend from tokenizer/model metadata rather than supplied directly
|
||||
/// by end users.
|
||||
#[serde(rename = "_eos_token_id")]
|
||||
pub eos_token_id: Option<u32>,
|
||||
/// Complete stop-token set used by engine-core for `min_tokens` masking.
|
||||
///
|
||||
/// This mirrors Python's internal `_all_stop_token_ids` field and should
|
||||
/// contain explicit `stop_token_ids` plus any frontend-derived EOS token
|
||||
/// IDs.
|
||||
#[serde(rename = "_all_stop_token_ids")]
|
||||
pub all_stop_token_ids: BTreeSet<u32>,
|
||||
/// Logit biases to apply during sampling.
|
||||
/// Keys are token IDs
|
||||
pub logit_bias: Option<HashMap<u32, f32>>,
|
||||
/// Restrict output to these token IDs only.
|
||||
pub allowed_token_ids: Option<Vec<u32>>,
|
||||
/// Tokenized bad words to avoid during generation.
|
||||
#[serde(rename = "_bad_words_token_ids")]
|
||||
pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
pub structured_outputs: Option<StructuredOutputsParams>,
|
||||
/// Specific token IDs for which log probabilities should be returned at
|
||||
/// each position.
|
||||
///
|
||||
/// When set, the engine returns logprobs for exactly these tokens in
|
||||
/// addition to the sampled/scored token. Mutually exclusive with the
|
||||
/// `logprobs` count field in practice.
|
||||
pub logprob_token_ids: Option<Vec<u32>>,
|
||||
/// If `Some(true)`, the request will not attempt to read from the prefix
|
||||
/// cache; newly computed blocks may still populate the cache. `None`
|
||||
/// defers to engine-core defaults.
|
||||
pub skip_reading_prefix_cache: Option<bool>,
|
||||
/// Additional request parameters for custom extensions (from `vllm_xargs`).
|
||||
pub extra_args: Option<HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl EngineCoreSamplingParams {
|
||||
/// Constructs a default sampling params for testing purposes only.
|
||||
pub fn for_test() -> Self {
|
||||
Self {
|
||||
temperature: 1.0,
|
||||
top_p: 1.0,
|
||||
top_k: 0,
|
||||
seed: None,
|
||||
max_tokens: 65536,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
repetition_penalty: 1.0,
|
||||
stop_token_ids: Vec::new(),
|
||||
eos_token_id: None,
|
||||
all_stop_token_ids: BTreeSet::new(),
|
||||
logit_bias: None,
|
||||
allowed_token_ids: None,
|
||||
bad_words_token_ids: None,
|
||||
structured_outputs: None,
|
||||
logprob_token_ids: None,
|
||||
skip_reading_prefix_cache: None,
|
||||
extra_args: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rmpv::Value;
|
||||
|
||||
use crate::protocol::decode_msgpack;
|
||||
use crate::protocol::request::EngineCoreRequest;
|
||||
|
||||
/// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
|
||||
/// must fall back to defaults. `python_compat` can't catch this since Rust
|
||||
/// encodes full maps (see `engine_core_request_serializes_as_full_array`).
|
||||
#[test]
|
||||
fn decodes_sampling_params_with_omitted_defaults() {
|
||||
let sampling_params = Value::Map(vec![
|
||||
(
|
||||
Value::from("stop_token_ids"),
|
||||
Value::Array(vec![Value::from(151643u32)]),
|
||||
),
|
||||
(Value::from("skip_reading_prefix_cache"), Value::from(false)),
|
||||
]);
|
||||
let request = Value::Array(vec![
|
||||
Value::from("req-omit-defaults"),
|
||||
Value::Array(vec![
|
||||
Value::from(1u32),
|
||||
Value::from(2u32),
|
||||
Value::from(3u32),
|
||||
]),
|
||||
Value::Nil,
|
||||
sampling_params,
|
||||
Value::Nil,
|
||||
Value::from(1.0f64),
|
||||
]);
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
rmpv::encode::write_value(&mut bytes, &request).unwrap();
|
||||
|
||||
let decoded: EngineCoreRequest = decode_msgpack(&bytes)
|
||||
.expect("a real omit_defaults request must decode (regression: missing field)");
|
||||
|
||||
assert_eq!(decoded.request_id, "req-omit-defaults");
|
||||
let sampling = decoded.sampling_params.expect("sampling params present");
|
||||
|
||||
assert_eq!(sampling.stop_token_ids, vec![151643]);
|
||||
assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
|
||||
|
||||
// Omitted fields -> Python defaults.
|
||||
assert_eq!(sampling.temperature, 1.0);
|
||||
assert_eq!(sampling.top_p, 1.0);
|
||||
assert_eq!(sampling.top_k, 0);
|
||||
assert_eq!(sampling.seed, None);
|
||||
assert_eq!(sampling.max_tokens, 16);
|
||||
assert_eq!(sampling.min_tokens, 0);
|
||||
assert_eq!(sampling.min_p, 0.0);
|
||||
assert_eq!(sampling.frequency_penalty, 0.0);
|
||||
assert_eq!(sampling.presence_penalty, 0.0);
|
||||
assert_eq!(sampling.repetition_penalty, 1.0);
|
||||
assert_eq!(sampling.logprobs, None);
|
||||
assert_eq!(sampling.prompt_logprobs, None);
|
||||
assert_eq!(sampling.eos_token_id, None);
|
||||
assert!(sampling.all_stop_token_ids.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Structured-output backend selected for EngineCore grammar compilation.
|
||||
///
|
||||
/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
|
||||
/// validation. The Rust frontend currently always lowers structured-output
|
||||
/// requests to guidance, while ignoring any user-supplied `_backend` value.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum StructuredOutputBackend {
|
||||
Xgrammar,
|
||||
#[default]
|
||||
Guidance,
|
||||
Outlines,
|
||||
LmFormatEnforcer,
|
||||
}
|
||||
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
///
|
||||
/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
|
||||
/// `json_object`, or `structural_tag`) should be set. The engine-core
|
||||
/// backend selects the appropriate grammar compiler based on which field
|
||||
/// is present.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct StructuredOutputsParams {
|
||||
/// JSON schema (as a dict/object or JSON string) constraining the output.
|
||||
pub json: Option<serde_json::Value>,
|
||||
/// Regular expression the output must match.
|
||||
pub regex: Option<String>,
|
||||
/// List of allowed output strings (the model must produce one of these).
|
||||
pub choice: Option<Vec<String>>,
|
||||
/// Context-free grammar (in EBNF-like notation) the output must conform to.
|
||||
pub grammar: Option<String>,
|
||||
/// When `true`, output must be valid JSON (free-form, no schema).
|
||||
pub json_object: Option<bool>,
|
||||
/// Disable any additional whitespace in guided JSON output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_any_whitespace: bool,
|
||||
/// Disable `additionalProperties` in JSON schema output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_additional_properties: bool,
|
||||
/// Custom whitespace pattern for guided JSON output.
|
||||
pub whitespace_pattern: Option<String>,
|
||||
/// Structural tag configuration (JSON-encoded string).
|
||||
pub structural_tag: Option<String>,
|
||||
/// Structured-output backend, mirroring Python's internal `_backend`.
|
||||
///
|
||||
/// User-supplied values are ignored during deserialization. This matches
|
||||
/// Python's request boundary, where `_backend` is set by validation rather
|
||||
/// than accepted as a request-level backend selector.
|
||||
#[serde(
|
||||
default,
|
||||
rename = "_backend",
|
||||
deserialize_with = "serde_with::rust::deserialize_ignore_any"
|
||||
)]
|
||||
pub backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_backend_ignores_deserialized_value() {
|
||||
let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
|
||||
"json_object": true,
|
||||
"_backend": "xgrammar",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
assert_eq!(value["_backend"], "guidance");
|
||||
}
|
||||
}
|
||||
@@ -22,13 +22,14 @@ use crate::protocol::multimodal::{
|
||||
MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,
|
||||
SliceSpec,
|
||||
};
|
||||
use crate::protocol::output::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, decode_engine_core_outputs,
|
||||
};
|
||||
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::tensor::WireTensor;
|
||||
use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
|
||||
use crate::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
|
||||
EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
|
||||
};
|
||||
use crate::test_utils::{
|
||||
IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
|
||||
setup_mock_engine_with_init, spawn_mock_engine_task,
|
||||
|
||||
@@ -18,9 +18,8 @@ use crate::error::{Error, Result, bail_unexpected_handshake_message};
|
||||
use crate::protocol::handshake::{
|
||||
EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
|
||||
};
|
||||
use crate::protocol::{
|
||||
EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack,
|
||||
};
|
||||
use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs};
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
|
||||
/// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the
|
||||
/// engine dies.
|
||||
|
||||
@@ -5,7 +5,7 @@ use clap::Parser;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm};
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ use futures::stream::FusedStream;
|
||||
use futures::{Stream, StreamExt as _, pin_mut};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use vllm_engine_core_client::protocol::logprobs::Logprobs;
|
||||
use vllm_engine_core_client::protocol::{EngineCoreFinishReason, StopReason};
|
||||
use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason};
|
||||
use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
@@ -4,9 +4,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use uuid::Uuid;
|
||||
use vllm_engine_core_client::protocol::lora::LoraRequest;
|
||||
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreRequest, EngineCoreSamplingParams, ReasoningParserKwargs,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs};
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
@@ -134,7 +133,8 @@ fn current_unix_timestamp_secs() -> f64 {
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use vllm_engine_core_client::protocol::{EngineCoreSamplingParams, ReasoningParserKwargs};
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
|
||||
use super::GenerateRequest;
|
||||
use crate::error::Error;
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use vllm_engine_core_client::protocol::output::{
|
||||
EngineCoreEvent, EngineCoreEventType, EngineCoreOutput,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType, EngineCoreOutput};
|
||||
use vllm_metrics::{
|
||||
EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
|
||||
};
|
||||
@@ -328,8 +330,8 @@ pub(crate) fn current_unix_timestamp_secs() -> f64 {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType};
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_engine_core_client::protocol::{EngineCoreEvent, EngineCoreEventType};
|
||||
|
||||
use super::{RequestMetricsTracker, diff_or_zero};
|
||||
|
||||
@@ -341,7 +343,7 @@ mod tests {
|
||||
2,
|
||||
10.0,
|
||||
100.2,
|
||||
&vllm_engine_core_client::protocol::EngineCoreOutput {
|
||||
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![1],
|
||||
finish_reason: None,
|
||||
@@ -369,7 +371,7 @@ mod tests {
|
||||
2,
|
||||
11.5,
|
||||
100.4,
|
||||
&vllm_engine_core_client::protocol::EngineCoreOutput {
|
||||
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![2, 3],
|
||||
finish_reason: None,
|
||||
|
||||
@@ -9,11 +9,13 @@ use uuid::Uuid;
|
||||
use vllm_engine_core_client::protocol::logprobs::{
|
||||
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
use vllm_engine_core_client::protocol::output::{
|
||||
EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
|
||||
EngineCoreOutputs, EngineCoreRequest, EngineCoreSamplingParams,
|
||||
EngineCoreOutputs,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
|
||||
use vllm_llm::{
|
||||
|
||||
@@ -11,12 +11,13 @@ use tokio::sync::mpsc;
|
||||
use tokio::task::yield_now;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
use vllm_engine_core_client::protocol::output::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::utility::{
|
||||
EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
|
||||
};
|
||||
|
||||
use super::Opt;
|
||||
|
||||
|
||||
@@ -4,10 +4,9 @@ use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
use vllm_engine_core_client::mock_engine::MockEngineDataSockets;
|
||||
use vllm_engine_core_client::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
|
||||
use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{decode_msgpack, encode_msgpack};
|
||||
use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage};
|
||||
|
||||
use crate::engine::{EngineInput, EngineOutput};
|
||||
|
||||
@@ -5,9 +5,9 @@ use anyhow::Result;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::test_utils::IpcNamespace;
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ expect-test.workspace = true
|
||||
futures.workspace = true
|
||||
openai-protocol.workspace = true
|
||||
tool-parser.workspace = true
|
||||
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
|
||||
|
||||
[[bench]]
|
||||
name = "deepseek_v3"
|
||||
|
||||
@@ -27,6 +27,10 @@ impl Tokenizer for BenchTokenizer {
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
Some(u32::MAX)
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
Some("\u{FFFD}".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Bench-only adapter that exposes a unified parser through the tool-parser
|
||||
|
||||
@@ -5,14 +5,13 @@
|
||||
|
||||
mod adapter;
|
||||
|
||||
pub(super) use adapter::UnifiedToolParserAdapter;
|
||||
use futures::FutureExt as _;
|
||||
use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
|
||||
use tool_parser::traits::ToolParser as ExternalToolParser;
|
||||
use vllm_parser::tool::test_utils::collect_stream;
|
||||
use vllm_parser::tool::{Tool, ToolParser};
|
||||
|
||||
pub(super) use adapter::UnifiedToolParserAdapter;
|
||||
|
||||
pub(super) fn openai_tools(tools: &[Tool]) -> Vec<OpenAiTool> {
|
||||
tools
|
||||
.iter()
|
||||
|
||||
@@ -231,9 +231,10 @@ fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
|
||||
Python::initialize();
|
||||
Python::attach(f)
|
||||
|
||||
@@ -49,11 +49,12 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::SeedOssReasoningParser;
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::ReasoningParser;
|
||||
use crate::reasoning::tests::{SEED_THINK_END_ID, SEED_THINK_START_ID, fake_tokenizer};
|
||||
|
||||
#[test]
|
||||
fn without_prompt_markers_expects_start_token() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
|
||||
@@ -66,10 +67,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn picks_up_prompt_start_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
|
||||
parser.initialize(&[10]).unwrap();
|
||||
// Prompt prefills `<seed:think>`, opening reasoning before the stream.
|
||||
parser.initialize(&[SEED_THINK_START_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("reason</seed:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
@@ -78,10 +79,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn respects_prompt_end_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt already closed reasoning with `</seed:think>` (id 11).
|
||||
parser.initialize(&[11]).unwrap();
|
||||
// Prompt already closed reasoning with `</seed:think>`.
|
||||
parser.initialize(&[SEED_THINK_END_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
@@ -91,7 +92,7 @@ mod tests {
|
||||
#[test]
|
||||
fn handles_explicit_start_token() {
|
||||
// An explicit start delimiter must not leak into reasoning text.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
|
||||
@@ -103,7 +104,7 @@ mod tests {
|
||||
fn streams_explicit_start_token_across_pushes() {
|
||||
// Start token, reasoning body, end token, and content arrive in separate
|
||||
// streaming deltas.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let mut reasoning = String::new();
|
||||
@@ -131,9 +132,9 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn handles_partial_delimiters_across_pushes() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[10]).unwrap();
|
||||
parser.initialize(&[SEED_THINK_START_ID]).unwrap();
|
||||
|
||||
// Closing delimiter `</seed:think>` arrives in two halves.
|
||||
let first = parser.push("reason</seed:").unwrap();
|
||||
|
||||
@@ -127,14 +127,15 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::Step3p5ReasoningParser;
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::ReasoningParser;
|
||||
use crate::reasoning::tests::{THINK_START_ID, fake_tokenizer};
|
||||
|
||||
#[test]
|
||||
fn picks_up_prompt_start_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
// Prompt prefills `<think>` (id 1), opening reasoning before the stream.
|
||||
parser.initialize(&[1]).unwrap();
|
||||
// Prompt prefills `<think>`, opening reasoning before the stream.
|
||||
parser.initialize(&[THINK_START_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
|
||||
assert_eq!(
|
||||
@@ -146,7 +147,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn handles_unterminated_reasoning() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let pushed = parser.push("<think>reason without end").unwrap();
|
||||
@@ -159,7 +160,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn handles_empty_input() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let pushed = parser.push("").unwrap();
|
||||
@@ -172,9 +173,9 @@ mod tests {
|
||||
fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
|
||||
// Only the immediately-adjacent framing `\n` is dropped on each side of
|
||||
// `</think>`; surrounding newlines remain part of reasoning/content.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[1]).unwrap();
|
||||
parser.initialize(&[THINK_START_ID]).unwrap();
|
||||
|
||||
let delta = parser
|
||||
.push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
|
||||
@@ -188,7 +189,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn drops_framing_newlines_in_single_push() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
|
||||
@@ -198,7 +199,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn drops_framing_newlines_across_pushes() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
// The trailing `\n` from the first push is held until we know whether
|
||||
@@ -219,7 +220,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn replays_held_newline_when_more_reasoning_follows() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason\n").unwrap();
|
||||
@@ -232,7 +233,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn finish_flushes_held_newline_in_unterminated_stream() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason\n").unwrap();
|
||||
@@ -245,7 +246,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn preserves_inner_newlines_in_reasoning() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
|
||||
@@ -257,7 +258,7 @@ mod tests {
|
||||
fn trims_only_one_trailing_reasoning_newline() {
|
||||
// Only the single framing newline immediately before `</think>` is
|
||||
// dropped; earlier newlines in the reasoning body are preserved.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
|
||||
@@ -269,7 +270,7 @@ mod tests {
|
||||
fn drops_only_first_content_newline_after_transition() {
|
||||
// The leading-`\n` drop applies only to the first content delta after
|
||||
// `</think>`; later deltas pass through untouched.
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let first = parser.push("<think>reason</think>").unwrap();
|
||||
@@ -288,7 +289,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn passes_through_clean_boundary_without_framing_newlines() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think>reason</think>tail").unwrap();
|
||||
@@ -298,7 +299,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn handles_empty_reasoning_section() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<think></think>answer").unwrap();
|
||||
|
||||
@@ -1,54 +1,42 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::{
|
||||
DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
|
||||
Qwen3ReasoningParser, ReasoningParser,
|
||||
};
|
||||
|
||||
pub(crate) struct FakeTokenizer;
|
||||
pub(crate) const THINK_START_ID: u32 = 256;
|
||||
pub(crate) const THINK_END_ID: u32 = 257;
|
||||
pub(crate) const START_THINKING_ID: u32 = 258;
|
||||
pub(crate) const END_THINKING_ID: u32 = 259;
|
||||
pub(crate) const MINIMAX_THINK_START_ID: u32 = 260;
|
||||
pub(crate) const MINIMAX_THINK_END_ID: u32 = 261;
|
||||
pub(crate) const SPECIAL_BOUNDARY_ID: u32 = 262;
|
||||
pub(crate) const MM_THINK_START_ID: u32 = 263;
|
||||
pub(crate) const MM_THINK_END_ID: u32 = 264;
|
||||
pub(crate) const SEED_THINK_START_ID: u32 = 265;
|
||||
pub(crate) const SEED_THINK_END_ID: u32 = 266;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.chars().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(token_ids
|
||||
.iter()
|
||||
.map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}'))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<think>" => Some(1),
|
||||
"</think>" => Some(2),
|
||||
"<|START_THINKING|>" => Some(3),
|
||||
"<|END_THINKING|>" => Some(4),
|
||||
"◁think▷" => Some(5),
|
||||
"◁/think▷" => Some(6),
|
||||
"<mm:think>" => Some(8),
|
||||
"</mm:think>" => Some(9),
|
||||
"<seed:think>" => Some(10),
|
||||
"</seed:think>" => Some(11),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_special_id(&self, token_id: u32) -> bool {
|
||||
token_id == 7
|
||||
}
|
||||
pub(crate) fn fake_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<think>", THINK_START_ID)
|
||||
.with_regular_token("</think>", THINK_END_ID)
|
||||
.with_regular_token("<|START_THINKING|>", START_THINKING_ID)
|
||||
.with_regular_token("<|END_THINKING|>", END_THINKING_ID)
|
||||
.with_regular_token("◁think▷", MINIMAX_THINK_START_ID)
|
||||
.with_regular_token("◁/think▷", MINIMAX_THINK_END_ID)
|
||||
.with_special_token("<special-boundary>", SPECIAL_BOUNDARY_ID)
|
||||
.with_regular_token("<mm:think>", MM_THINK_START_ID)
|
||||
.with_regular_token("</mm:think>", MM_THINK_END_ID)
|
||||
.with_regular_token("<seed:think>", SEED_THINK_START_ID)
|
||||
.with_regular_token("</seed:think>", SEED_THINK_END_ID)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delimited_content_only_stream() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser =
|
||||
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
|
||||
|
||||
@@ -60,7 +48,7 @@ fn delimited_content_only_stream() {
|
||||
|
||||
#[test]
|
||||
fn delimited_single_chunk_with_reasoning_and_content() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser =
|
||||
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
|
||||
|
||||
@@ -71,7 +59,7 @@ fn delimited_single_chunk_with_reasoning_and_content() {
|
||||
|
||||
#[test]
|
||||
fn delimited_partial_tokens_across_chunks() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser =
|
||||
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
|
||||
|
||||
@@ -83,10 +71,10 @@ fn delimited_partial_tokens_across_chunks() {
|
||||
|
||||
#[test]
|
||||
fn delimited_finish_flushes_buffer() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser =
|
||||
DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false).unwrap();
|
||||
parser.initialize(&[1]);
|
||||
parser.initialize(&[THINK_START_ID]);
|
||||
|
||||
let delta = parser.push("unfinished</thi");
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("unfinished"));
|
||||
@@ -96,7 +84,7 @@ fn delimited_finish_flushes_buffer() {
|
||||
|
||||
#[test]
|
||||
fn qwen3_without_prompt_markers_expects_start_token() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("reason</think>answer").unwrap();
|
||||
@@ -106,9 +94,9 @@ fn qwen3_without_prompt_markers_expects_start_token() {
|
||||
|
||||
#[test]
|
||||
fn qwen3_prompt_end_marker_starts_in_content() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[2]).unwrap();
|
||||
parser.initialize(&[THINK_END_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
@@ -117,7 +105,7 @@ fn qwen3_prompt_end_marker_starts_in_content() {
|
||||
|
||||
#[test]
|
||||
fn qwen3_tolerates_old_and_new_formats() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
|
||||
let mut old_parser = Qwen3ReasoningParser::new(tokenizer.clone()).unwrap();
|
||||
let old = old_parser.push("<think>reason</think>answer").unwrap();
|
||||
@@ -125,7 +113,7 @@ fn qwen3_tolerates_old_and_new_formats() {
|
||||
assert_eq!(old.content.as_deref(), Some("answer"));
|
||||
|
||||
let mut new_parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
|
||||
new_parser.initialize(&[1]).unwrap();
|
||||
new_parser.initialize(&[THINK_START_ID]).unwrap();
|
||||
let new = new_parser.push("reason</think>answer").unwrap();
|
||||
assert_eq!(new.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(new.content.as_deref(), Some("answer"));
|
||||
@@ -133,10 +121,10 @@ fn qwen3_tolerates_old_and_new_formats() {
|
||||
|
||||
#[test]
|
||||
fn qwen3_stops_scanning_at_last_special_token() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = Qwen3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
parser.initialize(&[1, 7]).unwrap();
|
||||
parser.initialize(&[THINK_START_ID, SPECIAL_BOUNDARY_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
@@ -145,7 +133,7 @@ fn qwen3_stops_scanning_at_last_special_token() {
|
||||
|
||||
#[test]
|
||||
fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("reason</think>answer").unwrap();
|
||||
@@ -155,10 +143,10 @@ fn deepseek_r1_defaults_to_reasoning_without_prompt_boundary() {
|
||||
|
||||
#[test]
|
||||
fn deepseek_r1_stops_scanning_at_last_special_token() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = DeepSeekR1ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
parser.initialize(&[2, 7]).unwrap();
|
||||
parser.initialize(&[THINK_END_ID, SPECIAL_BOUNDARY_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("reason</think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
@@ -167,7 +155,7 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_handles_explicit_think_delimiters() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
|
||||
@@ -177,7 +165,7 @@ fn minimax_m3_handles_explicit_think_delimiters() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_drops_leading_end_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("</mm:think>answer").unwrap();
|
||||
@@ -187,7 +175,7 @@ fn minimax_m3_drops_leading_end_marker() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_preserves_non_leading_end_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("XXX</mm:think>YYY").unwrap();
|
||||
@@ -197,7 +185,7 @@ fn minimax_m3_preserves_non_leading_end_marker() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_drops_split_leading_end_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
assert!(parser.push("</mm").unwrap().is_empty());
|
||||
@@ -208,9 +196,9 @@ fn minimax_m3_drops_split_leading_end_marker() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_uses_prompt_prefilled_start_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[8]).unwrap();
|
||||
parser.initialize(&[MM_THINK_START_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("reason</mm:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
@@ -219,9 +207,9 @@ fn minimax_m3_uses_prompt_prefilled_start_marker() {
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_uses_prompt_prefilled_end_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(fake_tokenizer());
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[9]).unwrap();
|
||||
parser.initialize(&[MM_THINK_END_ID]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, 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]
|
||||
|
||||
@@ -14,8 +14,6 @@ mod parameters;
|
||||
mod qwen_coder;
|
||||
#[cfg(any(test, feature = "test-util"))]
|
||||
pub mod test_utils;
|
||||
use crate::utils;
|
||||
|
||||
use std::collections::{BTreeMap, btree_map};
|
||||
|
||||
pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser};
|
||||
@@ -35,6 +33,8 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
pub use xgrammar_structural_tag::Model as StructuralTagModel;
|
||||
|
||||
use crate::utils;
|
||||
|
||||
/// One function-style tool made available to the model.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Tool {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user