Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6e714a103c | ||
|
|
c9951fd5c7 | ||
|
|
a78c15616f | ||
|
|
5c4db60f01 | ||
|
|
4e5ca89cfe | ||
|
|
a22e0dfc69 | ||
|
|
cc56379e28 | ||
|
|
024b06b0dc | ||
|
|
e7d0fcbc09 | ||
|
|
aa8bb5562e | ||
|
|
fa4bec9056 | ||
|
|
dee5da1dec | ||
|
|
ed41aa270a | ||
|
|
77a9c5ae28 | ||
|
|
f651a8a9a4 | ||
|
|
8f82be5705 | ||
|
|
a461070d1c | ||
|
|
4470ae84de | ||
|
|
697c34b97b | ||
|
|
5b431b905c | ||
|
|
89e99202f2 | ||
|
|
b446792306 | ||
|
|
c3b1f9e827 | ||
|
|
df802a87b7 | ||
|
|
93d8f834dd | ||
|
|
aeb35b90f0 | ||
|
|
9a08a5118e | ||
|
|
c5200d3565 | ||
|
|
3c1396bab6 | ||
|
|
9969466a59 | ||
|
|
3406e8f83d | ||
|
|
a264e41975 | ||
|
|
f098ee70c7 | ||
|
|
9294dd27eb | ||
|
|
b1190d03cc | ||
|
|
92c7fac640 | ||
|
|
ac521f6237 | ||
|
|
28242824e0 | ||
|
|
68294739d1 | ||
|
|
c8d2f3cb14 | ||
|
|
345b28ff2f | ||
|
|
248d1fbb71 | ||
|
|
11b26c5528 | ||
|
|
20434c472e | ||
|
|
c8f9c156a5 | ||
|
|
953bba488d | ||
|
|
3a9784b82c | ||
|
|
3cecee40f3 | ||
|
|
a7732537f4 | ||
|
|
727971f1c1 | ||
|
|
25671cb520 | ||
|
|
27d5f78b63 | ||
|
|
7a341fa109 | ||
|
|
f41e8ddc97 | ||
|
|
245888ff77 | ||
|
|
e840f0d3f5 | ||
|
|
fcaa84efa7 | ||
|
|
9e84ec8648 | ||
|
|
d8f483dc30 | ||
|
|
dc148dc4d7 | ||
|
|
7cf7cbcd95 | ||
|
|
c231d1f290 | ||
|
|
db808b3961 | ||
|
|
00ebf19cca | ||
|
|
ded6676458 | ||
|
|
7a327f0b4f | ||
|
|
1ab9522935 | ||
|
|
0fc2512094 | ||
|
|
62c7d8009f | ||
|
|
ab80b3dff4 | ||
|
|
91055efd36 | ||
|
|
3675bcff67 | ||
|
|
bdbd7278b6 | ||
|
|
5dc36a4fa5 | ||
|
|
aab7af0bcb |
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
+17
-1
@@ -5099,6 +5099,7 @@ dependencies = [
|
||||
"uuid",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-model-files",
|
||||
"vllm-parser",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
@@ -5236,6 +5237,20 @@ dependencies = [
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-model-files"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"hf-hub",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-parser"
|
||||
version = "0.1.0"
|
||||
@@ -5310,6 +5325,7 @@ dependencies = [
|
||||
"vllm-llm",
|
||||
"vllm-metrics",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
@@ -5323,7 +5339,6 @@ dependencies = [
|
||||
"enum-as-inner",
|
||||
"expect-test",
|
||||
"futures",
|
||||
"hf-hub",
|
||||
"itertools 0.14.0",
|
||||
"reqwest",
|
||||
"serde",
|
||||
@@ -5338,6 +5353,7 @@ dependencies = [
|
||||
"trait-set",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-llm",
|
||||
"vllm-model-files",
|
||||
"vllm-tokenizer",
|
||||
]
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ members = [
|
||||
"src/llm",
|
||||
"src/managed-engine",
|
||||
"src/metrics",
|
||||
"src/model-files",
|
||||
"src/mock-engine",
|
||||
"src/parser",
|
||||
"src/parser/python",
|
||||
@@ -126,6 +127,7 @@ vllm-chat = { path = "src/chat" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
vllm-managed-engine = { path = "src/managed-engine" }
|
||||
vllm-model-files = { path = "src/model-files" }
|
||||
vllm-metrics = { path = "src/metrics" }
|
||||
vllm-parser = { path = "src/parser" }
|
||||
vllm-server = { path = "src/server" }
|
||||
|
||||
@@ -32,6 +32,7 @@ trait-set.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-model-files.workspace = true
|
||||
vllm-parser.workspace = true
|
||||
vllm-text.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
@@ -50,6 +51,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(
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use thiserror::Error;
|
||||
use thiserror_ext::Macro;
|
||||
use vllm_model_files::Error as ModelFilesError;
|
||||
|
||||
type BoxedError = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
@@ -69,6 +70,8 @@ pub enum Error {
|
||||
#[error(transparent)]
|
||||
Text(#[from] vllm_text::Error),
|
||||
#[error(transparent)]
|
||||
ModelFiles(#[from] ModelFilesError),
|
||||
#[error(transparent)]
|
||||
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ use thiserror_ext::AsReport as _;
|
||||
use tracing::{info, trace, warn};
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::{
|
||||
HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config,
|
||||
ResolvedModelFiles, SpecialTokens, TokenizerConfig, load_tokenizer_config,
|
||||
};
|
||||
|
||||
use self::format::{
|
||||
@@ -42,7 +42,7 @@ pub struct HfChatRenderer {
|
||||
default_template: Option<CompiledChatTemplate>,
|
||||
default_template_kwargs: HashMap<String, JsonValue>,
|
||||
content_format: ContentFormatOption,
|
||||
special_tokens: Option<HfSpecialTokens>,
|
||||
special_tokens: Option<SpecialTokens>,
|
||||
multimodal: Option<MultimodalRenderInfo>,
|
||||
}
|
||||
|
||||
@@ -67,7 +67,7 @@ impl HfChatRenderer {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_special_tokens(mut self, special_tokens: Option<HfSpecialTokens>) -> Self {
|
||||
pub fn with_special_tokens(mut self, special_tokens: Option<SpecialTokens>) -> Self {
|
||||
self.special_tokens = special_tokens;
|
||||
self
|
||||
}
|
||||
@@ -83,7 +83,7 @@ impl HfChatRenderer {
|
||||
options: LoadModelBackendsOptions,
|
||||
multimodal: Option<MultimodalRenderInfo>,
|
||||
) -> Result<Self> {
|
||||
let HfTokenizerConfig {
|
||||
let TokenizerConfig {
|
||||
special_tokens,
|
||||
chat_template,
|
||||
..
|
||||
@@ -451,7 +451,7 @@ mod tests {
|
||||
use expect_test::expect;
|
||||
use serde_json::Value;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
|
||||
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
|
||||
|
||||
use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo};
|
||||
use crate::request::{
|
||||
@@ -675,7 +675,7 @@ mod tests {
|
||||
#[test]
|
||||
fn chat_template_injects_special_tokens_into_context() {
|
||||
let request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<bos>".to_string())),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
//!
|
||||
//! This module is inlined from SMG's tokenizer crate with local adaptations:
|
||||
//! - thinking-related detection/state is removed
|
||||
//! - special tokens are wired to `vllm_text::backends::hf::HfSpecialTokens`
|
||||
//! - special tokens are wired to `vllm_text::backends::hf::SpecialTokens`
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
@@ -11,7 +11,7 @@ use std::path::Path;
|
||||
use minijinja::Environment;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{self};
|
||||
use vllm_text::backend::hf::HfSpecialTokens;
|
||||
use vllm_text::backend::hf::SpecialTokens;
|
||||
|
||||
use super::error::TemplateError;
|
||||
use super::format::{
|
||||
@@ -46,7 +46,7 @@ pub(super) struct TemplateContext<'a> {
|
||||
pub(super) tools: Option<&'a [TemplateTool]>,
|
||||
pub(super) documents: Option<&'a [serde_json::Value]>,
|
||||
#[serde(flatten)]
|
||||
pub(super) special_tokens: Option<&'a HfSpecialTokens>,
|
||||
pub(super) special_tokens: Option<&'a SpecialTokens>,
|
||||
#[serde(flatten)]
|
||||
pub(super) template_kwargs: Option<&'a HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
@@ -133,7 +133,7 @@ mod tests {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::TempDir;
|
||||
use vllm_text::backend::hf::{HfSpecialTokens, NamedSpecialToken};
|
||||
use vllm_text::backend::hf::{NamedSpecialToken, SpecialTokens};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -170,7 +170,7 @@ mod tests {
|
||||
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
|
||||
.unwrap();
|
||||
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
|
||||
eos_token: Some(NamedSpecialToken::Text("</s>".to_string())),
|
||||
..Default::default()
|
||||
@@ -205,7 +205,7 @@ mod tests {
|
||||
CompiledChatTemplate::new(template.to_string(), ChatTemplateContentFormatOption::Auto)
|
||||
.unwrap();
|
||||
|
||||
let special_tokens = HfSpecialTokens {
|
||||
let special_tokens = SpecialTokens {
|
||||
bos_token: Some(NamedSpecialToken::Text("<s>".to_string())),
|
||||
eos_token: None,
|
||||
..Default::default()
|
||||
|
||||
+34
-103
@@ -21,15 +21,17 @@ use vllm_engine_core_client::protocol::{
|
||||
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 +160,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 +185,7 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: true,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -217,6 +193,7 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: false,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,13 +201,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 +265,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 +662,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 +676,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 +697,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>,
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "vllm-model-files"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
hf-hub.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,21 @@
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::json::read_json_file;
|
||||
|
||||
/// Minimal subset of `tokenizer_config.json` needed by tokenizer selection.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub(crate) struct TokenizerConfig {
|
||||
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
|
||||
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
|
||||
/// containing "Tiktoken" which can be used as a hint for backend
|
||||
/// selection.
|
||||
pub tokenizer_class: Option<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
use thiserror::Error as ThisError;
|
||||
|
||||
/// Error returned while resolving or reading model files.
|
||||
#[derive(Debug, ThisError)]
|
||||
#[error("model file error: {0}")]
|
||||
pub struct Error(String);
|
||||
|
||||
impl Error {
|
||||
pub(crate) fn new(message: impl Into<String>) -> Self {
|
||||
Self(message.into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Result type used by model-file discovery helpers.
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
@@ -0,0 +1,31 @@
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
use thiserror_ext::AsReport as _;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Read an optional JSON file into `T`, returning `T::default()` when absent.
|
||||
pub fn read_json_file<T>(path: Option<&Path>) -> Result<T>
|
||||
where
|
||||
T: for<'de> Deserialize<'de> + Default,
|
||||
{
|
||||
let Some(path) = path else {
|
||||
return Ok(T::default());
|
||||
};
|
||||
let content = fs::read_to_string(path).map_err(|error| {
|
||||
Error::new(format!(
|
||||
"failed to read {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})?;
|
||||
serde_json::from_str(&content).map_err(|error| {
|
||||
Error::new(format!(
|
||||
"failed to parse {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Hugging Face model file discovery shared by Rust frontend crates.
|
||||
|
||||
mod config;
|
||||
mod error;
|
||||
mod json;
|
||||
mod model_files;
|
||||
|
||||
pub use error::{Error, Result};
|
||||
pub use json::read_json_file;
|
||||
pub use model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
+70
-125
@@ -4,7 +4,7 @@ use hf_hub::Cache;
|
||||
use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
|
||||
use thiserror_ext::AsReport as _;
|
||||
|
||||
use super::config::{HfTokenizerConfig, load_tokenizer_config};
|
||||
use crate::config::{TokenizerConfig, load_tokenizer_config};
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
const HF_TOKEN_ENV: &str = "HF_TOKEN";
|
||||
@@ -26,6 +26,28 @@ pub enum TokenizerSource {
|
||||
}
|
||||
|
||||
impl TokenizerSource {
|
||||
/// Select a tokenizer source from a tokenizer file path.
|
||||
pub fn from_path(path: impl Into<PathBuf>) -> Result<Self> {
|
||||
let path = path.into();
|
||||
let file_name = path.file_name().and_then(|name| name.to_str()).ok_or_else(|| {
|
||||
Error::new(format!(
|
||||
"tokenizer path has no file name: {}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
match file_name {
|
||||
"tekken.json" => Ok(Self::Tekken(path)),
|
||||
"tokenizer.json" => Ok(Self::HuggingFace(path)),
|
||||
_ if is_tiktoken_file(&path) => Ok(Self::Tiktoken(path)),
|
||||
_ => Err(Error::new(format!(
|
||||
"unsupported tokenizer file '{}'",
|
||||
path.display()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the local filesystem path for this tokenizer source.
|
||||
pub fn path(&self) -> &Path {
|
||||
match self {
|
||||
Self::HuggingFace(path) | Self::Tiktoken(path) | Self::Tekken(path) => path,
|
||||
@@ -38,10 +60,15 @@ impl TokenizerSource {
|
||||
pub struct ResolvedModelFiles {
|
||||
/// The selected tokenizer source for this model.
|
||||
pub tokenizer: TokenizerSource,
|
||||
/// Path to `tokenizer_config.json` when present.
|
||||
pub tokenizer_config_path: Option<PathBuf>,
|
||||
/// Path to `generation_config.json` when present.
|
||||
pub generation_config_path: Option<PathBuf>,
|
||||
/// Path to `preprocessor_config.json` when present.
|
||||
pub preprocessor_config_path: Option<PathBuf>,
|
||||
/// Path to a discovered chat template file when present.
|
||||
pub chat_template_path: Option<PathBuf>,
|
||||
/// Path to `config.json` when present.
|
||||
pub config_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
@@ -76,10 +103,10 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
|
||||
}
|
||||
|
||||
async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
|
||||
let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
|
||||
let api = build_api().map_err(|error| Error::new(format!("{}", error.as_report())))?;
|
||||
let repo = api.model(model_id.to_string());
|
||||
let info = repo.info().await.map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"failed to fetch model '{model_id}': {}",
|
||||
error.as_report()
|
||||
))
|
||||
@@ -138,9 +165,10 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let model_dir = tokenizer.path().parent().ok_or_else(|| {
|
||||
Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
|
||||
})?;
|
||||
let model_dir = tokenizer
|
||||
.path()
|
||||
.parent()
|
||||
.ok_or_else(|| Error::new("resolved tokenizer file has no parent directory"))?;
|
||||
let generation_config_path = cache_repo.get("generation_config.json");
|
||||
let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
|
||||
let chat_template_path = discover_chat_template_in_dir(model_dir);
|
||||
@@ -162,113 +190,88 @@ async fn resolve_remote_tokenizer_source(
|
||||
siblings: &std::collections::BTreeSet<&str>,
|
||||
tokenizer_class: Option<&str>,
|
||||
) -> Result<TokenizerSource> {
|
||||
if let Some(tekken_path) = download_if_present(repo, model_id, siblings, "tekken.json").await? {
|
||||
return Ok(TokenizerSource::Tekken(tekken_path));
|
||||
}
|
||||
|
||||
let tokenizer_path = if siblings.contains("tokenizer.json") {
|
||||
let tokenizer_path = if siblings.contains("tekken.json") {
|
||||
download_known_file(repo, model_id, "tekken.json").await?
|
||||
} else if siblings.contains("tokenizer.json") {
|
||||
download_known_file(repo, model_id, "tokenizer.json").await?
|
||||
} else if let Some(tiktoken_name) = find_tiktoken_sibling(siblings) {
|
||||
download_known_file(repo, model_id, tiktoken_name).await?
|
||||
} else {
|
||||
return Err(Error::Tokenizer(format!(
|
||||
return Err(Error::new(format!(
|
||||
"model '{model_id}' does not expose a supported tokenizer file \
|
||||
(tokenizer.json, tiktoken.model, or *.tiktoken) on Hugging Face"
|
||||
)));
|
||||
};
|
||||
|
||||
Ok(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_class,
|
||||
None,
|
||||
))
|
||||
resolve_tokenizer_source(tokenizer_path, tokenizer_class)
|
||||
}
|
||||
|
||||
fn resolve_cached_tokenizer_source(
|
||||
cache_repo: &hf_hub::CacheRepo,
|
||||
tokenizer_config: &HfTokenizerConfig,
|
||||
tokenizer_config: &TokenizerConfig,
|
||||
) -> Result<Option<TokenizerSource>> {
|
||||
let tekken_path = cache_repo.get("tekken.json");
|
||||
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return Ok(Some(TokenizerSource::Tekken(tekken_path)));
|
||||
}
|
||||
|
||||
let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
|
||||
// tiktoken.model is the most common name, try it first.
|
||||
cache_repo.get("tiktoken.model").or_else(|| {
|
||||
// Scan for any *.tiktoken file in the cache snapshot directory.
|
||||
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
|
||||
discover_tiktoken_in_dir(&snapshot_dir)
|
||||
let Some(tokenizer_path) = cache_repo
|
||||
.get("tekken.json")
|
||||
.or_else(|| cache_repo.get("tokenizer.json"))
|
||||
.or_else(|| {
|
||||
// tiktoken.model is the most common name, try it first.
|
||||
cache_repo.get("tiktoken.model").or_else(|| {
|
||||
// Scan for any *.tiktoken file in the cache snapshot directory.
|
||||
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
|
||||
discover_tiktoken_in_dir(&snapshot_dir)
|
||||
})
|
||||
})
|
||||
}) else {
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_config.tokenizer_class.as_deref(),
|
||||
None,
|
||||
)))
|
||||
)?))
|
||||
}
|
||||
|
||||
fn resolve_local_tokenizer_source(
|
||||
model_dir: &Path,
|
||||
tokenizer_config: &HfTokenizerConfig,
|
||||
tokenizer_config: &TokenizerConfig,
|
||||
) -> Result<TokenizerSource> {
|
||||
let tekken_path = local_file_if_exists(model_dir, "tekken.json");
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return Ok(TokenizerSource::Tekken(tekken_path));
|
||||
}
|
||||
|
||||
let tokenizer_path = local_file_if_exists(model_dir, "tokenizer.json")
|
||||
let tokenizer_path = local_file_if_exists(model_dir, "tekken.json")
|
||||
.or_else(|| local_file_if_exists(model_dir, "tokenizer.json"))
|
||||
.or_else(|| local_file_if_exists(model_dir, "tiktoken.model"))
|
||||
.or_else(|| discover_tiktoken_in_dir(model_dir))
|
||||
.ok_or_else(|| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"local model directory '{}' does not contain a supported tokenizer file \
|
||||
(tokenizer.json, tiktoken.model, or *.tiktoken)",
|
||||
model_dir.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(resolve_tokenizer_source(
|
||||
tokenizer_path,
|
||||
tokenizer_config.tokenizer_class.as_deref(),
|
||||
None,
|
||||
))
|
||||
resolve_tokenizer_source(tokenizer_path, tokenizer_config.tokenizer_class.as_deref())
|
||||
}
|
||||
|
||||
/// Choose the tokenizer.
|
||||
///
|
||||
/// Selection order:
|
||||
/// 1. `tekken.json` — Mistral native tokenizer (preferred over HF `tokenizer.json` because the HF
|
||||
/// version has a known regex bug for Mistral models).
|
||||
/// 2. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
|
||||
/// 3. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
|
||||
/// 1. File extension — `.tiktoken` / `tiktoken.model` files use tiktoken from BPE data.
|
||||
/// 2. `tokenizer_class` in `tokenizer_config.json` — classes containing "Tiktoken" (case-
|
||||
/// insensitive) trigger tiktoken loading from a sibling BPE file.
|
||||
/// 4. Default — `tokenizer.json` in HuggingFace format.
|
||||
/// 3. Default — `tokenizer.json` in HuggingFace format.
|
||||
fn resolve_tokenizer_source(
|
||||
tokenizer_path: PathBuf,
|
||||
tokenizer_class: Option<&str>,
|
||||
tekken_path: Option<PathBuf>,
|
||||
) -> TokenizerSource {
|
||||
if let Some(tekken_path) = tekken_path {
|
||||
return TokenizerSource::Tekken(tekken_path);
|
||||
}
|
||||
) -> Result<TokenizerSource> {
|
||||
let tokenizer = TokenizerSource::from_path(tokenizer_path)?;
|
||||
|
||||
if is_tiktoken_file(&tokenizer_path) {
|
||||
return TokenizerSource::Tiktoken(tokenizer_path);
|
||||
}
|
||||
|
||||
if tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
|
||||
&& let Some(dir) = tokenizer_path.parent()
|
||||
if let TokenizerSource::HuggingFace(path) = &tokenizer
|
||||
&& tokenizer_class.is_some_and(|cls| cls.to_ascii_lowercase().contains("tiktoken"))
|
||||
&& let Some(dir) = path.parent()
|
||||
&& let Some(tiktoken_path) = discover_tiktoken_in_dir(dir)
|
||||
{
|
||||
return TokenizerSource::Tiktoken(tiktoken_path);
|
||||
return Ok(TokenizerSource::Tiktoken(tiktoken_path));
|
||||
}
|
||||
|
||||
TokenizerSource::HuggingFace(tokenizer_path)
|
||||
Ok(tokenizer)
|
||||
}
|
||||
|
||||
/// Download `filename` only if it exists in `siblings`.
|
||||
@@ -286,7 +289,7 @@ async fn download_if_present(
|
||||
|
||||
async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) -> Result<PathBuf> {
|
||||
repo.get(filename).await.map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
Error::new(format!(
|
||||
"failed to download '{filename}' for model '{model_id}': {}",
|
||||
error.as_report()
|
||||
))
|
||||
@@ -317,7 +320,7 @@ fn find_tiktoken_sibling<'a>(siblings: &std::collections::BTreeSet<&'a str>) ->
|
||||
}
|
||||
|
||||
/// Discover a tiktoken model file in a local directory.
|
||||
pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
|
||||
fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
|
||||
let tiktoken_model = dir.join("tiktoken.model");
|
||||
if tiktoken_model.exists() {
|
||||
return Some(tiktoken_model);
|
||||
@@ -337,7 +340,7 @@ pub(super) fn discover_tiktoken_in_dir(dir: &std::path::Path) -> Option<PathBuf>
|
||||
}
|
||||
|
||||
/// Returns `true` if `path` points to a tiktoken-format file (by name).
|
||||
pub(super) fn is_tiktoken_file(path: &std::path::Path) -> bool {
|
||||
fn is_tiktoken_file(path: &std::path::Path) -> bool {
|
||||
path.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.is_some_and(|name| name == "tiktoken.model" || name.ends_with(".tiktoken"))
|
||||
@@ -368,7 +371,6 @@ mod tests {
|
||||
use std::fs;
|
||||
|
||||
use tempfile::tempdir;
|
||||
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
|
||||
|
||||
use super::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
@@ -399,61 +401,4 @@ mod tests {
|
||||
Some(dir.path().join("tokenizer_config.json"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "too slow for CI and requires network access to Hugging Face"]
|
||||
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
|
||||
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
|
||||
.await
|
||||
.expect("resolve real Kimi K2.5 model files");
|
||||
|
||||
let tokenizer_path = match &files.tokenizer {
|
||||
TokenizerSource::Tiktoken(path) => path.clone(),
|
||||
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
|
||||
};
|
||||
|
||||
for backend in [
|
||||
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
|
||||
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
|
||||
] {
|
||||
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
|
||||
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
|
||||
let tool_section_id = backend
|
||||
.token_to_id("<|tool_calls_section_begin|>")
|
||||
.expect("resolve tool call section marker");
|
||||
let contraction_heavy_text =
|
||||
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
|
||||
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(think_id, end_think_id, tool_section_id),
|
||||
(163606, 163607, 163595)
|
||||
);
|
||||
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
|
||||
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
|
||||
assert_eq!(
|
||||
backend.decode(&[tool_section_id], true).unwrap(),
|
||||
"<|tool_calls_section_begin|>"
|
||||
);
|
||||
|
||||
// This demonstrates that we're using Kimi's custom BPE pattern.
|
||||
// With CL100K this will be 23 tokens instead.
|
||||
assert_eq!(
|
||||
contraction_heavy_ids,
|
||||
vec![
|
||||
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
|
||||
1573, 56229, 13922, 13,
|
||||
]
|
||||
);
|
||||
assert_eq!(contraction_heavy_ids.len(), 17);
|
||||
assert_eq!(
|
||||
backend.decode(&contraction_heavy_ids, false).unwrap(),
|
||||
contraction_heavy_text
|
||||
);
|
||||
|
||||
// Special-looking text that is not actually registered should fail gracefully.
|
||||
assert_eq!(backend.token_to_id("◁think▷"), None);
|
||||
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -49,11 +49,14 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::SeedOssReasoningParser;
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::{
|
||||
ReasoningParser,
|
||||
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 +69,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 +81,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 +94,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 +106,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 +134,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,17 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::Step3p5ReasoningParser;
|
||||
use crate::reasoning::{ReasoningParser, tests::FakeTokenizer};
|
||||
use crate::reasoning::{
|
||||
ReasoningParser,
|
||||
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 +149,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 +162,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 +175,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 +191,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 +201,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 +222,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 +235,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 +248,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 +260,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 +272,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 +291,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 +301,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]
|
||||
|
||||
@@ -685,7 +685,8 @@ mod tests {
|
||||
let mut parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
let error = parser.parse_chunk("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "\n<bad>\n</tool_call>": "#]]
|
||||
.assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -697,7 +698,7 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
|
||||
expect![[r#"tool parser parsing failed: near "\n<function=get_weather>\n<parameter=location>SF</function>\n</tool_call>": "#]].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -124,42 +124,17 @@ impl UnifiedParser for CombinedParser {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::CombinedParser;
|
||||
use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser};
|
||||
use crate::tool::{Qwen3XmlToolParser, Tool, ToolParser};
|
||||
use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
|
||||
|
||||
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 {
|
||||
"<think>" => Some(1),
|
||||
"</think>" => Some(2),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
fn tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<think>", 256)
|
||||
.with_regular_token("</think>", 257)
|
||||
}
|
||||
|
||||
fn test_tools() -> Vec<Tool> {
|
||||
@@ -273,7 +248,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn combined_parser_emits_reasoning_and_text() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let tokenizer = Arc::new(tokenizer());
|
||||
let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap();
|
||||
let mut parser = CombinedParser::new(Some(reasoning), None);
|
||||
|
||||
|
||||
@@ -515,7 +515,7 @@ mod tests {
|
||||
|
||||
use serde_json::{Value, json};
|
||||
use thiserror_ext::AsReport;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use winnow::combinator::{eof, terminated};
|
||||
use winnow::error::ErrMode;
|
||||
use winnow::prelude::*;
|
||||
@@ -527,66 +527,15 @@ mod tests {
|
||||
use crate::tool::Tool;
|
||||
use crate::unified::{UnifiedParserEvent, parsing_failed};
|
||||
|
||||
struct FakeTokenizer;
|
||||
const CHANNEL_START_ID: u32 = 256;
|
||||
const CHANNEL_END_ID: u32 = 257;
|
||||
const TURN_BOUNDARY_ID: u32 = 258;
|
||||
|
||||
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_START => Some(100),
|
||||
CHANNEL_END => Some(101),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_special_id(&self, token_id: u32) -> bool {
|
||||
matches!(token_id, 100..=105)
|
||||
}
|
||||
}
|
||||
|
||||
struct MissingTokenTokenizer;
|
||||
|
||||
impl Tokenizer for MissingTokenTokenizer {
|
||||
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
|
||||
}
|
||||
fn tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_special_token(CHANNEL_START, CHANNEL_START_ID)
|
||||
.with_special_token(CHANNEL_END, CHANNEL_END_ID)
|
||||
.with_special_token("<turn-boundary>", TURN_BOUNDARY_ID)
|
||||
}
|
||||
|
||||
trait UnifiedParserTestExt {
|
||||
@@ -716,12 +665,12 @@ mod tests {
|
||||
}
|
||||
|
||||
fn test_parser() -> Gemma4UnifiedParser {
|
||||
Gemma4UnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap()
|
||||
Gemma4UnifiedParser::new(&test_tools(), Arc::new(tokenizer())).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemma4_create_requires_channel_start_token() {
|
||||
let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(MissingTokenTokenizer)) {
|
||||
let error = match Gemma4UnifiedParser::new(&test_tools(), Arc::new(TestTokenizer::new())) {
|
||||
Ok(_) => panic!("expected missing token error"),
|
||||
Err(error) => error,
|
||||
};
|
||||
@@ -1046,7 +995,7 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() {
|
||||
let mut parser = test_parser();
|
||||
parser.initialize(&[100, 3000, 3001]).unwrap();
|
||||
parser.initialize(&[CHANNEL_START_ID, 3000, 3001]).unwrap();
|
||||
|
||||
let output = parser.parse_complete("reason<channel|>answer").unwrap();
|
||||
|
||||
@@ -1057,7 +1006,7 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_initialize_turn_prompt_starts_in_text() {
|
||||
let mut parser = test_parser();
|
||||
parser.initialize(&[104, 3000, 3001]).unwrap();
|
||||
parser.initialize(&[TURN_BOUNDARY_ID, 3000, 3001]).unwrap();
|
||||
|
||||
let output = parser.parse_complete("<|channel>thought\nreason<channel|>answer").unwrap();
|
||||
|
||||
@@ -1068,7 +1017,7 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_initialize_special_token_caps_boundary_scan() {
|
||||
let mut parser = test_parser();
|
||||
parser.initialize(&[100, 3000, 104, 3001]).unwrap();
|
||||
parser.initialize(&[CHANNEL_START_ID, 3000, TURN_BOUNDARY_ID, 3001]).unwrap();
|
||||
|
||||
let output = parser.parse_complete("answer").unwrap();
|
||||
|
||||
@@ -1079,7 +1028,7 @@ mod tests {
|
||||
#[test]
|
||||
fn gemma4_initialize_closed_channel_prompt_starts_in_text() {
|
||||
let mut parser = test_parser();
|
||||
parser.initialize(&[100, 3000, 3001, 101]).unwrap();
|
||||
parser.initialize(&[CHANNEL_START_ID, 3000, 3001, CHANNEL_END_ID]).unwrap();
|
||||
|
||||
let output = parser.parse_complete("answer").unwrap();
|
||||
|
||||
|
||||
@@ -395,9 +395,9 @@ pub fn parse_buffered_event<E>(
|
||||
Ok(event) => event,
|
||||
Err(ErrMode::Incomplete(_)) => return Ok(None),
|
||||
Err(ErrMode::Backtrack(e) | ErrMode::Cut(e)) => {
|
||||
// TODO: enrich context for error reporting
|
||||
let snippet = buffer.char_indices().nth(80).map_or(buffer, |(i, _)| &buffer[..i]);
|
||||
return Err(ToolParserError::ParsingFailed {
|
||||
message: e.to_string(),
|
||||
message: format!("near {snippet:?}: {e}"),
|
||||
});
|
||||
}
|
||||
};
|
||||
@@ -423,8 +423,9 @@ mod tests {
|
||||
use winnow::stream::{Offset, Partial, Stream};
|
||||
|
||||
use super::{
|
||||
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len,
|
||||
safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker,
|
||||
JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, parse_buffered_event,
|
||||
partial_prefix_len, safe_text_len, safe_text_len_mul, take_json_object, take_json_string,
|
||||
take_until_marker,
|
||||
};
|
||||
|
||||
#[test]
|
||||
@@ -832,4 +833,27 @@ mod tests {
|
||||
|
||||
assert!(matches!(error, ErrMode::Incomplete(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_buffered_event_error_includes_input_snippet() {
|
||||
let result = parse_buffered_event(" {\"x\":1}", |input| {
|
||||
take_json_object(input, &mut JsonObjectScanState::default())
|
||||
});
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("near \""), "error must include snippet");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_buffered_event_error_truncates_long_input() {
|
||||
let long_input = format!(" {}", "x".repeat(100));
|
||||
let result = parse_buffered_event(&long_input, |input| {
|
||||
take_json_object(input, &mut JsonObjectScanState::default())
|
||||
});
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(err.contains("near \""), "error must include snippet");
|
||||
assert!(
|
||||
!err.contains(&long_input),
|
||||
"snippet must be truncated for long input"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,8 +61,10 @@ expect-test.workspace = true
|
||||
rmp-serde.workspace = true
|
||||
serial_test.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
tower.workspace = true
|
||||
vllm-engine-core-client = { workspace = true, features = ["test-util"] }
|
||||
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
|
||||
zeromq.workspace = true
|
||||
|
||||
[lints]
|
||||
|
||||
@@ -4,21 +4,16 @@ mod convert;
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use futures::{Stream, StreamExt as _, stream};
|
||||
use futures::{Stream, StreamExt as _};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_openssl::SslStream;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tonic::transport::server::{Connected, TcpConnectInfo};
|
||||
use tonic::{Request, Response, Status};
|
||||
use tracing::info;
|
||||
use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _};
|
||||
|
||||
use self::convert::ResponseOpts;
|
||||
use crate::listener::{Listener, ListenerIo};
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Generated protobuf/gRPC types for the `vllm` package.
|
||||
@@ -31,78 +26,6 @@ pub use pb::generate_server::GenerateServer;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
/// Newtype over `tokio-openssl`'s `SslStream` so we can implement tonic's
|
||||
/// [`Connected`] on it (the orphan rule blocks doing so on the foreign type).
|
||||
pub(crate) struct GrpcTlsStream {
|
||||
inner: SslStream<ListenerIo>,
|
||||
}
|
||||
|
||||
impl GrpcTlsStream {
|
||||
pub(crate) fn new(inner: SslStream<ListenerIo>) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncRead for GrpcTlsStream {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut ReadBuf<'_>,
|
||||
) -> Poll<std::io::Result<()>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for GrpcTlsStream {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<std::io::Result<usize>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl Connected for GrpcTlsStream {
|
||||
type ConnectInfo = TcpConnectInfo;
|
||||
|
||||
fn connect_info(&self) -> TcpConnectInfo {
|
||||
self.inner.get_ref().connect_info()
|
||||
}
|
||||
}
|
||||
|
||||
/// Adapt the shared server listener into tonic's incoming stream shape.
|
||||
pub(crate) fn incoming(listener: Listener) -> impl Stream<Item = std::io::Result<ListenerIo>> {
|
||||
stream::unfold(listener, |mut listener| async move {
|
||||
let (io, _) = axum::serve::Listener::accept(&mut listener).await;
|
||||
Some((Ok(io), listener))
|
||||
})
|
||||
}
|
||||
|
||||
/// Wrap the gRPC listener so each accepted connection completes a TLS handshake
|
||||
/// before tonic serves it.
|
||||
pub(crate) fn tls_incoming(
|
||||
listener: Listener,
|
||||
context: openssl::ssl::SslContext,
|
||||
handshake_timeout: std::time::Duration,
|
||||
) -> impl Stream<Item = std::io::Result<GrpcTlsStream>> {
|
||||
tls_listener::builder(context)
|
||||
.handshake_timeout(handshake_timeout)
|
||||
.listen(listener)
|
||||
.map(|res| {
|
||||
res.map(|(inner, _addr)| GrpcTlsStream::new(inner))
|
||||
.map_err(std::io::Error::other)
|
||||
})
|
||||
}
|
||||
|
||||
/// gRPC Generate service implementation backed by the shared application state.
|
||||
pub struct GenerateServiceImpl {
|
||||
state: Arc<AppState>,
|
||||
|
||||
@@ -24,14 +24,15 @@ use vllm_engine_core_client::protocol::{
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::{Prompt, TextBackend};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
use super::pb::generate_client::GenerateClient;
|
||||
use super::{GenerateServer, GenerateServiceImpl, incoming, pb, tls_incoming};
|
||||
use crate::listener::Listener;
|
||||
use super::{GenerateServer, GenerateServiceImpl, pb};
|
||||
use crate::listener::{Listener, MaybeTlsListener};
|
||||
use crate::state::AppState;
|
||||
use crate::tls;
|
||||
use crate::tls_tests::{TestCerts, server_tls};
|
||||
@@ -155,37 +156,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
|
||||
#[derive(Clone, Debug)]
|
||||
struct FakeTextBackend;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FakeTokenizer;
|
||||
|
||||
impl Tokenizer for FakeTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.bytes().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
Ok(
|
||||
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
|
||||
.into_owned(),
|
||||
)
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
token.bytes().next().map(u32::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl TextBackend for FakeTextBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FakeTokenizer)
|
||||
Arc::new(TestTokenizer::new())
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
@@ -287,7 +260,7 @@ async fn grpc_test_server(
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = incoming(Listener::Tcp(listener));
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.serve_with_incoming(incoming)
|
||||
@@ -318,7 +291,7 @@ async fn grpc_tls_test_server(
|
||||
let addr = listener.local_addr().expect("local addr").to_string();
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = tls_incoming(Listener::Tcp(listener), context, tls::TLS_HANDSHAKE_TIMEOUT);
|
||||
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
|
||||
TonicServer::builder()
|
||||
.add_service(svc)
|
||||
.serve_with_incoming(incoming)
|
||||
@@ -413,7 +386,7 @@ async fn grpc_server_with_keepalive(
|
||||
}
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let incoming = incoming(Listener::Tcp(listener));
|
||||
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
|
||||
builder
|
||||
.add_service(svc)
|
||||
.serve_with_incoming(incoming)
|
||||
|
||||
+17
-62
@@ -27,7 +27,6 @@ pub use config::{
|
||||
ApiServerOptions, Config, CoordinatorMode, CorsConfig, DEFAULT_KEEP_ALIVE_TIMEOUT,
|
||||
HttpListenerMode, TlsConfig,
|
||||
};
|
||||
use futures::FutureExt as _;
|
||||
use hyper::body::Incoming;
|
||||
use hyper::server::conn::http1;
|
||||
use hyper_util::rt::{TokioIo, TokioTimer};
|
||||
@@ -45,7 +44,7 @@ use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::TextLlm;
|
||||
|
||||
use crate::listener::Listener;
|
||||
use crate::listener::{Listener, MaybeTlsListener};
|
||||
use crate::routes::build_router;
|
||||
use crate::server_info::ServerInfoSnapshot;
|
||||
use crate::state::AppState;
|
||||
@@ -179,7 +178,7 @@ where
|
||||
let listener = Listener::bind(&config.listener_mode)
|
||||
.await
|
||||
.context("failed to bind listener for OpenAI server")?;
|
||||
let bind_address = listener.local_addr()?;
|
||||
let bind_address = listener.local_addr_display()?;
|
||||
let model = state.primary_model_name().to_owned();
|
||||
let app = extend_router(build_router(state.clone()));
|
||||
|
||||
@@ -252,7 +251,6 @@ where
|
||||
// silent client cannot hold the connection open.
|
||||
let keep_alive_timeout = config.keep_alive_timeout;
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: tls::TLS_HANDSHAKE_TIMEOUT,
|
||||
header_read: if keep_alive_timeout.is_zero() {
|
||||
DEFAULT_KEEP_ALIVE_TIMEOUT
|
||||
} else {
|
||||
@@ -266,9 +264,15 @@ where
|
||||
let server_shutdown = server_shutdown.clone();
|
||||
let force_shutdown = force_shutdown.clone();
|
||||
async move {
|
||||
let listener = match tls_config {
|
||||
Some(context) => MaybeTlsListener::tls(listener, context),
|
||||
None => MaybeTlsListener::plain(listener),
|
||||
};
|
||||
let server = serve_connections(listener, app, shutdown.cancelled_owned(), timeouts);
|
||||
|
||||
let result = tokio::select! {
|
||||
result = serve_listener(listener, tls_config, app, shutdown.cancelled_owned(), timeouts) => {
|
||||
result
|
||||
result = server => {
|
||||
result.context("HTTP server failed")
|
||||
}
|
||||
_ = force_shutdown.cancelled() => {
|
||||
warn!("HTTP graceful shutdown deadline elapsed; aborting server");
|
||||
@@ -292,18 +296,11 @@ where
|
||||
shutdown.cancelled().await;
|
||||
return Ok(());
|
||||
};
|
||||
// Box to unify the TLS and plaintext arms' different stream types.
|
||||
let server = match grpc_tls {
|
||||
Some(context) => {
|
||||
let incoming =
|
||||
grpc::tls_incoming(grpc_listener, context, tls::TLS_HANDSHAKE_TIMEOUT);
|
||||
svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
|
||||
}
|
||||
None => {
|
||||
let incoming = grpc::incoming(grpc_listener);
|
||||
svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()).boxed()
|
||||
}
|
||||
let incoming = match grpc_tls {
|
||||
Some(context) => MaybeTlsListener::tls(grpc_listener, context),
|
||||
None => MaybeTlsListener::plain(grpc_listener),
|
||||
};
|
||||
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());
|
||||
|
||||
let result = tokio::select! {
|
||||
result = server => {
|
||||
@@ -333,61 +330,19 @@ where
|
||||
/// Per-connection timeouts applied while serving HTTP/HTTPS.
|
||||
#[derive(Clone, Copy)]
|
||||
pub(crate) struct ConnectionTimeouts {
|
||||
/// Max time for a client to complete the TLS handshake (TLS path only).
|
||||
pub(crate) handshake: Duration,
|
||||
/// HTTP/1 header-read timeout (bounds idle keep-alive and the head read).
|
||||
pub(crate) header_read: Duration,
|
||||
/// Whether HTTP/1 keep-alive is enabled; `false` closes after each response.
|
||||
pub(crate) keep_alive_enabled: bool,
|
||||
}
|
||||
|
||||
/// Apply optional TLS termination and per-connection HTTP timeouts, then serve
|
||||
/// `app`. Shared by [`serve_with_router_extension`] and the TLS tests.
|
||||
async fn serve_listener(
|
||||
listener: Listener,
|
||||
tls: Option<openssl::ssl::SslContext>,
|
||||
app: Router,
|
||||
shutdown: impl Future<Output = ()> + Send + 'static,
|
||||
timeouts: ConnectionTimeouts,
|
||||
) -> Result<()> {
|
||||
match tls {
|
||||
Some(context) => {
|
||||
// tls-listener terminates TLS (handshake + timeout); serve_connections
|
||||
// owns the HTTP keep-alive/idle bound that axum::serve cannot express.
|
||||
// Failed handshakes (incl. timeouts) log at ERROR via tls-listener.
|
||||
let listener = tls_listener::builder(context)
|
||||
.handshake_timeout(timeouts.handshake)
|
||||
.listen(listener);
|
||||
serve_connections(
|
||||
listener,
|
||||
app,
|
||||
shutdown,
|
||||
timeouts.header_read,
|
||||
timeouts.keep_alive_enabled,
|
||||
)
|
||||
.await
|
||||
.context("HTTPS server failed")
|
||||
}
|
||||
None => serve_connections(
|
||||
listener,
|
||||
app,
|
||||
shutdown,
|
||||
timeouts.header_read,
|
||||
timeouts.keep_alive_enabled,
|
||||
)
|
||||
.await
|
||||
.context("HTTP server failed"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Serve `app` per connection (HTTP/1) with a keep-alive idle timeout and
|
||||
/// graceful drain. Hand-rolled on hyper because [`axum::serve()`] takes no config.
|
||||
async fn serve_connections<L>(
|
||||
mut listener: L,
|
||||
app: Router,
|
||||
shutdown: impl Future<Output = ()> + Send,
|
||||
header_read: Duration,
|
||||
keep_alive_enabled: bool,
|
||||
timeouts: ConnectionTimeouts,
|
||||
) -> Result<()>
|
||||
where
|
||||
L: axum::serve::Listener,
|
||||
@@ -404,8 +359,8 @@ where
|
||||
app.clone().map_request(|req: Request<Incoming>| req.map(Body::new)),
|
||||
);
|
||||
let mut builder = http1::Builder::new();
|
||||
builder.timer(TokioTimer::new()).header_read_timeout(header_read);
|
||||
if !keep_alive_enabled {
|
||||
builder.timer(TokioTimer::new()).header_read_timeout(timeouts.header_read);
|
||||
if !timeouts.keep_alive_enabled {
|
||||
builder.keep_alive(false);
|
||||
}
|
||||
let connection = builder.serve_connection(TokioIo::new(io), service);
|
||||
|
||||
@@ -12,13 +12,14 @@ use std::pin::Pin;
|
||||
use std::task::{Context, Poll, ready};
|
||||
|
||||
use auto_enums::enum_derive;
|
||||
use openssl::ssl::SslContext;
|
||||
use socket2::Socket;
|
||||
use tls_listener::{AsyncAccept, AsyncListener};
|
||||
use tokio::net::{TcpListener, TcpStream, UnixListener, UnixStream};
|
||||
use tonic::transport::server::{Connected, TcpConnectInfo};
|
||||
use tracing::trace;
|
||||
|
||||
use crate::HttpListenerMode;
|
||||
use crate::{HttpListenerMode, tls};
|
||||
|
||||
/// Runtime listener type used by the OpenAI-compatible HTTP or gRPC server,
|
||||
/// which is either a TCP listener or a Unix-domain listener.
|
||||
@@ -61,7 +62,7 @@ impl Listener {
|
||||
|
||||
/// Return a log-friendly local address string for either TCP or Unix
|
||||
/// sockets.
|
||||
pub fn local_addr(&self) -> Result<String> {
|
||||
pub fn local_addr_display(&self) -> Result<String> {
|
||||
match self {
|
||||
Self::Tcp(listener) => Ok(listener.local_addr()?.to_string()),
|
||||
Self::Unix(listener) => Ok(match listener.local_addr()?.as_pathname() {
|
||||
@@ -92,7 +93,7 @@ impl Listener {
|
||||
}
|
||||
}
|
||||
|
||||
fn listener_addr(&self) -> Result<ListenerAddr> {
|
||||
fn local_addr(&self) -> Result<ListenerAddr> {
|
||||
match self {
|
||||
Self::Tcp(listener) => listener.local_addr().map(ListenerAddr::Tcp),
|
||||
Self::Unix(listener) => listener.local_addr().map(ListenerAddr::Unix),
|
||||
@@ -100,6 +101,7 @@ impl Listener {
|
||||
}
|
||||
}
|
||||
|
||||
/// Allow the unified listener to plug directly into tonic's gRPC server.
|
||||
impl Connected for ListenerIo {
|
||||
type ConnectInfo = TcpConnectInfo;
|
||||
|
||||
@@ -144,7 +146,7 @@ impl axum::serve::Listener for Listener {
|
||||
}
|
||||
|
||||
fn local_addr(&self) -> Result<Self::Addr> {
|
||||
self.listener_addr()
|
||||
self.local_addr()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -173,10 +175,98 @@ impl AsyncAccept for Listener {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncListener for Listener {
|
||||
fn local_addr(&self) -> Result<Self::Address> {
|
||||
self.listener_addr()
|
||||
self.local_addr()
|
||||
}
|
||||
}
|
||||
|
||||
/// A listener that may be either a plain TCP/UDS listener or a TLS listener over it.
|
||||
pub enum MaybeTlsListener {
|
||||
Plain(Listener),
|
||||
Tls(tls_listener::TlsListener<Listener, SslContext>),
|
||||
}
|
||||
|
||||
impl MaybeTlsListener {
|
||||
/// Create a plain listener without TLS.
|
||||
pub fn plain(listener: Listener) -> Self {
|
||||
Self::Plain(listener)
|
||||
}
|
||||
|
||||
/// Create a TLS listener over the given plain listener.
|
||||
pub fn tls(listener: Listener, context: SslContext) -> Self {
|
||||
Self::Tls(
|
||||
tls_listener::builder(context)
|
||||
.handshake_timeout(tls::TLS_HANDSHAKE_TIMEOUT)
|
||||
.listen(listener),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Listener I/O type that may be either a plain TCP/UDS stream or a TLS stream over it.
|
||||
#[derive(Debug)]
|
||||
#[enum_derive(tokio1::AsyncRead, tokio1::AsyncWrite)]
|
||||
pub enum MaybeTlsStream {
|
||||
Plain(ListenerIo),
|
||||
Tls(tokio_openssl::SslStream<ListenerIo>),
|
||||
}
|
||||
|
||||
/// Allow the maybe-TLS listener to plug directly into `axum::serve(...)`.
|
||||
impl axum::serve::Listener for MaybeTlsListener {
|
||||
type Addr = ListenerAddr;
|
||||
type Io = MaybeTlsStream;
|
||||
|
||||
async fn accept(&mut self) -> (Self::Io, Self::Addr) {
|
||||
match self {
|
||||
Self::Plain(listener) => {
|
||||
let (io, addr) = axum::serve::Listener::accept(listener).await;
|
||||
(MaybeTlsStream::Plain(io), addr)
|
||||
}
|
||||
Self::Tls(tls_listener) => {
|
||||
let (io, addr) = axum::serve::Listener::accept(tls_listener).await;
|
||||
(MaybeTlsStream::Tls(io), addr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn local_addr(&self) -> tokio::io::Result<Self::Addr> {
|
||||
match self {
|
||||
Self::Plain(listener) => listener.local_addr(),
|
||||
Self::Tls(tls_listener) => tls_listener.local_addr(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Allow the maybe-TLS listener to plug directly into tonic's gRPC server.
|
||||
impl Connected for MaybeTlsStream {
|
||||
type ConnectInfo = TcpConnectInfo;
|
||||
|
||||
fn connect_info(&self) -> TcpConnectInfo {
|
||||
match self {
|
||||
Self::Plain(stream) => stream.connect_info(),
|
||||
Self::Tls(stream) => stream.get_ref().connect_info(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Allow the maybe-TLS listener to be adaptable to tonic's incoming stream shape.
|
||||
impl futures::Stream for MaybeTlsListener {
|
||||
type Item = std::io::Result<MaybeTlsStream>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
match self.get_mut() {
|
||||
Self::Plain(listener) => {
|
||||
let listener = Pin::new(listener);
|
||||
let (io, _) = ready!(listener.poll_accept(cx))?;
|
||||
Poll::Ready(Some(Ok(MaybeTlsStream::Plain(io))))
|
||||
}
|
||||
Self::Tls(tls_listener) => {
|
||||
let tls_listener = Pin::new(tls_listener);
|
||||
let (io, _) =
|
||||
ready!(tls_listener.poll_accept(cx)).map_err(std::io::Error::other)?;
|
||||
Poll::Ready(Some(Ok(MaybeTlsStream::Tls(io))))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -24,8 +24,9 @@ use vllm_engine_core_client::protocol::{
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::{Prompt, TextBackend};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
@@ -151,37 +152,9 @@ fn test_llm(client: EngineCoreClient) -> Llm {
|
||||
#[derive(Clone, Debug)]
|
||||
struct FakeChatBackend;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FakeChatTokenizer;
|
||||
|
||||
impl Tokenizer for FakeChatTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.bytes().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
Ok(
|
||||
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
|
||||
.into_owned(),
|
||||
)
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
token.bytes().next().map(u32::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl TextBackend for FakeChatBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FakeChatTokenizer)
|
||||
Arc::new(TestTokenizer::new())
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
|
||||
@@ -194,11 +194,12 @@ mod tests {
|
||||
use axum::http::HeaderMap;
|
||||
use serde_json::json;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::prepare_completion_request;
|
||||
use crate::lora::LoraModelResolution;
|
||||
use crate::routes::openai::completions::types::CompletionRequest;
|
||||
use crate::routes::openai::utils::types::Normalizable;
|
||||
use crate::utils::{ResolvedRequestContext, resolve_request_context};
|
||||
|
||||
fn request_context(headers: &HeaderMap, request_id: Option<&str>) -> ResolvedRequestContext {
|
||||
@@ -212,32 +213,8 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TestTokenizer;
|
||||
|
||||
impl Tokenizer for TestTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
Ok(text.bytes().map(u32::from).collect())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
Ok(
|
||||
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
|
||||
.into_owned(),
|
||||
)
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
fn test_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
}
|
||||
|
||||
fn base_request_json() -> serde_json::Value {
|
||||
@@ -273,6 +250,28 @@ mod tests {
|
||||
assert!(request.ignore_eos);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_coerces_null_max_tokens_to_default() {
|
||||
// An absent `max_tokens` already gets the serde default.
|
||||
let absent: CompletionRequest =
|
||||
serde_json::from_value(base_request_json()).expect("parse request");
|
||||
assert_eq!(absent.max_tokens, Some(16));
|
||||
|
||||
// An explicit `null` deserializes to `None`, bypassing the default;
|
||||
// `normalize` must coerce it back to match Python vLLM.
|
||||
let mut request: CompletionRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hello",
|
||||
"stream": true,
|
||||
"max_tokens": null
|
||||
}))
|
||||
.expect("parse request");
|
||||
assert_eq!(request.max_tokens, None);
|
||||
|
||||
request.normalize();
|
||||
assert_eq!(request.max_tokens, Some(16));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_completion_request_maps_sampling_fields() {
|
||||
let request: CompletionRequest = serde_json::from_value(json!({
|
||||
@@ -297,7 +296,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -340,7 +339,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare")
|
||||
.text_request
|
||||
@@ -374,7 +373,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -399,7 +398,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -422,7 +421,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -446,7 +445,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -471,7 +470,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -498,7 +497,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -523,7 +522,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
@@ -549,7 +548,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
|
||||
@@ -574,7 +573,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
request_context(&headers, None),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
assert_eq!(prepared.text_request.data_parallel_rank, Some(3));
|
||||
@@ -593,7 +592,7 @@ mod tests {
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
&TestTokenizer,
|
||||
&test_tokenizer(),
|
||||
)
|
||||
.expect("prepare");
|
||||
assert_eq!(prepared.text_request.data_parallel_rank, None);
|
||||
|
||||
@@ -179,7 +179,17 @@ pub struct CompletionRequest {
|
||||
pub other: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl Normalizable for CompletionRequest {}
|
||||
impl Normalizable for CompletionRequest {
|
||||
/// Normalize the request by applying defaults.
|
||||
fn normalize(&mut self) {
|
||||
// An explicit `"max_tokens": null` deserializes to `None`, bypassing the
|
||||
// serde field default. Coerce it back to the default so it behaves like
|
||||
// an absent field, matching Python vLLM's `normalize_null_max_tokens`.
|
||||
if self.max_tokens.is_none() {
|
||||
self.max_tokens = default_completion_max_tokens();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `CompletionResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
|
||||
@@ -40,8 +40,9 @@ use vllm_engine_core_client::{
|
||||
};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_metrics::METRICS;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::{Prompt, TextBackend};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use zeromq::prelude::{SocketRecv, SocketSend};
|
||||
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
|
||||
|
||||
@@ -417,82 +418,20 @@ struct FakeChatBackend {
|
||||
}
|
||||
|
||||
/// Synthetic BOS id used when `add_special_tokens` is true in tests.
|
||||
const FAKE_BOS_TOKEN_ID: u32 = 1;
|
||||
const FAKE_BOS_TOKEN_ID: u32 = 256;
|
||||
const UNKNOWN_DECODE_TOKEN_ID: u32 = 10_000;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FakeChatTokenizer;
|
||||
|
||||
impl Tokenizer for FakeChatTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
let mut token_ids = Vec::new();
|
||||
if add_special_tokens {
|
||||
token_ids.push(FAKE_BOS_TOKEN_ID);
|
||||
}
|
||||
let mut rest = text;
|
||||
while !rest.is_empty() {
|
||||
if let Some(stripped) = rest.strip_prefix("<image>") {
|
||||
token_ids.push(999);
|
||||
rest = stripped;
|
||||
continue;
|
||||
}
|
||||
if let Some(stripped) = rest.strip_prefix("<|image_pad|>") {
|
||||
token_ids.push(151655);
|
||||
rest = stripped;
|
||||
continue;
|
||||
}
|
||||
|
||||
let ch = rest.chars().next().expect("rest is not empty");
|
||||
let mut buf = [0; 4];
|
||||
token_ids.extend(ch.encode_utf8(&mut buf).bytes().map(u32::from));
|
||||
rest = &rest[ch.len_utf8()..];
|
||||
}
|
||||
Ok(token_ids)
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
Ok(
|
||||
String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())
|
||||
.into_owned(),
|
||||
)
|
||||
}
|
||||
|
||||
fn token_to_id(&self, token: &str) -> Option<u32> {
|
||||
match token {
|
||||
"<image>" => Some(999),
|
||||
"<|image_pad|>" => Some(151655),
|
||||
"<think>" => Some(0xF001),
|
||||
"</think>" => Some(0xF002),
|
||||
"<|START_THINKING|>" => Some(0xF003),
|
||||
"<|END_THINKING|>" => Some(0xF004),
|
||||
"◁think▷" => Some(0xF005),
|
||||
"◁/think▷" => Some(0xF006),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn id_to_token(&self, id: u32) -> Option<String> {
|
||||
match id {
|
||||
FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
|
||||
999 => Some("<image>".to_string()),
|
||||
151655 => Some("<|image_pad|>".to_string()),
|
||||
0xF001 => Some("<think>".to_string()),
|
||||
0xF002 => Some("</think>".to_string()),
|
||||
0xF003 => Some("<|START_THINKING|>".to_string()),
|
||||
0xF004 => Some("<|END_THINKING|>".to_string()),
|
||||
0xF005 => Some("◁think▷".to_string()),
|
||||
0xF006 => Some("◁/think▷".to_string()),
|
||||
id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
fn fake_chat_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_bos_token("<bos>", FAKE_BOS_TOKEN_ID)
|
||||
.with_regular_token("<image>", 999)
|
||||
.with_regular_token("<|image_pad|>", 151655)
|
||||
.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 FakeChatBackend {
|
||||
@@ -530,7 +469,7 @@ impl fmt::Debug for FakeChatBackend {
|
||||
|
||||
impl TextBackend for FakeChatBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FakeChatTokenizer)
|
||||
Arc::new(fake_chat_tokenizer())
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
@@ -630,7 +569,7 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
|
||||
Some("qwen2_vl".to_string()),
|
||||
Some(&config_path),
|
||||
None,
|
||||
Arc::new(FakeChatTokenizer),
|
||||
Arc::new(fake_chat_tokenizer()),
|
||||
)
|
||||
.expect("load multimodal info")
|
||||
.expect("qwen multimodal info is registered");
|
||||
@@ -638,70 +577,6 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
|
||||
info
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct FailingDecodeChatBackend;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FailingDecodeTokenizer;
|
||||
|
||||
impl Tokenizer for FailingDecodeTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
text: &str,
|
||||
add_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<Vec<u32>> {
|
||||
FakeChatTokenizer.encode(text, add_special_tokens)
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
skip_special_tokens: bool,
|
||||
) -> vllm_text::tokenizer::Result<String> {
|
||||
if token_ids.contains(&(b'i' as u32)) {
|
||||
return Err(vllm_text::tokenizer::TokenizerError(
|
||||
"forced decode failure for streaming test".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 FailingDecodeChatBackend {
|
||||
fn tokenizer(&self) -> DynTokenizer {
|
||||
Arc::new(FailingDecodeTokenizer)
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
"test-model"
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatBackend for FailingDecodeChatBackend {
|
||||
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 FailingDecodeChatBackend {
|
||||
fn render(&self, request: &ChatRequest) -> vllm_chat::Result<vllm_chat::RenderedPrompt> {
|
||||
FakeChatBackend::new().render(request)
|
||||
}
|
||||
}
|
||||
|
||||
async fn test_models_with_engine_outputs_and_backend_inner(
|
||||
engine_id: impl Into<EngineId>,
|
||||
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
|
||||
@@ -2755,8 +2630,8 @@ async fn load_endpoint_resets_when_stream_response_is_dropped() {
|
||||
#[serial]
|
||||
async fn stream_error_is_returned_as_openai_error_sse() {
|
||||
let (app, engine_task) = test_app_with_backend_and_stream_output_specs(
|
||||
Arc::new(FailingDecodeChatBackend),
|
||||
default_stream_output_specs(),
|
||||
Arc::new(FakeChatBackend::new()),
|
||||
vec![(vec![UNKNOWN_DECODE_TOKEN_ID], None)],
|
||||
)
|
||||
.await;
|
||||
let response = app
|
||||
@@ -2789,7 +2664,9 @@ async fn stream_error_is_returned_as_openai_error_sse() {
|
||||
assert!(text.contains("\"role\":\"assistant\""), "{text}");
|
||||
assert!(text.contains("\"type\":\"server_error\""), "{text}");
|
||||
assert!(
|
||||
text.contains("forced decode failure for streaming test"),
|
||||
text.contains(&format!(
|
||||
"test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}"
|
||||
)),
|
||||
"{text}"
|
||||
);
|
||||
assert!(!text.contains("\"usage\":"), "{text}");
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
//! TLS tests: `build_server_config` unit checks plus end-to-end OpenSSL handshakes
|
||||
//! through the production `serve_listener` path, with a trivial router since TLS
|
||||
//! through the production listener/connection path, with a trivial router since TLS
|
||||
//! terminates below the app.
|
||||
|
||||
use std::pin::Pin;
|
||||
@@ -23,8 +23,8 @@ use tokio_openssl::SslStream;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::config::{HttpListenerMode, TlsConfig};
|
||||
use crate::listener::Listener;
|
||||
use crate::{ConnectionTimeouts, serve_listener, tls};
|
||||
use crate::listener::{Listener, MaybeTlsListener};
|
||||
use crate::{ConnectionTimeouts, serve_connections, tls};
|
||||
|
||||
// ============================================================================
|
||||
// Test infrastructure
|
||||
@@ -251,7 +251,6 @@ fn build_tls(certs: &TestCerts, cert: &str, key: Option<&str>) -> TlsConfig {
|
||||
|
||||
/// Generous per-connection timeouts that never fire during the fast tests.
|
||||
const TEST_TIMEOUTS: ConnectionTimeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_secs(5),
|
||||
keep_alive_enabled: true,
|
||||
};
|
||||
@@ -261,9 +260,8 @@ async fn spawn_server(tls_config: Option<TlsConfig>) -> (String, CancellationTok
|
||||
}
|
||||
|
||||
/// Bind an ephemeral listener and serve a trivial router via the production
|
||||
/// `serve_listener`, optionally with TLS. The listener is bound (and thus
|
||||
/// accepting into the backlog) before returning, so a client may connect
|
||||
/// immediately without a sleep.
|
||||
/// listener/connection path. The listener is bound (and thus accepting into the
|
||||
/// backlog) before returning, so a client may connect immediately without a sleep.
|
||||
async fn spawn_server_with_timeouts(
|
||||
tls_config: Option<TlsConfig>,
|
||||
timeouts: ConnectionTimeouts,
|
||||
@@ -274,7 +272,7 @@ async fn spawn_server_with_timeouts(
|
||||
})
|
||||
.await
|
||||
.expect("bind listener");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
let addr = listener.local_addr_display().expect("local addr");
|
||||
|
||||
let server_config =
|
||||
tls_config.map(|cfg| tls::build_server_config(&cfg).expect("build server config"));
|
||||
@@ -282,14 +280,11 @@ async fn spawn_server_with_timeouts(
|
||||
let shutdown = CancellationToken::new();
|
||||
let server_shutdown = shutdown.clone();
|
||||
tokio::spawn(async move {
|
||||
let _ = serve_listener(
|
||||
listener,
|
||||
server_config,
|
||||
app,
|
||||
server_shutdown.cancelled_owned(),
|
||||
timeouts,
|
||||
)
|
||||
.await;
|
||||
let listener = match server_config {
|
||||
Some(context) => MaybeTlsListener::tls(listener, context),
|
||||
None => MaybeTlsListener::plain(listener),
|
||||
};
|
||||
let _ = serve_connections(listener, app, server_shutdown.cancelled_owned(), timeouts).await;
|
||||
});
|
||||
(addr, shutdown)
|
||||
}
|
||||
@@ -521,20 +516,19 @@ async fn plain_http_serves_when_tls_is_disabled() {
|
||||
shutdown.cancel();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn tls_handshake_timeout_drops_silent_client() {
|
||||
// Silent client (no ClientHello) must be dropped at the handshake deadline.
|
||||
let certs = TestCerts::generate();
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_millis(150),
|
||||
header_read: Duration::from_secs(5),
|
||||
keep_alive_enabled: true,
|
||||
};
|
||||
let (addr, shutdown) = spawn_server_with_timeouts(Some(server_tls(&certs, 0)), timeouts).await;
|
||||
let (addr, shutdown) = spawn_server(Some(server_tls(&certs, 0))).await;
|
||||
|
||||
let mut tcp = TcpStream::connect(&addr).await.expect("connect");
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::advance(tls::TLS_HANDSHAKE_TIMEOUT + Duration::from_millis(1)).await;
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
let mut buf = [0u8; 1];
|
||||
let read = tokio::time::timeout(Duration::from_secs(5), tcp.read(&mut buf)).await;
|
||||
let read = tokio::time::timeout(Duration::from_secs(1), tcp.read(&mut buf)).await;
|
||||
assert!(
|
||||
matches!(read, Ok(Ok(0)) | Ok(Err(_))),
|
||||
"server must drop a stalled TLS handshake (expected close, got {read:?})"
|
||||
@@ -546,7 +540,6 @@ async fn tls_handshake_timeout_drops_silent_client() {
|
||||
async fn keep_alive_timeout_closes_idle_connection() {
|
||||
// Idle keep-alive connection must be closed at the deadline.
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_millis(150),
|
||||
keep_alive_enabled: true,
|
||||
};
|
||||
@@ -582,7 +575,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
|
||||
// still fires through tls-listener's post-handshake SslStream, not just plaintext.
|
||||
let certs = TestCerts::generate();
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_millis(150),
|
||||
keep_alive_enabled: true,
|
||||
};
|
||||
@@ -618,7 +610,6 @@ async fn keep_alive_timeout_closes_idle_tls_connection() {
|
||||
async fn idle_timeout_closes_silent_client() {
|
||||
// Silent client closed by the header-read timeout (http1-only arms it from byte 0).
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_millis(150),
|
||||
keep_alive_enabled: true,
|
||||
};
|
||||
@@ -638,7 +629,6 @@ async fn idle_timeout_closes_silent_client() {
|
||||
async fn keep_alive_zero_disables_keep_alive() {
|
||||
// 0 disables keep-alive (serve, then close), like uvicorn's timeout_keep_alive=0.
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_secs(5),
|
||||
keep_alive_enabled: false,
|
||||
};
|
||||
@@ -671,7 +661,6 @@ async fn disabled_keep_alive_still_closes_silent_client() {
|
||||
// Even with keep-alive off, the head read stays bounded, so a silent client
|
||||
// is dropped rather than held open.
|
||||
let timeouts = ConnectionTimeouts {
|
||||
handshake: Duration::from_secs(60),
|
||||
header_read: Duration::from_millis(150),
|
||||
keep_alive_enabled: false,
|
||||
};
|
||||
|
||||
@@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true
|
||||
easy-ext.workspace = true
|
||||
enum-as-inner.workspace = true
|
||||
futures.workspace = true
|
||||
hf-hub.workspace = true
|
||||
itertools.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
@@ -22,6 +21,7 @@ tracing.workspace = true
|
||||
trait-set.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-llm.workspace = true
|
||||
vllm-model-files.workspace = true
|
||||
vllm-tokenizer.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
@@ -31,6 +31,7 @@ serial_test.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio.workspace = true
|
||||
vllm-llm = { workspace = true, features = ["test-util"] }
|
||||
vllm-tokenizer = { workspace = true, features = ["test-utils"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -1,24 +1,18 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_model_files::read_json_file;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Minimal subset of `tokenizer_config.json` needed by chat/EOS handling.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct HfTokenizerConfig {
|
||||
pub struct TokenizerConfig {
|
||||
#[serde(flatten)]
|
||||
pub special_tokens: HfSpecialTokens,
|
||||
pub special_tokens: SpecialTokens,
|
||||
pub chat_template: Option<String>,
|
||||
/// The `tokenizer_class` field from HuggingFace tokenizer configs. Some
|
||||
/// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value
|
||||
/// containing "Tiktoken" which can be used as a hint for backend
|
||||
/// selection.
|
||||
pub tokenizer_class: Option<String>,
|
||||
}
|
||||
|
||||
/// Hugging Face named special tokens may be serialized as a string or an
|
||||
@@ -61,14 +55,14 @@ impl NamedSpecialToken {
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct HfSpecialTokens {
|
||||
pub struct SpecialTokens {
|
||||
pub bos_token: Option<NamedSpecialToken>,
|
||||
pub eos_token: Option<NamedSpecialToken>,
|
||||
pub unk_token: Option<NamedSpecialToken>,
|
||||
pub pad_token: Option<NamedSpecialToken>,
|
||||
}
|
||||
|
||||
impl HfSpecialTokens {
|
||||
impl SpecialTokens {
|
||||
/// Returns true if we don't discover any special tokens in the config.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.bos_token.is_none()
|
||||
@@ -234,42 +228,19 @@ impl ModelConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the tokenizer-side EOS metadata if a config file is present.
|
||||
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<HfTokenizerConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
|
||||
/// Load the generation-side EOS metadata if a config file is present.
|
||||
pub(super) fn load_generation_config(path: Option<&Path>) -> Result<GenerationConfig> {
|
||||
read_json_file(path)
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
/// Load the tokenizer-side EOS metadata if a config file is present.
|
||||
pub fn load_tokenizer_config(path: Option<&Path>) -> Result<TokenizerConfig> {
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
/// Load the model-side config (`config.json`) if present.
|
||||
pub fn load_model_config(path: Option<&Path>) -> Result<ModelConfig> {
|
||||
read_json_file(path)
|
||||
}
|
||||
|
||||
fn read_json_file<T>(path: Option<&Path>) -> Result<T>
|
||||
where
|
||||
T: for<'de> Deserialize<'de> + Default,
|
||||
{
|
||||
let Some(path) = path else {
|
||||
return Ok(T::default());
|
||||
};
|
||||
let content = fs::read_to_string(path).map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
"failed to read {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})?;
|
||||
serde_json::from_str(&content).map_err(|error| {
|
||||
Error::Tokenizer(format!(
|
||||
"failed to parse {}: {}",
|
||||
path.display(),
|
||||
error.as_report()
|
||||
))
|
||||
})
|
||||
Ok(read_json_file(path)?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
mod config;
|
||||
mod model_files;
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::Arc;
|
||||
@@ -9,12 +8,12 @@ use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, Tiktok
|
||||
|
||||
use self::config::{GenerationConfig, load_generation_config};
|
||||
pub use self::config::{
|
||||
HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config,
|
||||
ModelConfig, NamedSpecialToken, SpecialTokens, TokenizerConfig, load_model_config,
|
||||
load_tokenizer_config,
|
||||
};
|
||||
pub use self::model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
use crate::backend::{SamplingHints, TextBackend};
|
||||
use crate::error::Result;
|
||||
pub use vllm_model_files::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
fn load_tokenizer(tokenizer: &TokenizerSource) -> Result<DynTokenizer> {
|
||||
match tokenizer {
|
||||
@@ -125,3 +124,67 @@ impl TextBackend for HfTextBackend {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
|
||||
|
||||
use super::{ResolvedModelFiles, TokenizerSource};
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "too slow for CI and requires network access to Hugging Face"]
|
||||
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
|
||||
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
|
||||
.await
|
||||
.expect("resolve real Kimi K2.5 model files");
|
||||
|
||||
let tokenizer_path = match &files.tokenizer {
|
||||
TokenizerSource::Tiktoken(path) => path.clone(),
|
||||
other => panic!("expected tiktoken tokenizer source, got {other:?}"),
|
||||
};
|
||||
|
||||
for backend in [
|
||||
TiktokenTokenizer::new_riptoken(&tokenizer_path).expect("load riptoken backend"),
|
||||
TiktokenTokenizer::new_tiktoken_rs(&tokenizer_path).expect("load tiktoken-rs backend"),
|
||||
] {
|
||||
let think_id = backend.token_to_id("<think>").expect("resolve <think>");
|
||||
let end_think_id = backend.token_to_id("</think>").expect("resolve </think>");
|
||||
let tool_section_id = backend
|
||||
.token_to_id("<|tool_calls_section_begin|>")
|
||||
.expect("resolve tool call section marker");
|
||||
let contraction_heavy_text =
|
||||
"I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.";
|
||||
let contraction_heavy_ids = backend.encode(contraction_heavy_text, false).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(think_id, end_think_id, tool_section_id),
|
||||
(163606, 163607, 163595)
|
||||
);
|
||||
assert_eq!(backend.decode(&[think_id], true).unwrap(), "<think>");
|
||||
assert_eq!(backend.decode(&[end_think_id], true).unwrap(), "</think>");
|
||||
assert_eq!(
|
||||
backend.decode(&[tool_section_id], true).unwrap(),
|
||||
"<|tool_calls_section_begin|>"
|
||||
);
|
||||
|
||||
// This demonstrates that we're using Kimi's custom BPE pattern.
|
||||
// With CL100K this will be 23 tokens instead.
|
||||
assert_eq!(
|
||||
contraction_heavy_ids,
|
||||
vec![
|
||||
17172, 3287, 4643, 8201, 11, 996, 374, 8971, 3637, 20020, 8173, 473, 4643,
|
||||
1573, 56229, 13922, 13,
|
||||
]
|
||||
);
|
||||
assert_eq!(contraction_heavy_ids.len(), 17);
|
||||
assert_eq!(
|
||||
backend.decode(&contraction_heavy_ids, false).unwrap(),
|
||||
contraction_heavy_text
|
||||
);
|
||||
|
||||
// Special-looking text that is not actually registered should fail gracefully.
|
||||
assert_eq!(backend.token_to_id("◁think▷"), None);
|
||||
assert_eq!(backend.token_to_id("<|definitely_not_registered|>"), None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use thiserror::Error;
|
||||
use vllm_engine_core_client::Error as EngineCoreError;
|
||||
use vllm_llm::Error as LlmError;
|
||||
use vllm_model_files::Error as ModelFilesError;
|
||||
|
||||
pub use crate::lower::logprobs::LogprobsError;
|
||||
pub use crate::lower::token_ids::TokenIdsError;
|
||||
@@ -20,6 +21,8 @@ pub enum Error {
|
||||
Logprobs(#[from] LogprobsError),
|
||||
#[error(transparent)]
|
||||
TokenIds(#[from] TokenIdsError),
|
||||
#[error(transparent)]
|
||||
ModelFiles(#[from] ModelFilesError),
|
||||
#[error(
|
||||
"`min_tokens` must be less than or equal to `max_tokens`, \
|
||||
got min_tokens={min_tokens}, max_tokens={max_tokens}"
|
||||
|
||||
@@ -275,6 +275,7 @@ mod tests {
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
|
||||
use serial_test::file_serial;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::*;
|
||||
use crate::backend::hf::HfTextBackend;
|
||||
@@ -282,60 +283,8 @@ mod tests {
|
||||
use crate::error::{LogprobsError, TokenIdsError};
|
||||
use crate::request::{Prompt, TextRequest};
|
||||
|
||||
/// Stub tokenizer that returns empty token IDs — sufficient for tests that
|
||||
/// don't exercise bad-words tokenization.
|
||||
struct StubTokenizer;
|
||||
|
||||
impl Tokenizer for StubTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
_text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
_token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn stub_tokenizer() -> StubTokenizer {
|
||||
StubTokenizer
|
||||
}
|
||||
|
||||
struct FixedTokenizer {
|
||||
token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
impl Tokenizer for FixedTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
_text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
Ok(self.token_ids.clone())
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
_token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
fn stub_tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
}
|
||||
|
||||
fn sample_request() -> TextRequest {
|
||||
@@ -952,9 +901,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
|
||||
let tokenizer = FixedTokenizer {
|
||||
token_ids: vec![1999, 2000],
|
||||
};
|
||||
let tokenizer = TestTokenizer::new().with_regular_token("blocked", 2000);
|
||||
let error = lower_sampling_params(
|
||||
SamplingParams {
|
||||
bad_words: Some(vec!["blocked".to_string()]),
|
||||
|
||||
@@ -323,37 +323,11 @@ mod tests {
|
||||
use futures::{Stream, stream};
|
||||
use vllm_engine_core_client::AbortCause;
|
||||
use vllm_llm::GenerateOutput;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::*;
|
||||
use crate::output::TextOutputStreamExt as _;
|
||||
|
||||
/// Backend that treats each token ID as a raw byte, producing lossy UTF-8.
|
||||
struct ByteTokenizer;
|
||||
|
||||
impl Tokenizer for ByteTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
_text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
|
||||
Ok(String::from_utf8_lossy(&bytes).into_owned())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper: run `decoded_text_event_stream` to completion and return the
|
||||
/// collected output.
|
||||
async fn run_to_completion(
|
||||
@@ -366,7 +340,7 @@ mod tests {
|
||||
token_ids,
|
||||
Some(FinishReason::Length),
|
||||
))]);
|
||||
let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
|
||||
let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
|
||||
decoded_text_event_stream("test".into(), tokenizer, raw_stream, decode_options, false)
|
||||
.collect_output()
|
||||
.await
|
||||
@@ -419,7 +393,7 @@ mod tests {
|
||||
))),
|
||||
dropped_cause: Arc::clone(&dropped_cause),
|
||||
};
|
||||
let tokenizer: DynTokenizer = Arc::new(ByteTokenizer);
|
||||
let tokenizer: DynTokenizer = Arc::new(TestTokenizer::new());
|
||||
|
||||
let output = decoded_text_event_stream(
|
||||
"test".into(),
|
||||
|
||||
@@ -129,40 +129,13 @@ fn decode_position_logprobs<T: Tokenizer + ?Sized>(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use vllm_llm::{Logprobs, PositionLogprobs, TokenLogprob};
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ByteTokenizer;
|
||||
|
||||
impl vllm_tokenizer::Tokenizer for ByteTokenizer {
|
||||
fn encode(
|
||||
&self,
|
||||
_text: &str,
|
||||
_add_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<Vec<u32>> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
token_ids: &[u32],
|
||||
_skip_special_tokens: bool,
|
||||
) -> vllm_tokenizer::Result<String> {
|
||||
Ok(String::from_utf8_lossy(
|
||||
&token_ids.iter().map(|token_id| *token_id as u8).collect::<Vec<_>>(),
|
||||
)
|
||||
.into_owned())
|
||||
}
|
||||
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_logprobs_decodes_every_candidate_token() {
|
||||
let tokenizer = ByteTokenizer;
|
||||
let tokenizer = TestTokenizer::new();
|
||||
let logprobs = Logprobs {
|
||||
positions: vec![PositionLogprobs {
|
||||
entries: vec![
|
||||
@@ -205,7 +178,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn decode_prompt_logprobs_separates_first_prompt_token() {
|
||||
let tokenizer = ByteTokenizer;
|
||||
let tokenizer = TestTokenizer::new();
|
||||
let logprobs = Logprobs {
|
||||
positions: vec![PositionLogprobs {
|
||||
entries: vec![TokenLogprob {
|
||||
|
||||
@@ -4,6 +4,9 @@ version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[features]
|
||||
test-utils = []
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
fastokens.workspace = true
|
||||
|
||||
@@ -170,6 +170,10 @@ mod tests {
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -248,6 +252,10 @@ mod tests {
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -320,6 +328,10 @@ mod tests {
|
||||
fn token_to_id(&self, _token: &str) -> Option<u32> {
|
||||
unreachable!()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
unreachable!()
|
||||
}
|
||||
}
|
||||
|
||||
/// Without the char-boundary fix, this panics slicing mid-emoji.
|
||||
|
||||
@@ -8,6 +8,8 @@ mod error;
|
||||
mod hf;
|
||||
mod incremental;
|
||||
mod tekken;
|
||||
#[cfg(any(test, feature = "test-utils"))]
|
||||
pub mod test_utils;
|
||||
mod tiktoken;
|
||||
|
||||
pub use error::{Result, TokenizerError};
|
||||
@@ -28,11 +30,7 @@ pub trait Tokenizer: Send + Sync {
|
||||
fn token_to_id(&self, token: &str) -> Option<u32>;
|
||||
|
||||
/// Convert one token ID into the tokenizer's raw token string.
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
// TODO: remove default impl and require this to be implemented by all
|
||||
// tokenizers
|
||||
None
|
||||
}
|
||||
fn id_to_token(&self, id: u32) -> Option<String>;
|
||||
|
||||
/// Return the vocabulary size. Backends that cannot report it fall back to
|
||||
/// `usize::MAX`, an effectively unbounded value used only by test stubs.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user