forked from Karylab-cklius/vllm
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
95dcefaaa5 |
@@ -81,9 +81,7 @@ steps:
|
||||
'cd tests &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
set -o pipefail &&
|
||||
pytest -v -s lora/test_punica_ops.py::test_kernels &&
|
||||
pytest -v -s lora/test_punica_ops.py::test_kernels_hidden_size &&
|
||||
pytest -v -s lora/test_punica_ops.py::test_add_lora_fused_moe_early_exit'
|
||||
pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'
|
||||
|
||||
- label: LoRA Punica FP8/XPU Ops
|
||||
timeout_in_minutes: 45
|
||||
|
||||
@@ -103,31 +103,6 @@ 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 &&
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
|
||||
pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
|
||||
@@ -47,7 +47,8 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
|
||||
@@ -70,7 +71,8 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
|
||||
pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
|
||||
|
||||
@@ -94,7 +96,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av &&
|
||||
'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
|
||||
|
||||
@@ -119,7 +121,7 @@ steps:
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'pip install av matplotlib ftfy &&
|
||||
'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
|
||||
pip install open-clip-torch --no-deps &&
|
||||
cd tests &&
|
||||
pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
@@ -68,6 +68,7 @@ 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,20 +534,6 @@ 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 \
|
||||
@@ -555,7 +541,6 @@ 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}" \
|
||||
|
||||
+85
-34
@@ -8,6 +8,7 @@
|
||||
# Documentation
|
||||
# label(str): the name of the test. emojis allowed.
|
||||
# fast_check(bool): whether to run this on each commit on the fastcheck pipeline.
|
||||
# torch_nightly(bool): whether to run this on vllm against the torch nightly pipeline.
|
||||
# fast_check_only(bool): run this test on the fastcheck pipeline only
|
||||
# optional(bool): never run this test by default (i.e. need to unblock manually) unless it's a scheduled nightly run.
|
||||
# soft_fail(bool): allow this step to fail without failing the entire pipeline (useful for flaky or experimental tests).
|
||||
@@ -113,11 +114,32 @@ 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]
|
||||
agent_pool: mi250_1
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
@@ -191,6 +213,7 @@ steps:
|
||||
agent_pool: mi250_1
|
||||
no_gpu: true
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -229,12 +252,14 @@ steps:
|
||||
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/
|
||||
- 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
|
||||
|
||||
@@ -413,10 +438,11 @@ steps:
|
||||
#----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------#
|
||||
|
||||
- label: Basic Correctness # TBD
|
||||
timeout_in_minutes: 95
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -430,7 +456,7 @@ steps:
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
- label: Distributed Model Tests (2 GPUs) # TBD
|
||||
timeout_in_minutes: 110
|
||||
timeout_in_minutes: 65
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -472,25 +498,6 @@ steps:
|
||||
|
||||
#---------------------------------------------------------- mi300 · compile ----------------------------------------------------------#
|
||||
|
||||
- label: PyTorch Compilation Unit Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
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]
|
||||
@@ -548,6 +555,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
@@ -670,7 +678,7 @@ steps:
|
||||
- pytest -v -s distributed/test_eplb_spec_decode.py
|
||||
|
||||
- label: Distributed Tests (2xH100-2xMI300) # TBD
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -815,6 +823,7 @@ steps:
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -831,6 +840,7 @@ steps:
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -847,6 +857,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -862,6 +873,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -879,6 +891,7 @@ steps:
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -898,6 +911,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -911,6 +925,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -924,6 +939,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -937,6 +953,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -1205,7 +1222,7 @@ steps:
|
||||
#--------------------------------------------------------- mi300 · examples ----------------------------------------------------------#
|
||||
|
||||
- label: Examples # TBD
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1241,7 +1258,7 @@ steps:
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 100
|
||||
timeout_in_minutes: 55
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1275,7 +1292,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: 95
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1357,6 +1374,8 @@ 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
|
||||
@@ -1371,6 +1390,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
@@ -1392,6 +1412,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1408,6 +1429,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -1417,7 +1439,7 @@ steps:
|
||||
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
|
||||
|
||||
- label: Basic Models Tests (Other) # TBD
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1449,6 +1471,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -1462,6 +1485,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
torch_nightly: true
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1492,6 +1516,7 @@ 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
|
||||
|
||||
@@ -1505,6 +1530,7 @@ 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'
|
||||
|
||||
|
||||
@@ -1518,18 +1544,21 @@ 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
|
||||
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/
|
||||
- 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
|
||||
|
||||
@@ -1537,6 +1566,7 @@ steps:
|
||||
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:
|
||||
@@ -1544,6 +1574,7 @@ 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
|
||||
|
||||
@@ -1551,12 +1582,14 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- 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
|
||||
@@ -1572,6 +1605,7 @@ 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
|
||||
@@ -1587,6 +1621,7 @@ 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 -----------------------------------------------------#
|
||||
@@ -1868,7 +1903,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: 90
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -2084,7 +2119,7 @@ steps:
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs) # TBD
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 20
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
optional: true
|
||||
@@ -2237,7 +2272,7 @@ steps:
|
||||
#------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -2249,7 +2284,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: 75
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
@@ -2397,6 +2432,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -2426,12 +2462,14 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- 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
|
||||
|
||||
@@ -2518,6 +2556,7 @@ steps:
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2534,6 +2573,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -2549,6 +2589,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -2565,6 +2606,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -2585,6 +2627,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2598,6 +2641,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2611,6 +2655,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2783,10 +2828,9 @@ 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: 100
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
@@ -2913,6 +2957,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2934,6 +2979,7 @@ 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
|
||||
|
||||
@@ -2947,6 +2993,7 @@ 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
|
||||
@@ -2965,12 +3012,14 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- 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
|
||||
|
||||
@@ -2978,12 +3027,14 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- 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
|
||||
@@ -3126,7 +3177,7 @@ steps:
|
||||
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
|
||||
|
||||
- label: Weight Loading Multiple GPU # TBD
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 30
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_2
|
||||
num_gpus: 2
|
||||
@@ -3138,7 +3189,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: 75
|
||||
timeout_in_minutes: 30
|
||||
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.py
|
||||
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.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 (1xB200)
|
||||
key: lm-eval-small-models-1xb200
|
||||
- label: LM Eval Small Models (2xB200)
|
||||
key: lm-eval-small-models-2xb200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
@@ -65,10 +65,9 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Small Models Distributed (2xB200)
|
||||
key: lm-eval-small-models-distributed-2xb200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
- label: LM Eval Small Models (2xL4)
|
||||
key: lm-eval-small-models-tp
|
||||
timeout_in_minutes: 10
|
||||
num_devices: 2
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -18,7 +18,9 @@ steps:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
- pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
# This requires eager until we sort out CG correctness issues.
|
||||
# TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
|
||||
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
# Temporary hack filter to exclude ngram spec decoding based tests.
|
||||
|
||||
@@ -6,6 +6,7 @@ steps:
|
||||
key: basic-models-tests-initialization
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_initialization.py
|
||||
@@ -13,6 +14,8 @@ steps:
|
||||
commands:
|
||||
# Run a subset of model initialization tests
|
||||
- pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Basic Models Tests (Extra Initialization) %N
|
||||
device: h200_35gb
|
||||
@@ -28,6 +31,8 @@ steps:
|
||||
# test.) Also run if model initialization test file is modified
|
||||
- pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
parallelism: 2
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Basic Models Tests (Other)
|
||||
device: h200_35gb
|
||||
|
||||
@@ -14,6 +14,7 @@ steps:
|
||||
- pip freeze | grep -E 'torch'
|
||||
- pytest -v -s models/language -m 'core_model and (not slow_test)'
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
@@ -34,6 +35,7 @@ steps:
|
||||
- pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
parallelism: 2
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
@@ -65,6 +67,7 @@ steps:
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
parallelism: 2
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 90
|
||||
|
||||
@@ -10,6 +10,7 @@ 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:
|
||||
@@ -26,6 +27,7 @@ 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:
|
||||
@@ -42,6 +44,7 @@ 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:
|
||||
@@ -58,6 +61,7 @@ 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
|
||||
@@ -79,6 +83,7 @@ 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
|
||||
@@ -90,6 +95,7 @@ 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
|
||||
@@ -123,6 +129,7 @@ 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:
|
||||
@@ -139,6 +146,7 @@ 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)
|
||||
@@ -149,6 +157,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Pooling)
|
||||
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
runs-on: [self-hosted, linux, x64, vllm-runners]
|
||||
steps:
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
|
||||
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
# Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its
|
||||
|
||||
+30
-19
@@ -350,7 +350,9 @@ endif()
|
||||
if(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/torch_bindings.cpp"
|
||||
"csrc/custom_quickreduce.cu")
|
||||
"csrc/custom_quickreduce.cu"
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu")
|
||||
|
||||
message(STATUS "Enabling C extension.")
|
||||
define_extension_target(
|
||||
@@ -378,8 +380,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
#
|
||||
set(VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/activation_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/activation_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
|
||||
@@ -459,6 +459,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
@@ -1089,15 +1091,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
@@ -1107,6 +1108,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
elseif(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
|
||||
endif()
|
||||
|
||||
@@ -1143,8 +1150,7 @@ set(VLLM_MOE_EXT_SRC
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_MOE_EXT_SRC
|
||||
"csrc/libtorch_stable/moe/moe_wna16.cu"
|
||||
"csrc/libtorch_stable/moe/grouped_topk_kernels.cu"
|
||||
"csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels_fast.cu")
|
||||
"csrc/libtorch_stable/moe/grouped_topk_kernels.cu")
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -1315,20 +1321,25 @@ define_extension_target(
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
elseif(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _moe_C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
target_compile_definitions(_moe_C_stable_libtorch PRIVATE USE_ROCM)
|
||||
endif()
|
||||
|
||||
|
||||
@@ -132,8 +132,10 @@ def benchmark_function(
|
||||
reset_memory_stats()
|
||||
|
||||
# Benchmark
|
||||
start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
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)]
|
||||
|
||||
for i in range(benchmark_iters):
|
||||
logits_copy = logits.clone()
|
||||
|
||||
@@ -134,8 +134,8 @@ def benchmark_config(
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Benchmark
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.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.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
latencies: list[float] = []
|
||||
for _ in range(num_iters):
|
||||
start.record()
|
||||
|
||||
@@ -427,7 +427,6 @@ 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()
|
||||
|
||||
@@ -29,8 +29,7 @@ if(DEEPGEMM_SRC_DIR)
|
||||
else()
|
||||
# Keep in sync with tools/install_deepgemm.sh
|
||||
set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
|
||||
# NOTE: This is currently targeting nv-dev branch due to sm120 support
|
||||
set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
|
||||
set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
|
||||
|
||||
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _deepgemm_fc_root)
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG 2c839c33742309ec41e620bf837495ec9926c56e
|
||||
GIT_TAG b3964b1d8b95d8e8447435668ab169a2700bab65
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -126,18 +126,6 @@ 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);
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
// 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,8 +3,6 @@
|
||||
|
||||
#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>
|
||||
@@ -347,10 +345,6 @@ 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> {
|
||||
@@ -397,13 +391,6 @@ 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;
|
||||
@@ -510,15 +497,6 @@ 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,10 +298,6 @@ 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);
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#include <torch/all.h>
|
||||
#include <torch/cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// This function assumes that `cpu_tensor` is a CPU tensor,
|
||||
// and that UVA (Unified Virtual Addressing) is enabled.
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
|
||||
TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
|
||||
|
||||
// handle empty tensor
|
||||
if (cpu_tensor.numel() == 0) {
|
||||
return torch::empty(cpu_tensor.sizes(),
|
||||
cpu_tensor.options().device(torch::kCUDA));
|
||||
}
|
||||
|
||||
if (cpu_tensor.is_pinned()) {
|
||||
// If CPU tensor is pinned, directly get the device pointer.
|
||||
void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
|
||||
void* device_ptr = nullptr;
|
||||
cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
|
||||
return torch::from_blob(
|
||||
device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
|
||||
[base = cpu_tensor](void*) {}, // keep cpu tensor alive
|
||||
cpu_tensor.options().device(torch::kCUDA));
|
||||
}
|
||||
|
||||
// If CPU tensor is not pinned, allocate a new pinned memory buffer.
|
||||
torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
|
||||
size_t nbytes = contiguous_cpu.nbytes();
|
||||
|
||||
void* host_ptr = nullptr;
|
||||
cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
|
||||
if (err != cudaSuccess) {
|
||||
AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
|
||||
cudaMemcpyDefault);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
void* device_ptr = nullptr;
|
||||
err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
|
||||
|
||||
return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
|
||||
contiguous_cpu.strides(), deleter,
|
||||
contiguous_cpu.options().device(torch::kCUDA));
|
||||
}
|
||||
@@ -328,7 +328,7 @@ struct GmemLoaderB {
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
if (need_wait) {
|
||||
@@ -643,7 +643,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
|
||||
mma_computer.issue_mainloop();
|
||||
mma_computer.epi();
|
||||
}
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -100,7 +100,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
for (int ki = 0; ki < k_iterations; ki++) {
|
||||
@@ -146,7 +146,7 @@ __global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -249,7 +249,7 @@ __global__ void __launch_bounds__(1024)
|
||||
LamportComm<NRanks> comm(params.workspace, params.rank);
|
||||
int clear_access = comm.clear_size / kElemsPerAccess<DType>;
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
for (int idx = access_id; idx < tot_access;
|
||||
idx += access_stride, token_id += token_stride) {
|
||||
@@ -313,7 +313,7 @@ __global__ void __launch_bounds__(1024)
|
||||
}
|
||||
comm.update(params.size_q * NRanks);
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
@@ -384,7 +384,7 @@ __global__ void __launch_bounds__(1024)
|
||||
|
||||
DType norm_weight[kElemsPerAccess<DType>]{};
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
if (is_q) {
|
||||
if (is_valid_q) {
|
||||
@@ -596,7 +596,7 @@ __global__ void __launch_bounds__(1024)
|
||||
}
|
||||
} // end group loop
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
|
||||
int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
|
||||
@@ -804,6 +804,35 @@ 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,
|
||||
|
||||
@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
|
||||
}
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
// Process the GEMM in chunks
|
||||
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
|
||||
}
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
// Process the GEMM in chunks
|
||||
@@ -163,7 +163,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -48,8 +48,7 @@ static constexpr int NumTopGroupScores = 2;
|
||||
static constexpr int DefaultMaxNumTopExperts = 8;
|
||||
static constexpr int MaxSupportedTopExperts = 22;
|
||||
static constexpr int MaxNumTopGroups = 4;
|
||||
// The empirical value for small batch
|
||||
static constexpr int PDLEnableTokens = 16;
|
||||
|
||||
namespace warp_topk {
|
||||
|
||||
template <int size, typename T>
|
||||
@@ -565,8 +564,8 @@ __global__ void grouped_topk_fused_kernel(
|
||||
T* s_group_scores = reinterpret_cast<T*>(ptr_u);
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize(); // I think all prolog can be put before
|
||||
// acqbulk because it's ptr arithmetic
|
||||
asm volatile("griddepcontrol.wait;"); // I think all prolog can be put before
|
||||
// acqbulk because it's ptr arithmetic
|
||||
#endif
|
||||
|
||||
// phase 1: per-group scan
|
||||
@@ -610,7 +609,7 @@ __global__ void grouped_topk_fused_kernel(
|
||||
topk_values[i] = 1.0f / static_cast<float>(topk_i32);
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
@@ -671,7 +670,7 @@ __global__ void grouped_topk_fused_kernel(
|
||||
}
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
@@ -896,8 +895,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
|
||||
int64_t const num_experts, int64_t const n_group,
|
||||
int64_t const topk_group, int64_t const topk,
|
||||
bool const renormalize, double const routed_scaling_factor,
|
||||
const bool enable_pdl = false,
|
||||
cudaStream_t const stream = 0) {
|
||||
bool enable_pdl = false, cudaStream_t const stream = 0) {
|
||||
cudaLaunchConfig_t config;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
@@ -985,7 +983,7 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,
|
||||
int64_t const num_tokens, int64_t const num_experts, \
|
||||
int64_t const n_group, int64_t const topk_group, int64_t const topk, \
|
||||
bool const renormalize, double const routed_scaling_factor, \
|
||||
const bool enable_pdl, cudaStream_t const stream);
|
||||
bool enable_pdl, cudaStream_t const stream);
|
||||
|
||||
INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID);
|
||||
INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID);
|
||||
@@ -1039,7 +1037,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
|
||||
scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float);
|
||||
auto topk_indices = torch::stable::new_empty(
|
||||
scores, {num_tokens, topk}, torch::headeronly::ScalarType::Int);
|
||||
const bool pdl_flag = num_tokens <= vllm::moe::PDLEnableTokens;
|
||||
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(scores.get_device_index());
|
||||
auto const sf = static_cast<vllm::moe::ScoringFunc>(scoring_func);
|
||||
@@ -1054,7 +1052,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
|
||||
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
|
||||
reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
|
||||
num_experts, n_group, topk_group, topk, renormalize, \
|
||||
routed_scaling_factor, pdl_flag, stream); \
|
||||
routed_scaling_factor, false, stream); \
|
||||
break; \
|
||||
case vllm::moe::SCORING_SIGMOID: \
|
||||
vllm::moe::invokeNoAuxTc<T, BiasT, IdxT, vllm::moe::SCORING_SIGMOID>( \
|
||||
@@ -1063,7 +1061,7 @@ std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
|
||||
reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \
|
||||
reinterpret_cast<BiasT const*>(bias.data_ptr()), num_tokens, \
|
||||
num_experts, n_group, topk_group, topk, renormalize, \
|
||||
routed_scaling_factor, pdl_flag, stream); \
|
||||
routed_scaling_factor, false, stream); \
|
||||
break; \
|
||||
default: \
|
||||
STD_TORCH_CHECK(false, "Unsupported scoring_func"); \
|
||||
|
||||
@@ -74,22 +74,6 @@ void shuffle_rows(const torch::stable::Tensor& input_tensor,
|
||||
const torch::stable::Tensor& dst2src_map,
|
||||
torch::stable::Tensor& output_tensor);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// TRT-LLM fused sqrt-softplus routing gate (top-k, renormalize). Ported from
|
||||
// NVIDIA/TensorRT-LLM PR #15402 (customMoeRoutingKernels.cu). Signature mirrors
|
||||
// topk_softplus_sqrt for drop-in use: gating_output/topk_weights are float32,
|
||||
// topk_indices is int32, topk (=topk_weights.size(-1)) must be 6, n_experts in
|
||||
// {256, 384}, always renormalizes. Hash mode is inferred from tid2eid.
|
||||
void topk_softplus_sqrt_fast(
|
||||
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid);
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// DeepSeek V3 optimized router GEMM kernel for SM90+
|
||||
// Computes output = mat_a @ mat_b.T where:
|
||||
|
||||
@@ -173,7 +173,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
float row_chunk[VPT];
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaGridDependencySynchronize();
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
@@ -300,7 +300,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
return;
|
||||
} else {
|
||||
@@ -425,7 +425,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,263 +0,0 @@
|
||||
/*
|
||||
* Adapted from
|
||||
* https://github.com/NVIDIA/TensorRT-LLM/pull/15402
|
||||
* cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (gate_forward)
|
||||
* Copyright (c) 2026, The vLLM team.
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
#include <torch/headeronly/util/Exception.h>
|
||||
|
||||
#include "../../cuda_compat.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
// CUDA-only: depends on cub-based reduce_topk and a 32-lane warp layout.
|
||||
#ifndef USE_ROCM
|
||||
|
||||
#include <cooperative_groups.h>
|
||||
#include <cooperative_groups/reduce.h>
|
||||
|
||||
#include "moeTopKFuncs.cuh" // vllm::moe::reduce_topk::reduceTopK
|
||||
|
||||
namespace cg = cooperative_groups;
|
||||
|
||||
namespace vllm {
|
||||
namespace moe {
|
||||
|
||||
// CUDA kernel for gate forward.
|
||||
// Input: pre-computed scores from linear(x, weight) done outside the kernel.
|
||||
// Template parameters:
|
||||
// nExperts: number of experts
|
||||
// topK: number of top experts to select
|
||||
// hash: true for hash mode, false for topk mode
|
||||
// One warp per row (batch element).
|
||||
template <int nExperts, int topK, bool hash>
|
||||
__global__ void topk_softplus_sqrt_fast_kernel(
|
||||
float const* __restrict__ scores_in, // [batch_size, nExperts]
|
||||
float const* __restrict__ bias, // [nExperts] (only when hash=false)
|
||||
int const* __restrict__ input_ids, // [batch_size] (only when hash=true)
|
||||
int const* __restrict__ tid2eid, // [vocab_size, topK] (only hash=true)
|
||||
float* __restrict__ out_weights, // [batch_size, topK]
|
||||
int* __restrict__ out_indices, // [batch_size, topK]
|
||||
int batch_size, float route_scale) {
|
||||
// Compile-time constants
|
||||
constexpr int kExpertsPerThread = nExperts / WARP_SIZE;
|
||||
constexpr int kWarpsPerBlock = 4; // Adjust based on occupancy needs
|
||||
|
||||
// Shared memory for original scores (one array per warp in the block)
|
||||
__shared__ float smem_scores[kWarpsPerBlock][nExperts];
|
||||
|
||||
// One warp per batch element
|
||||
int const global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE;
|
||||
int const local_warp_id = (threadIdx.x / WARP_SIZE) % kWarpsPerBlock;
|
||||
int const lane_id = threadIdx.x % WARP_SIZE;
|
||||
|
||||
if (global_warp_id >= batch_size) return;
|
||||
|
||||
auto warp = cg::tiled_partition<WARP_SIZE>(cg::this_thread_block());
|
||||
|
||||
// Pointer to this warp's shared memory and input scores
|
||||
float* my_smem = smem_scores[local_warp_id];
|
||||
float const* scores_row = scores_in + global_warp_id * nExperts;
|
||||
|
||||
// Load scores, apply score function (softplus + sqrt), store to shared mem.
|
||||
#pragma unroll
|
||||
for (int e = 0; e < kExpertsPerThread; ++e) {
|
||||
int expert_id = lane_id + e * WARP_SIZE;
|
||||
float s = scores_row[expert_id];
|
||||
float sp = log1pf(expf(s));
|
||||
float score = sqrtf(sp);
|
||||
my_smem[expert_id] = score; // Store original score to shared memory
|
||||
}
|
||||
__syncwarp(); // Ensure all scores are written before reading
|
||||
|
||||
// Output: each of the first K lanes holds one value.
|
||||
float my_topk_value = 0.0f;
|
||||
int my_topk_index = 0;
|
||||
|
||||
if constexpr (hash) {
|
||||
// Hash mode: directly read from shared memory
|
||||
int token_id = input_ids[global_warp_id];
|
||||
int const* expert_ids = tid2eid + token_id * topK;
|
||||
|
||||
if (lane_id < topK) {
|
||||
int expert_id = expert_ids[lane_id];
|
||||
my_topk_index = expert_id;
|
||||
my_topk_value = my_smem[expert_id]; // Direct lookup from shared memory
|
||||
}
|
||||
} else {
|
||||
// Topk mode: load from shared memory, add bias in registers for topk.
|
||||
float scores[kExpertsPerThread];
|
||||
int indices[kExpertsPerThread];
|
||||
|
||||
#pragma unroll
|
||||
for (int e = 0; e < kExpertsPerThread; ++e) {
|
||||
int expert_id = lane_id + e * WARP_SIZE;
|
||||
indices[e] = expert_id;
|
||||
scores[e] = my_smem[expert_id] + bias[expert_id]; // biased for selection
|
||||
}
|
||||
|
||||
// Use reduceTopK to find the top-k experts (result broadcast to all lanes).
|
||||
float topk_values[topK];
|
||||
int32_t topk_indices[topK];
|
||||
constexpr float minValue = -1e30f;
|
||||
reduce_topk::reduceTopK<topK, float, kExpertsPerThread>(
|
||||
warp, topk_values, topk_indices, scores, indices, minValue, topK);
|
||||
|
||||
// Gather original weights (without bias) from shared memory.
|
||||
if (lane_id < topK) {
|
||||
int expert_id = topk_indices[lane_id];
|
||||
my_topk_index = expert_id;
|
||||
my_topk_value = my_smem[expert_id]; // original score (no bias)
|
||||
}
|
||||
}
|
||||
|
||||
// Reduce to get the sum (first K lanes have values, others have 0).
|
||||
float weight_sum = cg::reduce(warp, my_topk_value, cg::plus<float>{});
|
||||
|
||||
// Normalize weights and write output (first K lanes).
|
||||
if (lane_id < topK) {
|
||||
out_weights[global_warp_id * topK + lane_id] =
|
||||
(my_topk_value / weight_sum) * route_scale;
|
||||
out_indices[global_warp_id * topK + lane_id] = my_topk_index;
|
||||
}
|
||||
}
|
||||
|
||||
// C++ launcher (output tensors passed as parameters). All tensors are float32.
|
||||
template <int nExperts, int topK, bool hash>
|
||||
void launch_topk_softplus_sqrt_fast_kernel(float* scores_in, float* bias, int* input_ids,
|
||||
int* tid2eid, float* out_weights,
|
||||
int* out_indices, int batch_size,
|
||||
float route_scale, cudaStream_t stream) {
|
||||
constexpr int warps_per_block = 4;
|
||||
constexpr int threads_per_block = warps_per_block * WARP_SIZE;
|
||||
int const blocks = (batch_size + warps_per_block - 1) / warps_per_block;
|
||||
|
||||
topk_softplus_sqrt_fast_kernel<nExperts, topK, hash>
|
||||
<<<blocks, threads_per_block, 0, stream>>>(scores_in, bias, input_ids,
|
||||
tid2eid, out_weights,
|
||||
out_indices, batch_size,
|
||||
route_scale);
|
||||
}
|
||||
|
||||
// Dispatch over (n_experts, is_hash). topK and n_experts match the source:
|
||||
// DeepSeek-V4 uses top-k 6, and n_experts is 256 or 384.
|
||||
template <int topK>
|
||||
void topk_softplus_sqrt_fast_dispatch_experts(float* scores, float* bias, int* input_ids,
|
||||
int* tid2eid, float* weights, int* indices,
|
||||
int batch_size, int n_experts,
|
||||
float route_scale, bool is_hash,
|
||||
cudaStream_t stream) {
|
||||
switch (n_experts) {
|
||||
case 256:
|
||||
if (is_hash) {
|
||||
launch_topk_softplus_sqrt_fast_kernel<256, topK, true>(
|
||||
scores, nullptr, input_ids, tid2eid, weights, indices, batch_size,
|
||||
route_scale, stream);
|
||||
} else {
|
||||
launch_topk_softplus_sqrt_fast_kernel<256, topK, false>(
|
||||
scores, bias, nullptr, nullptr, weights, indices, batch_size,
|
||||
route_scale, stream);
|
||||
}
|
||||
break;
|
||||
case 384:
|
||||
if (is_hash) {
|
||||
launch_topk_softplus_sqrt_fast_kernel<384, topK, true>(
|
||||
scores, nullptr, input_ids, tid2eid, weights, indices, batch_size,
|
||||
route_scale, stream);
|
||||
} else {
|
||||
launch_topk_softplus_sqrt_fast_kernel<384, topK, false>(
|
||||
scores, bias, nullptr, nullptr, weights, indices, batch_size,
|
||||
route_scale, stream);
|
||||
}
|
||||
break;
|
||||
default:
|
||||
STD_TORCH_CHECK(false, "topk_softplus_sqrt_fast only supports n_experts 256 or 384, "
|
||||
"got ",
|
||||
n_experts);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe
|
||||
} // namespace vllm
|
||||
|
||||
// Stable-ABI entry. Signature mirrors `topk_softplus_sqrt` so it is a drop-in
|
||||
// replacement. gating_output/topk_weights are float32, topk_indices is int32.
|
||||
// topk is taken from topk_weights.size(-1) (must be 6), hash mode is inferred
|
||||
// from tid2eid, and route_scale = routed_scaling_factor. token_expert_indices
|
||||
// is accepted for ABI parity but unused (this kernel does not permute).
|
||||
// Selection: sqrt(softplus(score)) (+ bias for topk selection), top-k, then
|
||||
// renormalize by the sum of the selected (unbiased) weights and scale.
|
||||
void topk_softplus_sqrt_fast(
|
||||
torch::stable::Tensor& topk_weights, // [num_tokens, topk] fp32
|
||||
torch::stable::Tensor& topk_indices, // [num_tokens, topk] int32
|
||||
torch::stable::Tensor& token_expert_indices, // unused (ABI parity)
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, n_experts] fp32
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid) {
|
||||
const int n_experts = gating_output.size(-1);
|
||||
const int batch_size = gating_output.numel() / n_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
const bool is_hash = tid2eid.has_value();
|
||||
|
||||
STD_TORCH_CHECK(
|
||||
renormalize,
|
||||
"topk_softplus_sqrt_fast always renormalizes; renormalize must be true");
|
||||
STD_TORCH_CHECK(topk == 6, "topk_softplus_sqrt_fast only supports topk 6, got ",
|
||||
topk);
|
||||
STD_TORCH_CHECK(
|
||||
gating_output.scalar_type() == torch::headeronly::ScalarType::Float,
|
||||
"gating_output must be float32");
|
||||
STD_TORCH_CHECK(
|
||||
topk_weights.scalar_type() == torch::headeronly::ScalarType::Float,
|
||||
"topk_weights must be float32");
|
||||
STD_TORCH_CHECK(
|
||||
topk_indices.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"topk_indices must be int32");
|
||||
|
||||
float* bias_ptr = nullptr;
|
||||
int* input_ids_ptr = nullptr;
|
||||
int* tid2eid_ptr = nullptr;
|
||||
if (is_hash) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(), "hash mode requires input_ids");
|
||||
input_ids_ptr = input_ids.value().mutable_data_ptr<int>();
|
||||
tid2eid_ptr = tid2eid.value().mutable_data_ptr<int>();
|
||||
} else {
|
||||
STD_TORCH_CHECK(correction_bias.has_value(), "topk mode requires bias");
|
||||
bias_ptr = correction_bias.value().mutable_data_ptr<float>();
|
||||
}
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard guard(
|
||||
gating_output.get_device_index());
|
||||
const cudaStream_t stream =
|
||||
get_current_cuda_stream(gating_output.get_device_index());
|
||||
|
||||
float* scores = gating_output.mutable_data_ptr<float>();
|
||||
float* weights = topk_weights.mutable_data_ptr<float>();
|
||||
int* indices = topk_indices.mutable_data_ptr<int>();
|
||||
|
||||
vllm::moe::topk_softplus_sqrt_fast_dispatch_experts<6>(
|
||||
scores, bias_ptr, input_ids_ptr, tid2eid_ptr, weights, indices,
|
||||
batch_size, n_experts, static_cast<float>(routed_scaling_factor), is_hash,
|
||||
stream);
|
||||
}
|
||||
|
||||
#endif // USE_ROCM
|
||||
@@ -122,14 +122,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
"routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "
|
||||
"Tensor)");
|
||||
|
||||
// TRT-LLM fused sqrt-softplus routing gate. Signature mirrors
|
||||
// topk_softplus_sqrt for drop-in use (topk fixed at 6, always renormalizes).
|
||||
m.def(
|
||||
"topk_softplus_sqrt_fast(Tensor! topk_weights, Tensor! topk_indices, "
|
||||
"Tensor! token_expert_indices, Tensor gating_output, bool renormalize, "
|
||||
"float routed_scaling_factor, Tensor? bias, Tensor? input_ids, Tensor? "
|
||||
"tid2eid) -> ()");
|
||||
|
||||
// DeepSeek V3 optimized router GEMM for SM90+
|
||||
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
// conditionally compiled so impl registration is in source file
|
||||
@@ -149,7 +141,6 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
|
||||
m.impl("moe_wna16_gemm", TORCH_BOX(&moe_wna16_gemm));
|
||||
m.impl("shuffle_rows", TORCH_BOX(&shuffle_rows));
|
||||
m.impl("grouped_topk", TORCH_BOX(&grouped_topk));
|
||||
m.impl("topk_softplus_sqrt_fast", TORCH_BOX(&topk_softplus_sqrt_fast));
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -181,12 +181,13 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
// AllSpark ops: declarations are in the source files
|
||||
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
|
||||
#endif
|
||||
|
||||
// CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
|
||||
// TODO: Move this out once ROCm upgrade their torch to 2.11.
|
||||
// CPU tensor -> CUDA UVA view (shared CUDA)
|
||||
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
|
||||
torch::stable::Tensor& cpu_tensor);
|
||||
|
||||
#endif
|
||||
|
||||
// Attention kernels (shared CUDA/ROCm)
|
||||
void merge_attn_states(
|
||||
torch::stable::Tensor& output,
|
||||
@@ -287,6 +288,10 @@ 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,
|
||||
|
||||
@@ -29,10 +29,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"()");
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
|
||||
// Note about marlin kernel 'workspace' arguments:
|
||||
// Technically these should be mutable since they are modified by the kernel.
|
||||
// But since they are set back to zero once the kernel is finished we can
|
||||
@@ -448,6 +449,10 @@ 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, "
|
||||
@@ -700,6 +705,7 @@ 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",
|
||||
@@ -752,6 +758,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
|
||||
}
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#ifndef USE_ROCM
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
|
||||
ops.impl("get_cuda_view_from_cpu_tensor",
|
||||
TORCH_BOX(&get_cuda_view_from_cpu_tensor));
|
||||
@@ -770,6 +778,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
|
||||
TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
// These capability-check functions take only primitive args (no tensors), so
|
||||
// there is no device to dispatch on. CompositeExplicitAutograd makes them
|
||||
// available for all backends. This is the stable ABI equivalent of calling
|
||||
|
||||
+3
-2
@@ -35,8 +35,6 @@ 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);
|
||||
@@ -67,4 +65,7 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
|
||||
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t quant_level, bool cast_bf2half = false);
|
||||
int64_t qr_max_size();
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
#endif
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
// cache.h, which is no longer included here after cache ops moved to
|
||||
// _C_stable_libtorch).
|
||||
#include <torch/all.h>
|
||||
#include "cuda_utils.h"
|
||||
#include "ops.h"
|
||||
#include "core/registration.h"
|
||||
#include <torch/library.h>
|
||||
@@ -17,6 +18,19 @@
|
||||
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
|
||||
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations
|
||||
|
||||
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// vLLM custom ops
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// TODO: Remove this once we upgrade to torch 2.11.
|
||||
// ROCm still uses torch 2.10,
|
||||
// So we still need to use unstable torch ABI for now.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
|
||||
&get_cuda_view_from_cpu_tensor);
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef USE_ROCM
|
||||
TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
|
||||
// Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
|
||||
@@ -34,6 +48,20 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
|
||||
|
||||
custom_ar.def("qr_max_size", &qr_max_size);
|
||||
}
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
|
||||
// Cuda utils
|
||||
// Gets the specified device attribute.
|
||||
cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
|
||||
cuda_utils.impl("get_device_attribute", &get_device_attribute);
|
||||
|
||||
// Gets the maximum shared memory per block device attribute.
|
||||
cuda_utils.def(
|
||||
"get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
|
||||
cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
|
||||
&get_max_shared_memory_per_block_device_attribute);
|
||||
}
|
||||
#endif
|
||||
|
||||
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|
||||
|
||||
@@ -0,0 +1,326 @@
|
||||
#######
|
||||
#
|
||||
# THIS FILE IS DEPRECATED AND WILL BE REMOVED SHORTLY
|
||||
#
|
||||
# Please use the standard Dockerfile with PYTORCH_NIGHTLY=1 instead
|
||||
#
|
||||
#######
|
||||
|
||||
# The vLLM Dockerfile is used to construct vLLM image against torch nightly that can be directly used for testing
|
||||
|
||||
# for torch nightly, cuda >=12.6 is required,
|
||||
# use 12.8 due to FlashAttention issue with cuda 12.6 (https://github.com/vllm-project/vllm/issues/15435#issuecomment-2775924628)
|
||||
ARG CUDA_VERSION=12.8.0
|
||||
#
|
||||
#################### BASE BUILD IMAGE ####################
|
||||
# prepare basic build environment
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS base
|
||||
ARG CUDA_VERSION=12.8.0
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG TARGETPLATFORM
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
# Install Python and other dependencies
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y ccache software-properties-common git curl sudo \
|
||||
&& for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
done \
|
||||
&& apt-get update -y \
|
||||
&& apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
|
||||
&& update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
|
||||
&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
|
||||
&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
|
||||
&& curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
|
||||
&& python3 --version \
|
||||
&& python3 -m pip --version
|
||||
# Install uv for faster pip installs
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
python3 -m pip install uv
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
# Reference: https://github.com/astral-sh/uv/pull/1694
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# GCC >= 11.3 required for PyTorch C++20 headers (pytorch/pytorch#167929).
|
||||
RUN apt-get install -y gcc-11 g++-11
|
||||
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 --slave /usr/bin/g++ g++ /usr/bin/g++-11
|
||||
RUN <<EOF
|
||||
gcc --version
|
||||
EOF
|
||||
|
||||
# Workaround for https://github.com/openai/triton/issues/2507 and
|
||||
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
|
||||
# this won't be needed for future versions of this docker image
|
||||
# or future versions of triton.
|
||||
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# install build and runtime dependencies
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
|
||||
# install build and runtime dependencies without stable torch version
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
# install torch nightly
|
||||
ARG PINNED_TORCH_VERSION
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if [ -n "$PINNED_TORCH_VERSION" ]; then \
|
||||
pkgs="$PINNED_TORCH_VERSION"; \
|
||||
else \
|
||||
pkgs="torch torchaudio torchvision"; \
|
||||
fi && \
|
||||
uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system numba==0.65.0
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
|
||||
# build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
|
||||
# track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
|
||||
RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
|
||||
RUN cat torch_build_versions.txt
|
||||
|
||||
# cuda arch list used by torch
|
||||
# can be useful for `test`
|
||||
# explicitly set the list to avoid issues with torch 2.2
|
||||
# see https://github.com/pytorch/pytorch/pull/123243
|
||||
|
||||
#################### BASE BUILD IMAGE ####################
|
||||
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the wheel build
|
||||
# stage doesn't need the rust toolchain or protoc.
|
||||
FROM ubuntu:22.04 AS rust-build
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ca-certificates curl git build-essential unzip python3 python3-pip \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY tools/install_protoc.sh /tmp/install_protoc.sh
|
||||
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
|
||||
# by the wheel build stage.
|
||||
COPY rust rust
|
||||
COPY rust-toolchain.toml rust-toolchain.toml
|
||||
COPY tools/build_rust.py tools/build_rust.py
|
||||
COPY build_rust.sh build_rust.sh
|
||||
|
||||
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
|
||||
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
|
||||
ENV CARGO_BUILD_JOBS=4
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
|
||||
bash build_rust.sh
|
||||
|
||||
#################### RUST BUILD IMAGE ####################
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
FROM base AS build
|
||||
ARG TARGETPLATFORM
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
# Reference: https://github.com/astral-sh/uv/pull/1694
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
COPY . .
|
||||
|
||||
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
|
||||
# them and ships them as-is, skipping the local Rust build.
|
||||
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
|
||||
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/build/cuda.txt
|
||||
|
||||
ARG GIT_REPO_CHECK=0
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
if [ "$GIT_REPO_CHECK" != "0" ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
# Max jobs used by Ninja to build extensions
|
||||
ARG max_jobs=16
|
||||
ENV MAX_JOBS=${max_jobs}
|
||||
ARG nvcc_threads=2
|
||||
ENV NVCC_THREADS=$nvcc_threads
|
||||
|
||||
ARG USE_SCCACHE
|
||||
ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
|
||||
ARG SCCACHE_REGION_NAME=us-west-2
|
||||
ARG SCCACHE_S3_NO_CREDENTIALS=0
|
||||
|
||||
# if USE_SCCACHE is set, use sccache to speed up compilation
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
echo "Installing sccache..." \
|
||||
&& curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
|
||||
&& tar -xzf sccache.tar.gz \
|
||||
&& sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \
|
||||
&& rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \
|
||||
&& export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
|
||||
&& export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
|
||||
&& export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
|
||||
&& export SCCACHE_IDLE_TIMEOUT=0 \
|
||||
&& export CMAKE_BUILD_TYPE=Release \
|
||||
&& sccache --show-stats \
|
||||
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
|
||||
&& sccache --show-stats; \
|
||||
fi
|
||||
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
if [ "$USE_SCCACHE" != "1" ]; then \
|
||||
# Clean any existing CMake artifacts
|
||||
rm -rf .deps && \
|
||||
mkdir -p .deps && \
|
||||
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
|
||||
fi
|
||||
|
||||
#################### WHEEL BUILD IMAGE ####################
|
||||
|
||||
################### VLLM INSTALLED IMAGE ####################
|
||||
# Setup clean environment for vLLM and its dependencies for test and api server using ubuntu22.04 with AOT flashinfer
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS vllm-base
|
||||
# prepare for environment starts
|
||||
ARG CUDA_VERSION=12.8.0
|
||||
ARG PYTHON_VERSION=3.12
|
||||
WORKDIR /vllm-workspace
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ARG TARGETPLATFORM
|
||||
|
||||
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
|
||||
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
|
||||
|
||||
# Install Python and other dependencies
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
|
||||
&& apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
|
||||
&& for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
done \
|
||||
&& apt-get update -y \
|
||||
&& apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \
|
||||
&& update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
|
||||
&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
|
||||
&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
|
||||
&& curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
|
||||
&& python3 --version && python3 -m pip --version
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
python3 -m pip install uv
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
# Reference: https://github.com/astral-sh/uv/pull/1694
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Workaround for https://github.com/openai/triton/issues/2507 and
|
||||
# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
|
||||
# this won't be needed for future versions of this docker image
|
||||
# or future versions of triton.
|
||||
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
|
||||
|
||||
# get the nightly torch version used in the build to make sure the version is the same
|
||||
COPY --from=base /workspace/torch_build_versions.txt ./torch_build_versions.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system $(cat torch_build_versions.txt | xargs) --index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
|
||||
# install the vllm wheel
|
||||
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm-dist \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system vllm-dist/*.whl --verbose
|
||||
|
||||
ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
|
||||
|
||||
# install package for build flashinfer
|
||||
# see issue: https://github.com/flashinfer-ai/flashinfer/issues/738
|
||||
RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.post1
|
||||
|
||||
|
||||
# build flashinfer for torch nightly from source around 10 mins
|
||||
# release version: v0.6.13
|
||||
# todo(elainewy): cache flashinfer build result for faster build
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
echo "git clone flashinfer..." \
|
||||
&& git clone --depth 1 --branch v0.6.13 --recursive https://github.com/flashinfer-ai/flashinfer.git \
|
||||
&& cd flashinfer \
|
||||
&& git submodule update --init --recursive \
|
||||
&& echo "finish git clone flashinfer..." \
|
||||
&& rm -rf build \
|
||||
&& export TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list} \
|
||||
&& FLASHINFER_ENABLE_AOT=1 python3 setup.py bdist_wheel --dist-dir=../flashinfer-dist --verbose \
|
||||
&& cd .. \
|
||||
&& rm -rf flashinfer
|
||||
|
||||
# install flashinfer
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system flashinfer-dist/*.whl --verbose
|
||||
|
||||
# install common packages
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
|
||||
COPY examples examples
|
||||
COPY benchmarks benchmarks
|
||||
COPY ./vllm/collect_env.py .
|
||||
|
||||
RUN python3 use_existing_torch.py
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
|
||||
################### VLLM INSTALLED IMAGE ####################
|
||||
|
||||
|
||||
#################### UNITTEST IMAGE #############################
|
||||
FROM vllm-base as test
|
||||
COPY tests/ tests/
|
||||
|
||||
# install build and runtime dependencies without stable torch version
|
||||
COPY requirements/test/nightly-torch.txt requirements/test/nightly-torch.txt
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
# Reference: https://github.com/astral-sh/uv/pull/1694
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# install development dependencies (for testing)
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -e tests/vllm_test_utils
|
||||
|
||||
# enable fast downloads from hf (for testing)
|
||||
ENV HF_XET_HIGH_PERFORMANCE 1
|
||||
|
||||
# increase timeout for hf downloads (for testing)
|
||||
ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/test/nightly-torch.txt
|
||||
|
||||
# Logging to confirm the torch versions
|
||||
RUN pip freeze | grep -E 'torch|vllm|flashinfer'
|
||||
|
||||
# Logging to confirm all the packages are installed
|
||||
RUN pip freeze
|
||||
|
||||
#################### UNITTEST IMAGE #############################
|
||||
@@ -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 |
|
||||
|
||||
+4
-25
@@ -30,7 +30,7 @@ or just on the low or high end.
|
||||
| [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
|
||||
| [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
|
||||
| [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always |
|
||||
| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm (+ FP8 quant) → 1 kernel | O1 (ROCm/AITER only) | 1-2% | No | Always |
|
||||
| [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm → single kernel | O1 (ROCm/AITER only) | ~2% | No | Always |
|
||||
|
||||
## Support Matrix
|
||||
|
||||
@@ -381,32 +381,11 @@ q_normed, kv_normed = fused_mla_dual_rms_norm(
|
||||
Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
|
||||
when AITER is available.
|
||||
|
||||
**FP8 attention variant (per-token quant).** With a per-token FP8 `q_b_proj`,
|
||||
only the *q* latent is FP8-quantized while *kv* stays bf16.
|
||||
`RocmAiterRMSNormQuantFusionPass` first folds the q side into
|
||||
`rocm_aiter_rmsnorm_fused_dynamic_quant`, leaving kv a plain
|
||||
`rms_norm` — breaking the symmetric pattern above. The same pass then matches
|
||||
this asymmetric pair and lowers it to `fused_mla_dual_rms_norm_per_token_quant`.
|
||||
|
||||
```text
|
||||
# Unfused (q norm+quant fused; kv still plain rms_norm):
|
||||
q_c, kv_lora = split(projected, [q_dim, kv_dim])
|
||||
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
|
||||
q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8)
|
||||
kv_normed = rms_norm(kv_c, kv_weight, eps) # bf16
|
||||
|
||||
# Fused:
|
||||
q_c, kv_lora = split(projected, [q_dim, kv_dim])
|
||||
kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
|
||||
q_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant(
|
||||
q_c, q_weight, kv_c, kv_weight, eps1, eps2)
|
||||
```
|
||||
|
||||
**Code locations.**
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`, `MLADualRMSPerTokenQuantPattern`)
|
||||
- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`, `fused_mla_dual_rms_norm_per_token_quant`)
|
||||
- AITER kernels: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442), `fused_qk_rmsnorm_per_token_quant`
|
||||
- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)
|
||||
- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)
|
||||
- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)
|
||||
|
||||
## See Also
|
||||
|
||||
|
||||
@@ -86,7 +86,6 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
|
||||
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
|
||||
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
|
||||
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
|
||||
| `use_heterogeneous_vocab` | `boolean` | `false` | Allow draft and target models with different vocabularies. Builds a token-level intersection at initialisation and constrains draft logits to shared tokens only. Only compatible with `method=draft_model`. Probabilistic draft sampling (`draft_sample_method='probabilistic'`) is not yet supported when this option is enabled. |
|
||||
|
||||
!!! note
|
||||
Gemma 4 assistant checkpoints are handled as Gemma 4 MTP speculators, not
|
||||
@@ -143,33 +142,6 @@ vllm serve <target-model> \
|
||||
}'
|
||||
```
|
||||
|
||||
#### Cross-Vocabulary Draft Models (TLI)
|
||||
|
||||
By default, vLLM requires the draft and target models to share the same
|
||||
vocabulary. Setting `use_heterogeneous_vocab: true` enables the
|
||||
**Token-Level Intersection (TLI)** algorithm, which allows draft models
|
||||
from a different model family with a different tokenizer.
|
||||
|
||||
At initialisation, vLLM builds a mapping between the two vocabularies by
|
||||
normalising token strings and computing their intersection. Draft logits are
|
||||
constrained to the shared tokens before sampling, and the sampled token IDs
|
||||
are translated to the target vocabulary before rejection sampling.
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
speculative_config={
|
||||
"method": "draft_model",
|
||||
"model": "HuggingFaceTB/SmolLM2-135M-Instruct",
|
||||
"num_speculative_tokens": 3,
|
||||
"use_heterogeneous_vocab": True,
|
||||
},
|
||||
gpu_memory_utilization=0.5,
|
||||
)
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--speculative-config` expects a JSON object on the CLI. In YAML config
|
||||
@@ -181,7 +153,6 @@ vllm serve <target-model> \
|
||||
- Internal fields such as `target_model_config`, `draft_model_config`,
|
||||
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
|
||||
are populated by vLLM and are not intended to be set by users.
|
||||
- `use_heterogeneous_vocab` currently supports greedy draft sampling only. Probabilistic acceptance (temperature > 0 draft sampling) is not yet supported and will be added in a future release.
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
|
||||
@@ -76,34 +76,6 @@ The code used to request as completions as a client remains unchanged:
|
||||
print(completion)
|
||||
```
|
||||
|
||||
## Draft Model Method with heterogeneous vocabs
|
||||
|
||||
By default, vLLM requires the draft and target models to share the same vocabulary. Setting `use_heterogeneous_vocab: true` enables the **Token-Level Intersection (TLI)** algorithm, which allows draft models from a different model family with a different tokenizer.
|
||||
|
||||
Currently,`use_heterogeneous_vocab` currently requires `draft_sample_method='greedy'` (the default). Probabilistic draft sampling is not yet supported and will be added in a
|
||||
future release.
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
speculative_config={
|
||||
"method": "draft_model",
|
||||
"model": "HuggingFaceTB/SmolLM2-135M-Instruct",
|
||||
"num_speculative_tokens": 3,
|
||||
"use_heterogeneous_vocab": True,
|
||||
},
|
||||
gpu_memory_utilization=0.5,
|
||||
)
|
||||
outputs = llm.generate(prompts,sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative-config` to set all configurations related
|
||||
to speculative decoding. The previous method of specifying the model
|
||||
|
||||
@@ -405,6 +405,7 @@ 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` | ✅︎ | ✅︎ |
|
||||
@@ -476,6 +477,7 @@ 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. | ✅︎ | ✅︎ |
|
||||
@@ -488,9 +490,7 @@ 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,6 +532,7 @@ 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. | ✅︎ | ✅︎ |
|
||||
@@ -578,10 +579,9 @@ 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`, `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
|
||||
| `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. | ✅︎ | ✅︎ |
|
||||
| `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. | | ✅︎ |
|
||||
| `LlavaOnevision2ForConditionalGeneration` | LLaVA-OneVision-2 | T + I<sup>+</sup> + V<sup>+</sup> | `lmms-lab-encoder/LLaVA-OneVision-2-8B-Instruct` | | |
|
||||
| `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. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
@@ -594,6 +594,7 @@ 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` | ✅︎ | ✅︎ |
|
||||
@@ -625,6 +626,8 @@ 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. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -677,6 +680,9 @@ 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/). 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.
|
||||
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.
|
||||
|
||||
### Mid Level `reload_weights` API
|
||||
|
||||
|
||||
@@ -17,7 +17,6 @@ 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
|
||||
|
||||
@@ -42,7 +41,7 @@ vllm serve my-model \
|
||||
--weight-transfer-config '{"backend": "nccl"}'
|
||||
```
|
||||
|
||||
The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
|
||||
The `backend` field accepts `"nccl"` (default) or `"ipc"`.
|
||||
|
||||
## API Endpoints
|
||||
|
||||
@@ -70,7 +69,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()
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
|
||||
# 3. Send weights to inference workers
|
||||
EngineClass.trainer_send_weights(
|
||||
|
||||
@@ -11,23 +11,15 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data
|
||||
|
||||
### Abstract Methods
|
||||
|
||||
Subclasses must implement these methods:
|
||||
Subclasses must implement these four methods:
|
||||
|
||||
| Method | Side | Description |
|
||||
| ------ | ---- | ----------- |
|
||||
| `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
|
||||
| `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` |
|
||||
| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
|
||||
| `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.
|
||||
@@ -89,7 +81,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo):
|
||||
### 2. Implement the Engine
|
||||
|
||||
```python
|
||||
from collections.abc import Iterator
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Any
|
||||
import torch
|
||||
|
||||
@@ -101,25 +93,18 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
|
||||
# Set up connection to trainer using init_info.endpoint, etc.
|
||||
...
|
||||
|
||||
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 = []
|
||||
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
|
||||
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)
|
||||
weights.append((name, weight))
|
||||
self.model.load_weights(weights)
|
||||
load_weights([(name, weight)])
|
||||
|
||||
def shutdown(self) -> None:
|
||||
# Clean up resources
|
||||
@@ -136,6 +121,9 @@ 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
|
||||
@@ -159,7 +147,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`, `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.
|
||||
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.
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
||||
@@ -167,8 +155,7 @@ from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
|
||||
# Create an engine from config
|
||||
engine = WeightTransferEngineFactory.create_engine(
|
||||
config=weight_transfer_config,
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
parallel_config=parallel_config,
|
||||
model=model,
|
||||
)
|
||||
```
|
||||
|
||||
@@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
|
||||
llm_handle=llm_actor_handle,
|
||||
)
|
||||
# start
|
||||
ray.get(llm_actor_handle.start_weight_update.remote())
|
||||
ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
|
||||
# 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={}, timeout=60)
|
||||
response = requests.post(url, json={"is_checkpoint_format": True}, 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 the weights.
|
||||
2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
|
||||
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()
|
||||
llm.start_weight_update(is_checkpoint_format=True)
|
||||
|
||||
# 2. Receive weights (can be called multiple times for chunked transfers)
|
||||
llm.update_weights(
|
||||
@@ -116,22 +116,19 @@ 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 NCCL engine receives checkpoint-format weights and applies
|
||||
layerwise reload processing automatically inside `start_weight_update` /
|
||||
`finish_weight_update`.
|
||||
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).
|
||||
|
||||
## 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.
|
||||
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`.
|
||||
|
||||
## 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 `backend="sparse_nccl"` 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 `start_weight_update(is_checkpoint_format=False)` 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,27 +326,6 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the
|
||||
|
||||
**Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).
|
||||
|
||||
## gRPC Interface
|
||||
|
||||
vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server.
|
||||
|
||||
**Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network.
|
||||
|
||||
### Security Implications
|
||||
|
||||
An attacker who can reach the gRPC port can:
|
||||
|
||||
1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials
|
||||
2. **Consume GPU and compute resources** by submitting unbounded generation requests
|
||||
3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.
|
||||
|
||||
### Recommendations
|
||||
|
||||
- Only enable `--grpc-port` when you have a specific need for gRPC-based inference
|
||||
- Ensure the gRPC port is only accessible from trusted hosts or services
|
||||
- Use firewall rules to block external access to the gRPC port
|
||||
- Consider deploying the gRPC interface on a dedicated internal network interface
|
||||
|
||||
## Cache Directory Security
|
||||
|
||||
vLLM assumes that its cache directories are **private and trusted**. Cache contents are loaded without cryptographic integrity verification, including formats that support arbitrary code execution. If an untrusted user or process can write to vLLM's cache directories, they may be able to crash vLLM or cause it to execute arbitrary code.
|
||||
|
||||
@@ -327,9 +327,6 @@ async def handle_request(api: str, request: Request):
|
||||
session, decode_response = await decode_request_task
|
||||
stream_generator = stream_decode_response(session, decode_response, request_id)
|
||||
response = await make_response(stream_generator)
|
||||
response.headers["Content-Type"] = decode_response.headers.get(
|
||||
"Content-Type", "application/json"
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
logger.exception("An error occurred while handling the request: %s", e)
|
||||
|
||||
@@ -72,7 +72,6 @@ def parse_args():
|
||||
parser.add_argument("--max-num-seqs", type=int, default=None)
|
||||
parser.add_argument("--parallel-drafting", action="store_true")
|
||||
parser.add_argument("--allowed-local-media-path", type=str, default="")
|
||||
parser.add_argument("--use-heterogeneous-vocab", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -136,7 +135,6 @@ def main(args):
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"max_model_len": args.max_model_len,
|
||||
"parallel_drafting": args.parallel_drafting,
|
||||
"use_heterogeneous_vocab": args.use_heterogeneous_vocab,
|
||||
}
|
||||
elif args.method == "mtp":
|
||||
speculative_config = {
|
||||
|
||||
@@ -91,6 +91,44 @@ 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"
|
||||
@@ -527,6 +565,7 @@ 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,6 +74,39 @@ 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"
|
||||
|
||||
@@ -1242,6 +1275,55 @@ 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"
|
||||
@@ -1387,6 +1469,7 @@ 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,
|
||||
@@ -1424,6 +1507,8 @@ 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,6 +68,28 @@ 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"
|
||||
@@ -1355,6 +1377,28 @@ 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"]
|
||||
@@ -2303,8 +2347,68 @@ 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,
|
||||
@@ -2345,6 +2449,7 @@ 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,
|
||||
@@ -2374,6 +2479,8 @@ 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())
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
inference_handle = llm.update_weights.remote(
|
||||
WeightTransferUpdateRequest(
|
||||
|
||||
@@ -80,10 +80,14 @@ def init_weight_transfer_engine(base_url: str) -> None:
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def start_weight_update(base_url: str) -> None:
|
||||
def start_weight_update(
|
||||
base_url: str,
|
||||
is_checkpoint_format: bool = True,
|
||||
) -> None:
|
||||
"""Start a weight update via HTTP endpoint."""
|
||||
url = f"{base_url}/start_weight_update"
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
payload = {"is_checkpoint_format": is_checkpoint_format}
|
||||
response = requests.post(url, json=payload, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
@@ -166,7 +170,7 @@ def main():
|
||||
pause_generation(BASE_URL)
|
||||
|
||||
# Start weight update, broadcast via IPC, then finish
|
||||
start_weight_update(BASE_URL)
|
||||
start_weight_update(BASE_URL, is_checkpoint_format=False)
|
||||
|
||||
print("Broadcasting weights via CUDA IPC (HTTP)...")
|
||||
trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
|
||||
|
||||
@@ -83,10 +83,14 @@ def init_weight_transfer_engine(
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
def start_weight_update(base_url: str) -> None:
|
||||
def start_weight_update(
|
||||
base_url: str,
|
||||
is_checkpoint_format: bool = True,
|
||||
) -> None:
|
||||
"""Start a weight update via HTTP endpoint."""
|
||||
url = f"{base_url}/start_weight_update"
|
||||
response = requests.post(url, json={}, timeout=60)
|
||||
payload = {"is_checkpoint_format": is_checkpoint_format}
|
||||
response = requests.post(url, json=payload, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
@@ -219,7 +223,7 @@ def main():
|
||||
shapes.append(list(p.shape))
|
||||
|
||||
# Start weight update
|
||||
start_weight_update(BASE_URL)
|
||||
start_weight_update(BASE_URL, is_checkpoint_format=True)
|
||||
|
||||
# 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())
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
ray.get(train_model.broadcast_weights.remote(llm))
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
|
||||
|
||||
@@ -277,8 +277,15 @@ class DataParallelInferenceEngine:
|
||||
]
|
||||
)
|
||||
|
||||
def start_weight_update(self):
|
||||
ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
|
||||
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 finish_weight_update(self):
|
||||
ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
|
||||
@@ -385,7 +392,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())
|
||||
ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
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())
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
# 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()
|
||||
await engine.start_weight_update(is_checkpoint_format=True)
|
||||
|
||||
print("[sync] Broadcasting weights from FSDP → vLLM...")
|
||||
broadcast_handles = [
|
||||
|
||||
@@ -44,14 +44,11 @@ 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"
|
||||
@@ -247,6 +244,7 @@ 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
|
||||
|
||||
@@ -273,7 +271,7 @@ class TrainModel:
|
||||
raise RuntimeError("Sparse patch has not been prepared")
|
||||
|
||||
start = time.perf_counter()
|
||||
SparseNCCLWeightTransferEngine.trainer_send_weights(
|
||||
NCCLWeightTransferEngine.trainer_send_sparse_weights(
|
||||
iter(self.pending_sparse_patches),
|
||||
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
|
||||
)
|
||||
@@ -284,7 +282,6 @@ class TrainModel:
|
||||
|
||||
def launch_llm(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
backend: str = "nccl",
|
||||
):
|
||||
return ray.remote(
|
||||
num_cpus=0,
|
||||
@@ -296,7 +293,7 @@ def launch_llm(
|
||||
tensor_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
gpu_memory_utilization=0.7,
|
||||
weight_transfer_config=WeightTransferConfig(backend=backend),
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
)
|
||||
|
||||
|
||||
@@ -335,7 +332,7 @@ def run_dense_phase(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference, backend="nccl")
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
dense_before = collect_vllm_generations(llm)
|
||||
|
||||
@@ -354,7 +351,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())
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
dense_update_info, dense_payload_bytes = ray.get(
|
||||
train_model.get_dense_update_info.remote()
|
||||
@@ -394,7 +391,7 @@ def run_sparse_phase(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference, backend="sparse_nccl")
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
sparse_before = collect_vllm_generations(llm)
|
||||
|
||||
@@ -413,7 +410,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())
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
|
||||
|
||||
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
|
||||
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
|
||||
|
||||
@@ -8,15 +8,14 @@ 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.1.0
|
||||
PyNvVideoCodec==2.0.4
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
flashinfer-python==0.6.13
|
||||
flashinfer-cubin==0.6.13
|
||||
apache-tvm-ffi==0.1.9
|
||||
tilelang==0.1.9
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
# Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
|
||||
nvtx==0.2.15
|
||||
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.3.2
|
||||
|
||||
|
||||
Generated
-1
@@ -5310,7 +5310,6 @@ dependencies = [
|
||||
"vllm-llm",
|
||||
"vllm-metrics",
|
||||
"vllm-text",
|
||||
"vllm-tokenizer",
|
||||
"zeromq",
|
||||
]
|
||||
|
||||
|
||||
@@ -50,7 +50,6 @@ 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,8 +154,7 @@ mod tests {
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_text::Prompt;
|
||||
use vllm_text::backend::hf::TokenizerSource;
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
|
||||
use super::HfChatBackend;
|
||||
use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions};
|
||||
@@ -197,8 +196,32 @@ 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::new())
|
||||
Arc::new(TestTokenizer)
|
||||
}
|
||||
|
||||
fn backend_for_selection(
|
||||
|
||||
@@ -50,8 +50,7 @@ mod request;
|
||||
mod stream;
|
||||
|
||||
use vllm_engine_core_client::EngineCoreClient;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
|
||||
use vllm_engine_core_client::protocol::{ModelDtype, ReasoningParserKwargs};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::{Prompt, TextLlm, TextRequest};
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ use llm_multimodal::{
|
||||
TrackedMedia,
|
||||
};
|
||||
use tracing::warn;
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::{
|
||||
MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargsItem,
|
||||
MmSharedField, MmSlice, PlaceholderRange, SliceSpec,
|
||||
@@ -563,7 +563,7 @@ mod tests {
|
||||
|
||||
use llm_multimodal::TokenId;
|
||||
use vllm_engine_core_client::protocol::tensor::WireArrayData;
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use vllm_text::tokenizer::{IncrementalDecoder, Tokenizer, TokenizerError};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -574,14 +574,60 @@ mod tests {
|
||||
const LLAMA4_TILE_X_SEPARATOR_ID: u32 = 200093;
|
||||
const LLAMA4_TILE_Y_SEPARATOR_ID: u32 = 200094;
|
||||
|
||||
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)
|
||||
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 test_info(model_type: &str, config: serde_json::Value) -> MultimodalModelInfo {
|
||||
@@ -589,7 +635,7 @@ mod tests {
|
||||
model_id: format!("{model_type}-test"),
|
||||
model_type: Some(model_type.to_string()),
|
||||
config,
|
||||
tokenizer: TokenizerResolver(Arc::new(llama4_tokenizer())),
|
||||
tokenizer: TokenizerResolver(Arc::new(TestTokenizer)),
|
||||
};
|
||||
let spec = context
|
||||
.resolve_model_spec()
|
||||
|
||||
@@ -2,7 +2,7 @@ use std::collections::HashMap;
|
||||
|
||||
use half::{bf16, f16};
|
||||
use llm_multimodal::{ModelSpecificValue, PreprocessedImages};
|
||||
use vllm_engine_core_client::protocol::dtype::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::ModelDtype;
|
||||
use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
|
||||
use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
|
||||
|
||||
|
||||
@@ -189,19 +189,46 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
use super::DefaultChatOutputProcessor;
|
||||
use crate::Error;
|
||||
use crate::parser::ParserSelection;
|
||||
use crate::request::ChatRequest;
|
||||
|
||||
fn tokenizer() -> Arc<TestTokenizer> {
|
||||
Arc::new(
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<|channel>", 256)
|
||||
.with_regular_token("<channel|>", 257),
|
||||
)
|
||||
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)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
//! Applies xgrammar structural-tag constraints for strict tool calling.
|
||||
|
||||
use thiserror_ext::AsReport;
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
|
||||
use vllm_parser::tool::StructuralTagModel;
|
||||
use xgrammar_structural_tag::{
|
||||
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
|
||||
@@ -47,8 +45,9 @@ pub(super) fn apply_structural_tag_constraint(
|
||||
|
||||
// Overwrite any existing structured output settings with the structural tag constraint.
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
structural_tag: Some(structural_tag),
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::structural_tag(structural_tag)
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
Ok(())
|
||||
@@ -77,9 +76,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option<StructuralTagTool
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::{StructuredOutputBackend, StructuredOutputsParams};
|
||||
use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser};
|
||||
|
||||
use super::*;
|
||||
@@ -118,11 +115,10 @@ mod tests {
|
||||
.as_ref()
|
||||
.expect("structured outputs should be set");
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Xgrammar);
|
||||
let structural_tag = params
|
||||
.constraint
|
||||
.as_structural_tag()
|
||||
.expect("structured output constraint should be structural_tag");
|
||||
serde_json::from_str(structural_tag).expect("structural_tag should be valid JSON")
|
||||
serde_json::from_str(
|
||||
params.structural_tag.as_deref().expect("structural_tag should be set"),
|
||||
)
|
||||
.expect("structural_tag should be valid JSON")
|
||||
}
|
||||
|
||||
fn structured_outputs(request: &ChatRequest) -> &StructuredOutputsParams {
|
||||
@@ -161,8 +157,9 @@ mod tests {
|
||||
fn auto_strict_tool_choice_overwrites_existing_json_guidance() {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
json: Some(json!({"type": "object"})),
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json(json!({"type": "object"}))
|
||||
..Default::default()
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
@@ -170,7 +167,8 @@ mod tests {
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_structural_tag());
|
||||
assert!(params.json.is_none());
|
||||
assert!(params.structural_tag.is_some());
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
@@ -193,8 +191,9 @@ mod tests {
|
||||
fn required_tool_choice_overwrites_existing_json_object_guidance() {
|
||||
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
json_object: Some(true),
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json_object()
|
||||
..Default::default()
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
@@ -202,7 +201,8 @@ mod tests {
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_structural_tag());
|
||||
assert!(params.json_object.is_none());
|
||||
assert!(params.structural_tag.is_some());
|
||||
let tag = structural_tag_value(&request);
|
||||
assert_eq!(tag["type"], "structural_tag");
|
||||
assert!(tag.to_string().contains("search"));
|
||||
@@ -241,8 +241,9 @@ mod tests {
|
||||
fn none_tool_choice_preserves_existing_json_object_guidance() {
|
||||
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
json_object: Some(true),
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
..StructuredOutputsParams::json_object()
|
||||
..Default::default()
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
@@ -250,6 +251,7 @@ mod tests {
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
assert!(params.constraint.is_json_object());
|
||||
assert_eq!(params.json_object, Some(true));
|
||||
assert!(params.structural_tag.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,14 @@
|
||||
//! 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::*;
|
||||
@@ -85,6 +91,7 @@ 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(
|
||||
@@ -120,6 +127,7 @@ 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();
|
||||
@@ -145,6 +153,7 @@ 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(
|
||||
@@ -180,6 +189,7 @@ 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."),
|
||||
@@ -207,6 +217,7 @@ 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"),
|
||||
@@ -238,6 +249,7 @@ 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",
|
||||
@@ -290,6 +302,7 @@ 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(
|
||||
@@ -340,6 +353,7 @@ 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,9 +1,32 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
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();
|
||||
@@ -84,7 +107,7 @@ fn factory_resolves_minimax_m3_before_generic_minimax() {
|
||||
|
||||
#[test]
|
||||
fn factory_rejects_unknown_parser_names() {
|
||||
let tokenizer = Arc::new(TestTokenizer::new());
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let factory = ReasoningParserFactory::new();
|
||||
let error = match factory.create("missing", tokenizer) {
|
||||
Ok(_) => panic!("expected parser lookup to fail"),
|
||||
|
||||
@@ -75,14 +75,39 @@ impl UnifiedParserFactory {
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_tokenizer::test_utils::TestTokenizer;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
use super::{UnifiedParserFactory, names};
|
||||
|
||||
fn tokenizer() -> TestTokenizer {
|
||||
TestTokenizer::new()
|
||||
.with_regular_token("<|channel>", 256)
|
||||
.with_regular_token("<channel|>", 257)
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -94,6 +119,6 @@ mod tests {
|
||||
factory.resolve_name_for_model("google/gemma-4-27b-it"),
|
||||
Some(names::GEMMA4)
|
||||
);
|
||||
factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap();
|
||||
factory.create(names::GEMMA4, &[], Arc::new(FakeTokenizer)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
+127
-70
@@ -15,24 +15,21 @@ use vllm_chat::{
|
||||
use vllm_engine_core_client::protocol::logprobs::{
|
||||
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::output::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, StopReason,
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
|
||||
use vllm_llm::Llm;
|
||||
use vllm_text::tokenizer::DynTokenizer;
|
||||
use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
|
||||
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,
|
||||
@@ -161,18 +158,45 @@ async fn connect_chat_llm_with_ipc(
|
||||
struct FakeChatBackend {
|
||||
has_template: bool,
|
||||
model_id: String,
|
||||
tokenizer: DynTokenizer,
|
||||
}
|
||||
|
||||
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)
|
||||
#[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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for FakeChatBackend {
|
||||
@@ -186,7 +210,6 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: true,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -194,7 +217,6 @@ impl FakeChatBackend {
|
||||
Self {
|
||||
has_template: false,
|
||||
model_id: "test-model".to_string(),
|
||||
tokenizer: Arc::new(fake_chat_tokenizer()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -202,19 +224,13 @@ 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::clone(&self.tokenizer)
|
||||
Arc::new(FakeChatTokenizer)
|
||||
}
|
||||
|
||||
fn model_id(&self) -> &str {
|
||||
@@ -266,6 +282,65 @@ 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>>
|
||||
@@ -342,7 +417,7 @@ async fn chat_streams_text_events() {
|
||||
);
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output("chat-1", vec![b'H' as u32], None, None),
|
||||
request_output(
|
||||
@@ -354,8 +429,7 @@ async fn chat_streams_text_events() {
|
||||
],
|
||||
finished_requests: Some(BTreeSet::from(["chat-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -455,7 +529,7 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output("chat-utf8", bytes_to_token_ids(&[0xe4]), None, None),
|
||||
request_output(
|
||||
@@ -467,8 +541,7 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
|
||||
],
|
||||
finished_requests: Some(BTreeSet::from(["chat-utf8".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -544,7 +617,7 @@ async fn chat_stream_flushes_held_text_on_finish() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output(
|
||||
"chat-final-flush",
|
||||
bytes_to_token_ids(b"ok st"),
|
||||
@@ -553,8 +626,7 @@ async fn chat_stream_flushes_held_text_on_finish() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["chat-final-flush".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -665,24 +737,19 @@ async fn chat_stream_reports_decode_failure_as_error_event() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
outputs: vec![request_output(
|
||||
"chat-4",
|
||||
vec![UNKNOWN_DECODE_TOKEN_ID],
|
||||
None,
|
||||
None,
|
||||
)],
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output("chat-4", vec![b'i' as u32], None, None)],
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
},
|
||||
);
|
||||
|
||||
let backend: Arc<dyn ChatTextBackend> =
|
||||
Arc::new(FakeChatBackend::new().with_tokenizer(Arc::new(TestTokenizer::new())));
|
||||
let backend: Arc<dyn ChatTextBackend> = Arc::new(FailingDecodeBackend {
|
||||
inner: FakeChatBackend::new(),
|
||||
});
|
||||
let chat = connect_chat_llm_with_ipc(
|
||||
EngineCoreClientConfig::new_single(handshake_address),
|
||||
&ipc,
|
||||
@@ -702,10 +769,7 @@ 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,
|
||||
format!("test tokenizer cannot decode unknown token id {UNKNOWN_DECODE_TOKEN_ID}")
|
||||
);
|
||||
assert_eq!(message, "decode failed");
|
||||
}
|
||||
other => panic!("unexpected event after close: {other:?}"),
|
||||
}
|
||||
@@ -729,7 +793,7 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output(
|
||||
"chat-include-stop",
|
||||
vec![b'H' as u32, b'i' as u32, b'!' as u32],
|
||||
@@ -738,8 +802,7 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["chat-include-stop".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -817,7 +880,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output(
|
||||
"chat-reasoning",
|
||||
@@ -846,8 +909,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
|
||||
],
|
||||
finished_requests: Some(BTreeSet::from(["chat-reasoning".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -961,7 +1023,7 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output(
|
||||
"chat-collect",
|
||||
bytes_to_token_ids(b"<think>inner</think>outer"),
|
||||
@@ -970,8 +1032,7 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["chat-collect".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1025,7 +1086,7 @@ async fn chat_explicitly_disables_reasoning_parser() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output(
|
||||
"chat-reasoning-disabled",
|
||||
@@ -1056,8 +1117,7 @@ async fn chat_explicitly_disables_reasoning_parser() {
|
||||
"chat-reasoning-disabled".to_string()
|
||||
])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1104,7 +1164,7 @@ async fn chat_stream_parses_tool_calls_automatically() {
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output(
|
||||
"chat-tool",
|
||||
@@ -1129,8 +1189,7 @@ async fn chat_stream_parses_tool_calls_automatically() {
|
||||
],
|
||||
finished_requests: Some(BTreeSet::from(["chat-tool".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1213,7 +1272,7 @@ async fn chat_collect_message_preserves_tool_call_arguments_in_final_only_mode()
|
||||
let _ = recv_engine_message(dealer).await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output(
|
||||
"chat-final-only-tool",
|
||||
@@ -1240,8 +1299,7 @@ async fn chat_collect_message_preserves_tool_call_arguments_in_final_only_mode()
|
||||
"chat-final-only-tool".to_string()
|
||||
])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1292,7 +1350,7 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
|
||||
let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output_with_logprobs(
|
||||
&request.request_id,
|
||||
@@ -1313,8 +1371,7 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
|
||||
],
|
||||
finished_requests: Some(BTreeSet::from([request.request_id])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
//! Roundtrip tests for the real chat-template and output-processor pairing.
|
||||
//! Text-level 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.
|
||||
//! assistant-completion text.
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
@@ -18,10 +18,6 @@ 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)]
|
||||
@@ -195,28 +191,14 @@ 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_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
|
||||
($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
|
||||
paste::paste! {
|
||||
$(
|
||||
$(
|
||||
#[tokio::test]
|
||||
$(#[$fixture_attr])*
|
||||
#[file_serial([<hf_ $case>])]
|
||||
async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
|
||||
[<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
|
||||
@@ -235,9 +217,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.
|
||||
@@ -439,10 +421,10 @@ struct RoundtripResult {
|
||||
parsed_message: AssistantMessage,
|
||||
/// Assistant-completion suffix cut from rendering the expected assistant as
|
||||
/// history.
|
||||
closed_completion: Prompt,
|
||||
closed_completion: String,
|
||||
/// Assistant-completion suffix cut after rendering the parsed assistant
|
||||
/// back as history.
|
||||
rerendered_closed_completion: Prompt,
|
||||
rerendered_closed_completion: String,
|
||||
}
|
||||
|
||||
/// Render, parse, and rerender one assistant turn through the production
|
||||
@@ -454,59 +436,60 @@ async fn run_roundtrip(
|
||||
assistant: AssistantMessage,
|
||||
) -> Result<RoundtripResult> {
|
||||
let renderer = backends.chat_backend.chat_renderer();
|
||||
let rendered = render_closed_completion(renderer.as_ref(), request, &assistant)?;
|
||||
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 parsed_message = parse_completion(case, backends, request, &rendered).await?;
|
||||
let rerendered = render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
|
||||
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)?;
|
||||
|
||||
Ok(RoundtripResult {
|
||||
parsed_message,
|
||||
closed_completion: rendered.completion,
|
||||
rerendered_closed_completion: rerendered.completion,
|
||||
closed_completion: closed_completion_text,
|
||||
rerendered_closed_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<RenderedTurn> {
|
||||
) -> Result<(String, String)> {
|
||||
let mut prompt_request = base_request.clone();
|
||||
prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
|
||||
let prompt = renderer.render(&prompt_request).context("failed to render prompt")?.prompt;
|
||||
let prompt = render_text(renderer, &prompt_request).context("failed to render 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 = renderer.render(&full_request).context("failed to render full prompt")?.prompt;
|
||||
let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
|
||||
|
||||
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:?}"),
|
||||
};
|
||||
ensure!(
|
||||
full.starts_with(&prompt),
|
||||
"full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
|
||||
);
|
||||
let completion = full[prompt.len()..].to_string();
|
||||
|
||||
Ok(RenderedTurn { prompt, completion })
|
||||
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:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed one rendered assistant completion body into the real output processor
|
||||
@@ -515,15 +498,13 @@ async fn parse_completion(
|
||||
case: &RoundtripCase,
|
||||
backends: &vllm_chat::LoadedModelBackends,
|
||||
base_request: &ChatRequest,
|
||||
rendered: &RenderedTurn,
|
||||
prompt: &str,
|
||||
completion_body: &str,
|
||||
) -> Result<AssistantMessage> {
|
||||
let tokenizer = backends.text_backend.tokenizer();
|
||||
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 prompt_token_ids = tokenizer
|
||||
.encode(prompt, base_request.add_special_tokens)
|
||||
.context("failed to encode rendered prompt")?;
|
||||
|
||||
let mut request = base_request.clone();
|
||||
let processor = backends.chat_backend.new_chat_output_processor(
|
||||
@@ -534,12 +515,7 @@ async fn parse_completion(
|
||||
},
|
||||
)?;
|
||||
|
||||
let decoded = decoded_completion_stream(
|
||||
tokenizer.as_ref(),
|
||||
prompt_token_ids,
|
||||
&rendered.completion,
|
||||
case.assistant_stop_suffix,
|
||||
)?;
|
||||
let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
|
||||
let mut events = processor.process(decoded)?;
|
||||
|
||||
while let Some(event) = events.next().await {
|
||||
@@ -562,46 +538,16 @@ 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: &Prompt,
|
||||
assistant_stop_suffix: &str,
|
||||
) -> Result<Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>>> {
|
||||
completion_body: &str,
|
||||
) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
|
||||
let prompt_token_count = prompt_token_ids.len();
|
||||
let mut events = vec![DecodedTextEvent::Start {
|
||||
prompt_token_ids: Arc::from(prompt_token_ids.clone().into_boxed_slice()),
|
||||
prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
|
||||
prompt_logprobs: None,
|
||||
}];
|
||||
|
||||
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,
|
||||
)?
|
||||
}
|
||||
};
|
||||
|
||||
let chunks = split_by_chars(completion_body, 7);
|
||||
if chunks.is_empty() {
|
||||
events.push({
|
||||
DecodedTextEvent::TextDelta {
|
||||
@@ -609,7 +555,11 @@ fn decoded_completion_stream(
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished: Some(Finished {
|
||||
usage: Default::default(),
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: 0,
|
||||
output_token_count: 0,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
@@ -619,26 +569,24 @@ 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: Default::default(),
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count,
|
||||
output_token_count: completion_body.chars().count(),
|
||||
cached_token_count: 0,
|
||||
},
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
});
|
||||
events.push(DecodedTextEvent::TextDelta {
|
||||
delta: chunk.delta,
|
||||
token_ids: chunk.token_ids,
|
||||
delta: chunk,
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finished,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(stream::iter(events).map(Ok).boxed())
|
||||
}
|
||||
|
||||
/// One decoded completion chunk fed into the output processor.
|
||||
struct DecodedCompletionChunk {
|
||||
delta: String,
|
||||
token_ids: Vec<u32>,
|
||||
stream::iter(events).map(Ok).boxed()
|
||||
}
|
||||
|
||||
/// Split text into chunks containing at most `chunk_chars` Unicode scalar
|
||||
@@ -664,49 +612,6 @@ 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>,
|
||||
|
||||
+1
-47
@@ -13,8 +13,8 @@ use std::time::Duration;
|
||||
|
||||
use clap::{Args, Parser, Subcommand};
|
||||
use educe::Educe;
|
||||
use serde::Deserialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use serde_with::{DefaultOnNull, OneOrMany, serde_as};
|
||||
use thiserror_ext::AsReport as _;
|
||||
@@ -290,15 +290,6 @@ pub struct SharedRuntimeArgs {
|
||||
#[serde(default)]
|
||||
pub ssl_ciphers: Option<String>,
|
||||
|
||||
/// Profiler configuration forwarded by the Python supervisor.
|
||||
///
|
||||
/// When set with a non-null `profiler` type, the Rust frontend registers
|
||||
/// the `/start_profile` and `/stop_profile` routes and forwards calls to
|
||||
/// the engine via the `"profile"` utility RPC.
|
||||
#[arg(long, value_parser = parse_json::<ProfilerConfig>, value_name = "JSON")]
|
||||
#[serde(default)]
|
||||
pub profiler_config: Option<ProfilerConfig>,
|
||||
|
||||
/// Unsupported Python vLLM frontend arguments recognized but not yet
|
||||
/// implemented in Rust.
|
||||
#[educe(Debug(ignore))]
|
||||
@@ -326,20 +317,6 @@ impl SharedRuntimeArgs {
|
||||
.map_or(DEFAULT_KEEP_ALIVE_TIMEOUT, Duration::from_secs)
|
||||
}
|
||||
|
||||
/// Return the configured profiler mode, when profiling is enabled.
|
||||
pub fn profiler(&self) -> Option<String> {
|
||||
self.profiler_config.as_ref().and_then(|c| c.profiler.clone())
|
||||
}
|
||||
|
||||
/// Return the profiler config JSON for managed Python engine forwarding.
|
||||
pub fn profiler_config_json(&self) -> Option<String> {
|
||||
self.profiler_config
|
||||
.as_ref()
|
||||
.map(serde_json::to_string)
|
||||
.transpose()
|
||||
.expect("profiler config serialization should not fail")
|
||||
}
|
||||
|
||||
/// Apply fallback logic for API key configuration from env variables.
|
||||
fn apply_env_api_key_fallback(&mut self) {
|
||||
if self.api_key.is_empty()
|
||||
@@ -368,7 +345,6 @@ impl SharedRuntimeArgs {
|
||||
let api_server_options = self.api_server_options();
|
||||
let cors = self.cors_config();
|
||||
let tls = self.tls_config();
|
||||
let profiler = self.profiler();
|
||||
|
||||
Config {
|
||||
transport_mode: TransportMode::Bootstrapped {
|
||||
@@ -401,7 +377,6 @@ impl SharedRuntimeArgs {
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
keep_alive_timeout,
|
||||
profiler,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -422,7 +397,6 @@ impl SharedRuntimeArgs {
|
||||
let api_server_options = self.api_server_options();
|
||||
let cors = self.cors_config();
|
||||
let tls = self.tls_config();
|
||||
let profiler = self.profiler();
|
||||
|
||||
Config {
|
||||
transport_mode: TransportMode::HandshakeOwner {
|
||||
@@ -453,7 +427,6 @@ impl SharedRuntimeArgs {
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
keep_alive_timeout,
|
||||
profiler,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -504,23 +477,6 @@ fn default_py_bootstrap_parser_selection() -> ParserSelection {
|
||||
ParserSelection::None
|
||||
}
|
||||
|
||||
/// Minimal profiler configuration parsed from `--profiler-config`.
|
||||
///
|
||||
/// Only the `profiler` field is inspected by the Rust frontend to decide
|
||||
/// whether to register the `/start_profile` and `/stop_profile` routes.
|
||||
/// All other fields are accepted but ignored — they are consumed by the
|
||||
/// Python engine layer.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize, Serialize)]
|
||||
pub struct ProfilerConfig {
|
||||
/// Profiler backend type (e.g. `"torch"`, `"cuda"`). When `null` or
|
||||
/// absent, profiling is disabled.
|
||||
#[serde(default)]
|
||||
pub profiler: Option<String>,
|
||||
/// Additional Python profiler config fields consumed by the engine layer.
|
||||
#[serde(flatten)]
|
||||
pub extra: serde_json::Map<String, Value>,
|
||||
}
|
||||
|
||||
fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
|
||||
serde_json::from_str(value).map_err(|e| format!("invalid JSON object: {}", e.as_report()))
|
||||
}
|
||||
@@ -647,13 +603,11 @@ impl ServeArgs {
|
||||
pub fn to_managed_engine_config(&self, handshake_port: u16) -> ManagedEngineConfig {
|
||||
let reasoning_parser =
|
||||
effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model);
|
||||
let profiler_config = self.runtime.profiler_config_json();
|
||||
|
||||
self.managed_engine.clone().into_config(
|
||||
self.runtime.model.clone(),
|
||||
self.runtime.max_model_len,
|
||||
self.runtime.max_logprobs,
|
||||
profiler_config,
|
||||
reasoning_parser.as_deref(),
|
||||
self.runtime.language_model_only,
|
||||
self.runtime.disable_log_stats,
|
||||
|
||||
@@ -71,7 +71,6 @@ fn serve_args_forward_python_flags_with_separator() {
|
||||
ssl_ca_certs: None,
|
||||
ssl_cert_reqs: 0,
|
||||
ssl_ciphers: None,
|
||||
profiler_config: None,
|
||||
},
|
||||
managed_engine: ManagedEngineArgs {
|
||||
python: "../vllm/.venv/bin/python",
|
||||
@@ -181,40 +180,6 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() {
|
||||
.assert_debug_eq(&config.python_args);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_forward_profiler_config_to_managed_engine() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"serve",
|
||||
"Qwen/Qwen3-0.6B",
|
||||
"--profiler-config",
|
||||
r#"{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Serve(args) = cli.command else {
|
||||
panic!("expected serve args");
|
||||
};
|
||||
assert_eq!(args.runtime.profiler().as_deref(), Some("torch"));
|
||||
|
||||
let config = args.to_managed_engine_config(5555);
|
||||
let profiler_flag_index = config
|
||||
.python_args
|
||||
.iter()
|
||||
.position(|arg| arg == "--profiler-config")
|
||||
.expect("profiler config flag");
|
||||
let profiler_config: serde_json::Value =
|
||||
serde_json::from_str(&config.python_args[profiler_flag_index + 1])
|
||||
.expect("profiler config json");
|
||||
assert_eq!(
|
||||
profiler_config,
|
||||
serde_json::json!({
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "/tmp/profile",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
|
||||
let cli = Cli::try_parse_from([
|
||||
@@ -767,7 +732,6 @@ fn frontend_args_accept_json() {
|
||||
ssl_ca_certs: None,
|
||||
ssl_cert_reqs: 0,
|
||||
ssl_ciphers: None,
|
||||
profiler_config: None,
|
||||
},
|
||||
},
|
||||
),
|
||||
@@ -1289,7 +1253,6 @@ fn serve_args_accept_handshake_aliases() {
|
||||
ssl_ca_certs: None,
|
||||
ssl_cert_reqs: 0,
|
||||
ssl_ciphers: None,
|
||||
profiler_config: None,
|
||||
},
|
||||
managed_engine: ManagedEngineArgs {
|
||||
python: "python3",
|
||||
@@ -1429,7 +1392,6 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
keep_alive_timeout: 5s,
|
||||
profiler: None,
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&Config {
|
||||
@@ -1513,7 +1475,6 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
keep_alive_timeout: 5s,
|
||||
profiler: None,
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&config);
|
||||
@@ -1615,7 +1576,6 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
keep_alive_timeout: 5s,
|
||||
profiler: None,
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&config);
|
||||
@@ -1644,75 +1604,3 @@ fn serve_frontend_config_uses_unix_listener_when_uds_is_present() {
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_enables_profiling_when_profiler_config_set() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":"torch","torch_profiler_dir":"/tmp/profile"}}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
assert_eq!(args.runtime.profiler().as_deref(), Some("torch"));
|
||||
let config = args.into_config();
|
||||
assert_eq!(config.profiler.as_deref(), Some("torch"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_disables_profiling_when_profiler_config_absent() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
assert_eq!(args.runtime.profiler(), None);
|
||||
let config = args.into_config();
|
||||
assert_eq!(config.profiler, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_disables_profiling_when_profiler_type_is_null() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","profiler_config":{"profiler":null}}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
assert_eq!(args.runtime.profiler(), None);
|
||||
let config = args.into_config();
|
||||
assert_eq!(config.profiler, None);
|
||||
}
|
||||
|
||||
@@ -5,9 +5,9 @@ use clap::Parser;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use vllm_engine_core_client::protocol::output::EngineCoreFinishReason;
|
||||
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
|
||||
};
|
||||
use vllm_engine_core_client::{
|
||||
EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
|
||||
};
|
||||
|
||||
@@ -11,11 +11,10 @@ use tracing::{debug, info, trace};
|
||||
use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop};
|
||||
use crate::coordinator::CoordinatorHandle;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
use crate::protocol::handshake::EngineCoreReadyResponse;
|
||||
use crate::protocol::lora::LoraRequest;
|
||||
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
|
||||
use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
|
||||
use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
|
||||
use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
|
||||
use crate::transport::{self, ConnectedEngine};
|
||||
|
||||
@@ -738,18 +737,6 @@ impl EngineCoreClient {
|
||||
self.call_utility_consensus("is_scheduler_paused", ()).await
|
||||
}
|
||||
|
||||
/// Start profiling the engine.
|
||||
pub async fn start_profile(&self, profile_prefix: Option<&str>) -> Result<()> {
|
||||
self.call_utility::<(), _>("profile", (true, profile_prefix)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop profiling the engine.
|
||||
pub async fn stop_profile(&self, profile_prefix: Option<&str>) -> Result<()> {
|
||||
self.call_utility::<(), _>("profile", (false, profile_prefix)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Shut down local client tasks and close transport state.
|
||||
pub async fn shutdown(self) -> Result<()> {
|
||||
let Self {
|
||||
|
||||
@@ -16,11 +16,12 @@ use crate::client::stream::EngineCoreStreamOutput;
|
||||
use crate::client::{AbortCause, AbortRequest};
|
||||
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
|
||||
use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
|
||||
use crate::protocol::encode_msgpack;
|
||||
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
|
||||
use crate::protocol::request::EngineCoreRequestType;
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
use crate::protocol::{
|
||||
ClassifiedEngineCoreOutputs, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequestType,
|
||||
encode_msgpack,
|
||||
};
|
||||
use crate::transport::{ConnectedEngine, EngineId};
|
||||
use crate::{Error, Result, transport};
|
||||
|
||||
@@ -351,8 +352,8 @@ pub(crate) async fn run_output_dispatcher_loop(
|
||||
)),
|
||||
}?;
|
||||
|
||||
match outputs {
|
||||
EngineCoreOutputs::RequestBatch(batch) => {
|
||||
match outputs.classify() {
|
||||
ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
|
||||
let senders = inner.take_senders_for_outputs(&batch.outputs);
|
||||
for (output, sender) in batch.outputs.into_iter().zip(senders) {
|
||||
let request_id = output.request_id.clone();
|
||||
@@ -403,7 +404,7 @@ pub(crate) async fn run_output_dispatcher_loop(
|
||||
let (running, waiting) = inner.lora_adapter_states();
|
||||
lora_info.update(&METRICS.scheduler, running, waiting);
|
||||
}
|
||||
EngineCoreOutputs::Utility(utility) => {
|
||||
ClassifiedEngineCoreOutputs::Utility(utility) => {
|
||||
let call_id = utility.output.call_id;
|
||||
if inner.resolve_utility_output(utility.output) {
|
||||
trace!(
|
||||
@@ -419,7 +420,8 @@ pub(crate) async fn run_output_dispatcher_loop(
|
||||
);
|
||||
}
|
||||
}
|
||||
other => {
|
||||
other @ (ClassifiedEngineCoreOutputs::DpControl { .. }
|
||||
| ClassifiedEngineCoreOutputs::Other(_)) => {
|
||||
Err::<(), _>(unexpected_dispatcher_output!(
|
||||
"received unexpected output on main dispatcher path: {other:?}"
|
||||
))?;
|
||||
|
||||
@@ -7,9 +7,9 @@ use tracing::trace;
|
||||
use crate::EngineId;
|
||||
use crate::client::stream::EngineCoreStreamOutput;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
use crate::protocol::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::transport::ConnectedEngine;
|
||||
|
||||
pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
|
||||
@@ -452,7 +452,7 @@ mod tests {
|
||||
EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry,
|
||||
};
|
||||
use crate::mock_engine::default_ready_response;
|
||||
use crate::protocol::output::{
|
||||
use crate::protocol::{
|
||||
EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput,
|
||||
};
|
||||
use crate::transport::ConnectedEngine;
|
||||
|
||||
@@ -10,7 +10,7 @@ use tracing::{debug, error, warn};
|
||||
|
||||
use crate::client::AbortRequest;
|
||||
use crate::client::state::OutputReceiver;
|
||||
use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
|
||||
use crate::{AbortCause, Error, Result};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
|
||||
@@ -10,9 +10,10 @@ use zeromq::{XPubSocket, ZmqMessage};
|
||||
use crate::client::imp::ClientInner;
|
||||
use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState};
|
||||
use crate::error::{Error, Result, bail_unexpected_coordinator_output};
|
||||
use crate::protocol::encode_msgpack;
|
||||
use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs};
|
||||
use crate::protocol::request::EngineCoreRequestType;
|
||||
use crate::protocol::{
|
||||
ClassifiedEngineCoreOutputs, DpControlMessage, EngineCoreOutputs, EngineCoreRequestType,
|
||||
encode_msgpack,
|
||||
};
|
||||
|
||||
/// Coordinator-to-engine `START_DP_WAVE` control payload encoded on the
|
||||
/// engine-facing coordinator socket.
|
||||
@@ -109,19 +110,19 @@ impl InProcCoordinatorRunner {
|
||||
/// Apply one engine-originated control output to the coordinator state
|
||||
/// machine.
|
||||
async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> {
|
||||
match outputs {
|
||||
EngineCoreOutputs::RequestBatch(batch)
|
||||
match outputs.classify() {
|
||||
ClassifiedEngineCoreOutputs::RequestBatch(batch)
|
||||
if batch.outputs.is_empty() && batch.finished_requests.is_none() =>
|
||||
{
|
||||
// Stats-only output for coordinator.
|
||||
// Ignore since the Rust coordinator doesn't track stats for
|
||||
// routing decisions.
|
||||
}
|
||||
EngineCoreOutputs::DpControl(DpControlOutput {
|
||||
ClassifiedEngineCoreOutputs::DpControl {
|
||||
engine_index,
|
||||
control,
|
||||
..
|
||||
}) => match control {
|
||||
} => match control {
|
||||
// The engines signals they completed the current wave and are now paused.
|
||||
// Advance the current wave and mark the state as paused.
|
||||
DpControlMessage::WaveComplete(wave) => {
|
||||
|
||||
@@ -25,8 +25,6 @@ pub enum Error {
|
||||
ValueDecode(#[from] rmpv::decode::Error),
|
||||
#[error("messagepack ext value decode failed: {message}")]
|
||||
ExtValueDecode { message: String },
|
||||
#[error("invalid structured outputs params: {message}")]
|
||||
InvalidStructuredOutputsParams { message: String },
|
||||
#[error("io error")]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("transport error")]
|
||||
|
||||
@@ -8,9 +8,8 @@ use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
|
||||
|
||||
use crate::EngineId;
|
||||
use crate::error::{Error, Result, bail_unexpected_handshake_message};
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
|
||||
|
||||
/// Default model length advertised by reusable mock engine helpers.
|
||||
pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
|
||||
|
||||
@@ -0,0 +1,252 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use enum_as_inner::EnumAsInner;
|
||||
|
||||
use super::utility::UtilityOutput;
|
||||
use super::{EngineCoreOutput, EngineCoreOutputs};
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
|
||||
/// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum DpControlMessage {
|
||||
WaveComplete(u32),
|
||||
StartWave(u32),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct RequestBatchOutputs {
|
||||
pub engine_index: u32,
|
||||
pub outputs: Vec<EngineCoreOutput>,
|
||||
pub scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
pub timestamp: f64,
|
||||
pub finished_requests: Option<BTreeSet<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct UtilityCallOutput {
|
||||
pub engine_index: u32,
|
||||
pub timestamp: f64,
|
||||
pub output: UtilityOutput,
|
||||
}
|
||||
|
||||
/// Semantic classification of a raw `EngineCoreOutputs` message.
|
||||
///
|
||||
/// Python currently uses one product-shaped wire struct for several distinct
|
||||
/// output families. This enum exposes those families more explicitly without
|
||||
/// changing the wire format.
|
||||
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
|
||||
pub enum ClassifiedEngineCoreOutputs {
|
||||
RequestBatch(RequestBatchOutputs),
|
||||
Utility(UtilityCallOutput),
|
||||
DpControl {
|
||||
engine_index: u32,
|
||||
timestamp: f64,
|
||||
control: DpControlMessage,
|
||||
},
|
||||
/// Fallback for wire-shape combinations that do not map cleanly onto the
|
||||
/// current semantic families.
|
||||
Other(EngineCoreOutputs),
|
||||
}
|
||||
|
||||
impl EngineCoreOutputs {
|
||||
/// Classify the raw wire message into a more semantic Rust enum.
|
||||
pub fn classify(self) -> ClassifiedEngineCoreOutputs {
|
||||
let has_request_payload = !self.outputs.is_empty()
|
||||
|| self.scheduler_stats.is_some()
|
||||
|| self.finished_requests.is_some();
|
||||
|
||||
match (
|
||||
has_request_payload,
|
||||
&self.utility_output,
|
||||
&self.wave_complete,
|
||||
&self.start_wave,
|
||||
) {
|
||||
(true, None, None, None) => {
|
||||
ClassifiedEngineCoreOutputs::RequestBatch(RequestBatchOutputs {
|
||||
engine_index: self.engine_index,
|
||||
outputs: self.outputs,
|
||||
scheduler_stats: self.scheduler_stats,
|
||||
timestamp: self.timestamp,
|
||||
finished_requests: self.finished_requests,
|
||||
})
|
||||
}
|
||||
(false, Some(_), None, None) => {
|
||||
ClassifiedEngineCoreOutputs::Utility(UtilityCallOutput {
|
||||
engine_index: self.engine_index,
|
||||
timestamp: self.timestamp,
|
||||
output: self.utility_output.unwrap(),
|
||||
})
|
||||
}
|
||||
(false, None, Some(_), None) => ClassifiedEngineCoreOutputs::DpControl {
|
||||
engine_index: self.engine_index,
|
||||
timestamp: self.timestamp,
|
||||
control: DpControlMessage::WaveComplete(self.wave_complete.unwrap()),
|
||||
},
|
||||
(false, None, None, Some(_)) => ClassifiedEngineCoreOutputs::DpControl {
|
||||
engine_index: self.engine_index,
|
||||
timestamp: self.timestamp,
|
||||
control: DpControlMessage::StartWave(self.start_wave.unwrap()),
|
||||
},
|
||||
_ => ClassifiedEngineCoreOutputs::Other(self),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::EngineCoreOutput;
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_request_batch() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![7],
|
||||
..Default::default()
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
RequestBatch(
|
||||
RequestBatchOutputs {
|
||||
engine_index: 0,
|
||||
outputs: [
|
||||
EngineCoreOutput {
|
||||
request_id: "req-1",
|
||||
new_token_ids: [
|
||||
7,
|
||||
],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: None,
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
},
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
finished_requests: Some(
|
||||
{
|
||||
"req-1",
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&outputs.classify());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_utility() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 42_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
Utility(
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
call_id: 42,
|
||||
failure_message: None,
|
||||
result: None,
|
||||
},
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&outputs.classify());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_control() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
start_wave: Some(3),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
DpControl {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
control: StartWave(
|
||||
3,
|
||||
),
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&outputs.classify());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_mixed_shape_as_raw() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![7],
|
||||
..Default::default()
|
||||
}],
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 1_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
Other(
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: [
|
||||
EngineCoreOutput {
|
||||
request_id: "req-1",
|
||||
new_token_ids: [
|
||||
7,
|
||||
],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: None,
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
},
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
utility_output: Some(
|
||||
UtilityOutput {
|
||||
call_id: 1,
|
||||
failure_message: None,
|
||||
result: None,
|
||||
},
|
||||
),
|
||||
finished_requests: None,
|
||||
wave_complete: None,
|
||||
start_wave: None,
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&outputs.classify());
|
||||
}
|
||||
}
|
||||
@@ -2,8 +2,7 @@ use std::collections::BTreeMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::protocol::OpaqueValue;
|
||||
use crate::protocol::dtype::ModelDtype;
|
||||
use crate::protocol::{ModelDtype, OpaqueValue};
|
||||
|
||||
/// Decoded engine startup-handshake payload sent on the handshake socket.
|
||||
///
|
||||
|
||||
@@ -9,7 +9,8 @@ use enum_as_inner::EnumAsInner;
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
|
||||
use self::wire::*;
|
||||
use crate::error::{Error, Result, bail_ext_value_decode};
|
||||
use super::{EngineCoreOutput, EngineCoreOutputs, decode_msgpack};
|
||||
use crate::error::{Error, Result, bail_ext_value_decode, ext_value_decode};
|
||||
use crate::protocol::tensor::{WireArrayData, WireNdArray};
|
||||
|
||||
/// One token candidate and its logprob metadata for a single sequence position.
|
||||
@@ -159,7 +160,7 @@ impl Serialize for MaybeWireLogprobs {
|
||||
impl MaybeWireLogprobs {
|
||||
/// Resolve the wire representation into decoded logprobs by looking up aux
|
||||
/// frames and decoding raw views as needed.
|
||||
pub(super) fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
|
||||
fn resolve<Frame>(self, frames: &[Frame], field_prefix: &str) -> Result<Self>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
@@ -170,6 +171,37 @@ impl MaybeWireLogprobs {
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineCoreOutputs {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
for output in &mut self.outputs {
|
||||
output.resolve_in_place(frames)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
self.new_logprobs = (self.new_logprobs.take())
|
||||
.map(|value| value.resolve(frames, "new_logprobs"))
|
||||
.transpose()?;
|
||||
self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
|
||||
.map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl WireLogprobs {
|
||||
/// Convert semantic per-position logprobs into the Python wire tuple shape.
|
||||
///
|
||||
@@ -283,3 +315,16 @@ impl WireLogprobs {
|
||||
Ok(Logprobs { positions })
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one ordinary or multipart engine-core output message into the strong
|
||||
/// typed public protocol shape.
|
||||
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
|
||||
|
||||
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
|
||||
outputs.resolve_in_place(frames)?;
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
@@ -3,8 +3,8 @@ use std::collections::BTreeSet;
|
||||
use bytes::Bytes;
|
||||
use rmpv::Value;
|
||||
|
||||
use super::{Logprobs, PositionLogprobs, TokenLogprob};
|
||||
use crate::protocol::output::{EngineCoreFinishReason, decode_engine_core_outputs};
|
||||
use super::{Logprobs, PositionLogprobs, TokenLogprob, decode_engine_core_outputs};
|
||||
use crate::protocol::EngineCoreFinishReason;
|
||||
|
||||
fn encode_value(value: &Value) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
@@ -183,7 +183,7 @@ fn decodes_inline_new_logprobs() {
|
||||
Some(inline_logprobs_value()),
|
||||
None,
|
||||
)))];
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap();
|
||||
|
||||
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
|
||||
assert_eq!(logprobs, expected_sample_logprobs());
|
||||
@@ -214,7 +214,7 @@ fn decodes_multipart_new_logprobs() {
|
||||
]),
|
||||
Bytes::from_static(&[1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0]),
|
||||
];
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap();
|
||||
|
||||
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
|
||||
assert_eq!(logprobs, expected_sample_logprobs());
|
||||
@@ -226,7 +226,7 @@ fn decodes_inline_prompt_logprobs() {
|
||||
None,
|
||||
Some(inline_prompt_logprobs_value()),
|
||||
)))];
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap();
|
||||
|
||||
let logprobs = decoded.outputs[0]
|
||||
.new_prompt_logprobs_tensors
|
||||
@@ -252,7 +252,7 @@ fn decodes_big_endian_payloads() {
|
||||
])),
|
||||
None,
|
||||
)))];
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
|
||||
let decoded = decode_engine_core_outputs(&frames).unwrap();
|
||||
let logprobs = decoded.outputs[0].new_logprobs.clone().unwrap().into_direct().unwrap();
|
||||
assert_eq!(
|
||||
logprobs,
|
||||
|
||||
@@ -1,11 +1,28 @@
|
||||
use std::any::type_name;
|
||||
use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::Bytes;
|
||||
use rmpv::Value;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_repr::{Deserialize_repr, Serialize_repr};
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::protocol::logprobs::MaybeWireLogprobs;
|
||||
use crate::protocol::multimodal::MmFeatures;
|
||||
use crate::protocol::stats::{PrefillStats, SchedulerStats};
|
||||
use crate::protocol::utility::UtilityOutput;
|
||||
|
||||
// TODO: This module currently mixes reusable frontend-facing semantic types
|
||||
// (for example `FinishReason`, `StopReason`, `RequestOutputKind`, and future
|
||||
// cleaned-up frontend sampling types) with engine-core-specific wire DTOs and
|
||||
// handshake/control messages. While the Rust frontend is still evolving
|
||||
// quickly, keep them co-located here for iteration speed. Once the higher-level
|
||||
// API boundary stabilizes, move the truly reusable semantic types into a
|
||||
// lower-level common crate and keep the engine transport/wire messages here.
|
||||
|
||||
/// Dynamic msgpack value used for schema positions that are preserved but not
|
||||
/// yet strongly typed in the early-stage Rust client.
|
||||
@@ -19,18 +36,499 @@ fn is_false(v: &bool) -> bool {
|
||||
!v
|
||||
}
|
||||
|
||||
fn default_top_p() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repetition_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_temperature() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
16
|
||||
}
|
||||
|
||||
mod classified_outputs;
|
||||
pub mod dtype;
|
||||
pub mod handshake;
|
||||
pub mod logprobs;
|
||||
pub mod lora;
|
||||
pub mod multimodal;
|
||||
pub mod output;
|
||||
pub mod request;
|
||||
pub mod sampling;
|
||||
pub mod stats;
|
||||
pub mod structured_outputs;
|
||||
pub mod tensor;
|
||||
pub mod utility;
|
||||
pub use classified_outputs::{
|
||||
ClassifiedEngineCoreOutputs, DpControlMessage, RequestBatchOutputs, UtilityCallOutput,
|
||||
};
|
||||
pub use dtype::ModelDtype;
|
||||
pub use logprobs::decode_engine_core_outputs;
|
||||
|
||||
/// Request types are encoded as single-byte protocol constants so they can be
|
||||
/// sent over the ZMQ socket without an extra encoding step.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreRequestType {
|
||||
Add = 0,
|
||||
Abort = 1,
|
||||
StartDpWave = 2,
|
||||
Utility = 3,
|
||||
}
|
||||
|
||||
impl EngineCoreRequestType {
|
||||
/// Decode the single-byte request type frame used on the engine input
|
||||
/// socket. Returns `None` for unrecognized values.
|
||||
pub fn from_frame(frame: &[u8]) -> Option<Self> {
|
||||
let [value] = frame else {
|
||||
return None;
|
||||
};
|
||||
|
||||
match value {
|
||||
0 => Some(Self::Add),
|
||||
1 => Some(Self::Abort),
|
||||
2 => Some(Self::StartDpWave),
|
||||
3 => Some(Self::Utility),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode the request type as the single-byte frame used on the engine
|
||||
/// input socket.
|
||||
pub fn to_frame(self) -> Bytes {
|
||||
Bytes::from_static(match self {
|
||||
Self::Add => b"\x00",
|
||||
Self::Abort => b"\x01",
|
||||
Self::StartDpWave => b"\x02",
|
||||
Self::Utility => b"\x03",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Reason a request finished: stop, length, abort, error, or repetition.
|
||||
///
|
||||
/// This mirrors the Python enum and uses integer encoding for compact wire
|
||||
/// representation.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreFinishReason {
|
||||
/// A stop string was emitted.
|
||||
Stop = 0,
|
||||
/// `max_tokens` or `max_model_len` was reached.
|
||||
Length = 1,
|
||||
/// The request was aborted by the client.
|
||||
Abort = 2,
|
||||
/// A retryable request-level internal error occurred.
|
||||
Error = 3,
|
||||
/// A repetitive token pattern was detected.
|
||||
Repetition = 4,
|
||||
}
|
||||
|
||||
/// Event types emitted by engine-core for one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreEventType {
|
||||
Queued = 1,
|
||||
Scheduled = 2,
|
||||
Preempted = 3,
|
||||
}
|
||||
|
||||
/// A timestamped engine-core event associated with one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EngineCoreEvent {
|
||||
pub r#type: EngineCoreEventType,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
/// Controls how intermediate outputs are returned to the frontend.
|
||||
///
|
||||
/// `Cumulative = 0` is intentionally not supported in Rust frontend.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L146-L152>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum RequestOutputKind {
|
||||
/// Return only token deltas in each update.
|
||||
#[default]
|
||||
Delta = 1,
|
||||
/// Suppress intermediate updates and return only the final output.
|
||||
FinalOnly = 2,
|
||||
}
|
||||
|
||||
/// Structured-output backend selected for EngineCore grammar compilation.
|
||||
///
|
||||
/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
|
||||
/// validation. The Rust frontend currently always lowers structured-output
|
||||
/// requests to guidance, while ignoring any user-supplied `_backend` value.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum StructuredOutputBackend {
|
||||
Xgrammar,
|
||||
#[default]
|
||||
Guidance,
|
||||
Outlines,
|
||||
LmFormatEnforcer,
|
||||
}
|
||||
|
||||
/// The stop reason associated with a finished output.
|
||||
///
|
||||
/// Python models this as the union-typed `stop_reason: int | str | None`
|
||||
/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
|
||||
///
|
||||
/// Original Python field:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum StopReason {
|
||||
TokenId(u32),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
///
|
||||
/// Exactly one constraint field (`json`, `regex`, `choice`, `grammar`,
|
||||
/// `json_object`, or `structural_tag`) should be set. The engine-core
|
||||
/// backend selects the appropriate grammar compiler based on which field
|
||||
/// is present.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct StructuredOutputsParams {
|
||||
/// JSON schema (as a dict/object or JSON string) constraining the output.
|
||||
pub json: Option<serde_json::Value>,
|
||||
/// Regular expression the output must match.
|
||||
pub regex: Option<String>,
|
||||
/// List of allowed output strings (the model must produce one of these).
|
||||
pub choice: Option<Vec<String>>,
|
||||
/// Context-free grammar (in EBNF-like notation) the output must conform to.
|
||||
pub grammar: Option<String>,
|
||||
/// When `true`, output must be valid JSON (free-form, no schema).
|
||||
pub json_object: Option<bool>,
|
||||
/// Disable any additional whitespace in guided JSON output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_any_whitespace: bool,
|
||||
/// Disable `additionalProperties` in JSON schema output.
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
pub disable_additional_properties: bool,
|
||||
/// Custom whitespace pattern for guided JSON output.
|
||||
pub whitespace_pattern: Option<String>,
|
||||
/// Structural tag configuration (JSON-encoded string).
|
||||
pub structural_tag: Option<String>,
|
||||
/// Structured-output backend, mirroring Python's internal `_backend`.
|
||||
///
|
||||
/// User-supplied values are ignored during deserialization. This matches
|
||||
/// Python's request boundary, where `_backend` is set by validation rather
|
||||
/// than accepted as a request-level backend selector.
|
||||
#[serde(
|
||||
default,
|
||||
rename = "_backend",
|
||||
deserialize_with = "serde_with::rust::deserialize_ignore_any"
|
||||
)]
|
||||
pub backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
/// Engine-core-facing sampling parameters for text generation.
|
||||
///
|
||||
/// This is the normalized southbound subset used by the Rust frontend when it
|
||||
/// talks to Python engine-core over the wire. User-facing request semantics
|
||||
/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
|
||||
/// intentionally handled by higher layers before values reach this DTO.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
|
||||
// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
|
||||
// default-valued keys; default the whole struct. Per-field fns cover the
|
||||
// non-zero defaults.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
|
||||
#[serde(default)]
|
||||
pub struct EngineCoreSamplingParams {
|
||||
/// Controls randomness. Lower values are more deterministic; zero means
|
||||
/// greedy sampling.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: f32,
|
||||
/// Cumulative probability threshold for nucleus sampling.
|
||||
#[serde(default = "default_top_p")]
|
||||
pub top_p: f32,
|
||||
/// Maximum number of top tokens to consider. `0` means all tokens.
|
||||
pub top_k: u32,
|
||||
/// Random seed used by the sampler when present.
|
||||
pub seed: Option<i64>,
|
||||
/// Maximum number of tokens to generate per output sequence.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: u32,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` means unlimited; the
|
||||
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
|
||||
/// reaching this DTO, so only non-negative values are sent. Enforced
|
||||
/// engine-side (and only when a reasoning parser is configured).
|
||||
pub thinking_token_budget: Option<u64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
pub logprobs: Option<i32>,
|
||||
/// Number of log probabilities to return per prompt token.
|
||||
///
|
||||
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
|
||||
pub prompt_logprobs: Option<i32>,
|
||||
/// Minimum probability threshold for token sampling.
|
||||
pub min_p: f32,
|
||||
/// Frequency penalty applied by the sampler.
|
||||
pub frequency_penalty: f32,
|
||||
/// Presence penalty applied by the sampler.
|
||||
pub presence_penalty: f32,
|
||||
/// Repetition penalty applied by the sampler.
|
||||
#[serde(default = "default_repetition_penalty")]
|
||||
pub repetition_penalty: f32,
|
||||
/// Token IDs that stop generation.
|
||||
pub stop_token_ids: Vec<u32>,
|
||||
/// Primary EOS token ID used by engine-core's dedicated EOS stop path.
|
||||
///
|
||||
/// This mirrors Python's internal `_eos_token_id` field and is derived by
|
||||
/// the frontend from tokenizer/model metadata rather than supplied directly
|
||||
/// by end users.
|
||||
#[serde(rename = "_eos_token_id")]
|
||||
pub eos_token_id: Option<u32>,
|
||||
/// Complete stop-token set used by engine-core for `min_tokens` masking.
|
||||
///
|
||||
/// This mirrors Python's internal `_all_stop_token_ids` field and should
|
||||
/// contain explicit `stop_token_ids` plus any frontend-derived EOS token
|
||||
/// IDs.
|
||||
#[serde(rename = "_all_stop_token_ids")]
|
||||
pub all_stop_token_ids: BTreeSet<u32>,
|
||||
/// Logit biases to apply during sampling.
|
||||
/// Keys are token IDs
|
||||
pub logit_bias: Option<HashMap<u32, f32>>,
|
||||
/// Restrict output to these token IDs only.
|
||||
pub allowed_token_ids: Option<Vec<u32>>,
|
||||
/// Tokenized bad words to avoid during generation.
|
||||
#[serde(rename = "_bad_words_token_ids")]
|
||||
pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
pub structured_outputs: Option<StructuredOutputsParams>,
|
||||
/// Specific token IDs for which log probabilities should be returned at
|
||||
/// each position.
|
||||
///
|
||||
/// When set, the engine returns logprobs for exactly these tokens in
|
||||
/// addition to the sampled/scored token. Mutually exclusive with the
|
||||
/// `logprobs` count field in practice.
|
||||
pub logprob_token_ids: Option<Vec<u32>>,
|
||||
/// If `Some(true)`, the request will not attempt to read from the prefix
|
||||
/// cache; newly computed blocks may still populate the cache. `None`
|
||||
/// defers to engine-core defaults.
|
||||
pub skip_reading_prefix_cache: Option<bool>,
|
||||
/// Additional request parameters for custom extensions (from `vllm_xargs`).
|
||||
pub extra_args: Option<HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl EngineCoreSamplingParams {
|
||||
/// Constructs a default sampling params for testing purposes only.
|
||||
pub fn for_test() -> Self {
|
||||
Self {
|
||||
temperature: 1.0,
|
||||
top_p: 1.0,
|
||||
top_k: 0,
|
||||
seed: None,
|
||||
max_tokens: 65536,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
repetition_penalty: 1.0,
|
||||
stop_token_ids: Vec::new(),
|
||||
eos_token_id: None,
|
||||
all_stop_token_ids: BTreeSet::new(),
|
||||
logit_bias: None,
|
||||
allowed_token_ids: None,
|
||||
bad_words_token_ids: None,
|
||||
structured_outputs: None,
|
||||
logprob_token_ids: None,
|
||||
skip_reading_prefix_cache: None,
|
||||
extra_args: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extra kwargs consumed by engine-side reasoning parsers.
|
||||
///
|
||||
/// Original Python construction point:
|
||||
/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ReasoningParserKwargs {
|
||||
/// Effective kwargs visible to the chat template for this request.
|
||||
pub chat_template_kwargs: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Engine-core add-request payload sent from frontend to engine.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreRequest {
|
||||
pub request_id: String,
|
||||
pub prompt_token_ids: Option<Vec<u32>>,
|
||||
/// Multimodal features attached to the request.
|
||||
pub mm_features: Option<MmFeatures>,
|
||||
pub sampling_params: Option<EngineCoreSamplingParams>,
|
||||
/// Pooling parameters are preserved in the schema but not yet strongly
|
||||
/// typed.
|
||||
pub pooling_params: Option<OpaqueValue>,
|
||||
pub arrival_time: f64,
|
||||
#[serde(default)]
|
||||
pub lora_request: Option<lora::LoraRequest>,
|
||||
#[serde(default)]
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub data_parallel_rank: Option<u32>,
|
||||
/// Unsupported in the first-stage Rust client because Python uses a custom
|
||||
/// tensor/aux-frame encoding path for this field.
|
||||
#[serde(default)]
|
||||
pub prompt_embeds: Option<OpaqueValue>,
|
||||
/// Per-position mask for mixed-mode inputs (e.g. chat completion with
|
||||
/// `prompt_embeds` content parts). `Some(true)` means real token id;
|
||||
/// `Some(false)` means the position uses a pre-computed entry from
|
||||
/// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
|
||||
#[serde(default)]
|
||||
pub prompt_is_token_ids: Option<Vec<bool>>,
|
||||
/// Index of the client, used to ensure outputs are sent back to the same
|
||||
/// client when scaling out the frontend.
|
||||
#[serde(default)]
|
||||
pub client_index: u32,
|
||||
/// In DP mode, indicates which wave this request is expected to belong to.
|
||||
#[serde(default)]
|
||||
pub current_wave: u32,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<BTreeMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub resumable: bool,
|
||||
/// Original user-provided request ID, used for output reporting and aborts.
|
||||
#[serde(default)]
|
||||
pub external_req_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub reasoning_ended: Option<bool>,
|
||||
/// Reasoning-parser kwargs forwarded from the frontend to the
|
||||
/// structured-output backend.
|
||||
#[serde(default)]
|
||||
pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
|
||||
/// If `true`, the request should be added to the scheduler's waiting queue
|
||||
/// and immediately aborted, so connector-side cleanup runs via the
|
||||
/// standard `request_finished` hook.
|
||||
#[serde(default)]
|
||||
pub abort_immediately: bool,
|
||||
}
|
||||
|
||||
impl EngineCoreRequest {
|
||||
/// Validate fields intentionally not supported in the first-stage client.
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.prompt_embeds.is_some() {
|
||||
return Err(Error::UnsupportedField {
|
||||
context: "EngineCoreRequest",
|
||||
field: "prompt_embeds",
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Engine-core output for a single request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutput {
|
||||
pub request_id: String,
|
||||
pub new_token_ids: Vec<u32>,
|
||||
/// Decoded sample logprobs for the newly generated positions in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_logprobs: Option<MaybeWireLogprobs>,
|
||||
/// Decoded prompt logprobs for the scored prompt positions emitted in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
|
||||
#[serde(default)]
|
||||
pub pooling_output: Option<OpaqueValue>,
|
||||
#[serde(default)]
|
||||
pub finish_reason: Option<EngineCoreFinishReason>,
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<StopReason>,
|
||||
#[serde(default)]
|
||||
pub events: Option<Vec<EngineCoreEvent>>,
|
||||
#[serde(default)]
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<OpaqueValue>,
|
||||
/// Breakdown of the scheduled prefill computation, set on the first output
|
||||
/// of a newly scheduled prefill and elided for subsequent decode outputs.
|
||||
#[serde(default)]
|
||||
pub prefill_stats: Option<PrefillStats>,
|
||||
#[serde(default)]
|
||||
pub routed_experts: Option<OpaqueValue>,
|
||||
/// Number of NaNs seen in logits. Values above zero indicate corruption.
|
||||
#[serde(default)]
|
||||
pub num_nans_in_logits: u32,
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Returns whether this output is terminal for the request.
|
||||
pub fn finished(&self) -> bool {
|
||||
self.finish_reason.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch of engine-core outputs returned to a frontend client.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutputs {
|
||||
#[serde(default)]
|
||||
pub engine_index: u32,
|
||||
/// Outputs grouped for this client in the current engine tick.
|
||||
#[serde(default)]
|
||||
pub outputs: Vec<EngineCoreOutput>,
|
||||
#[serde(default)]
|
||||
pub scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
#[serde(default)]
|
||||
pub timestamp: f64,
|
||||
#[serde(default)]
|
||||
pub utility_output: Option<UtilityOutput>,
|
||||
#[serde(default)]
|
||||
pub finished_requests: Option<BTreeSet<String>>,
|
||||
/// In DP mode, signals that the current wave finished and engines are
|
||||
/// paused.
|
||||
#[serde(default)]
|
||||
pub wave_complete: Option<u32>,
|
||||
/// In DP mode, signals that a request arrived for an old wave and the next
|
||||
/// wave needs to start in other engines.
|
||||
#[serde(default)]
|
||||
pub start_wave: Option<u32>,
|
||||
}
|
||||
|
||||
/// Encode a Rust value into msgpack using the protocol crate's serde model.
|
||||
pub fn encode_msgpack<T>(value: &T) -> Result<Vec<u8>>
|
||||
@@ -66,17 +564,81 @@ where
|
||||
})
|
||||
}
|
||||
|
||||
/// Decode a msgpack payload into a dynamic value for diagnostics and tests.
|
||||
pub fn decode_value(bytes: &[u8]) -> Result<Value> {
|
||||
Ok(rmpv::decode::read_value(&mut Cursor::new(bytes))?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn engine_core_request_serializes_as_full_array() {
|
||||
let request = EngineCoreRequest {
|
||||
request_id: "req-1".to_string(),
|
||||
prompt_token_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: Some(EngineCoreSamplingParams {
|
||||
max_tokens: 8,
|
||||
..EngineCoreSamplingParams::for_test()
|
||||
}),
|
||||
arrival_time: 1234.5,
|
||||
client_index: 7,
|
||||
..EngineCoreRequest::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&request).unwrap();
|
||||
let value = decode_value(&encoded).unwrap();
|
||||
let array = match value {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected array, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(array.len(), 20);
|
||||
assert_eq!(array[0], Value::from("req-1"));
|
||||
assert_eq!(array[2], Value::Nil);
|
||||
assert_eq!(array[4], Value::Nil);
|
||||
assert_eq!(array[10], Value::Nil);
|
||||
assert_eq!(array[11], Value::from(7));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_roundtrip_finished_fields() {
|
||||
let outputs = EngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![42],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(EngineCoreFinishReason::Length),
|
||||
stop_reason: Some(StopReason::Text("stop".to_string())),
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&outputs).unwrap();
|
||||
let decoded: EngineCoreOutputs = decode_msgpack(&encoded).unwrap();
|
||||
|
||||
assert_eq!(decoded.outputs.len(), 1);
|
||||
assert_eq!(
|
||||
decoded.outputs[0].finish_reason,
|
||||
Some(EngineCoreFinishReason::Length)
|
||||
);
|
||||
assert_eq!(
|
||||
decoded.finished_requests,
|
||||
Some(BTreeSet::from(["req-1".to_string()]))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_msgpack_includes_type_name_and_value_fallback() {
|
||||
let error = decode_msgpack::<u64>(
|
||||
@@ -86,4 +648,72 @@ mod tests {
|
||||
|
||||
expect_test::expect![[r#"messagepack decode failed for u64: wrong msgpack marker FixMap(1); value fallback: {"status": "READY"}"#]].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_backend_ignores_deserialized_value() {
|
||||
let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
|
||||
"json_object": true,
|
||||
"_backend": "xgrammar",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
assert_eq!(value["_backend"], "guidance");
|
||||
}
|
||||
|
||||
/// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
|
||||
/// must fall back to defaults. `python_compat` can't catch this since Rust
|
||||
/// encodes full maps (see `engine_core_request_serializes_as_full_array`).
|
||||
#[test]
|
||||
fn decodes_sampling_params_with_omitted_defaults() {
|
||||
let sampling_params = Value::Map(vec![
|
||||
(
|
||||
Value::from("stop_token_ids"),
|
||||
Value::Array(vec![Value::from(151643u32)]),
|
||||
),
|
||||
(Value::from("skip_reading_prefix_cache"), Value::from(false)),
|
||||
]);
|
||||
let request = Value::Array(vec![
|
||||
Value::from("req-omit-defaults"),
|
||||
Value::Array(vec![
|
||||
Value::from(1u32),
|
||||
Value::from(2u32),
|
||||
Value::from(3u32),
|
||||
]),
|
||||
Value::Nil,
|
||||
sampling_params,
|
||||
Value::Nil,
|
||||
Value::from(1.0f64),
|
||||
]);
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
rmpv::encode::write_value(&mut bytes, &request).unwrap();
|
||||
|
||||
let decoded: EngineCoreRequest = decode_msgpack(&bytes)
|
||||
.expect("a real omit_defaults request must decode (regression: missing field)");
|
||||
|
||||
assert_eq!(decoded.request_id, "req-omit-defaults");
|
||||
let sampling = decoded.sampling_params.expect("sampling params present");
|
||||
|
||||
assert_eq!(sampling.stop_token_ids, vec![151643]);
|
||||
assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
|
||||
|
||||
// Omitted fields -> Python defaults.
|
||||
assert_eq!(sampling.temperature, 1.0);
|
||||
assert_eq!(sampling.top_p, 1.0);
|
||||
assert_eq!(sampling.top_k, 0);
|
||||
assert_eq!(sampling.seed, None);
|
||||
assert_eq!(sampling.max_tokens, 16);
|
||||
assert_eq!(sampling.min_tokens, 0);
|
||||
assert_eq!(sampling.min_p, 0.0);
|
||||
assert_eq!(sampling.frequency_penalty, 0.0);
|
||||
assert_eq!(sampling.presence_penalty, 0.0);
|
||||
assert_eq!(sampling.repetition_penalty, 1.0);
|
||||
assert_eq!(sampling.logprobs, None);
|
||||
assert_eq!(sampling.prompt_logprobs, None);
|
||||
assert_eq!(sampling.eos_token_id, None);
|
||||
assert!(sampling.all_stop_token_ids.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,518 +0,0 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use enum_as_inner::EnumAsInner;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_repr::{Deserialize_repr, Serialize_repr};
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
|
||||
use super::utility::UtilityOutput;
|
||||
use crate::error::{Error, Result, ext_value_decode};
|
||||
use crate::protocol::logprobs::MaybeWireLogprobs;
|
||||
use crate::protocol::stats::{PrefillStats, SchedulerStats};
|
||||
use crate::protocol::{OpaqueValue, decode_msgpack};
|
||||
|
||||
/// The stop reason associated with a finished output.
|
||||
///
|
||||
/// Python models this as the union-typed `stop_reason: int | str | None`
|
||||
/// field on `EngineCoreOutput`; the Rust client narrows it into a tagged enum.
|
||||
///
|
||||
/// Original Python field:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L155>
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum StopReason {
|
||||
TokenId(u32),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Reason a request finished: stop, length, abort, error, or repetition.
|
||||
///
|
||||
/// This mirrors the Python enum and uses integer encoding for compact wire
|
||||
/// representation.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L41-L63>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreFinishReason {
|
||||
/// A stop string was emitted.
|
||||
Stop = 0,
|
||||
/// `max_tokens` or `max_model_len` was reached.
|
||||
Length = 1,
|
||||
/// The request was aborted by the client.
|
||||
Abort = 2,
|
||||
/// A retryable request-level internal error occurred.
|
||||
Error = 3,
|
||||
/// A repetitive token pattern was detected.
|
||||
Repetition = 4,
|
||||
}
|
||||
|
||||
/// Event types emitted by engine-core for one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L113-L118>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreEventType {
|
||||
Queued = 1,
|
||||
Scheduled = 2,
|
||||
Preempted = 3,
|
||||
}
|
||||
|
||||
/// A timestamped engine-core event associated with one request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L121-L130>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct EngineCoreEvent {
|
||||
pub r#type: EngineCoreEventType,
|
||||
pub timestamp: f64,
|
||||
}
|
||||
|
||||
/// Engine-core output for a single request.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/d3af8c18317c0dc008d42e4367fbb9045cfb7bf6/vllm/v1/engine/__init__.py#L154-L184>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreOutput {
|
||||
pub request_id: String,
|
||||
pub new_token_ids: Vec<u32>,
|
||||
/// Decoded sample logprobs for the newly generated positions in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_logprobs: Option<MaybeWireLogprobs>,
|
||||
/// Decoded prompt logprobs for the scored prompt positions emitted in this
|
||||
/// output.
|
||||
#[serde(default)]
|
||||
pub new_prompt_logprobs_tensors: Option<MaybeWireLogprobs>,
|
||||
#[serde(default)]
|
||||
pub pooling_output: Option<OpaqueValue>,
|
||||
#[serde(default)]
|
||||
pub finish_reason: Option<EngineCoreFinishReason>,
|
||||
#[serde(default)]
|
||||
pub stop_reason: Option<StopReason>,
|
||||
#[serde(default)]
|
||||
pub events: Option<Vec<EngineCoreEvent>>,
|
||||
#[serde(default)]
|
||||
pub kv_transfer_params: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<OpaqueValue>,
|
||||
/// Breakdown of the scheduled prefill computation, set on the first output
|
||||
/// of a newly scheduled prefill and elided for subsequent decode outputs.
|
||||
#[serde(default)]
|
||||
pub prefill_stats: Option<PrefillStats>,
|
||||
#[serde(default)]
|
||||
pub routed_experts: Option<OpaqueValue>,
|
||||
/// Number of NaNs seen in logits. Values above zero indicate corruption.
|
||||
#[serde(default)]
|
||||
pub num_nans_in_logits: u32,
|
||||
}
|
||||
|
||||
impl EngineCoreOutput {
|
||||
/// Returns whether this output is terminal for the request.
|
||||
pub fn finished(&self) -> bool {
|
||||
self.finish_reason.is_some()
|
||||
}
|
||||
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
self.new_logprobs = (self.new_logprobs.take())
|
||||
.map(|value| value.resolve(frames, "new_logprobs"))
|
||||
.transpose()?;
|
||||
self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take())
|
||||
.map(|value| value.resolve(frames, "new_prompt_logprobs_tensors"))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw Python/msgpack engine-core output envelope.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L186-L214>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
struct WireEngineCoreOutputs {
|
||||
#[serde(default)]
|
||||
engine_index: u32,
|
||||
/// Outputs grouped for this client in the current engine tick.
|
||||
#[serde(default)]
|
||||
outputs: Vec<EngineCoreOutput>,
|
||||
#[serde(default)]
|
||||
scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
#[serde(default)]
|
||||
timestamp: f64,
|
||||
#[serde(default)]
|
||||
utility_output: Option<UtilityOutput>,
|
||||
#[serde(default)]
|
||||
finished_requests: Option<BTreeSet<String>>,
|
||||
/// In DP mode, signals that the current wave finished and engines are
|
||||
/// paused.
|
||||
#[serde(default)]
|
||||
wave_complete: Option<u32>,
|
||||
/// In DP mode, signals that a request arrived for an old wave and the next
|
||||
/// wave needs to start in other engines.
|
||||
#[serde(default)]
|
||||
start_wave: Option<u32>,
|
||||
}
|
||||
|
||||
/// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum DpControlMessage {
|
||||
WaveComplete(u32),
|
||||
StartWave(u32),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default)]
|
||||
pub struct RequestBatchOutputs {
|
||||
pub engine_index: u32,
|
||||
pub outputs: Vec<EngineCoreOutput>,
|
||||
pub scheduler_stats: Option<Box<SchedulerStats>>,
|
||||
pub timestamp: f64,
|
||||
pub finished_requests: Option<BTreeSet<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default)]
|
||||
pub struct UtilityCallOutput {
|
||||
pub engine_index: u32,
|
||||
pub timestamp: f64,
|
||||
pub output: UtilityOutput,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DpControlOutput {
|
||||
pub engine_index: u32,
|
||||
pub timestamp: f64,
|
||||
pub control: DpControlMessage,
|
||||
}
|
||||
|
||||
/// Semantic engine-core output families.
|
||||
///
|
||||
/// Python currently uses one product-shaped wire struct. The Rust protocol
|
||||
/// exposes the finite semantic families while preserving the same msgpack shape
|
||||
/// for serialization.
|
||||
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
|
||||
pub enum EngineCoreOutputs {
|
||||
RequestBatch(RequestBatchOutputs),
|
||||
Utility(UtilityCallOutput),
|
||||
DpControl(DpControlOutput),
|
||||
}
|
||||
|
||||
impl From<RequestBatchOutputs> for EngineCoreOutputs {
|
||||
fn from(outputs: RequestBatchOutputs) -> Self {
|
||||
Self::RequestBatch(outputs)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<UtilityCallOutput> for EngineCoreOutputs {
|
||||
fn from(output: UtilityCallOutput) -> Self {
|
||||
Self::Utility(output)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<DpControlOutput> for EngineCoreOutputs {
|
||||
fn from(output: DpControlOutput) -> Self {
|
||||
Self::DpControl(output)
|
||||
}
|
||||
}
|
||||
|
||||
impl EngineCoreOutputs {
|
||||
/// Resolve all wire-format fields in-place by looking up aux frames and
|
||||
/// decoding raw-view payloads as needed.
|
||||
fn resolve_in_place<Frame>(&mut self, frames: &[Frame]) -> Result<()>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
if let Self::RequestBatch(batch) = self {
|
||||
for output in &mut batch.outputs {
|
||||
output.resolve_in_place(frames)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Classify the raw wire message into a more semantic Rust enum.
|
||||
impl TryFrom<WireEngineCoreOutputs> for EngineCoreOutputs {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(value: WireEngineCoreOutputs) -> Result<Self> {
|
||||
let has_request_payload = !value.outputs.is_empty()
|
||||
|| value.scheduler_stats.is_some()
|
||||
|| value.finished_requests.is_some();
|
||||
|
||||
match (
|
||||
has_request_payload,
|
||||
&value.utility_output,
|
||||
&value.wave_complete,
|
||||
&value.start_wave,
|
||||
) {
|
||||
(true, None, None, None) => Ok(RequestBatchOutputs {
|
||||
engine_index: value.engine_index,
|
||||
outputs: value.outputs,
|
||||
scheduler_stats: value.scheduler_stats,
|
||||
timestamp: value.timestamp,
|
||||
finished_requests: value.finished_requests,
|
||||
}
|
||||
.into()),
|
||||
(false, Some(_), None, None) => Ok(UtilityCallOutput {
|
||||
engine_index: value.engine_index,
|
||||
timestamp: value.timestamp,
|
||||
output: value.utility_output.unwrap(),
|
||||
}
|
||||
.into()),
|
||||
(false, None, Some(_), None) => Ok(DpControlOutput {
|
||||
engine_index: value.engine_index,
|
||||
timestamp: value.timestamp,
|
||||
control: DpControlMessage::WaveComplete(value.wave_complete.unwrap()),
|
||||
}
|
||||
.into()),
|
||||
(false, None, None, Some(_)) => Ok(DpControlOutput {
|
||||
engine_index: value.engine_index,
|
||||
timestamp: value.timestamp,
|
||||
control: DpControlMessage::StartWave(value.start_wave.unwrap()),
|
||||
}
|
||||
.into()),
|
||||
|
||||
_ => Err(Error::Decode {
|
||||
target_type: "EngineCoreOutputs",
|
||||
message: "invalid wire shape".to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<EngineCoreOutputs> for WireEngineCoreOutputs {
|
||||
fn from(value: EngineCoreOutputs) -> Self {
|
||||
match value {
|
||||
EngineCoreOutputs::RequestBatch(batch) => Self {
|
||||
engine_index: batch.engine_index,
|
||||
outputs: batch.outputs,
|
||||
scheduler_stats: batch.scheduler_stats,
|
||||
timestamp: batch.timestamp,
|
||||
finished_requests: batch.finished_requests,
|
||||
..Default::default()
|
||||
},
|
||||
EngineCoreOutputs::Utility(utility) => Self {
|
||||
engine_index: utility.engine_index,
|
||||
timestamp: utility.timestamp,
|
||||
utility_output: Some(utility.output),
|
||||
..Default::default()
|
||||
},
|
||||
EngineCoreOutputs::DpControl(control) => {
|
||||
let (wave_complete, start_wave) = match control.control {
|
||||
DpControlMessage::WaveComplete(wave) => (Some(wave), None),
|
||||
DpControlMessage::StartWave(wave) => (None, Some(wave)),
|
||||
};
|
||||
Self {
|
||||
engine_index: control.engine_index,
|
||||
timestamp: control.timestamp,
|
||||
wave_complete,
|
||||
start_wave,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for EngineCoreOutputs {
|
||||
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
WireEngineCoreOutputs::from(self.clone()).serialize(serializer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for EngineCoreOutputs {
|
||||
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
WireEngineCoreOutputs::deserialize(deserializer)?
|
||||
.try_into()
|
||||
.map_err(serde::de::Error::custom)
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode one ordinary or multipart engine-core output message into the strong
|
||||
/// typed public protocol shape.
|
||||
pub fn decode_engine_core_outputs<Frame>(frames: &[Frame]) -> Result<EngineCoreOutputs>
|
||||
where
|
||||
Frame: AsRef<[u8]>,
|
||||
{
|
||||
let first_frame = frames.first().ok_or_else(|| ext_value_decode!("missing output frame"))?;
|
||||
|
||||
let mut outputs: EngineCoreOutputs = decode_msgpack(first_frame.as_ref())?;
|
||||
outputs.resolve_in_place(frames)?;
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::output::EngineCoreOutput;
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_roundtrip_finished_fields() {
|
||||
let outputs = WireEngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![42],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(EngineCoreFinishReason::Length),
|
||||
stop_reason: Some(StopReason::Text("stop".to_string())),
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&outputs).unwrap();
|
||||
let decoded: WireEngineCoreOutputs = decode_msgpack(&encoded).unwrap();
|
||||
|
||||
assert_eq!(decoded.outputs.len(), 1);
|
||||
assert_eq!(
|
||||
decoded.outputs[0].finish_reason,
|
||||
Some(EngineCoreFinishReason::Length)
|
||||
);
|
||||
assert_eq!(
|
||||
decoded.finished_requests,
|
||||
Some(BTreeSet::from(["req-1".to_string()]))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_request_batch() {
|
||||
let outputs = WireEngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![7],
|
||||
..Default::default()
|
||||
}],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
RequestBatch(
|
||||
RequestBatchOutputs {
|
||||
engine_index: 0,
|
||||
outputs: [
|
||||
EngineCoreOutput {
|
||||
request_id: "req-1",
|
||||
new_token_ids: [
|
||||
7,
|
||||
],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: None,
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
},
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
finished_requests: Some(
|
||||
{
|
||||
"req-1",
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_utility() {
|
||||
let outputs = WireEngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 42_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
Utility(
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
call_id: 42,
|
||||
failure_message: None,
|
||||
result: None,
|
||||
},
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_classify_control() {
|
||||
let outputs = WireEngineCoreOutputs {
|
||||
start_wave: Some(3),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
expect_test::expect![[r#"
|
||||
DpControl(
|
||||
DpControlOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
control: StartWave(
|
||||
3,
|
||||
),
|
||||
},
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&EngineCoreOutputs::try_from(outputs).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn engine_core_outputs_rejects_mixed_shape() {
|
||||
let outputs = WireEngineCoreOutputs {
|
||||
outputs: vec![EngineCoreOutput {
|
||||
request_id: "req-1".to_string(),
|
||||
new_token_ids: vec![7],
|
||||
..Default::default()
|
||||
}],
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 1_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let error = EngineCoreOutputs::try_from(outputs).unwrap_err();
|
||||
expect_test::expect![[
|
||||
r#"messagepack decode failed for EngineCoreOutputs: invalid wire shape"#
|
||||
]]
|
||||
.assert_eq(&error.to_string());
|
||||
}
|
||||
}
|
||||
@@ -1,175 +0,0 @@
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
use bytes::Bytes;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
|
||||
use crate::protocol::multimodal::MmFeatures;
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::{OpaqueValue, lora};
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Request types are encoded as single-byte protocol constants so they can be
|
||||
/// sent over the ZMQ socket without an extra encoding step.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/v1/engine/__init__.py#L217-L228>
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[repr(u8)]
|
||||
pub enum EngineCoreRequestType {
|
||||
Add = 0,
|
||||
Abort = 1,
|
||||
StartDpWave = 2,
|
||||
Utility = 3,
|
||||
}
|
||||
|
||||
impl EngineCoreRequestType {
|
||||
/// Decode the single-byte request type frame used on the engine input
|
||||
/// socket. Returns `None` for unrecognized values.
|
||||
pub fn from_frame(frame: &[u8]) -> Option<Self> {
|
||||
let [value] = frame else {
|
||||
return None;
|
||||
};
|
||||
|
||||
match value {
|
||||
0 => Some(Self::Add),
|
||||
1 => Some(Self::Abort),
|
||||
2 => Some(Self::StartDpWave),
|
||||
3 => Some(Self::Utility),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode the request type as the single-byte frame used on the engine
|
||||
/// input socket.
|
||||
pub fn to_frame(self) -> Bytes {
|
||||
Bytes::from_static(match self {
|
||||
Self::Add => b"\x00",
|
||||
Self::Abort => b"\x01",
|
||||
Self::StartDpWave => b"\x02",
|
||||
Self::Utility => b"\x03",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Extra kwargs consumed by engine-side reasoning parsers.
|
||||
///
|
||||
/// Original Python construction point:
|
||||
/// <https://github.com/vllm-project/vllm/blob/cec2ec11760f9f3beabd4c90451936078bf91533/vllm/entrypoints/openai/chat_completion/serving.py#L367-L369>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ReasoningParserKwargs {
|
||||
/// Effective kwargs visible to the chat template for this request.
|
||||
pub chat_template_kwargs: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Engine-core add-request payload sent from frontend to engine.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/3f5bd482f5c1a5dbdffbbf68d624e20bb7032013/vllm/v1/engine/__init__.py#L80-L129>
|
||||
#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple, DefaultFromSerde)]
|
||||
pub struct EngineCoreRequest {
|
||||
pub request_id: String,
|
||||
pub prompt_token_ids: Option<Vec<u32>>,
|
||||
/// Multimodal features attached to the request.
|
||||
pub mm_features: Option<MmFeatures>,
|
||||
pub sampling_params: Option<EngineCoreSamplingParams>,
|
||||
/// Pooling parameters are preserved in the schema but not yet strongly
|
||||
/// typed.
|
||||
pub pooling_params: Option<OpaqueValue>,
|
||||
pub arrival_time: f64,
|
||||
#[serde(default)]
|
||||
pub lora_request: Option<lora::LoraRequest>,
|
||||
#[serde(default)]
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub data_parallel_rank: Option<u32>,
|
||||
/// Unsupported in the first-stage Rust client because Python uses a custom
|
||||
/// tensor/aux-frame encoding path for this field.
|
||||
#[serde(default)]
|
||||
pub prompt_embeds: Option<OpaqueValue>,
|
||||
/// Per-position mask for mixed-mode inputs (e.g. chat completion with
|
||||
/// `prompt_embeds` content parts). `Some(true)` means real token id;
|
||||
/// `Some(false)` means the position uses a pre-computed entry from
|
||||
/// `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
|
||||
#[serde(default)]
|
||||
pub prompt_is_token_ids: Option<Vec<bool>>,
|
||||
/// Index of the client, used to ensure outputs are sent back to the same
|
||||
/// client when scaling out the frontend.
|
||||
#[serde(default)]
|
||||
pub client_index: u32,
|
||||
/// In DP mode, indicates which wave this request is expected to belong to.
|
||||
#[serde(default)]
|
||||
pub current_wave: u32,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
#[serde(default)]
|
||||
pub trace_headers: Option<BTreeMap<String, String>>,
|
||||
#[serde(default)]
|
||||
pub resumable: bool,
|
||||
/// Original user-provided request ID, used for output reporting and aborts.
|
||||
#[serde(default)]
|
||||
pub external_req_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub reasoning_ended: Option<bool>,
|
||||
/// Reasoning-parser kwargs forwarded from the frontend to the
|
||||
/// structured-output backend.
|
||||
#[serde(default)]
|
||||
pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,
|
||||
/// If `true`, the request should be added to the scheduler's waiting queue
|
||||
/// and immediately aborted, so connector-side cleanup runs via the
|
||||
/// standard `request_finished` hook.
|
||||
#[serde(default)]
|
||||
pub abort_immediately: bool,
|
||||
}
|
||||
|
||||
impl EngineCoreRequest {
|
||||
/// Validate fields intentionally not supported in the first-stage client.
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.prompt_embeds.is_some() {
|
||||
return Err(Error::UnsupportedField {
|
||||
context: "EngineCoreRequest",
|
||||
field: "prompt_embeds",
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rmpv::Value;
|
||||
|
||||
use super::*;
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::{decode_value, encode_msgpack};
|
||||
|
||||
#[test]
|
||||
fn engine_core_request_serializes_as_full_array() {
|
||||
let request = EngineCoreRequest {
|
||||
request_id: "req-1".to_string(),
|
||||
prompt_token_ids: Some(vec![1, 2, 3]),
|
||||
sampling_params: Some(EngineCoreSamplingParams {
|
||||
max_tokens: 8,
|
||||
..EngineCoreSamplingParams::for_test()
|
||||
}),
|
||||
arrival_time: 1234.5,
|
||||
client_index: 7,
|
||||
..EngineCoreRequest::default()
|
||||
};
|
||||
|
||||
let encoded = encode_msgpack(&request).unwrap();
|
||||
let value = decode_value(&encoded).unwrap();
|
||||
let array = match value {
|
||||
Value::Array(array) => array,
|
||||
other => panic!("expected array, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(array.len(), 20);
|
||||
assert_eq!(array[0], Value::from("req-1"));
|
||||
assert_eq!(array[2], Value::Nil);
|
||||
assert_eq!(array[4], Value::Nil);
|
||||
assert_eq!(array[10], Value::Nil);
|
||||
assert_eq!(array[11], Value::from(7));
|
||||
}
|
||||
}
|
||||
@@ -1,211 +0,0 @@
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_default::DefaultFromSerde;
|
||||
|
||||
use crate::protocol::structured_outputs::StructuredOutputsParams;
|
||||
|
||||
fn default_top_p() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_repetition_penalty() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_temperature() -> f32 {
|
||||
1.0
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
16
|
||||
}
|
||||
|
||||
/// Engine-core-facing sampling parameters for text generation.
|
||||
///
|
||||
/// This is the normalized southbound subset used by the Rust frontend when it
|
||||
/// talks to Python engine-core over the wire. User-facing request semantics
|
||||
/// such as `stop` strings, `n`, `ignore_eos`, and output aggregation mode are
|
||||
/// intentionally handled by higher layers before values reach this DTO.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
|
||||
// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
|
||||
// default-valued keys; default the whole struct. Per-field fns cover the
|
||||
// non-zero defaults.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
|
||||
#[serde(default)]
|
||||
pub struct EngineCoreSamplingParams {
|
||||
/// Controls randomness. Lower values are more deterministic; zero means
|
||||
/// greedy sampling.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: f32,
|
||||
/// Cumulative probability threshold for nucleus sampling.
|
||||
#[serde(default = "default_top_p")]
|
||||
pub top_p: f32,
|
||||
/// Maximum number of top tokens to consider. `0` means all tokens.
|
||||
pub top_k: u32,
|
||||
/// Random seed used by the sampler when present.
|
||||
pub seed: Option<i64>,
|
||||
/// Maximum number of tokens to generate per output sequence.
|
||||
#[serde(default = "default_max_tokens")]
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: u32,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` means unlimited; the
|
||||
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
|
||||
/// reaching this DTO, so only non-negative values are sent. Enforced
|
||||
/// engine-side (and only when a reasoning parser is configured).
|
||||
pub thinking_token_budget: Option<u64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
pub logprobs: Option<i32>,
|
||||
/// Number of log probabilities to return per prompt token.
|
||||
///
|
||||
/// `None` disables prompt logprobs. `-1` requests the full vocabulary.
|
||||
pub prompt_logprobs: Option<i32>,
|
||||
/// Minimum probability threshold for token sampling.
|
||||
pub min_p: f32,
|
||||
/// Frequency penalty applied by the sampler.
|
||||
pub frequency_penalty: f32,
|
||||
/// Presence penalty applied by the sampler.
|
||||
pub presence_penalty: f32,
|
||||
/// Repetition penalty applied by the sampler.
|
||||
#[serde(default = "default_repetition_penalty")]
|
||||
pub repetition_penalty: f32,
|
||||
/// Token IDs that stop generation.
|
||||
pub stop_token_ids: Vec<u32>,
|
||||
/// Primary EOS token ID used by engine-core's dedicated EOS stop path.
|
||||
///
|
||||
/// This mirrors Python's internal `_eos_token_id` field and is derived by
|
||||
/// the frontend from tokenizer/model metadata rather than supplied directly
|
||||
/// by end users.
|
||||
#[serde(rename = "_eos_token_id")]
|
||||
pub eos_token_id: Option<u32>,
|
||||
/// Complete stop-token set used by engine-core for `min_tokens` masking.
|
||||
///
|
||||
/// This mirrors Python's internal `_all_stop_token_ids` field and should
|
||||
/// contain explicit `stop_token_ids` plus any frontend-derived EOS token
|
||||
/// IDs.
|
||||
#[serde(rename = "_all_stop_token_ids")]
|
||||
pub all_stop_token_ids: BTreeSet<u32>,
|
||||
/// Logit biases to apply during sampling.
|
||||
/// Keys are token IDs
|
||||
pub logit_bias: Option<HashMap<u32, f32>>,
|
||||
/// Restrict output to these token IDs only.
|
||||
pub allowed_token_ids: Option<Vec<u32>>,
|
||||
/// Tokenized bad words to avoid during generation.
|
||||
#[serde(rename = "_bad_words_token_ids")]
|
||||
pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
pub structured_outputs: Option<StructuredOutputsParams>,
|
||||
/// Specific token IDs for which log probabilities should be returned at
|
||||
/// each position.
|
||||
///
|
||||
/// When set, the engine returns logprobs for exactly these tokens in
|
||||
/// addition to the sampled/scored token. Mutually exclusive with the
|
||||
/// `logprobs` count field in practice.
|
||||
pub logprob_token_ids: Option<Vec<u32>>,
|
||||
/// If `Some(true)`, the request will not attempt to read from the prefix
|
||||
/// cache; newly computed blocks may still populate the cache. `None`
|
||||
/// defers to engine-core defaults.
|
||||
pub skip_reading_prefix_cache: Option<bool>,
|
||||
/// Additional request parameters for custom extensions (from `vllm_xargs`).
|
||||
pub extra_args: Option<HashMap<String, serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl EngineCoreSamplingParams {
|
||||
/// Constructs a default sampling params for testing purposes only.
|
||||
pub fn for_test() -> Self {
|
||||
Self {
|
||||
temperature: 1.0,
|
||||
top_p: 1.0,
|
||||
top_k: 0,
|
||||
seed: None,
|
||||
max_tokens: 65536,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
repetition_penalty: 1.0,
|
||||
stop_token_ids: Vec::new(),
|
||||
eos_token_id: None,
|
||||
all_stop_token_ids: BTreeSet::new(),
|
||||
logit_bias: None,
|
||||
allowed_token_ids: None,
|
||||
bad_words_token_ids: None,
|
||||
structured_outputs: None,
|
||||
logprob_token_ids: None,
|
||||
skip_reading_prefix_cache: None,
|
||||
extra_args: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rmpv::Value;
|
||||
|
||||
use crate::protocol::decode_msgpack;
|
||||
use crate::protocol::request::EngineCoreRequest;
|
||||
|
||||
/// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
|
||||
/// must fall back to defaults. `python_compat` can't catch this since Rust
|
||||
/// encodes full maps (see `engine_core_request_serializes_as_full_array`).
|
||||
#[test]
|
||||
fn decodes_sampling_params_with_omitted_defaults() {
|
||||
let sampling_params = Value::Map(vec![
|
||||
(
|
||||
Value::from("stop_token_ids"),
|
||||
Value::Array(vec![Value::from(151643u32)]),
|
||||
),
|
||||
(Value::from("skip_reading_prefix_cache"), Value::from(false)),
|
||||
]);
|
||||
let request = Value::Array(vec![
|
||||
Value::from("req-omit-defaults"),
|
||||
Value::Array(vec![
|
||||
Value::from(1u32),
|
||||
Value::from(2u32),
|
||||
Value::from(3u32),
|
||||
]),
|
||||
Value::Nil,
|
||||
sampling_params,
|
||||
Value::Nil,
|
||||
Value::from(1.0f64),
|
||||
]);
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
rmpv::encode::write_value(&mut bytes, &request).unwrap();
|
||||
|
||||
let decoded: EngineCoreRequest = decode_msgpack(&bytes)
|
||||
.expect("a real omit_defaults request must decode (regression: missing field)");
|
||||
|
||||
assert_eq!(decoded.request_id, "req-omit-defaults");
|
||||
let sampling = decoded.sampling_params.expect("sampling params present");
|
||||
|
||||
assert_eq!(sampling.stop_token_ids, vec![151643]);
|
||||
assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
|
||||
|
||||
// Omitted fields -> Python defaults.
|
||||
assert_eq!(sampling.temperature, 1.0);
|
||||
assert_eq!(sampling.top_p, 1.0);
|
||||
assert_eq!(sampling.top_k, 0);
|
||||
assert_eq!(sampling.seed, None);
|
||||
assert_eq!(sampling.max_tokens, 16);
|
||||
assert_eq!(sampling.min_tokens, 0);
|
||||
assert_eq!(sampling.min_p, 0.0);
|
||||
assert_eq!(sampling.frequency_penalty, 0.0);
|
||||
assert_eq!(sampling.presence_penalty, 0.0);
|
||||
assert_eq!(sampling.repetition_penalty, 1.0);
|
||||
assert_eq!(sampling.logprobs, None);
|
||||
assert_eq!(sampling.prompt_logprobs, None);
|
||||
assert_eq!(sampling.eos_token_id, None);
|
||||
assert!(sampling.all_stop_token_ids.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,305 +0,0 @@
|
||||
use enum_as_inner::EnumAsInner;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Structured-output backend selected for EngineCore grammar compilation.
|
||||
///
|
||||
/// Python vLLM stores this in `StructuredOutputsParams._backend` after request
|
||||
/// validation. The Rust frontend currently always lowers structured-output
|
||||
/// requests to guidance, while ignoring any user-supplied `_backend` value.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum StructuredOutputBackend {
|
||||
Xgrammar,
|
||||
#[default]
|
||||
Guidance,
|
||||
Outlines,
|
||||
LmFormatEnforcer,
|
||||
}
|
||||
|
||||
/// The single structured-output constraint selected for a request.
|
||||
#[derive(Debug, Clone, PartialEq, EnumAsInner)]
|
||||
pub enum StructuredOutputConstraint {
|
||||
/// JSON schema (as a dict/object or JSON string) constraining the output.
|
||||
Json(Value),
|
||||
/// Regular expression the output must match.
|
||||
Regex(String),
|
||||
/// List of allowed output strings (the model must produce one of these).
|
||||
Choice(Vec<String>),
|
||||
/// Context-free grammar (in EBNF-like notation) the output must conform to.
|
||||
Grammar(String),
|
||||
/// Output must be valid JSON (free-form, no schema).
|
||||
JsonObject,
|
||||
/// Structural tag configuration (JSON-encoded string).
|
||||
StructuralTag(String),
|
||||
}
|
||||
|
||||
/// Additional structured-output options that do not select the constraint mode.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct StructuredOutputOptions {
|
||||
/// Disable any additional whitespace in guided JSON output.
|
||||
pub disable_any_whitespace: bool,
|
||||
/// Disable `additionalProperties` in JSON schema output.
|
||||
pub disable_additional_properties: bool,
|
||||
/// Custom whitespace pattern for guided JSON output.
|
||||
pub whitespace_pattern: Option<String>,
|
||||
}
|
||||
|
||||
/// Parameters for configuring structured outputs (guided decoding).
|
||||
///
|
||||
/// This is the semantic Rust representation: exactly one constraint mode is
|
||||
/// always selected. The Python/msgpack product-shaped representation is kept in
|
||||
/// the private wire type below and used only at serde boundaries.
|
||||
///
|
||||
/// Original Python definition:
|
||||
/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L36-L107>
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct StructuredOutputsParams {
|
||||
pub constraint: StructuredOutputConstraint,
|
||||
pub options: StructuredOutputOptions,
|
||||
/// Structured-output backend, mirroring Python's internal `_backend`.
|
||||
///
|
||||
/// User-supplied values are ignored during deserialization. This matches
|
||||
/// Python's request boundary, where `_backend` is set by validation rather
|
||||
/// than accepted as a request-level backend selector.
|
||||
pub backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
impl StructuredOutputsParams {
|
||||
pub fn json(json: Value) -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::Json(json))
|
||||
}
|
||||
|
||||
pub fn regex(regex: impl Into<String>) -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::Regex(regex.into()))
|
||||
}
|
||||
|
||||
pub fn choice(choice: Vec<String>) -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::Choice(choice))
|
||||
}
|
||||
|
||||
pub fn grammar(grammar: impl Into<String>) -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::Grammar(grammar.into()))
|
||||
}
|
||||
|
||||
pub fn json_object() -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::JsonObject)
|
||||
}
|
||||
|
||||
pub fn structural_tag(structural_tag: impl Into<String>) -> Self {
|
||||
Self::from_constraint(StructuredOutputConstraint::StructuralTag(
|
||||
structural_tag.into(),
|
||||
))
|
||||
}
|
||||
|
||||
fn from_constraint(constraint: StructuredOutputConstraint) -> Self {
|
||||
Self {
|
||||
constraint,
|
||||
options: StructuredOutputOptions::default(),
|
||||
backend: StructuredOutputBackend::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wire-compatible structured-output payload used by Python engine-core.
|
||||
///
|
||||
/// Python models `StructuredOutputsParams` as a product-shaped dataclass with
|
||||
/// several optional constraint fields, then validates that exactly one of those
|
||||
/// fields is present. Rust exposes [`StructuredOutputsParams`] as an enum-backed
|
||||
/// domain type instead, while using this private wire type for ser/de.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
struct WireStructuredOutputsParams {
|
||||
json: Option<Value>,
|
||||
regex: Option<String>,
|
||||
choice: Option<Vec<String>>,
|
||||
grammar: Option<String>,
|
||||
json_object: Option<bool>,
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
disable_any_whitespace: bool,
|
||||
#[serde(skip_serializing_if = "crate::protocol::is_false")]
|
||||
disable_additional_properties: bool,
|
||||
whitespace_pattern: Option<String>,
|
||||
structural_tag: Option<String>,
|
||||
#[serde(
|
||||
default,
|
||||
rename = "_backend",
|
||||
deserialize_with = "serde_with::rust::deserialize_ignore_any"
|
||||
)]
|
||||
backend: StructuredOutputBackend,
|
||||
}
|
||||
|
||||
impl TryFrom<WireStructuredOutputsParams> for StructuredOutputsParams {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(raw: WireStructuredOutputsParams) -> Result<Self> {
|
||||
use StructuredOutputConstraint::*;
|
||||
|
||||
let mut constraint = None;
|
||||
|
||||
macro_rules! insert_constraint {
|
||||
($name:literal, $value:expr) => {
|
||||
if let Some(value) = $value {
|
||||
if let Some((existing, _)) = constraint {
|
||||
return Err(Error::InvalidStructuredOutputsParams {
|
||||
message: format!(
|
||||
"multiple structured output constraints specified: {existing}, {}",
|
||||
$name
|
||||
),
|
||||
});
|
||||
}
|
||||
constraint = Some(($name, value));
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
insert_constraint!("json", raw.json.map(Json));
|
||||
insert_constraint!("regex", raw.regex.map(Regex));
|
||||
insert_constraint!("choice", raw.choice.map(Choice));
|
||||
insert_constraint!("grammar", raw.grammar.map(Grammar));
|
||||
match raw.json_object {
|
||||
Some(true) => {
|
||||
insert_constraint!("json_object", Some(JsonObject))
|
||||
}
|
||||
Some(false) => {
|
||||
return Err(Error::InvalidStructuredOutputsParams {
|
||||
message: "structured_outputs.json_object must be true if set; omit structured_outputs to disable structured outputs".to_string(),
|
||||
});
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag));
|
||||
|
||||
Ok(Self {
|
||||
constraint: constraint.map(|(_, c)| c).ok_or_else(|| {
|
||||
Error::InvalidStructuredOutputsParams {
|
||||
message: "missing structured output constraint".to_string(),
|
||||
}
|
||||
})?,
|
||||
options: StructuredOutputOptions {
|
||||
disable_any_whitespace: raw.disable_any_whitespace,
|
||||
disable_additional_properties: raw.disable_additional_properties,
|
||||
whitespace_pattern: raw.whitespace_pattern,
|
||||
},
|
||||
backend: raw.backend,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl From<StructuredOutputsParams> for WireStructuredOutputsParams {
|
||||
fn from(params: StructuredOutputsParams) -> Self {
|
||||
let mut raw = Self {
|
||||
disable_any_whitespace: params.options.disable_any_whitespace,
|
||||
disable_additional_properties: params.options.disable_additional_properties,
|
||||
whitespace_pattern: params.options.whitespace_pattern,
|
||||
backend: params.backend,
|
||||
..Self::default()
|
||||
};
|
||||
|
||||
match params.constraint {
|
||||
StructuredOutputConstraint::Json(json) => raw.json = Some(json),
|
||||
StructuredOutputConstraint::Regex(regex) => raw.regex = Some(regex),
|
||||
StructuredOutputConstraint::Choice(choice) => raw.choice = Some(choice),
|
||||
StructuredOutputConstraint::Grammar(grammar) => raw.grammar = Some(grammar),
|
||||
StructuredOutputConstraint::JsonObject => raw.json_object = Some(true),
|
||||
StructuredOutputConstraint::StructuralTag(structural_tag) => {
|
||||
raw.structural_tag = Some(structural_tag);
|
||||
}
|
||||
}
|
||||
|
||||
raw
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for StructuredOutputsParams {
|
||||
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
WireStructuredOutputsParams::from(self.clone()).serialize(serializer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for StructuredOutputsParams {
|
||||
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
WireStructuredOutputsParams::deserialize(deserializer)?
|
||||
.try_into()
|
||||
.map_err(serde::de::Error::custom)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_backend_ignores_deserialized_value() {
|
||||
let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({
|
||||
"json_object": true,
|
||||
"_backend": "xgrammar",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(params.backend, StructuredOutputBackend::Guidance);
|
||||
assert_eq!(params.constraint, StructuredOutputConstraint::JsonObject);
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
assert_eq!(value["_backend"], "guidance");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_rejects_missing_constraint() {
|
||||
let error =
|
||||
serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({})).unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("missing structured output constraint"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_rejects_multiple_constraints() {
|
||||
let error = serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({
|
||||
"json": {"type": "object"},
|
||||
"regex": ".*",
|
||||
}))
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("multiple structured output constraints specified: json, regex")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_rejects_json_object_false() {
|
||||
let error = serde_json::from_value::<StructuredOutputsParams>(serde_json::json!({
|
||||
"json_object": false,
|
||||
}))
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("json_object must be true"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_outputs_serializes_through_raw_shape() {
|
||||
let params = StructuredOutputsParams {
|
||||
constraint: StructuredOutputConstraint::StructuralTag(
|
||||
r#"{"type":"structural_tag"}"#.to_string(),
|
||||
),
|
||||
options: StructuredOutputOptions::default(),
|
||||
backend: StructuredOutputBackend::Xgrammar,
|
||||
};
|
||||
|
||||
let value = serde_json::to_value(params).unwrap();
|
||||
|
||||
assert_eq!(value["structural_tag"], r#"{"type":"structural_tag"}"#);
|
||||
assert_eq!(value["_backend"], "xgrammar");
|
||||
assert!(value.get("json").is_none());
|
||||
}
|
||||
}
|
||||
@@ -22,15 +22,13 @@ use crate::protocol::multimodal::{
|
||||
MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmSlice, PlaceholderRange,
|
||||
SliceSpec,
|
||||
};
|
||||
use crate::protocol::output::{
|
||||
DpControlMessage, DpControlOutput, EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs,
|
||||
RequestBatchOutputs, UtilityCallOutput, decode_engine_core_outputs,
|
||||
};
|
||||
use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
|
||||
use crate::protocol::sampling::EngineCoreSamplingParams;
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::protocol::tensor::WireTensor;
|
||||
use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
|
||||
use crate::protocol::{
|
||||
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
|
||||
EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
|
||||
};
|
||||
use crate::test_utils::{
|
||||
IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
|
||||
setup_mock_engine_with_init, spawn_mock_engine_task,
|
||||
@@ -572,7 +570,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-1",
|
||||
vec![],
|
||||
@@ -580,19 +579,17 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
send_outputs(
|
||||
&mut coordinator.output_push,
|
||||
DpControlOutput {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
control: DpControlMessage::WaveComplete(0),
|
||||
}
|
||||
.into(),
|
||||
wave_complete: Some(0),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -607,7 +604,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-3",
|
||||
vec![],
|
||||
@@ -615,8 +613,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-3".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -653,7 +650,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 1,
|
||||
outputs: vec![request_output(
|
||||
"req-2",
|
||||
@@ -662,8 +659,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -769,12 +765,11 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
|
||||
|
||||
send_outputs(
|
||||
&mut coordinator.output_push,
|
||||
DpControlOutput {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 1,
|
||||
timestamp: 0.0,
|
||||
control: DpControlMessage::StartWave(4),
|
||||
}
|
||||
.into(),
|
||||
start_wave: Some(4),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -825,16 +820,15 @@ async fn coordinator_accepts_stats_only_outputs() {
|
||||
|
||||
send_outputs(
|
||||
&mut coordinator.output_push,
|
||||
RequestBatchOutputs {
|
||||
outputs: Vec::new(),
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
scheduler_stats: Some(Box::new(SchedulerStats {
|
||||
num_running_reqs: 1,
|
||||
current_wave: 0,
|
||||
..Default::default()
|
||||
})),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -845,7 +839,8 @@ async fn coordinator_accepts_stats_only_outputs() {
|
||||
|
||||
send_outputs(
|
||||
&mut data_socket.push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-stats",
|
||||
vec![],
|
||||
@@ -853,8 +848,7 @@ async fn coordinator_accepts_stats_only_outputs() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-stats".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -915,34 +909,30 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 1_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
},
|
||||
}),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
send_outputs(
|
||||
push,
|
||||
DpControlOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
control: DpControlMessage::StartWave(3),
|
||||
}
|
||||
.into(),
|
||||
EngineCoreOutputs {
|
||||
start_wave: Some(3),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output("req-1", vec![999], None)],
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -994,6 +984,91 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn client_fail_closes_when_main_output_path_receives_mixed_shape_output() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-0".to_vec();
|
||||
|
||||
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
|
||||
handshake_address.clone(),
|
||||
engine_id.clone(),
|
||||
|dealer, push| {
|
||||
Box::pin(async move {
|
||||
let add_1 = recv_engine_message(dealer).await;
|
||||
assert_eq!(add_1[0].as_ref(), &[0x00]);
|
||||
let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
|
||||
assert_eq!(request_1.client_index, 7);
|
||||
assert_eq!(request_1.request_id, "req-1");
|
||||
|
||||
let add_2 = recv_engine_message(dealer).await;
|
||||
assert_eq!(add_2[0].as_ref(), &[0x00]);
|
||||
let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
|
||||
assert_eq!(request_2.client_index, 7);
|
||||
assert_eq!(request_2.request_id, "req-2");
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: 1_u64.into(),
|
||||
failure_message: None,
|
||||
result: None,
|
||||
}),
|
||||
outputs: vec![request_output("req-1", vec![999], None)],
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
})
|
||||
},
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
1,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
7,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(client.engine_identities()[0], b"engine-0");
|
||||
assert!(client.ready_responses()[0].max_model_len > 0);
|
||||
|
||||
let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
|
||||
let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();
|
||||
|
||||
let error_2 = timeout(Duration::from_secs(1), stream_2.next())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap_err();
|
||||
assert!(is_unexpected_dispatcher_output(&error_2));
|
||||
|
||||
let error_1 = timeout(Duration::from_secs(1), stream_1.next())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap_err();
|
||||
assert!(is_unexpected_dispatcher_output(&error_1));
|
||||
|
||||
assert!(matches!(
|
||||
client.health_error().as_deref(),
|
||||
Some(error) if is_unexpected_dispatcher_output(error)
|
||||
));
|
||||
|
||||
let _ = shutdown_tx.send(());
|
||||
engine_task.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
|
||||
init_tracing();
|
||||
@@ -1015,7 +1090,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output(
|
||||
"req-1",
|
||||
vec![],
|
||||
@@ -1023,8 +1098,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1083,12 +1157,10 @@ async fn finished_requests_without_final_output_is_treated_as_unexpected_close()
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
outputs: Vec::new(),
|
||||
EngineCoreOutputs {
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1144,11 +1216,10 @@ async fn dropping_a_live_stream_triggers_abort() {
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![request_output("req-1", vec![99], None)],
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1203,15 +1274,14 @@ async fn dropping_multiple_live_streams_aborts_all_in_a_burst() {
|
||||
}
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output("req-1", vec![99], None),
|
||||
request_output("req-2", vec![99], None),
|
||||
request_output("req-3", vec![99], None),
|
||||
],
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1365,16 +1435,14 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(true)),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1422,16 +1490,14 @@ async fn call_utility_failure_message_surfaces_as_error() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: Some("boom".to_string()),
|
||||
result: None,
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1748,7 +1814,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
finish_req_1_rx.await.unwrap();
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
&request_1.request_id,
|
||||
vec![10],
|
||||
@@ -1756,8 +1823,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from([request_1.request_id.clone()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1769,7 +1835,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
finish_req_3_rx.await.unwrap();
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
&request_3.request_id,
|
||||
vec![30],
|
||||
@@ -1777,8 +1844,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from([request_3.request_id.clone()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1797,7 +1863,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
finish_req_2_rx.await.unwrap();
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 1,
|
||||
outputs: vec![request_output(
|
||||
&request_2.request_id,
|
||||
@@ -1806,8 +1872,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from([request_2.request_id.clone()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1924,16 +1989,14 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
assert_eq!(array[2], Value::from("is_sleeping"));
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(true)),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1948,7 +2011,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
assert_eq!(aborted_ids, vec!["req-1".to_string()]);
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-1",
|
||||
vec![],
|
||||
@@ -1956,8 +2020,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -1980,16 +2043,14 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
assert_eq!(array[2], Value::from("is_sleeping"));
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(true)),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -2004,7 +2065,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
assert_eq!(aborted_ids, vec!["req-2".to_string()]);
|
||||
send_outputs(
|
||||
push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 1,
|
||||
outputs: vec![request_output(
|
||||
"req-2",
|
||||
@@ -2013,8 +2074,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-2".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -2094,16 +2154,14 @@ async fn collective_rpc_flattens_results_from_all_engines() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(vec!["engine-0-worker"])),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -2127,16 +2185,14 @@ async fn collective_rpc_flattens_results_from_all_engines() {
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(vec!["engine-1-worker"])),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -2207,16 +2263,14 @@ fn spawn_mock_utility_engine(
|
||||
assert_eq!(array[3], expected_args, "unexpected utility args");
|
||||
send_outputs(
|
||||
push,
|
||||
UtilityCallOutput {
|
||||
engine_index: 0,
|
||||
timestamp: 0.0,
|
||||
output: UtilityOutput {
|
||||
EngineCoreOutputs {
|
||||
utility_output: Some(UtilityOutput {
|
||||
call_id: call_id.into(),
|
||||
failure_message: None,
|
||||
result: Some(utility_result_value(result)),
|
||||
},
|
||||
}
|
||||
.into(),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
@@ -2497,40 +2551,41 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
|
||||
let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
|
||||
expect_test::expect![[r#"
|
||||
RequestBatch(
|
||||
RequestBatchOutputs {
|
||||
engine_index: 0,
|
||||
outputs: [
|
||||
EngineCoreOutput {
|
||||
request_id: "req-1",
|
||||
new_token_ids: [
|
||||
7,
|
||||
8,
|
||||
],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(
|
||||
Length,
|
||||
),
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
},
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
finished_requests: Some(
|
||||
{
|
||||
"req-1",
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: [
|
||||
EngineCoreOutput {
|
||||
request_id: "req-1",
|
||||
new_token_ids: [
|
||||
7,
|
||||
8,
|
||||
],
|
||||
new_logprobs: None,
|
||||
new_prompt_logprobs_tensors: None,
|
||||
pooling_output: None,
|
||||
finish_reason: Some(
|
||||
Length,
|
||||
),
|
||||
stop_reason: None,
|
||||
events: None,
|
||||
kv_transfer_params: None,
|
||||
trace_headers: None,
|
||||
prefill_stats: None,
|
||||
routed_experts: None,
|
||||
num_nans_in_logits: 0,
|
||||
},
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
utility_output: None,
|
||||
finished_requests: Some(
|
||||
{
|
||||
"req-1",
|
||||
},
|
||||
),
|
||||
wave_complete: None,
|
||||
start_wave: None,
|
||||
}
|
||||
"#]]
|
||||
.assert_debug_eq(&decoded_outputs);
|
||||
|
||||
@@ -2543,7 +2598,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
let inline_logprobs =
|
||||
decode_engine_core_outputs(&decode_frames(inline_logprobs_frames)).unwrap();
|
||||
expect_sample_logprobs(
|
||||
inline_logprobs.as_request_batch().unwrap().outputs[0]
|
||||
inline_logprobs.outputs[0]
|
||||
.new_logprobs
|
||||
.as_ref()
|
||||
.expect("inline logprobs decoded"),
|
||||
@@ -2552,7 +2607,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
let multipart_logprobs =
|
||||
decode_engine_core_outputs(&decode_frames(multipart_logprobs_frames)).unwrap();
|
||||
expect_sample_logprobs(
|
||||
multipart_logprobs.as_request_batch().unwrap().outputs[0]
|
||||
multipart_logprobs.outputs[0]
|
||||
.new_logprobs
|
||||
.as_ref()
|
||||
.expect("multipart logprobs decoded"),
|
||||
@@ -2560,7 +2615,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
|
||||
let inline_prompt = decode_engine_core_outputs(&decode_frames(inline_prompt_frames)).unwrap();
|
||||
expect_prompt_logprobs(
|
||||
inline_prompt.as_request_batch().unwrap().outputs[0]
|
||||
inline_prompt.outputs[0]
|
||||
.new_prompt_logprobs_tensors
|
||||
.as_ref()
|
||||
.expect("inline prompt logprobs decoded"),
|
||||
@@ -2569,7 +2624,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
let multipart_prompt =
|
||||
decode_engine_core_outputs(&decode_frames(multipart_prompt_frames)).unwrap();
|
||||
expect_prompt_logprobs(
|
||||
multipart_prompt.as_request_batch().unwrap().outputs[0]
|
||||
multipart_prompt.outputs[0]
|
||||
.new_prompt_logprobs_tensors
|
||||
.as_ref()
|
||||
.expect("multipart prompt logprobs decoded"),
|
||||
@@ -2889,7 +2944,8 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
|
||||
|
||||
send_outputs(
|
||||
&mut push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-1",
|
||||
vec![],
|
||||
@@ -2897,8 +2953,7 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -2964,7 +3019,8 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
|
||||
|
||||
send_outputs(
|
||||
&mut push,
|
||||
RequestBatchOutputs {
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![request_output(
|
||||
"req-1",
|
||||
vec![],
|
||||
@@ -2972,8 +3028,7 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() {
|
||||
)],
|
||||
finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
|
||||
..Default::default()
|
||||
}
|
||||
.into(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -18,8 +18,9 @@ use crate::error::{Error, Result, bail_unexpected_handshake_message};
|
||||
use crate::protocol::handshake::{
|
||||
EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage,
|
||||
};
|
||||
use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs};
|
||||
use crate::protocol::{decode_msgpack, encode_msgpack};
|
||||
use crate::protocol::{
|
||||
EngineCoreOutputs, decode_engine_core_outputs, decode_msgpack, encode_msgpack,
|
||||
};
|
||||
|
||||
/// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the
|
||||
/// engine dies.
|
||||
|
||||
@@ -5,7 +5,7 @@ use clap::Parser;
|
||||
use futures::StreamExt as _;
|
||||
use tokio::time::timeout;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
|
||||
use vllm_llm::{FinishReason, GenerateOutputStream, GenerateRequest, Llm};
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user