Compare commits

..
Author SHA1 Message Date
dependabot[bot]andGitHub 95dcefaaa5 Bump actions/setup-python from 6.1.0 to 6.3.0
Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6.1.0 to 6.3.0.
- [Release notes](https://github.com/actions/setup-python/releases)
- [Commits](https://github.com/actions/setup-python/compare/83679a892e2d95755f2dac6acb0bfd1e9ac5d548...ece7cb06caefa5fff74198d8649806c4678c61a1)

---
updated-dependencies:
- dependency-name: actions/setup-python
  dependency-version: 6.3.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
2026-06-30 12:18:49 +00:00
486 changed files with 13634 additions and 19544 deletions
+1 -3
View File
@@ -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
-25
View File
@@ -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
+1
View File
@@ -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
View File
@@ -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"
+1 -1
View File
@@ -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
+5 -6
View File
@@ -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:
+3 -1
View File
@@ -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.
+5
View File
@@ -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)
+1 -1
View File
@@ -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
View File
@@ -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()
+4 -2
View File
@@ -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()
+2 -2
View File
@@ -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()
-1
View File
@@ -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()
+1 -2
View File
@@ -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
-12
View File
@@ -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);
-128
View File
@@ -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
-22
View File
@@ -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));
-4
View File
@@ -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);
+60
View File
@@ -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));
}
+2 -2
View File
@@ -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
}
+2 -2
View File
@@ -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"); \
-16
View File
@@ -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
}
+8 -3
View File
@@ -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,
+12 -2
View File
@@ -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
View File
@@ -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
+28
View File
@@ -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)
+326
View File
@@ -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 #############################
+1 -1
View File
@@ -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
View File
@@ -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(promptssampling_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
+10 -4
View File
@@ -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>
+1 -1
View File
@@ -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
+2 -3
View File
@@ -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(
+15 -28
View File
@@ -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,
)
```
+2 -2
View File
@@ -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 -14
View File
@@ -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
-21
View File
@@ -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,
}
+1 -1
View File
@@ -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(
+7 -3
View File
@@ -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)
+7 -3
View File
@@ -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
+1 -1
View File
@@ -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())
+10 -3
View File
@@ -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(
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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 = [
+8 -11
View File
@@ -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))
+2 -3
View File
@@ -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
-1
View File
@@ -5310,7 +5310,6 @@ dependencies = [
"vllm-llm",
"vllm-metrics",
"vllm-text",
"vllm-tokenizer",
"zeromq",
]
-1
View File
@@ -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]
+26 -3
View File
@@ -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(
+1 -2
View File
@@ -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};
+57 -11
View File
@@ -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()
+1 -1
View File
@@ -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};
+34 -7
View File
@@ -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());
}
}
+14
View File
@@ -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();
+25 -2
View File
@@ -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"),
+31 -6
View File
@@ -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
View File
@@ -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;
}
+63 -158
View File
@@ -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
View File
@@ -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,
-112
View File
@@ -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,
};
+1 -14
View File
@@ -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) => {
-2
View File
@@ -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,
+636 -6
View File
@@ -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());
}
}
+222 -167
View File
@@ -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;
+3 -2
View File
@@ -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