Compare commits

...
Author SHA1 Message Date
Tyler Michael SmithandClaude 5b4f6d5284 Use cumem allocator for the KV cache by default
Enable the cumem (CUDA VMM) allocator by default on CUDA and ROCm
platforms so users get stable physical pages for KV cache without
needing to set --enable-cumem-allocator explicitly. This is required
for MNNVL KV transfers.

Changes:
- ModelConfig.enable_cumem_allocator defaults to None, resolved to
  True when the cumem C extension is available
- Cached fabric handle probe in csrc/cumem_allocator.cpp uses a real
  cuMemCreate probe instead of trusting cuDeviceGetAttribute, with
  fallback to POSIX FD when fabric handles aren't available
- Safety-net fallback: if fabric was probed as available but a real
  allocation still fails, updates the cache and retries with POSIX FD
- Weight loading bypasses cumem pool when sleep mode is off (cumem is
  only needed for KV cache stability in that case)
- Guard against stale engine allocations corrupting the singleton

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-10 15:25:35 -04:00
Andreas KaratzasandGitHub b697119800 [ROCm][CI] Stabilize ModernBERT token-classification parity against Hugging Face (#44040)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-09 16:52:36 +01:00
Kunshang JiandGitHub b4c6dc6454 [WIP][XPU] upgrade torch-xpu to 2.12 (#42262)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
2026-06-09 15:51:39 +00:00
Raushan TurganbayandGitHub 2ee5106372 Remove raw_inputs from transformers backend (#39425)
Signed-off-by: raushan <raushan@huggingface.co>
2026-06-09 15:01:04 +00:00
Jiangyun ZhuandGitHub 7a89b72564 [Perf] fuse qk rmsnorm rope gate for qwen3.5 (#44176)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-06-09 22:12:17 +08:00
Jee Jee LiandGitHub dc10e467a9 [Bugfix] Fix minimax_qk_norm_fusion (#44983) 2026-06-09 06:43:46 -07:00
Terrence ZhaoandGitHub ee4d7df2b5 [Cohere] Cohere2 moe parser fix (#44907)
Signed-off-by: Terrencezzj <terrence@cohere.ai>
2026-06-09 06:32:18 -07:00
Terrence ZhaoandGitHub 3e8afdf785 [Cohere] Fix Cohere2MoE weight loading when using Transformers ≥5.10 (#44747)
Signed-off-by: Terrencezzj <terrence@cohere.ai>
2026-06-09 06:27:40 -07:00
Nicolò LucchesiandGitHub 6690a0c4de [PD][Bugfix] Fix KV Cache sharing with HMA (#44629)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-06-09 06:10:06 -07:00
Maria GuevaraandGitHub 1c23c42030 [Rust Frontend] Support Kimi K2 tool call IDs (#44901) 2026-06-09 05:31:26 -07:00
xiangdongandGitHub b12e42d132 [XPU][CI] Refine docker image build and pull/create lock mechanism in Intel GPU CI (#44481)
Signed-off-by: zengxian <xiangdong.zeng@intel.com>
2026-06-09 20:20:32 +08:00
69fdaffbcd [Rust Frontend] Add /tokenize and /detokenize endpoints (#44222)
Signed-off-by: Tan Ngoc Do <darkknightkhtn2008@gmail.com>
Signed-off-by: TanNgocDo <darkknightkhtn2008@gmail.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
2026-06-09 05:11:37 -07:00
80e2c4462d [ROCm][Compile] Fuse AR + RMSNorm + per-group FP8 quant (+ DSv3.2 indexer fan-out) (#42864)
Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>
Co-authored-by: Frida Andersson <fanderss@amd.com>
2026-06-09 12:06:56 +00:00
SageandGitHub 5b3807e862 [KV Events] Switch event structs from array to map encoding (#42892)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-06-09 11:39:52 +00:00
Qiuyang YueandGitHub 59401ac9f1 [Kernel][Perf] Tune fused_moe FP8 config for Qwen3-Next-80B tp=4 on H100 (+25% at batch 96-512) (#44830)
Signed-off-by: Qiuyang Yue <yueqiuyang1389@gmail.com>
2026-06-09 04:15:51 -07:00
Chao-Ju ChenGitHubBugen Zhaomergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d841386d27 [Rust Frontend] Support API key authentication (#44321)
Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-09 10:15:20 +00:00
Mohammad Miadh AngkadandGitHub fff9210b2a [CI/Docs] Remove stale disagg prefill links (#44918)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
2026-06-09 03:05:53 -07:00
Ma JianandGitHub 70db1488c5 [DSV4][XPU] Add MHC fused_post_pre support (#44144)
Signed-off-by: Ma Jian <jian1.ma@intel.com>
2026-06-09 17:23:17 +08:00
Andreas KaratzasandGitHub 2385e140d6 [ROCm][CI] Stabilize sleep-mode memory release (#43022)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-09 16:51:12 +08:00
Nicolò LucchesiandGitHub dab60fc658 [Bugfix][CI] Fix test_offloading_connector.py::test_fs_tiering_offloading (#44903)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-06-09 00:57:34 -07:00
wang.yuqiandGitHub 996222f4bf [CI] Reorganize entrypoints CI (#44947)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-09 00:46:11 -07:00
Parth Ashwin JainGitHubParth Ashwin JainCyrus Leungmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
e6fc848d4f [Bugfix][MiniCPM-o] Fix cuda/cpu device mismatch in Resampler2_5 pos_embed (#43844)
Signed-off-by: Parth Ashwin Jain <parthash@amd.com>
Co-authored-by: Parth Ashwin Jain <parthash@amd.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-08 23:28:26 -07:00
Andreas KaratzasandGitHub f843ac1a1c [Bugfix][CI] Gemma3 Transformers multimodal encoder profiling and build prompt-embedding fixtures (#44952)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-09 05:49:30 +00:00
jeff.yeGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
7c2aa3108a fix: prevent MM cache hang from stale LRU order keys (#43595)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-08 22:48:31 -07:00
ebf53ba373 [Bugfix][Rust Frontend] Set a structured-output backend so requests do not 500 (#44729)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
2026-06-08 22:30:54 -07:00
baacbfcebf [ROCm][MLA][Bugfix] Reserve FP8 prefill workspace before lock for Kimi-K2.5 (#42978)
Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-06-08 22:25:52 -07:00
d8218b1ee7 [Bugfix] Propagate ImportError from load_audio_pyav when vllm[audio] … (#44750)
Signed-off-by: littlecircle0730 <littlecircle0730@gmail.com>
Signed-off-by: Hsiao-Yuan Chen <hy.c@Hsiao-YuandeMacBook-Pro.local>
Co-authored-by: Hsiao-Yuan Chen <hy.c@Hsiao-YuandeMacBook-Pro.local>
2026-06-09 04:24:52 +00:00
9f153aa781 [MM][Perf][CG] Support ViT full CUDA graph for glm4_1v image and video inference (#40576)
Signed-off-by: grYe99 <guorongye99@gmail.com>
Co-authored-by: grYe99 <guorongye99@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-09 11:13:56 +08:00
Kunshang JiandGitHub d3de61502f [XPU][CI] fix test case path (#44940)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-08 20:02:31 -07:00
Lanze LiuandGitHub 540aaf2140 [Bugfix][Model] Qwen3-Omni: move cu_seqlens to GPU before VIT attention (#44264)
Signed-off-by: Lanze Liu <lanzetech@gmail.com>
2026-06-08 20:02:27 -07:00
Lanze LiuandGitHub 4128605ad4 [Docs] Remove broken link to deleted disaggregated_prefill.sh (#44929)
Signed-off-by: Lanze Liu <lanzetech@gmail.com>
2026-06-09 01:40:06 +00:00
e2f993dc41 [WideEP] Integrate DeepEP v2 (#41183)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-06-08 18:07:29 -07:00
Andreas KaratzasandGitHub 05cb606cad [ROCm][CI] Re-route NixlConnector jobs (#44809)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-08 18:57:11 -05:00
3f627ebef7 [Misc] usage_stats: report more engine, spec-decode, and EP config (#44595)
Signed-off-by: Zach Xi <zachary.xi@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-08 15:20:00 -07:00
Bugen ZhaoandGitHub bc941f375d [Rust Frontend] [Refactor] Refine utility call interfaces (#44856)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-08 15:13:08 -07:00
Michael GoinandGitHub 6afa25000c [Bugfix] Canonicalize FP8 weight layout to (K, N) at the source (#44735)
Signed-off-by: mgoin <mgoin64@gmail.com>
2026-06-08 14:37:36 -06:00
Mohammad Miadh AngkadandGitHub 823a0ab754 [Bugfix][MoE] Fix fused MoE expert mapping helper call sites (#44897)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
2026-06-08 13:35:04 -07:00
Wentao YeandGitHub 2c27c294c0 [Model Runner V2] Fix mrv2 mm lora issue (#44450)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-08 14:30:09 -04:00
ba94a3b998 [Attention] Extract KV-cache update from CPU attention backend (#40470)
Signed-off-by: Diego Maniloff <diego.maniloff@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-06-08 15:43:05 +00:00
bnellnmandGitHub dc68bd8c41 [MoE Refactor] FusedMoE/MoERunner inversion refactor (#41184)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-06-08 10:42:58 -04:00
Walter Beller-MoralesGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
753e9d55e6 [Quantization] add online fp8 ptpc (#44132)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-08 22:42:11 +08:00
akii96andGitHub ac3409d162 [Benchmark] Auto-detect and correct client/server tokenizer mismatch for random dataset (#44708) 2026-06-08 06:10:20 -07:00
wang.yuqiandGitHub 93ee4cd47f [CI] Consolidate multimodal entrypoint tests. (#44819)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-08 04:48:08 -07:00
Li, JiangandGitHub 980796cd07 [CI/Build][CPU] Fix flaky CI image build failure and unexpected warnings (#44852)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-06-08 11:10:06 +00:00
Nicolò LucchesiandGitHub 5add018beb [Connector] Remove P2pNcclConnector (#44854)
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-06-08 18:58:29 +08:00
Jonathan MamouGitHubLi, Jiang <bigpyj64@gmail.com>Li, Jiang <jiang1.li@intel.com>
d5fe994e79 [CPU][Spec Decode] Warn about throughput loss when libiomp5 is not preloaded (#44419)
Signed-off-by: jmamou <jonathan.mamou@intel.com>
Signed-off-by: Jonathan Mamou <jonathan.mamou@intel.com>
Co-authored-by: Li, Jiang <bigpyj64@gmail.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-06-08 03:45:08 -07:00
Chaojun ZhangandGitHub fa662b1a8b [XPU] Cap topk/topp Triton BLOCK_SIZE to 4096 to fix Top-p mask difference failures (#44470)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
2026-06-08 09:36:51 +00:00
3c0b4432be [Rust Frontend] Add /pause, /resume, /is_paused endpoints (#44499)
Signed-off-by: Sahil Singh <sahiilsiingh37@gmail.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
2026-06-08 17:28:37 +08:00
Sungjae LeeandGitHub 469f3dcf1d [BugFix] Use served model name in gemma4 audio-tower error message (#44828)
Signed-off-by: Sungjae Lee <33976427+llsj14@users.noreply.github.com>
Signed-off-by: Sungjae Lee <sung-jae.lee@navercorp.com>
2026-06-08 06:58:31 +00:00
xiangdongandGitHub 94fcdd007f [XPU][CI] Add more test cases in Intel GPU CI (#43663)
Signed-off-by: zengxian <xiangdong.zeng@intel.com>
2026-06-08 06:21:24 +00:00
Andreas KaratzasandGitHub d9ff7e4e9a [ROCm][CI] Stabilizing teardown and timeout of flaky tests to prevent rare OOMs (#44761)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-08 14:11:17 +08:00
268 changed files with 9762 additions and 5921 deletions
@@ -0,0 +1,23 @@
group: Expert Parallelism
depends_on:
- image-build-xpu
steps:
- label: EPLB Algorithm
key: eplb-algorithm
timeout_in_minutes: 45
device: intel_gpu
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/eplb
- tests/distributed/test_eplb_algo.py
- tests/distributed/test_eplb_utils.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s distributed/test_eplb_algo.py'
+111 -2
View File
@@ -38,7 +38,17 @@ steps:
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/inputs/
- vllm/logger.py
- vllm/model_executor/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/sample
- tests/v1/logits_processors
- tests/v1/test_oracle.py
@@ -52,7 +62,8 @@ steps:
pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
pytest -v -s v1/test_oracle.py &&
pytest -v -s v1/test_request.py &&
pytest -v -s v1/test_outputs.py'
pytest -v -s v1/test_outputs.py &&
pytest -v -s v1/sample/test_topk_topp_sampler.py'
- label: XPU CPU Offload
timeout_in_minutes: 60
@@ -76,3 +87,101 @@ steps:
cd tests &&
pytest -v -s v1/kv_offload &&
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
- label: Regression
key: regression
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/inputs/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/test_regression
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install modelscope &&
cd tests &&
pytest -v -s test_regression.py'
- label: Metrics, Tracing (2 GPUs)
key: metrics-tracing-2-gpus
timeout_in_minutes: 30
num_devices: 2
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/inputs/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/tracing/
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/v1/tracing
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 &&
cd tests &&
pytest -v -s v1/tracing'
- label: Async Engine, Inputs, Utils, Worker
key: async-engine-inputs-utils-worker
timeout_in_minutes: 30
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/assets/
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/inputs/
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/sampling_params.py
- vllm/tokenizers/
- vllm/transformers_utils/
- vllm/utils/
- vllm/v1/
- tests/detokenizer
- tests/multimodal
- tests/utils_
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pip install av &&
pytest -v -s detokenizer &&
pytest -v -s -m "not cpu_test" ./multimodal &&
pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
@@ -0,0 +1,111 @@
group: Models - Multimodal
depends_on:
- image-build-xpu
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/multimodal
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'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'
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/multimodal
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'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"
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/multimodal
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'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'
- label: "Multi-Modal Models (Standard) 4: other + whisper"
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/multimodal
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'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'
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 45
device: intel_gpu
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/multimodal
- tests/models/registry.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
pip install open-clip-torch --no-deps &&
cd tests &&
pytest -v -s models/multimodal/processing/test_tensor_schema.py
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
--deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
parallelism: 4
+1 -1
View File
@@ -85,5 +85,5 @@ steps:
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'pip install av &&
cd tests &&
pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
pytest -v -s benchmarks/test_serve_cli.py'
@@ -324,23 +324,6 @@ IMAGE="${IMAGE_TAG_XPU:-${image_name}}"
echo "Using image: ${IMAGE}"
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
echo "Image already exists locally, skipping pull"
else
echo "Image not found locally, waiting for lock..."
flock /tmp/docker-pull.lock bash -c "
if docker image inspect '${IMAGE}' >/dev/null 2>&1; then
echo 'Image already pulled by another runner'
else
echo 'Pulling image...'
timeout 900 docker pull '${IMAGE}'
fi
"
echo "Pull step completed"
fi
remove_docker_container() {
docker rm -f "${container_name}" || true
}
@@ -357,9 +340,12 @@ export HF_TOKEN ZE_AFFINITY_MASK
{
flock 9
if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
echo 'Image missing before container creation, pulling again...'
if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
echo "Image already exists locally, skipping pull"
else
echo "Image not found locally, pulling image..."
timeout 900 docker pull "${IMAGE}"
echo "Pull step completed"
fi
docker create \
@@ -372,6 +358,8 @@ export HF_TOKEN ZE_AFFINITY_MASK
--entrypoint='' \
-e HF_TOKEN \
-e ZE_AFFINITY_MASK \
-e BUILDKITE_PARALLEL_JOB \
-e BUILDKITE_PARALLEL_JOB_COUNT \
-e CMDS \
--name "${container_name}" \
"${IMAGE}" \
+146 -115
View File
@@ -88,16 +88,16 @@
# - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function. #
# * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock. #
# * [Plugin Tests (2 GPUs)]: #
# - {`pytest -v -s entrypoints/openai/test_oot_registration.py`}: It needs a clean process #
# - {`pytest -v -s models/test_oot_registration.py`}: It needs a clean process #
# - {`pytest -v -s plugins/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
# - {`pytest -v -s plugins_tests/test_oot_registration_online.py`}: It needs a clean process #
# - {`pytest -v -s plugins_tests/test_oot_registration_offline.py`}: It needs a clean process #
# - {`pytest -v -s plugins_tests/lora_resolvers`}: Unit tests for in-tree lora resolver plugins #
# * [LoRA TP (Distributed)]: #
# - There is some Tensor Parallelism related processing logic in LoRA that requires multi-GPU testing for validation. #
# - {`pytest -v -s -x lora/test_gptoss_tp.py`}: Disabled for now because MXFP4 backend on non-cuda platform doesn't support #
# LoRA yet. #
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Don't test llama model here, it seems hf implementation is buggy. See: #
# https://github.com/vllm-project/vllm/pull/5689 #
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
# * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293 #
# in favor of new tests in fusions_e2e. We avoid replicating the new jobs in #
# this file as it's deprecated. #
# #
@@ -617,9 +617,9 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s models/test_oot_registration.py
- pytest -v -s plugins/lora_resolvers
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
#------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------#
@@ -804,7 +804,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: V1 e2e (2 GPUs) # TBD
timeout_in_minutes: 180
@@ -830,7 +830,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
#------------------------------------------------------------- mi250 · misc ------------------------------------------------------------#
@@ -1187,88 +1187,20 @@ steps:
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
- label: Entrypoints Integration (API Server 2) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
- label: Entrypoints Unit Tests # TBD
timeout_in_minutes: 180
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:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- vllm/entrypoints
- tests/entrypoints/unit_tests
- tests/entrypoints/weight_transfer
- vllm/platforms/rocm.py
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
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/
- tests/entrypoints/speech_to_text
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- pytest -v -s entrypoints/unit_tests
- pytest -v -s entrypoints/weight_transfer
- label: Entrypoints Integration (LLM) # TBD
timeout_in_minutes: 180
@@ -1287,19 +1219,75 @@ steps:
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
- label: Entrypoints Integration (Pooling) # TBD
- label: Entrypoints Integration (API Server) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/pooling
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/pooling
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
timeout_in_minutes: 180
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:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
timeout_in_minutes: 180
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:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: Entrypoints Integration (API Server Generate) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
fast_check: true
torch_nightly: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/tool_use
- tests/entrypoints/tool_parsers
- tests/entrypoints/anthropic
- tests/entrypoints/generate
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use
- pytest -v -s entrypoints/tool_parsers
- pytest -v -s entrypoints/generate
- pytest -v -s entrypoints/anthropic
- label: Entrypoints Integration (Responses API) # TBD
timeout_in_minutes: 180
@@ -1315,19 +1303,47 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/responses
- label: Entrypoints Unit Tests # TBD
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
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/entrypoints
- tests/entrypoints/
- vllm/platforms/rocm.py
- vllm/
- tests/entrypoints/speech_to_text
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Multimodal)
timeout_in_minutes: 180
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/
- tests/entrypoints/multimodal
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/multimodal
- label: Entrypoints Integration (Pooling) # TBD
timeout_in_minutes: 180
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/
- tests/entrypoints/pooling
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/pooling
- label: OpenAI API correctness # TBD
timeout_in_minutes: 180
@@ -1338,7 +1354,6 @@ steps:
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
- vllm/model_executor/layers/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
@@ -2391,7 +2406,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed DP Tests (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -2427,7 +2442,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -2441,7 +2456,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -2455,7 +2470,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: V1 e2e (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -2783,7 +2798,7 @@ steps:
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
- label: Entrypoints Integration (API Server 2) # TBD
- label: Entrypoints Integration (API Server) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
@@ -2799,7 +2814,7 @@ steps:
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
@@ -2813,9 +2828,9 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
@@ -2827,29 +2842,31 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
- label: Entrypoints Integration (API Server Generate) # TBD
timeout_in_minutes: 180
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:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/tool_use
- tests/entrypoints/tool_parsers
- tests/entrypoints/anthropic
- tests/entrypoints/generate
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- pytest -v -s tool_use
- pytest -v -s entrypoints/tool_parsers
- pytest -v -s entrypoints/generate
- pytest -v -s entrypoints/anthropic
- label: Entrypoints Integration (Speech to Text) # TBD
timeout_in_minutes: 180
@@ -2865,6 +2882,20 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Multimodal)
timeout_in_minutes: 180
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/
- tests/entrypoints/multimodal
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/multimodal
- label: Entrypoints Integration (Pooling) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3042,7 +3073,7 @@ steps:
- vllm/_aiter_ops.py
commands:
- rocm-smi
- python3 examples/basic/offline_inference/chat.py
- python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
- pytest -v -s tests/kernels/attention/test_attention_selector.py
- label: Kernels Attention Test %N # TBD
@@ -3406,7 +3437,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -3421,7 +3452,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
timeout_in_minutes: 180
@@ -3436,7 +3467,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
#------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------#
+87 -75
View File
@@ -8,10 +8,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/entrypoints
- tests/entrypoints/
- tests/entrypoints/unit_tests
- tests/entrypoints/weight_transfer
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- pytest -v -s entrypoints/unit_tests
- pytest -v -s entrypoints/weight_transfer
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
@@ -31,68 +32,9 @@ steps:
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server openai - Part 1)
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
mirror:
amd:
device: mi300_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server openai - Part 2)
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
mirror:
amd:
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server openai - Part 3)
key: entrypoints-integration-api-server-openai-part-3
timeout_in_minutes: 50
device: h200_18gb
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
mirror:
amd:
device: mi325_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server 2)
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
key: entrypoints-integration-api-server-2
timeout_in_minutes: 130
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -108,6 +50,75 @@ steps:
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server OpenAI - Part 1)
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
mirror:
amd:
device: mi300_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server OpenAI - Part 2)
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/chat_completion
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
mirror:
amd:
device: mi300_1
timeout_in_minutes: 80
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server Generate)
key: entrypoints-integration-api-server-generate
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/tool_use
- tests/entrypoints/tool_parsers
- tests/entrypoints/anthropic
- tests/entrypoints/generate
commands:
- pytest -v -s tool_use
- pytest -v -s entrypoints/tool_parsers
- pytest -v -s entrypoints/generate
- pytest -v -s entrypoints/anthropic
mirror:
amd:
device: mi300_1
timeout_in_minutes: 60
depends_on:
- image-build-amd
- label: Entrypoints Integration (Responses API)
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai/responses
commands:
- pytest -v -s entrypoints/openai/responses
- label: Entrypoints Integration (Speech to Text)
device: h200_35gb
key: entrypoints-integration-speech_to_text
@@ -120,6 +131,18 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/speech_to_text
- label: Entrypoints Integration (Multimodal)
device: h200_35gb
key: entrypoints-integration-multimodal
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/multimodal
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/multimodal
- label: Entrypoints Integration (Pooling)
key: entrypoints-integration-pooling
timeout_in_minutes: 50
@@ -131,16 +154,6 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/pooling
- label: Entrypoints Integration (Responses API)
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/openai/responses
commands:
- pytest -v -s entrypoints/openai/responses
- label: OpenAI API Correctness
key: openai-api-correctness
timeout_in_minutes: 30
@@ -158,7 +171,6 @@ steps:
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
- vllm/model_executor/layers/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
+3 -1
View File
@@ -21,8 +21,9 @@ steps:
- csrc/
- tests/kernels/core
- tests/kernels/test_concat_mla_q.py
- tests/kernels/test_fused_qk_norm_rope_gate.py
commands:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
key: kernels-minimax-reduce-rms-test-2-gpus
@@ -299,3 +300,4 @@ steps:
- vllm/config
commands:
- pytest -v -s kernels/moe/test_moe_layer.py
- pytest -v -s kernels/moe/test_deepep_v2_moe.py
+3 -3
View File
@@ -37,6 +37,6 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
- pytest -v -s models/test_oot_registration.py # it needs a clean process
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
- pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
@@ -1,143 +0,0 @@
#!/bin/bash
# benchmark the overhead of disaggregated prefill.
# methodology:
# - send all request to prefill vLLM instance. It will buffer KV cache.
# - then send all request to decode instance.
# - The TTFT of decode instance is the overhead.
set -ex
kill_gpu_processes() {
# kill all processes on GPU.
pgrep pt_main_thread | xargs -r kill -9
pgrep python3 | xargs -r kill -9
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
pgrep VLLM | xargs -r kill -9
sleep 10
# remove vllm config file
rm -rf ~/.config/vllm
# Print the GPU memory usage
# so that we know if all GPU processes are killed.
gpu_memory_usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0)
# The memory usage should be 0 MB.
echo "GPU 0 Memory Usage: $gpu_memory_usage MB"
}
wait_for_server() {
# wait for vllm server to start
# return 1 if vllm server crashes
local port=$1
timeout 1200 bash -c "
until curl -s localhost:${port}/v1/completions > /dev/null; do
sleep 1
done" && return 0 || return 1
}
benchmark() {
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
# compare chunked prefill with disaggregated prefill
results_folder="./results"
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
dataset_name="sonnet"
dataset_path="../sonnet_4x.txt"
num_prompts=10
qps=$1
prefix_len=50
input_len=2048
output_len=$2
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
--port 8100 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
--port 8200 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
wait_for_server 8100
wait_for_server 8200
# let the prefill instance finish prefill
vllm bench serve \
--backend vllm \
--model $model \
--dataset-name $dataset_name \
--dataset-path $dataset_path \
--sonnet-input-len $input_len \
--sonnet-output-len "$output_len" \
--sonnet-prefix-len $prefix_len \
--num-prompts $num_prompts \
--port 8100 \
--save-result \
--result-dir $results_folder \
--result-filename disagg_prefill_tp1.json \
--request-rate "inf"
# send the request to decode.
# The TTFT of this command will be the overhead of disagg prefill impl.
vllm bench serve \
--backend vllm \
--model $model \
--dataset-name $dataset_name \
--dataset-path $dataset_path \
--sonnet-input-len $input_len \
--sonnet-output-len "$output_len" \
--sonnet-prefix-len $prefix_len \
--num-prompts $num_prompts \
--port 8200 \
--save-result \
--result-dir $results_folder \
--result-filename disagg_prefill_tp1_overhead.json \
--request-rate "$qps"
kill_gpu_processes
}
main() {
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
(which jq) || (apt-get -y install jq)
(which socat) || (apt-get -y install socat)
pip install quart httpx datasets
cd "$(dirname "$0")"
cd ..
# create sonnet-4x.txt
echo "" > sonnet_4x.txt
for _ in {1..4}
do
cat sonnet.txt >> sonnet_4x.txt
done
cd disagg_benchmarks
rm -rf results
mkdir results
default_qps=1
default_output_len=1
benchmark $default_qps $default_output_len
}
main "$@"
@@ -1,157 +0,0 @@
#!/bin/bash
# Requirement: 2x GPUs.
# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
# Resource: 2x GPU
# Approaches:
# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
# Prefilling instance: max_output_token=1
# Decoding instance: force the input tokens be the same across requests to bypass prefilling
set -ex
kill_gpu_processes() {
# kill all processes on GPU.
pgrep pt_main_thread | xargs -r kill -9
pgrep python3 | xargs -r kill -9
# vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
pgrep VLLM | xargs -r kill -9
for port in 8000 8100 8200; do lsof -t -i:$port | xargs -r kill -9; done
sleep 1
}
wait_for_server() {
# wait for vllm server to start
# return 1 if vllm server crashes
local port=$1
timeout 1200 bash -c "
until curl -s localhost:${port}/v1/completions > /dev/null; do
sleep 1
done" && return 0 || return 1
}
launch_chunked_prefill() {
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
# disagg prefill
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
--port 8100 \
--max-model-len 10000 \
--enable-chunked-prefill \
--gpu-memory-utilization 0.6 &
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
--port 8200 \
--max-model-len 10000 \
--enable-chunked-prefill \
--gpu-memory-utilization 0.6 &
wait_for_server 8100
wait_for_server 8200
python3 round_robin_proxy.py &
sleep 1
}
launch_disagg_prefill() {
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
# disagg prefill
CUDA_VISIBLE_DEVICES=0 vllm serve $model \
--port 8100 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
CUDA_VISIBLE_DEVICES=1 vllm serve $model \
--port 8200 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
wait_for_server 8100
wait_for_server 8200
python3 disagg_prefill_proxy_server.py &
sleep 1
}
benchmark() {
results_folder="./results"
model="meta-llama/Meta-Llama-3.1-8B-Instruct"
dataset_name="sonnet"
dataset_path="../sonnet_4x.txt"
num_prompts=100
qps=$1
prefix_len=50
input_len=1024
output_len=$2
tag=$3
vllm bench serve \
--backend vllm \
--model $model \
--dataset-name $dataset_name \
--dataset-path $dataset_path \
--sonnet-input-len $input_len \
--sonnet-output-len "$output_len" \
--sonnet-prefix-len $prefix_len \
--num-prompts $num_prompts \
--port 8000 \
--save-result \
--result-dir $results_folder \
--result-filename "$tag"-qps-"$qps".json \
--request-rate "$qps"
sleep 2
}
main() {
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
(which jq) || (apt-get -y install jq)
(which socat) || (apt-get -y install socat)
(which lsof) || (apt-get -y install lsof)
pip install quart httpx matplotlib aiohttp datasets
cd "$(dirname "$0")"
cd ..
# create sonnet-4x.txt so that we can sample 2048 tokens for input
echo "" > sonnet_4x.txt
for _ in {1..4}
do
cat sonnet.txt >> sonnet_4x.txt
done
cd disagg_benchmarks
rm -rf results
mkdir results
default_output_len=6
export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
launch_chunked_prefill
for qps in 2 4 6 8; do
benchmark $qps $default_output_len chunked_prefill
done
kill_gpu_processes
launch_disagg_prefill
for qps in 2 4 6 8; do
benchmark $qps $default_output_len disagg_prefill
done
kill_gpu_processes
python3 visualize_benchmark_results.py
}
main "$@"
@@ -1,260 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import asyncio
import logging
import os
import time
import uuid
from urllib.parse import urlparse
import aiohttp
from quart import Quart, Response, make_response, request
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def parse_args():
"""parse command line arguments"""
parser = argparse.ArgumentParser(description="vLLM P/D disaggregation proxy server")
# Add args
parser.add_argument(
"--timeout",
type=float,
default=6 * 60 * 60,
help="Timeout for backend service requests in seconds (default: 21600)",
)
parser.add_argument(
"--port",
type=int,
default=8000,
help="Port to run the server on (default: 8000)",
)
parser.add_argument(
"--prefill-url",
type=str,
default="http://localhost:8100",
help="Prefill service base URL (protocol + host[:port])",
)
parser.add_argument(
"--decode-url",
type=str,
default="http://localhost:8200",
help="Decode service base URL (protocol + host[:port])",
)
parser.add_argument(
"--kv-host",
type=str,
default="localhost",
help="Hostname or IP used by KV transfer (default: localhost)",
)
parser.add_argument(
"--prefill-kv-port",
type=int,
default=14579,
help="Prefill KV port (default: 14579)",
)
parser.add_argument(
"--decode-kv-port",
type=int,
default=14580,
help="Decode KV port (default: 14580)",
)
return parser.parse_args()
def main():
"""parse command line arguments"""
args = parse_args()
# Initialize configuration using command line parameters
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
PREFILL_SERVICE_URL = args.prefill_url
DECODE_SERVICE_URL = args.decode_url
PORT = args.port
PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
logger.info(
"Proxy resolved KV addresses -> prefill: %s, decode: %s",
PREFILL_KV_ADDR,
DECODE_KV_ADDR,
)
app = Quart(__name__)
# Attach the configuration object to the application instance so helper
# coroutines can read the resolved backend URLs and timeouts without using
# globals.
app.config.update(
{
"AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
"PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
"DECODE_SERVICE_URL": DECODE_SERVICE_URL,
"PREFILL_KV_ADDR": PREFILL_KV_ADDR,
"DECODE_KV_ADDR": DECODE_KV_ADDR,
}
)
def _normalize_base_url(url: str) -> str:
"""Remove any trailing slash so path joins behave predictably."""
return url.rstrip("/")
def _get_host_port(url: str) -> str:
"""Return the hostname:port portion for logging and KV headers."""
parsed = urlparse(url)
host = parsed.hostname or "localhost"
port = parsed.port
if port is None:
port = 80 if parsed.scheme == "http" else 443
return f"{host}:{port}"
PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
def _build_headers(request_id: str) -> dict[str, str]:
"""Construct the headers expected by vLLM's P2P disagg connector."""
headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
api_key = os.environ.get("OPENAI_API_KEY")
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
async def _run_prefill(
request_path: str,
payload: dict,
headers: dict[str, str],
request_id: str,
):
url = f"{PREFILL_BASE}{request_path}"
start_ts = time.perf_counter()
logger.info("[prefill] start request_id=%s url=%s", request_id, url)
try:
async with (
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
session.post(url=url, json=payload, headers=headers) as resp,
):
if resp.status != 200:
error_text = await resp.text()
raise RuntimeError(
f"Prefill backend error {resp.status}: {error_text}"
)
await resp.read()
logger.info(
"[prefill] done request_id=%s status=%s elapsed=%.2fs",
request_id,
resp.status,
time.perf_counter() - start_ts,
)
except asyncio.TimeoutError as exc:
raise RuntimeError(f"Prefill service timeout at {url}") from exc
except aiohttp.ClientError as exc:
raise RuntimeError(f"Prefill service unavailable at {url}") from exc
async def _stream_decode(
request_path: str,
payload: dict,
headers: dict[str, str],
request_id: str,
):
url = f"{DECODE_BASE}{request_path}"
# Stream tokens from the decode service once the prefill stage has
# materialized KV caches on the target workers.
logger.info("[decode] start request_id=%s url=%s", request_id, url)
try:
async with (
aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
session.post(url=url, json=payload, headers=headers) as resp,
):
if resp.status != 200:
error_text = await resp.text()
logger.error(
"Decode backend error %s - %s", resp.status, error_text
)
err_msg = (
'{"error": "Decode backend error ' + str(resp.status) + '"}'
)
yield err_msg.encode()
return
logger.info(
"[decode] streaming response request_id=%s status=%s",
request_id,
resp.status,
)
async for chunk_bytes in resp.content.iter_chunked(1024):
yield chunk_bytes
logger.info("[decode] finished streaming request_id=%s", request_id)
except asyncio.TimeoutError:
logger.error("Decode service timeout at %s", url)
yield b'{"error": "Decode service timeout"}'
except aiohttp.ClientError as exc:
logger.error("Decode service error at %s: %s", url, exc)
yield b'{"error": "Decode service unavailable"}'
async def process_request():
"""Process a single request through prefill and decode stages"""
try:
original_request_data = await request.get_json()
# Create prefill request (max_tokens=1)
prefill_request = original_request_data.copy()
prefill_request["max_tokens"] = 1
if "max_completion_tokens" in prefill_request:
prefill_request["max_completion_tokens"] = 1
# Execute prefill stage
# The request id encodes both KV socket addresses so the backend can
# shuttle tensors directly via NCCL once the prefill response
# completes.
request_id = (
f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
)
headers = _build_headers(request_id)
await _run_prefill(request.path, prefill_request, headers, request_id)
# Execute decode stage and stream response
# Pass the unmodified user request so the decode phase can continue
# sampling with the already-populated KV cache.
generator = _stream_decode(
request.path, original_request_data, headers, request_id
)
response = await make_response(generator)
response.timeout = None # Disable timeout for streaming response
return response
except Exception:
logger.exception("Error processing request")
return Response(
response=b'{"error": "Internal server error"}',
status=500,
content_type="application/json",
)
@app.route("/v1/completions", methods=["POST"])
async def handle_request():
"""Handle incoming API requests with concurrency and rate limiting"""
try:
return await process_request()
except asyncio.CancelledError:
logger.warning("Request cancelled")
return Response(
response=b'{"error": "Request cancelled"}',
status=503,
content_type="application/json",
)
# Start the Quart server with host can be set to 0.0.0.0
app.run(port=PORT)
if __name__ == "__main__":
main()
@@ -1,63 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
import itertools
import aiohttp
from aiohttp import web
class RoundRobinProxy:
def __init__(self, target_ports):
self.target_ports = target_ports
self.port_cycle = itertools.cycle(self.target_ports)
async def handle_request(self, request):
target_port = next(self.port_cycle)
target_url = f"http://localhost:{target_port}{request.path_qs}"
async with aiohttp.ClientSession() as session:
try:
# Forward the request
async with session.request(
method=request.method,
url=target_url,
headers=request.headers,
data=request.content,
) as response:
# Start sending the response
resp = web.StreamResponse(
status=response.status, headers=response.headers
)
await resp.prepare(request)
# Stream the response content
async for chunk in response.content.iter_any():
await resp.write(chunk)
await resp.write_eof()
return resp
except Exception as e:
return web.Response(text=f"Error: {str(e)}", status=500)
async def main():
proxy = RoundRobinProxy([8100, 8200])
app = web.Application()
app.router.add_route("*", "/{path:.*}", proxy.handle_request)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "localhost", 8000)
await site.start()
print("Proxy server started on http://localhost:8000")
# Keep the server running
await asyncio.Event().wait()
if __name__ == "__main__":
asyncio.run(main())
@@ -1,47 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import matplotlib.pyplot as plt
import pandas as pd
if __name__ == "__main__":
data = []
for name in ["disagg_prefill", "chunked_prefill"]:
for qps in [2, 4, 6, 8]:
with open(f"results/{name}-qps-{qps}.json") as f:
x = json.load(f)
x["name"] = name
x["qps"] = qps
data.append(x)
df = pd.DataFrame.from_dict(data)
dis_df = df[df["name"] == "disagg_prefill"]
chu_df = df[df["name"] == "chunked_prefill"]
plt.style.use("bmh")
plt.rcParams["font.size"] = 20
for key in [
"mean_ttft_ms",
"median_ttft_ms",
"p99_ttft_ms",
"mean_itl_ms",
"median_itl_ms",
"p99_itl_ms",
]:
fig, ax = plt.subplots(figsize=(11, 7))
plt.plot(
dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4
)
plt.plot(
chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4
)
ax.legend()
ax.set_xlabel("QPS")
ax.set_ylabel(key)
ax.set_ylim(bottom=0)
fig.savefig(f"results/{key}.png")
plt.close(fig)
+1 -1
View File
@@ -250,7 +250,7 @@ def benchmark_config(
num_experts=num_experts,
experts_per_token=topk,
hidden_dim=hidden_size,
intermediate_size_per_partition=shard_intermediate_size,
intermediate_size=shard_intermediate_size,
num_local_experts=num_experts,
num_logical_experts=num_experts,
activation=MoEActivation.SILU,
+98 -13
View File
@@ -1,6 +1,7 @@
// A CUDAPluggableAllocator based on cumem* APIs.
// Important: allocation size, CUdeviceptr and CUmemGenericAllocationHandle*
// need to be unsigned long long
#include <atomic>
#include <iostream>
#include "cumem_allocator_compat.h"
@@ -9,6 +10,7 @@
static const char* PYARGS_PARSE = "KKKK";
#else
#include <cstdlib>
#include <cstdint>
#include <cerrno>
#include <climits>
@@ -46,6 +48,29 @@ static inline unsigned long long my_min(unsigned long long a,
return a < b ? a : b;
}
static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
size_t alignment) {
CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);
if (status == CUresult(0) || alignment == 0) {
return status;
}
// Some ROCm stacks can report OOM while reserving VA with an explicit
// alignment even when physical VRAM is free. Let HIP choose the default
// alignment, then verify that the returned address still satisfies the
// requested alignment before accepting it.
status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
if (status != CUresult(0)) {
return status;
}
if (((std::uintptr_t)(*d_mem) % alignment) == 0) {
return status;
}
(void)cuMemAddressFree(*d_mem, size);
return hipErrorNotSupported;
}
static const char* PYARGS_PARSE = "KKKO";
#endif
@@ -92,6 +117,59 @@ void ensure_context(unsigned long long device) {
}
}
// ---------------------------------------------------------------------------
// Cached fabric handle probe (CUDA 12.4+, NVIDIA only):
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040
// Per-device cache: 0 = not probed, 1 = supported, 2 = not supported
static constexpr int MAX_DEVICES = 32;
static std::atomic<int> fabric_support[MAX_DEVICES] = {};
static bool probe_fabric_support(unsigned long long device) {
if (device >= MAX_DEVICES) return false;
int cached = fabric_support[device].load(std::memory_order_acquire);
if (cached != 0) return cached == 1;
int fab_flag = 0;
CUresult r = cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
if (r != CUDA_SUCCESS || !fab_flag) {
fabric_support[device].store(2, std::memory_order_release);
return false;
}
// Attribute says supported — verify with a real allocation.
// cuDeviceGetAttribute can report supported even when IMEX is not
// configured, so we need a real probe.
CUmemAllocationProp probe_prop = {};
probe_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
probe_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
probe_prop.location.id = device;
probe_prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
size_t granularity;
r = cuMemGetAllocationGranularity(&granularity, &probe_prop,
CU_MEM_ALLOC_GRANULARITY_MINIMUM);
if (r != CUDA_SUCCESS) {
fabric_support[device].store(2, std::memory_order_release);
return false;
}
CUmemGenericAllocationHandle test_handle;
r = cuMemCreate(&test_handle, granularity, &probe_prop, 0);
if (r == CUDA_SUCCESS) {
cuMemRelease(test_handle);
fabric_support[device].store(1, std::memory_order_release);
return true;
}
fabric_support[device].store(2, std::memory_order_release);
return false;
}
#endif
// ---------------------------------------------------------------------------
void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
#ifndef USE_ROCM
CUmemGenericAllocationHandle* p_memHandle) {
@@ -112,32 +190,40 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
CUresult rdma_result = cuDeviceGetAttribute(
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
device);
if (rdma_result == CUDA_SUCCESS &&
flag) { // support GPUDirect RDMA if possible
if (rdma_result == CUDA_SUCCESS && flag) {
prop.allocFlags.gpuDirectRDMACapable = 1;
}
int fab_flag = 0;
CUresult fab_result = cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
if (fab_result == CUDA_SUCCESS &&
fab_flag) { // support fabric handle if possible
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
if (probe_fabric_support(device)) {
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
} else {
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
}
#else
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
#endif
#endif
#ifndef USE_ROCM
// Allocate memory using cuMemCreate
CUresult ret = (CUresult)cuMemCreate(p_memHandle, size, &prop, 0);
if (ret) {
if (fab_flag &&
#if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
// Safety net: if fabric was probed as available but this allocation
// still fails, fall back to POSIX FD and update the cache.
if (device < MAX_DEVICES &&
fabric_support[device].load(std::memory_order_acquire) == 1 &&
(ret == CUDA_ERROR_NOT_PERMITTED || ret == CUDA_ERROR_NOT_SUPPORTED)) {
// Fabric allocation may fail without multi-node nvlink,
// fallback to POSIX file descriptor
fabric_support[device].store(2, std::memory_order_release);
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));
} else {
CUDA_CHECK(ret);
}
#else
CUDA_CHECK(ret);
#endif
}
if (error_code != 0) {
return;
@@ -302,14 +388,13 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
// first allocation, align the size, and reserve an address, and also allocate
// a CUmemGenericAllocationHandle
// Define memory allocation properties
// No handle type here; create_and_map sets fabric/POSIX as needed.
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = device;
prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
// Check if the allocation is supported
size_t granularity;
CUDA_CHECK(cuMemGetAllocationGranularity(&granularity, &prop,
CU_MEM_ALLOC_GRANULARITY_MINIMUM));
@@ -325,7 +410,7 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
return nullptr;
}
#else
CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity));
if (error_code != 0) {
return nullptr;
}
+3 -3
View File
@@ -118,9 +118,9 @@ ENV CARGO_BUILD_JOBS=4
# Build the release binary. Cache cargo registry/git and target/, but copy the
# binary out of the target/ cache mount so it persists into the image layer
# for later COPY --from=rust-build.
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
--mount=type=cache,target=/workspace/rust/target,sharing=locked \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
######################### BUILD IMAGE #########################
+92 -79
View File
@@ -89,8 +89,8 @@ RUN mkdir neo && \
ENV PATH="/root/.local/bin:$PATH"
ENV VIRTUAL_ENV="/opt/venv"
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
&& uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
@@ -99,8 +99,8 @@ RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
rm "${ONECCL_INSTALLER}" && \
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
RUN rm -f /opt/intel/oneapi/ccl/latest && \
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
rm -f /opt/intel/oneapi/ccl/latest && \
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
SHELL ["bash", "-c"]
@@ -119,19 +119,98 @@ ENV UV_LINK_MODE="copy"
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
uv pip install --upgrade pip && \
uv pip install -r requirements/xpu.txt && \
uv pip install grpcio-tools protobuf nanobind && \
source /opt/intel/oneapi/setvars.sh --force && \
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt
uv pip install --upgrade pip
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
CMD ["/bin/bash"]
######################### UCX + NIXL BUILD STAGE #########################
# Build UCX and NIXL in a dedicated stage so compiler/autotools layers are
# never included in the final runtime image (mirrors ROCm's build_rixl stage).
FROM vllm-base AS ucx-nixl-build
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
ARG NIXL_VERSION=0.7.0
# Build-time only: compiler, autotools, and verbs dev headers
RUN apt-get update -y && apt-get install -y --no-install-recommends \
build-essential \
autoconf \
automake \
libtool \
pkg-config \
libibverbs-dev \
librdmacm-dev \
&& rm -rf /var/lib/apt/lists/*
# Build UCX and produce a NIXL wheel so the final image needs no compiler.
# patchelf (installed via uv) is used by the NIXL wheel build to rewrite
# RPATH entries, making the wheel portable across stages.
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
bash autogen.sh && \
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
uv pip install --upgrade meson pybind11 patchelf && \
uv pip install -r requirements.txt && \
PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
LD_LIBRARY_PATH=/tmp/ucx_install/lib \
python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
rm -rf /tmp/ucx_source /tmp/nixl_source
FROM vllm-base AS vllm-openai
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
ARG NIXL_VERSION=0.7.0
# Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
# No compiler or autotools are installed in this stage.
COPY --from=ucx-nixl-build /tmp/ucx_install /tmp/ucx_install
COPY --from=ucx-nixl-build /tmp/nixl_wheels /tmp/nixl_wheels
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
# Install RDMA runtime libraries (no build tools) and the pre-built NIXL wheel.
# Do not uninstall/reinstall large Python packages here to avoid extra layer
# churn; final package resolution remains in the later app install step.
RUN --mount=type=cache,target=/root/.cache/uv \
apt-get update -y && apt-get install -y --no-install-recommends \
rdma-core \
libibverbs1 \
librdmacm1 \
libibumad3 \
libibmad5 \
libmlx5-1 \
libmlx4-1 \
ibverbs-providers \
librdmacm1t64 \
&& rm -rf /var/lib/apt/lists/* \
&& uv pip install --no-deps /tmp/nixl_wheels/nixl-*.whl \
&& rm -rf /tmp/nixl_wheels
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
--mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
--mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \
uv pip install grpcio-tools protobuf nanobind && \
uv pip install -r /workspace/vllm/requirements/xpu.txt && \
uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.7.1 && \
uv pip uninstall oneccl oneccl-devel && \
source /opt/intel/oneapi/setvars.sh --force && \
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
# Keep source-dependent layers near the end so frequent code-only changes
# don't invalidate heavy dependency and UCX/NIXL layers.
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
@@ -147,75 +226,9 @@ ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=.git,target=.git \
uv pip install --no-build-isolation .
uv pip install --no-build-isolation --no-deps .
CMD ["/bin/bash"]
FROM vllm-base AS vllm-openai
# install development dependencies (for testing)
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -e tests/vllm_test_utils
# install NIXL and UCX from source code
ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
ARG NIXL_VERSION=0.7.0
RUN apt-get update && apt-get install -y \
pciutils \
net-tools \
iproute2 \
hwloc \
numactl \
wget \
curl \
git \
build-essential \
autoconf \
automake \
libtool \
pkg-config \
rdma-core \
libibverbs-dev \
ibverbs-utils \
libibverbs1 \
librdmacm-dev \
librdmacm1 \
libibumad-dev \
libibumad3 \
libibmad-dev \
libibmad5 \
infiniband-diags \
perftest \
ibutils \
libmlx5-1 \
libmlx4-1 \
ibverbs-providers \
librdmacm1t64
ENV PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig:${PKG_CONFIG_PATH}
ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/uv \
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
bash autogen.sh && \
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
cd /tmp/nixl_source && \
uv pip install --upgrade meson pybind11 patchelf && \
uv pip install -r requirements.txt && \
uv pip install . && \
rm -rf /tmp/ucx_source /tmp/nixl_source
# FIX triton
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.7.0
# remove torch bundled oneccl to avoid conflicts
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall oneccl oneccl-devel
ENTRYPOINT ["vllm", "serve"]
+1
View File
@@ -88,6 +88,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |
!!! note
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
+1 -1
View File
@@ -60,7 +60,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod]
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.UnquantizedFusedMoEMethod]
## Fused Experts Kernels
-319
View File
@@ -1,319 +0,0 @@
# P2P NCCL Connector
An implementation of xPyD with dynamic scaling based on point-to-point communication, partly inspired by Dynamo.
## Detailed Design
### Overall Process
As shown in Figure 1, the overall process of this **PD disaggregation** solution is described through a request flow:
1. The client sends an HTTP request to the Proxy/Router's `/v1/completions` interface.
2. The Proxy/Router selects a **1P1D (1 Prefill instance + 1 Decode instance)** through either through round-robin or random selection, generates a `request_id` (rules to be introduced later), modifies the `max_tokens` in the HTTP request message to **1**, and then forwards the request to the **P instance**.
3. Immediately afterward, the Proxy/Router forwards the **original HTTP request** to the **D instance**.
4. The **P instance** performs **Prefill** and then **actively sends the generated KV cache** to the D instance (using **PUT_ASYNC** mode). The D instance's `zmq_addr` can be resolved through the `request_id`.
5. The **D instance** has a **dedicated thread** for receiving the KV cache (to avoid blocking the main process). The received KV cache is saved into the **GPU memory buffer**, the size of which is determined by the vLLM startup parameter `kv_buffer_size`. When the GPU buffer is full, the KV cache is stored in the **local Tensor memory pool**.
6. During the **Decode**, the D instance's main process retrieves the KV cache (transmitted by the P instance) from either the **GPU buffer** or the **memory pool**, thereby **skipping Prefill**.
7. After completing **Decode**, the D instance returns the result to the **Proxy/Router**, which then forwards it to the **client**.
![image1](https://github.com/user-attachments/assets/fb01bde6-755b-49f7-ad45-48a94b1e10a7)
### Proxy/Router (Demo)
A simple HTTP service acts as the entry point for client requests and starts a background thread to listen for P/D instances reporting their HTTP IP and PORT, as well as ZMQ IP and PORT. It maintains a dictionary of `http_addr -> zmq_addr`. The `http_addr` is the IP:PORT for the vLLM instance's request, while the `zmq_addr` is the address for KV cache handshake and metadata reception.
The Proxy/Router is responsible for selecting 1P1D based on the characteristics of the client request, such as the prompt, and generating a corresponding `request_id`, for example:
```text
cmpl-___prefill_addr_10.0.1.2:21001___decode_addr_10.0.1.3:22001_93923d63113b4b338973f24d19d4bf11-0
```
Currently, to quickly verify whether xPyD can work, a round-robin selection of 1P1D is used. In the future, it is planned to use a trie combined with the load status of instances to select appropriate P and D.
Each P/D instance periodically sends a heartbeat packet to the Proxy/Router (currently every 3 seconds) to register (i.e., report `http_addr -> zmq_addr`) and keep the connection alive. If an instance crashes and fails to send a ping for a certain period of time, the Proxy/Router will remove the timed-out instance (this feature has not yet been developed).
### KV Cache Transfer Methods
There are three methods for KVCache transfer: PUT, GET, and PUT_ASYNC. These methods can be specified using the `--kv-transfer-config` and `kv_connector_extra_config` parameters, specifically through the `send_type` field. Both PUT and PUT_ASYNC involve the P instance actively sending KVCache to the D instance. The difference is that PUT is a synchronous transfer method that blocks the main process, while PUT_ASYNC is an asynchronous transfer method. PUT_ASYNC uses a dedicated thread for sending KVCache, which means it does not block the main process. In contrast, the GET method involves the P instance saving the KVCache to the memory buffer after computing the prefill. The D instance then actively retrieves the computed KVCache from the P instance once it has allocated space for the KVCache.
Experimental results have shown that the performance of these methods, from highest to lowest, is as follows: PUT_ASYNC → GET → PUT.
### P2P Communication via ZMQ & NCCL
As long as the address of the counterpart is known, point-to-point KV cache transfer (using NCCL) can be performed, without being constrained by rank and world size. To support dynamic scaling (expansion and contraction) of instances with PD disaggregation. This means that adding or removing P/D instances does not require a full system restart.
Each P/D instance only needs to create a single `P2pNcclEngine` instance. This instance maintains a ZMQ Server, which runs a dedicated thread to listen on the `zmq_addr` address and receive control flow requests from other instances. These requests include requests to establish an NCCL connection and requests to send KVCache metadata (such as tensor shapes and data types). However, it does not actually transmit the KVCache data itself.
When a P instance and a D instance transmit KVCache for the first time, they need to establish a ZMQ connection and an NCCL group. For subsequent KVCache transmissions, this ZMQ connection and NCCL group are reused. The NCCL group consists of only two ranks, meaning the world size is equal to 2. This design is intended to support dynamic scaling, which means that adding or removing P/D instances does not require a full system restart. As long as the address of the counterpart is known, point-to-point KVCache transmission can be performed, without being restricted by rank or world size.
### NCCL Group Topology
Currently, only symmetric TP (Tensor Parallelism) methods are supported for KVCache transmission. Asymmetric TP and PP (Pipeline Parallelism) methods will be supported in the future. Figure 2 illustrates the 1P2D setup, where each instance has a TP (Tensor Parallelism) degree of 2. There are a total of 7 NCCL groups: three vLLM instances each have one NCCL group with TP=2. Additionally, the 0th GPU card of the P instance establishes an NCCL group with the 0th GPU card of each D instance. Similarly, the 1st GPU card of the P instance establishes an NCCL group with the 1st GPU card of each D instance.
![image2](https://github.com/user-attachments/assets/837e61d6-365e-4cbf-8640-6dd7ab295b36)
Each NCCL group occupies a certain amount of GPU memory buffer for communication, the size of which is primarily influenced by the `NCCL_MAX_NCHANNELS` environment variable. When `NCCL_MAX_NCHANNELS=16`, an NCCL group typically occupies 100MB, while when `NCCL_MAX_NCHANNELS=8`, it usually takes up 52MB. For large-scale xPyD configurations—such as DeepSeek's 96P144D—this implementation is currently not feasible. Moving forward, we are considering using RDMA for point-to-point communication and are also keeping an eye on UCCL.
### GPU Memory Buffer and Tensor Memory Pool
The trade-off in the size of the memory buffer is as follows: For P instances, the memory buffer is not required in PUT and PUT_ASYNC modes, but it is necessary in GET mode. For D instances, a memory buffer is needed in all three modes. The memory buffer for D instances should not be too large. Similarly, for P instances in GET mode, the memory buffer should also not be too large. The memory buffer of D instances is used to temporarily store KVCache sent by P instances. If it is too large, it will reduce the KVCache space available for normal inference by D instances, thereby decreasing the inference batch size and ultimately leading to a reduction in output throughput. The size of the memory buffer is configured by the parameter `kv_buffer_size`, measured in bytes, and is typically set to 5%10% of the memory size.
If the `--max-num-seqs` parameter for P instances is set to a large value, due to the large batch size, P instances will generate a large amount of KVCache simultaneously. This may exceed the capacity of the memory buffer of D instances, resulting in KVCache loss. Once KVCache is lost, D instances need to recompute Prefill, which is equivalent to performing Prefill twice. Consequently, the time-to-first-token (TTFT) will significantly increase, leading to degraded performance.
To address the above issues, I have designed and developed a local Tensor memory pool for storing KVCache, inspired by the buddy system used in Linux memory modules. Since the memory is sufficiently large, typically in the TB range on servers, there is no need to consider prefix caching or using block-based designs to reuse memory, thereby saving space. When the memory buffer is insufficient, KVCache can be directly stored in the Tensor memory pool, and D instances can subsequently retrieve KVCache from it. The read and write speed is that of PCIe, with PCIe 4.0 having a speed of approximately 21 GB/s, which is usually faster than the Prefill speed. Otherwise, solutions like Mooncake and lmcache would not be necessary. The Tensor memory pool acts as a flood diversion area, typically unused except during sudden traffic surges. In the worst-case scenario, my solution performs no worse than the normal situation with a Cache store.
## Install vLLM
```shell
pip install "vllm>=0.9.2"
```
## Run xPyD
### Instructions
- The following examples are run on an A800 (80GB) device, using the Meta-Llama-3.1-8B-Instruct model.
- Pay attention to the setting of the `kv_buffer_size` (in bytes). The empirical value is 10% of the GPU memory size. This is related to the kvcache size. If it is too small, the GPU memory buffer for temporarily storing the received kvcache will overflow, causing the kvcache to be stored in the tensor memory pool, which increases latency. If it is too large, the kvcache available for inference will be reduced, leading to a smaller batch size and decreased throughput.
- For Prefill instances, when using non-GET mode, the `kv_buffer_size` can be set to 1, as Prefill currently does not need to receive kvcache. However, when using GET mode, a larger `kv_buffer_size` is required because it needs to store the kvcache sent to the D instance.
- You may need to modify the `kv_buffer_size` and `port` in the following commands (if there is a conflict).
- `PUT_ASYNC` offers the best performance and should be prioritized.
- The `--port` must be consistent with the `http_port` in the `--kv-transfer-config`.
- The `disagg_proxy_p2p_nccl_xpyd.py` script will use port 10001 (for receiving client requests) and port 30001 (for receiving service discovery from P and D instances).
- The node running the proxy must have `quart` installed.
- Supports multiple nodes; you just need to modify the `proxy_ip` and `proxy_port` in `--kv-transfer-config`.
- In the following examples, it is assumed that **the proxy's IP is 10.0.1.1**.
### Run 1P3D
#### Proxy (e.g. 10.0.1.1)
```shell
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
python3 disagg_proxy_p2p_nccl_xpyd.py &
```
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20001 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
```
#### Decode1 (e.g. 10.0.1.3 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20002 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
```
#### Decode2 (e.g. 10.0.1.4 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20003 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
```
#### Decode3 (e.g. 10.0.1.5 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20004 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
```
### Run 3P1D
#### Proxy (e.g. 10.0.1.1)
```shell
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
python3 disagg_proxy_p2p_nccl_xpyd.py &
```
#### Prefill1 (e.g. 10.0.1.2 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=0 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20001 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"21001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20001"}}' > /var/vllm.log 2>&1 &
```
#### Prefill2 (e.g. 10.0.1.3 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=1 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20002 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"22001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20002"}}' > /var/vllm.log 2>&1 &
```
#### Prefill3 (e.g. 10.0.1.4 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=2 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20003 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_buffer_size":"1e1","kv_port":"23001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20003"}}' > /var/vllm.log 2>&1 &
```
#### Decode1 (e.g. 10.0.1.5 or 10.0.1.1)
??? console "Command"
```shell
CUDA_VISIBLE_DEVICES=3 vllm serve {your model directory} \
--host 0.0.0.0 \
--port 20004 \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name base_model \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_buffer_size":"8e9","kv_port":"24001","kv_connector_extra_config":{"proxy_ip":"10.0.1.1","proxy_port":"30001","http_port":"20004"}}' > /var/vllm.log 2>&1 &
```
## Single request
```shell
curl -X POST -s http://10.0.1.1:10001/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "base_model",
"prompt": "San Francisco is a",
"max_tokens": 10,
"temperature": 0
}'
```
## Benchmark
??? console "Command"
```shell
vllm bench serve \
--backend vllm \
--model base_model \
--tokenizer meta-llama/Llama-3.1-8B-Instruct \
--dataset-name "random" \
--host 10.0.1.1 \
--port 10001 \
--random-input-len 1024 \
--random-output-len 1024 \
--ignore-eos \
--burstiness 100 \
--percentile-metrics "ttft,tpot,itl,e2el" \
--metric-percentiles "90,95,99" \
--seed $(date +%s) \
--trust-remote-code \
--request-rate 3 \
--num-prompts 1000
```
## Shut down
```shell
pgrep python | xargs kill -9 && pkill -f python
```
## Test data
### **Scenario**: 1K input & 200 output tokens, E2E P99 latency ~2s
![testdata](https://github.com/user-attachments/assets/cef0953b-4567-4bf9-b940-405b92a28eb1)
-7
View File
@@ -17,8 +17,6 @@ Two main reasons:
## Usage example
Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
Now supports 9 types of connectors:
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
@@ -29,7 +27,6 @@ Now supports 9 types of connectors:
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
```
- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
- **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
@@ -50,10 +47,6 @@ Now supports 9 types of connectors:
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
```
## Benchmarks
Please refer to [benchmarks/disagg_benchmarks](../../benchmarks/disagg_benchmarks) for disaggregated prefilling benchmarks.
## Development
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
@@ -41,12 +41,12 @@ pip install -v -r requirements/xpu.txt
```bash
pip uninstall -y triton triton-xpu
pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu
pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
```
!!! note
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
- Finally, build and install vLLM XPU backend:
@@ -1,127 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of disaggregated prefilling
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
and then transfer the KV cache between them.
"""
import os
import time
from multiprocessing import Event, Process
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
def run_prefill(prefill_done):
# We use GPU 0 for prefill node.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# The prefill node receives two requests, while the decode node receives
# three requests. So the decode node will only receive the KV Cache for
# requests 1 and 3. The decode node will use the KV Cache of requests 1
# and 3 and do prefilling on request 2.
prompts = [
"Hello, my name is",
"Hi, your name is",
# The decode node will actually "prefill" this request.
"Tell me a very long story",
]
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
# This instance is the prefill node (kv_producer, rank 0).
# The number of parallel instances for KV cache transfer is set to 2,
# as required for P2pNcclConnector.
ktc = KVTransferConfig(
kv_connector="P2pNcclConnector",
kv_role="kv_producer",
kv_rank=0,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
# memory. You may need to adjust the value to fit your GPU.
llm = LLM(
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
)
llm.generate(prompts, sampling_params)
print("Prefill node is finished.")
prefill_done.set()
# To keep the prefill node running in case the decode node is not done;
# otherwise, the script might exit prematurely, causing incomplete decoding.
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
print("Script stopped by user.")
def run_decode(prefill_done):
# We use GPU 1 for decode node.
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
prompts = [
"Hello, my name is",
"Hi, your name is",
"Tell me a very long story",
]
sampling_params = SamplingParams(temperature=0, top_p=0.95)
# Using P2pNcclConnector to transmit KV caches between vLLM instances.
# This instance is the decode node (kv_consumer, rank 1).
# The number of parallel instances for KV cache transfer is set to 2,
# as required for P2pNcclConnector.
ktc = KVTransferConfig(
kv_connector="P2pNcclConnector",
kv_role="kv_consumer",
kv_rank=1,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A6000 GPU with 40GB
# memory. You may need to adjust the value to fit your GPU.
llm = LLM(
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
)
# Wait for the producer to start the pipe
print("Waiting for prefill node to finish...")
prefill_done.wait()
# At this point when the prefill_done is set, the kv-cache should have been
# transferred to this decode node, so we can start decoding.
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
def main():
prefill_done = Event()
prefill_process = Process(target=run_prefill, args=(prefill_done,))
decode_process = Process(target=run_decode, args=(prefill_done,))
# Start prefill node
prefill_process.start()
# Start decode node
decode_process.start()
# Terminate the prefill node when decode is finished
decode_process.join()
prefill_process.terminate()
if __name__ == "__main__":
main()
@@ -1,125 +0,0 @@
#!/bin/bash
# This file demonstrates the example usage of disaggregated prefilling
# We will launch 2 vllm instances (1 for prefill and 1 for decode),
# and then transfer the KV cache between them.
set -xe
echo "🚧🚧 Warning: The usage of disaggregated prefill is experimental and subject to change 🚧🚧"
sleep 1
# meta-llama/Meta-Llama-3.1-8B-Instruct or deepseek-ai/DeepSeek-V2-Lite
MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
# Trap the SIGINT signal (triggered by Ctrl+C)
trap 'cleanup' INT
# Cleanup function
cleanup() {
echo "Caught Ctrl+C, cleaning up..."
# Cleanup commands
pgrep python | xargs kill -9
pkill -f python
echo "Cleanup complete. Exiting."
exit 0
}
if [[ -z "${VLLM_HOST_IP:-}" ]]; then
export VLLM_HOST_IP=127.0.0.1
echo "Using default VLLM_HOST_IP=127.0.0.1 (override by exporting VLLM_HOST_IP before running this script)"
else
echo "Using provided VLLM_HOST_IP=${VLLM_HOST_IP}"
fi
# install quart first -- required for disagg prefill proxy serve
if python3 -c "import quart" &> /dev/null; then
echo "Quart is already installed."
else
echo "Quart is not installed. Installing..."
python3 -m pip install quart
fi
# a function that waits vLLM server to start
wait_for_server() {
local port=$1
timeout 1200 bash -c "
until curl -i localhost:${port}/v1/models > /dev/null; do
sleep 1
done" && return 0 || return 1
}
# You can also adjust --kv-ip and --kv-port for distributed inference.
# prefilling instance, which is the KV producer
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
--host 0.0.0.0 \
--port 8100 \
--max-model-len 100 \
--gpu-memory-utilization 0.8 \
--trust-remote-code \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
# decoding instance, which is the KV consumer
CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
--host 0.0.0.0 \
--port 8200 \
--max-model-len 100 \
--gpu-memory-utilization 0.8 \
--trust-remote-code \
--kv-transfer-config \
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":"1e10","kv_port":"14580","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8200","send_type":"PUT_ASYNC"}}' &
# wait until prefill and decode instances are ready
wait_for_server 8100
wait_for_server 8200
# launch a proxy server that opens the service at port 8000
# the workflow of this proxy:
# - send the request to prefill vLLM instance (port 8100), change max_tokens
# to 1
# - after the prefill vLLM finishes prefill, send the request to decode vLLM
# instance
# NOTE: the usage of this API is subject to change --- in the future we will
# introduce "vllm connect" to connect between prefill and decode instances
python3 ../../benchmarks/disagg_benchmarks/disagg_prefill_proxy_server.py &
sleep 1
# serve two example requests
output1=$(curl -X POST -s http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "'"$MODEL_NAME"'",
"prompt": "San Francisco is a",
"max_tokens": 10,
"temperature": 0
}')
output2=$(curl -X POST -s http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "'"$MODEL_NAME"'",
"prompt": "Santa Clara is a",
"max_tokens": 10,
"temperature": 0
}')
# Cleanup commands
pgrep python | xargs kill -9
pkill -f python
echo ""
sleep 1
# Print the outputs of the curl requests
echo ""
echo "Output of first request: $output1"
echo "Output of second request: $output2"
echo "🎉🎉 Successfully finished 2 test requests! 🎉🎉"
echo ""
@@ -1,245 +0,0 @@
#!/bin/bash
# =============================================================================
# vLLM Disaggregated Serving Script - P2P NCCL XpYd Architecture
# =============================================================================
# This script demonstrates disaggregated prefill and decode serving using
# P2P NCCL communication. The architecture supports various XpYd configurations:
#
# - 1P3D: 1 Prefill server + 3 Decode servers (current default)
# - 3P1D: 3 Prefill servers + 1 Decode server
# - etc.
#
# Configuration can be customized via environment variables:
# MODEL: Model to serve
# PREFILL_GPUS: Comma-separated GPU IDs for prefill servers
# DECODE_GPUS: Comma-separated GPU IDs for decode servers
# PREFILL_PORTS: Comma-separated ports for prefill servers
# DECODE_PORTS: Comma-separated ports for decode servers
# PROXY_PORT: Proxy server port used to setup XpYd connection.
# TIMEOUT_SECONDS: Server startup timeout
# =============================================================================
# Configuration - can be overridden via environment variables
MODEL=${MODEL:-meta-llama/Llama-3.1-8B-Instruct}
TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200}
PROXY_PORT=${PROXY_PORT:-30001}
# Default 1P3D configuration (1 Prefill + 3 Decode)
PREFILL_GPUS=${PREFILL_GPUS:-0}
DECODE_GPUS=${DECODE_GPUS:-1,2,3}
PREFILL_PORTS=${PREFILL_PORTS:-20003}
DECODE_PORTS=${DECODE_PORTS:-20005,20007,20009}
echo "Warning: P2P NCCL disaggregated prefill XpYd support for vLLM v1 is experimental and subject to change."
echo ""
echo "Architecture Configuration:"
echo " Model: $MODEL"
echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS"
echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS"
echo " Proxy Port: $PROXY_PORT"
echo " Timeout: ${TIMEOUT_SECONDS}s"
echo ""
PIDS=()
# Switch to the directory of the current script
cd "$(dirname "${BASH_SOURCE[0]}")"
check_required_files() {
local files=("disagg_proxy_p2p_nccl_xpyd.py")
for file in "${files[@]}"; do
if [[ ! -f "$file" ]]; then
echo "Required file $file not found in $(pwd)"
exit 1
fi
done
}
check_hf_token() {
if [ -z "$HF_TOKEN" ]; then
echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
echo "Example: export HF_TOKEN=your_token_here"
exit 1
fi
if [[ "$HF_TOKEN" != hf_* ]]; then
echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
exit 1
fi
echo "HF_TOKEN is set and valid."
}
check_num_gpus() {
# Check if the number of GPUs are >=2 via nvidia-smi
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
if [ "$num_gpus" -lt 2 ]; then
echo "You need at least 2 GPUs to run disaggregated prefill."
exit 1
else
echo "Found $num_gpus GPUs."
fi
}
ensure_python_library_installed() {
echo "Checking if $1 is installed..."
if ! python3 -c "import $1" > /dev/null 2>&1; then
echo "$1 is not installed. Please install it via pip install $1."
exit 1
else
echo "$1 is installed."
fi
}
cleanup() {
echo "Stopping everything…"
trap - INT TERM # prevent re-entrancy
pkill -9 -f "disagg_proxy_p2p_nccl_xpyd.py"
kill -- -$$ # negative PID == "this whole process-group"
wait # reap children so we don't leave zombies
exit 0
}
wait_for_server() {
local port=$1
local timeout_seconds=$TIMEOUT_SECONDS
local start_time=$(date +%s)
echo "Waiting for server on port $port..."
while true; do
if curl -s "localhost:${port}/v1/completions" > /dev/null; then
echo "Server on port $port is ready."
return 0
fi
local now=$(date +%s)
if (( now - start_time >= timeout_seconds )); then
echo "Timeout waiting for server on port $port"
return 1
fi
sleep 1
done
}
main() {
check_required_files
check_hf_token
check_num_gpus
ensure_python_library_installed pandas
ensure_python_library_installed datasets
ensure_python_library_installed vllm
ensure_python_library_installed quart
trap cleanup INT
trap cleanup USR1
trap cleanup TERM
echo "Launching disaggregated serving components..."
echo "Please check the log files for detailed output:"
echo " - prefill*.log: Prefill server logs"
echo " - decode*.log: Decode server logs"
echo " - proxy.log: Proxy server log"
# =============================================================================
# Launch Proxy Server
# =============================================================================
echo ""
echo "Starting proxy server on port $PROXY_PORT..."
python3 disagg_proxy_p2p_nccl_xpyd.py &
PIDS+=($!)
# Parse GPU and port arrays
IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS"
IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS"
IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS"
IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
# =============================================================================
# Launch Prefill Servers (X Producers)
# =============================================================================
echo ""
echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..."
for i in "${!PREFILL_GPU_ARRAY[@]}"; do
local gpu_id=${PREFILL_GPU_ARRAY[$i]}
local port=${PREFILL_PORT_ARRAY[$i]}
local kv_port=$((21001 + i))
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
--enforce-eager \
--host 0.0.0.0 \
--port "$port" \
--tensor-parallel-size 1 \
--seed 1024 \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
"{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"1e1\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > prefill$((i+1)).log 2>&1 &
PIDS+=($!)
done
# =============================================================================
# Launch Decode Servers (Y Decoders)
# =============================================================================
echo ""
echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..."
for i in "${!DECODE_GPU_ARRAY[@]}"; do
local gpu_id=${DECODE_GPU_ARRAY[$i]}
local port=${DECODE_PORT_ARRAY[$i]}
local kv_port=$((22001 + i))
echo " Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
--enforce-eager \
--host 0.0.0.0 \
--port "$port" \
--tensor-parallel-size 1 \
--seed 1024 \
--dtype float16 \
--max-model-len 10000 \
--max-num-batched-tokens 10000 \
--max-num-seqs 256 \
--trust-remote-code \
--gpu-memory-utilization 0.7 \
--kv-transfer-config \
"{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > decode$((i+1)).log 2>&1 &
PIDS+=($!)
done
# =============================================================================
# Wait for All Servers to Start
# =============================================================================
echo ""
echo "Waiting for all servers to start..."
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
if ! wait_for_server "$port"; then
echo "Failed to start server on port $port"
cleanup
# shellcheck disable=SC2317
exit 1
fi
done
echo ""
echo "All servers are up. Starting benchmark..."
# =============================================================================
# Run Benchmark
# =============================================================================
cd ../../../benchmarks/
vllm bench serve --port 10001 --seed "$(date +%s)" \
--model "$MODEL" \
--dataset-name random --random-input-len 7500 --random-output-len 200 \
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
echo "Benchmarking done. Cleaning up..."
cleanup
}
main
@@ -1,190 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import socket
import threading
import time
import uuid
from typing import Any
import aiohttp
import msgpack
import zmq
from quart import Quart, make_response, request
count = 0
prefill_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
decode_instances: dict[str, Any] = {} # http_address: (zmq_address, stamp)
prefill_cv = threading.Condition()
decode_cv = threading.Condition()
DEFAULT_PING_SECONDS = 5
def _remove_oldest_instances(instances: dict[str, Any]) -> None:
oldest_key = next(iter(instances), None)
while oldest_key is not None:
value = instances[oldest_key]
if value[1] > time.time():
break
print(f"🔴Remove [HTTP:{oldest_key}, ZMQ:{value[0]}, stamp:{value[1]}]")
instances.pop(oldest_key, None)
oldest_key = next(iter(instances), None)
def _listen_for_register(poller, router_socket):
while True:
socks = dict(poller.poll())
if router_socket in socks:
remote_address, message = router_socket.recv_multipart()
# data: {"type": "P", "http_address": "ip:port",
# "zmq_address": "ip:port"}
data = msgpack.loads(message)
if data["type"] == "P":
global prefill_instances
global prefill_cv
with prefill_cv:
node = prefill_instances.get(data["http_address"], None)
prefill_instances[data["http_address"]] = (
data["zmq_address"],
time.time() + DEFAULT_PING_SECONDS,
)
_remove_oldest_instances(prefill_instances)
elif data["type"] == "D":
global decode_instances
global decode_cv
with decode_cv:
node = decode_instances.get(data["http_address"], None)
decode_instances[data["http_address"]] = (
data["zmq_address"],
time.time() + DEFAULT_PING_SECONDS,
)
_remove_oldest_instances(decode_instances)
else:
print(
"Unexpected, Received message from %s, data: %s",
remote_address,
data,
)
return
if node is None:
print(f"🔵Add [HTTP:{data['http_address']}, ZMQ:{data['zmq_address']}]")
def start_service_discovery(hostname, port):
if not hostname:
hostname = socket.gethostname()
if port == 0:
raise ValueError("Port cannot be 0")
context = zmq.Context()
router_socket = context.socket(zmq.ROUTER)
router_socket.bind(f"tcp://{hostname}:{port}")
poller = zmq.Poller()
poller.register(router_socket, zmq.POLLIN)
_listener_thread = threading.Thread(
target=_listen_for_register, args=[poller, router_socket], daemon=True
)
_listener_thread.start()
return _listener_thread
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
app = Quart(__name__)
def random_uuid() -> str:
return str(uuid.uuid4().hex)
async def forward_request(url, data, request_id):
async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
"X-Request-Id": request_id,
}
async with session.post(url=url, json=data, headers=headers) as response:
if response.status == 200:
if True:
async for chunk_bytes in response.content.iter_chunked(1024):
yield chunk_bytes
else:
content = await response.read()
yield content
@app.route("/v1/completions", methods=["POST"])
@app.route("/v1/chat/completions", methods=["POST"])
async def handle_request():
try:
original_request_data = await request.get_json()
prefill_request = original_request_data.copy()
# change max_tokens = 1 to let it only do prefill
prefill_request["max_tokens"] = 1
if "max_completion_tokens" in prefill_request:
prefill_request["max_completion_tokens"] = 1
global count
global prefill_instances
global prefill_cv
with prefill_cv:
prefill_list = list(prefill_instances.items())
prefill_addr, prefill_zmq_addr = prefill_list[count % len(prefill_list)]
prefill_zmq_addr = prefill_zmq_addr[0]
global decode_instances
global decode_cv
with decode_cv:
decode_list = list(decode_instances.items())
decode_addr, decode_zmq_addr = decode_list[count % len(decode_list)]
decode_zmq_addr = decode_zmq_addr[0]
print(
f"handle_request count: {count}, [HTTP:{prefill_addr}, "
f"ZMQ:{prefill_zmq_addr}] 👉 [HTTP:{decode_addr}, "
f"ZMQ:{decode_zmq_addr}]"
)
count += 1
request_id = (
f"___prefill_addr_{prefill_zmq_addr}___decode_addr_"
f"{decode_zmq_addr}_{random_uuid()}"
)
# finish prefill
async for _ in forward_request(
f"http://{prefill_addr}{request.path}", prefill_request, request_id
):
continue
# return decode
generator = forward_request(
f"http://{decode_addr}{request.path}", original_request_data, request_id
)
response = await make_response(generator)
response.timeout = None
return response
except Exception as e:
import sys
import traceback
exc_info = sys.exc_info()
print("Error occurred in disagg prefill proxy server")
print(e)
print("".join(traceback.format_exception(*exc_info)))
if __name__ == "__main__":
t = start_service_discovery("0.0.0.0", 30001)
app.run(host="0.0.0.0", port=10001)
t.join()
@@ -17,9 +17,7 @@ class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
events: list[Any]
class KVCacheEvent(
msgspec.Struct, array_like=True, omit_defaults=True, gc=False, tag=True
):
class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
"""Base class for all KV cache-related events"""
@@ -2562,6 +2562,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
"qwen3_5",
"qwen3_5_moe",
"stepvl",
"glm4_1v",
]
+1 -1
View File
@@ -31,7 +31,7 @@ schemathesis
jiwer
bm25s
pystemmer
mteb[bm25s]
mteb[bm25s]>=2, <3 # required for mteb test
num2words
pqdm
+3 -3
View File
@@ -646,7 +646,7 @@ tokenizers==0.22.2
# via
# -c requirements/common.txt
# transformers
torch==2.11.0+xpu
torch==2.12.0+xpu
# via
# -c requirements/xpu.txt
# accelerate
@@ -654,7 +654,7 @@ torch==2.11.0+xpu
# sentence-transformers
# timm
# torchvision
torchvision==0.26.0+xpu
torchvision==0.27.0+xpu
# via timm
tqdm==4.67.3
# via
@@ -672,7 +672,7 @@ transformers==5.5.3
# via
# -c requirements/common.txt
# sentence-transformers
triton-xpu==3.7.0
triton-xpu==3.7.1
# via torch
typepy==1.3.4
# via
+2 -2
View File
@@ -12,9 +12,9 @@ jinja2>=3.1.6
datasets # for benchmark scripts
numba == 0.65.0 # Required for N-gram speculative decoding
--extra-index-url=https://download.pytorch.org/whl/xpu
torch==2.11.0
torch==2.12.0
torchaudio
torchvision
auto_round_lib>=0.13.0
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9/vllm_xpu_kernels-0.1.9-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
+3
View File
@@ -5799,6 +5799,7 @@ dependencies = [
"axum",
"bytes",
"clap",
"educe",
"expect-test",
"futures",
"http-body",
@@ -5813,7 +5814,9 @@ dependencies = [
"serde_json",
"serde_with",
"serial_test",
"sha2",
"socket2",
"subtle",
"thiserror-ext",
"tokio",
"tokio-stream",
+2
View File
@@ -75,8 +75,10 @@ serde_repr = "0.1.20"
serde_tuple = "1.1.3"
serde_with = "3.18.0"
serial_test = { version = "3.2.0", features = ["file_locks"] }
sha2 = "0.10.9"
socket2 = "0.6.3"
subenum = "1.1.3"
subtle = "2.6"
task-local = "0.1.1"
tekken = { package = "tekken-rs", version = "0.1.1", default-features = false }
tempfile = "3.23.0"
+2
View File
@@ -66,6 +66,8 @@ pub enum Error {
ToolCallStreamInvariant { message: String },
#[error(transparent)]
Text(#[from] vllm_text::Error),
#[error(transparent)]
Tokenizer(#[from] vllm_tokenizer::TokenizerError),
}
pub type Result<T> = std::result::Result<T, Error>;
+28 -1
View File
@@ -52,7 +52,7 @@ mod stream;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::ModelDtype;
use vllm_llm::Llm;
use vllm_text::{TextLlm, TextRequest};
use vllm_text::{Prompt, TextLlm, TextRequest};
/// Validate explicit parser override names without starting request processing.
pub fn validate_parser_overrides(
@@ -198,6 +198,33 @@ impl ChatLlm {
Ok(ChatEventStream::new(request.request_id, structured_stream))
}
/// Render through the chat template and tokenize, without submitting to the engine.
///
/// Same render → [`multimodal::finalize_rendered_prompt`] → encode pipeline as
/// [`Self::chat`], but stops after token IDs so `/tokenize` counts match what
/// generation would see. Used by `POST /tokenize` (chat form).
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
request.validate()?;
let rendered = self.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let tokenizer = self.text.tokenizer();
let token_ids = match prompt {
// Rendered string from the template (usual chat path).
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
// Already tokenized (e.g. multimodal path); pass through unchanged.
Prompt::TokenIds(ids) => ids,
};
Ok(token_ids)
}
/// Shut down the underlying LLM client and its background tasks.
pub async fn shutdown(self) -> Result<()> {
self.text.shutdown().await?;
+103 -1
View File
@@ -121,7 +121,11 @@ impl ToolState {
None => true,
};
if is_new_tool {
let id = generate_tool_call_id();
let id = self
.parser
.tool_call_id(item.tool_index)
.map(str::to_string)
.unwrap_or_else(generate_tool_call_id);
self.open_call_index = Some(item.tool_index);
events.push(AssistantEvent::ToolCallStart { id, name });
}
@@ -291,6 +295,11 @@ mod tests {
buffered: String,
}
struct IdScriptedParser {
output: ToolParserOutput,
tool_call_id: Option<String>,
}
impl ToolParser for FailingParser {
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
where
@@ -351,6 +360,35 @@ mod tests {
}
}
impl ToolParser for IdScriptedParser {
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self {
output: ToolParserOutput::default(),
tool_call_id: None,
}))
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
(tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten()
}
fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
output.append(std::mem::take(&mut self.output));
Ok(())
}
fn finish(&mut self) -> Result<ToolParserOutput> {
Ok(ToolParserOutput::default())
}
fn reset(&mut self) -> String {
String::new()
}
}
impl ToolParser for PartialThenFailParser {
fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
where
@@ -505,6 +543,70 @@ mod tests {
assert!(matches!(events[3], AssistantEvent::Done { .. }));
}
#[tokio::test]
async fn tool_stream_preserves_parser_provided_tool_call_id() {
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
kind: AssistantBlockKind::Text,
delta: "ignored".to_string(),
})]);
let parser = IdScriptedParser {
output: ToolParserOutput {
normal_text: String::new(),
calls: vec![crate::parser::tool::ToolCallDelta {
tool_index: 0,
name: Some("get_weather".to_string()),
arguments: "{}".to_string(),
}],
},
tool_call_id: Some("functions.get_weather:0".to_string()),
};
let events = tool_event_stream(events, Some(Box::new(parser)))
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<crate::Result<Vec<_>>>()
.unwrap();
assert!(matches!(
&events[0],
AssistantEvent::ToolCallStart { id, name }
if id == "functions.get_weather:0" && name == "get_weather"
));
}
#[tokio::test]
async fn tool_stream_generates_tool_call_id_when_parser_omits_one() {
let events = stream::iter(vec![Ok(ContentEvent::TextDelta {
kind: AssistantBlockKind::Text,
delta: "ignored".to_string(),
})]);
let parser = IdScriptedParser {
output: ToolParserOutput {
normal_text: String::new(),
calls: vec![crate::parser::tool::ToolCallDelta {
tool_index: 0,
name: Some("get_weather".to_string()),
arguments: "{}".to_string(),
}],
},
tool_call_id: None,
};
let events = tool_event_stream(events, Some(Box::new(parser)))
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<crate::Result<Vec<_>>>()
.unwrap();
assert!(matches!(
&events[0],
AssistantEvent::ToolCallStart { id, name }
if id.starts_with("call_") && name == "get_weather"
));
}
#[tokio::test]
async fn real_buffered_parser_error_matches_streaming_and_non_streaming() {
let prefix = "I will check both.\n";
-2
View File
@@ -128,8 +128,6 @@ trait_set! {
/// Generate the northbound tool-call ID using the OpenAI-style `call_<id>`
/// format.
// TODO: support other ID scheme like Kimi-K2's
// `functions.{name}:{global_index}`.
pub(crate) fn generate_tool_call_id() -> String {
format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])
}
+1 -2
View File
@@ -170,8 +170,7 @@ roundtrip_tests! {
glm47 => [reasoning_and_content, tool_call_mix],
// Note: Kimi K2.5 strips the reasoning content in history.
// TODO: we don't respect model-generated tool call id now so `tool_call_mix` cannot pass.
// kimi_k25 => [tool_call_mix],
kimi_k25 => [tool_call_mix],
}
/// Run the fixed reasoning+content fixture for one model/parser case.
+25 -1
View File
@@ -16,6 +16,7 @@ use educe::Educe;
use serde::Deserialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use serde_with::{DefaultOnNull, OneOrMany, serde_as};
use thiserror_ext::AsReport as _;
use uuid::Uuid;
use vllm_engine_core_client::TransportMode;
@@ -84,6 +85,7 @@ pub enum Command {
}
/// Runtime arguments shared by the external-engine and managed-engine paths.
#[serde_as]
#[derive(Educe, Clone, Args, PartialEq, Eq, Deserialize)]
#[educe(Debug)]
pub struct SharedRuntimeArgs {
@@ -178,6 +180,14 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub enable_request_id_headers: bool,
/// If provided, the server will require one of these keys to be presented
/// in the Authorization header.
#[educe(Debug(ignore))]
#[arg(long, env = "VLLM_API_KEY", value_delimiter = ' ')]
#[serde_as(as = "DefaultOnNull<OneOrMany<_>>")]
#[serde(default)]
pub api_key: Vec<String>,
/// Disable periodic logging of engine statistics (throughput, queue depth,
/// cache usage).
#[arg(long)]
@@ -215,6 +225,15 @@ impl SharedRuntimeArgs {
Duration::from_secs(self.shutdown_timeout)
}
/// Apply fallback logic for API key configuration from env variables.
fn apply_env_api_key_fallback(&mut self) {
if self.api_key.is_empty()
&& let Ok(api_key) = std::env::var("VLLM_API_KEY")
{
self.api_key.push(api_key);
}
}
/// Build the OpenAI-server config for the Python-bootstrap worker contract.
///
/// The resulting config binds the Python-supplied transport addresses and
@@ -253,6 +272,7 @@ impl SharedRuntimeArgs {
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
api_keys: self.api_key,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
@@ -295,6 +315,7 @@ impl SharedRuntimeArgs {
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
api_keys: self.api_key,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
@@ -311,8 +332,11 @@ fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
}
fn parse_runtime_args_json(value: &str) -> Result<SharedRuntimeArgs, String> {
let args: SharedRuntimeArgs = serde_json::from_str(value)
let mut args: SharedRuntimeArgs = serde_json::from_str(value)
.map_err(|e| format!("invalid JSON arguments: {}", e.as_report()))?;
// --args-json is parsed with serde, so clap's env support does not run for
// the Python-supervised frontend path.
args.apply_env_api_key_fallback();
args.unsupported.check()?;
Ok(args)
}
+78 -4
View File
@@ -169,6 +169,76 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() {
assert!(config.enable_request_id_headers);
}
#[test]
fn serve_passes_api_keys_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--api-key",
"secret-a",
"--api-key",
"secret-b",
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
let debug = format!("{config:#?}");
assert!(debug.contains("api_keys: [<redacted>; 2]"));
assert!(!debug.contains("secret-a"));
assert!(!debug.contains("secret-b"));
}
#[test]
fn frontend_args_json_accepts_api_key_string() {
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","api_key":"secret"}"#,
])
.unwrap();
let Command::Frontend(args) = cli.command else {
panic!("expected frontend args");
};
let config = args.into_config();
assert_eq!(config.api_keys, vec!["secret"]);
}
#[test]
fn frontend_args_json_accepts_api_key_list() {
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","api_key":["secret-a","secret-b"]}"#,
])
.unwrap();
let Command::Frontend(args) = cli.command else {
panic!("expected frontend args");
};
let config = args.into_config();
assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
}
#[test]
fn serve_args_reject_unknown_renderer_value() {
let error = Cli::try_parse_from([
@@ -446,20 +516,21 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() {
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}"#,
r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}"#,
])
.unwrap_err();
let actual = error.to_string().replace(": \n", ":\n");
expect![[r#"
error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}' for '--args-json <JSON>':
error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>':
The following arguments are not implemented in Rust frontend yet:
- allow_credentials
- api_key
- ssl_keyfile
Remove these arguments to continue.
For more information, try '--help'.
"#]].assert_eq(&error.to_string());
"#]].assert_eq(&actual);
}
#[test]
@@ -793,6 +864,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_keys: [],
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -857,6 +929,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_keys: [],
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -936,6 +1009,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_keys: [],
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
-5
View File
@@ -564,11 +564,6 @@ pub struct ServerUnsupportedArgs {
#[arg(long)]
pub allowed_headers: Option<Unsupported>,
/// If provided, the server will require one of these keys to be presented
/// in the header.
#[arg(long)]
pub api_key: Option<Unsupported>,
/// The file path to the SSL key file.
#[arg(long)]
pub ssl_keyfile: Option<Unsupported>,
@@ -3,6 +3,7 @@ use std::time::Duration;
use anyhow::{Context, Result, bail};
use clap::Parser;
use tracing_subscriber::EnvFilter;
use vllm_engine_core_client::protocol::utility::PauseMode;
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
#[derive(Debug, Parser)]
@@ -32,8 +33,8 @@ struct Args {
reset_external: bool,
#[arg(long, default_value_t = 1)]
sleep_level: u32,
#[arg(long, default_value = "abort")]
sleep_mode: String,
#[arg(long, default_value_t = PauseMode::Abort)]
sleep_mode: PauseMode,
#[arg(
long,
default_value_t = false,
@@ -106,7 +107,7 @@ async fn main() -> Result<()> {
if args.skip_sleep_wake {
println!("sleep_wake=skipped");
} else {
client.sleep(args.sleep_level, &args.sleep_mode).await.with_context(|| {
client.sleep(args.sleep_level, args.sleep_mode).await.with_context(|| {
format!(
"failed to call sleep utility with level={} mode={}",
args.sleep_level, args.sleep_mode
+45 -34
View File
@@ -2,6 +2,7 @@ use std::sync::Arc;
use std::time::Duration;
use futures::future::{join_all, try_join_all};
use itertools::Itertools;
use serde::Serialize;
use tokio::sync::mpsc;
use tokio_util::task::AbortOnDropHandle;
@@ -12,7 +13,7 @@ use crate::coordinator::CoordinatorHandle;
use crate::error::{Error, Result};
use crate::protocol::handshake::EngineCoreReadyResponse;
use crate::protocol::lora::LoraRequest;
use crate::protocol::utility::EngineCoreUtilityRequest;
use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
use crate::transport::{self, ConnectedEngine};
@@ -571,6 +572,26 @@ impl EngineCoreClient {
try_join_all(futures).await
}
/// Call a utility method on all connected engines and return the shared
/// result if every engine agrees.
pub async fn call_utility_consensus<T, A>(&self, method: &str, args: A) -> Result<T>
where
T: serde::de::DeserializeOwned + std::fmt::Debug + PartialEq,
A: serde::Serialize + std::fmt::Debug,
{
let results: Vec<T> = self.call_utility(method, args).await?;
if results.iter().all_equal() {
// `engine_count >= 1` is enforced during startup handshake so `results` must be non-empty.
Ok(results.into_iter().next().unwrap())
} else {
Err(Error::InconsistentUtilityResults {
method: method.to_string(),
values: format!("{results:?}"),
})
}
}
/// Execute `collective_rpc` on all engines and flatten all engine results
/// into one list.
pub async fn collective_rpc<A, K>(
@@ -599,27 +620,8 @@ impl EngineCoreClient {
}
/// Return whether the engine is currently sleeping at any level.
///
/// Under data parallel, all engines should agree on the sleep state: a
/// divergence signals a control-plane bug. Returns
/// `Error::InconsistentUtilityResults` if engines disagree.
pub async fn is_sleeping(&self) -> Result<bool> {
let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
// `engine_count >= 1` is enforced during startup handshake, so `results`
// is normally non-empty; fall back to a fail-loud error rather than
// indexing in case that invariant is ever bypassed.
let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
method: "is_sleeping".to_string(),
values: "[]".to_string(),
})?;
if results.iter().all(|&v| v == first) {
Ok(first)
} else {
Err(Error::InconsistentUtilityResults {
method: "is_sleeping".to_string(),
values: format!("{results:?}"),
})
}
self.call_utility_consensus("is_sleeping", ()).await
}
/// Reset the multi-modal cache.
@@ -643,22 +645,14 @@ impl EngineCoreClient {
reset_running_requests: bool,
reset_connector: bool,
) -> Result<bool> {
let results: Vec<bool> = self
Ok(self
.call_utility(
"reset_prefix_cache",
(reset_running_requests, reset_connector),
)
.await?;
// `engine_count >= 1` is enforced during startup handshake, so `results`
// is normally non-empty; fail loud rather than reporting a vacuous
// success (`[].all() == true`) in case that invariant is ever bypassed.
if results.is_empty() {
return Err(Error::InconsistentUtilityResults {
method: "reset_prefix_cache".to_string(),
values: "[]".to_string(),
});
}
Ok(results.into_iter().all(|ok| ok))
.await?
.into_iter()
.all(|reset| reset))
}
/// Load or refresh one LoRA adapter on every connected engine.
@@ -680,7 +674,7 @@ impl EngineCoreClient {
}
/// Put the engine to sleep.
pub async fn sleep(&self, level: u32, mode: &str) -> Result<()> {
pub async fn sleep(&self, level: u32, mode: PauseMode) -> Result<()> {
self.call_utility::<(), _>("sleep", (level, mode)).await?;
Ok(())
}
@@ -692,6 +686,23 @@ impl EngineCoreClient {
Ok(())
}
/// Pause the scheduler so generation can be halted
pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {
self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;
Ok(())
}
/// Resume the scheduler after a pause
pub async fn resume_scheduler(&self) -> Result<()> {
self.call_utility::<(), _>("resume_scheduler", ()).await?;
Ok(())
}
/// Return whether the scheduler is currently in any pause state.
pub async fn is_scheduler_paused(&self) -> Result<bool> {
self.call_utility_consensus("is_scheduler_paused", ()).await
}
/// Shut down local client tasks and close transport state.
pub async fn shutdown(self) -> Result<()> {
let Self {
@@ -162,6 +162,21 @@ pub enum RequestOutputKind {
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`
@@ -209,6 +224,17 @@ pub struct StructuredOutputsParams {
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.
@@ -600,4 +626,18 @@ 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");
}
}
@@ -1,15 +1,63 @@
use std::any::type_name;
use std::fmt;
use std::{fmt, str::FromStr};
use rmpv::Value;
use serde::{Deserialize, Serialize};
use serde_default::DefaultFromSerde;
use serde_tuple::{Deserialize_tuple, Serialize_tuple};
use serde_with::{DeserializeFromStr, SerializeDisplay};
use thiserror_ext::AsReport;
use super::{OpaqueValue, default_opaque_value_nil};
use crate::error::{Error, Result};
/// How pause/sleep utility calls handle in-flight requests.
///
/// Use display/from-str serde so MessagePack utility args stay as Python
/// literal strings instead of serde enum variant tuples.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, SerializeDisplay, DeserializeFromStr)]
pub enum PauseMode {
/// Abort all in-flight requests immediately.
#[default]
Abort,
/// Wait for in-flight requests to complete.
Wait,
/// Freeze queued requests so they can resume later.
Keep,
}
impl PauseMode {
/// Return the Python literal used on the utility-call wire.
pub fn as_str(self) -> &'static str {
match self {
Self::Abort => "abort",
Self::Wait => "wait",
Self::Keep => "keep",
}
}
}
impl fmt::Display for PauseMode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl FromStr for PauseMode {
type Err = String;
fn from_str(value: &str) -> std::result::Result<Self, Self::Err> {
match value {
"abort" => Ok(Self::Abort),
"wait" => Ok(Self::Wait),
"keep" => Ok(Self::Keep),
other => Err(format!(
"invalid pause mode `{other}`; expected one of: abort, wait, keep"
)),
}
}
}
/// Utility call id as carried on the engine-core MessagePack wire.
///
/// Python emits utility ids as MessagePack integers, including values that may
@@ -212,7 +260,7 @@ mod tests {
use rmpv::Value;
use serde::Serialize;
use super::{EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope};
use super::{EngineCoreUtilityRequest, PauseMode, UtilityOutput, UtilityResultEnvelope};
use crate::Error;
use crate::protocol::{decode_msgpack, decode_value, encode_msgpack};
@@ -241,6 +289,26 @@ mod tests {
assert_eq!(array[3], Value::Array(Vec::new()));
}
#[test]
fn pause_mode_serializes_as_python_literal() {
let request =
EngineCoreUtilityRequest::new(7, 42, "pause_scheduler", (PauseMode::Abort, true))
.unwrap();
let encoded = encode_msgpack(&request).unwrap();
let value = decode_value(&encoded).unwrap();
let array = match value {
Value::Array(array) => array,
other => panic!("expected utility request array, got {other:?}"),
};
assert_eq!(array[2], Value::from("pause_scheduler"));
assert_eq!(
array[3],
Value::Array(vec![Value::from("abort"), Value::from(true)])
);
}
#[test]
fn utility_output_decodes_typed_result() {
let output = UtilityOutput {
+3
View File
@@ -8,6 +8,7 @@ license.workspace = true
anyhow.workspace = true
asynk-strim-attr.workspace = true
axum.workspace = true
educe.workspace = true
futures.workspace = true
http-body.workspace = true
itertools.workspace = true
@@ -19,7 +20,9 @@ rmpv.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_with.workspace = true
sha2.workspace = true
socket2.workspace = true
subtle.workspace = true
thiserror-ext.workspace = true
tokio.workspace = true
tokio-stream.workspace = true
@@ -70,6 +70,7 @@ async fn main() -> Result<()> {
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
enable_log_requests: false,
enable_request_id_headers: false,
api_keys: Vec::new(),
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: Duration::ZERO,
+24 -1
View File
@@ -1,7 +1,9 @@
use std::collections::HashMap;
use std::fmt;
use std::time::Duration;
use anyhow::Result;
use educe::Educe;
use serde::Serialize;
use serde_json::Value;
use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
@@ -33,7 +35,8 @@ pub enum CoordinatorMode {
}
/// Normalized runtime configuration for the minimal OpenAI-compatible server.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[derive(Educe, Clone, PartialEq, Eq, Serialize)]
#[educe(Debug)]
pub struct Config {
/// Frontend-to-engine transport setup.
pub transport_mode: TransportMode,
@@ -67,6 +70,10 @@ pub struct Config {
pub enable_log_requests: bool,
/// When `true`, set `X-Request-Id` on every HTTP response.
pub enable_request_id_headers: bool,
/// API keys accepted as bearer tokens for guarded routes.
#[serde(skip_serializing)]
#[educe(Debug(method(fmt_redacted_api_keys)))]
pub api_keys: Vec<String>,
/// When `true`, suppress periodic stats logging (throughput, queue depth,
/// cache usage).
pub disable_log_stats: bool,
@@ -114,3 +121,19 @@ impl Config {
}
}
}
struct RedactedApiKeys<'a>(&'a [String]);
impl fmt::Debug for RedactedApiKeys<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.0.is_empty() {
f.debug_list().finish()
} else {
write!(f, "[<redacted>; {}]", self.0.len())
}
}
}
fn fmt_redacted_api_keys(api_keys: &[String], f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&RedactedApiKeys(api_keys), f)
}
+2 -1
View File
@@ -92,7 +92,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
AppState::new(served_model_names, chat)
.with_log_requests(config.enable_log_requests)
.with_request_id_headers(config.enable_request_id_headers)
.with_server_info(ServerInfoSnapshot::from_config(config)),
.with_server_info(ServerInfoSnapshot::from_config(config))
.with_api_keys(config.api_keys.clone()),
))
}
+91
View File
@@ -0,0 +1,91 @@
use std::sync::Arc;
use axum::Json;
use axum::extract::{Request, State};
use axum::http::header::AUTHORIZATION;
use axum::http::{HeaderValue, Method, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use serde_json::json;
use crate::state::{ApiKeyHash, AppState, hash_api_key};
const GUARDED_PREFIXES: &[&str] = &["/v1", "/v2", "/inference"];
/// Authenticate guarded HTTP routes with an OpenAI-compatible bearer token.
///
/// Mirrors Python `AuthenticationMiddleware`: OPTIONS requests and non-guarded
/// helper endpoints such as `/health` are allowed through without a token.
pub async fn authenticate_api_key(
State(state): State<Arc<AppState>>,
req: Request,
next: Next,
) -> Response {
if req.method() == Method::OPTIONS || !requires_auth(req.uri().path()) {
return next.run(req).await;
}
if verify_token(req.headers().get(AUTHORIZATION), state.api_key_hashes()) {
return next.run(req).await;
}
(
StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
)
.into_response()
}
fn requires_auth(path: &str) -> bool {
GUARDED_PREFIXES.iter().any(|prefix| path.starts_with(prefix))
}
fn verify_token(authorization: Option<&HeaderValue>, api_key_hashes: &[ApiKeyHash]) -> bool {
let Some(authorization) = authorization else {
return false;
};
let Ok(authorization) = authorization.to_str() else {
return false;
};
let Some((scheme, token)) = authorization.split_once(' ') else {
return false;
};
if !scheme.eq_ignore_ascii_case("bearer") {
return false;
}
let token_hash = hash_api_key(token);
let mut token_match = false;
for api_key_hash in api_key_hashes {
token_match |= constant_time_eq(&token_hash, api_key_hash);
}
token_match
}
fn constant_time_eq(left: &ApiKeyHash, right: &ApiKeyHash) -> bool {
use subtle::ConstantTimeEq;
bool::from(left.ct_eq(right))
}
#[cfg(test)]
mod tests {
use super::constant_time_eq;
use crate::state::hash_api_key;
#[test]
fn constant_time_eq_checks_sha256_digests() {
assert!(constant_time_eq(
&hash_api_key("secret"),
&hash_api_key("secret")
));
assert!(!constant_time_eq(
&hash_api_key("secret"),
&hash_api_key("secrex")
));
assert!(!constant_time_eq(
&hash_api_key("secret"),
&hash_api_key("secret-more")
));
}
}
+2
View File
@@ -1,7 +1,9 @@
mod auth;
mod load;
mod metrics;
mod request_id;
pub use auth::authenticate_api_key;
pub use load::track_server_load;
pub use metrics::track_http_metrics;
pub use request_id::set_request_id_header;
+25 -4
View File
@@ -6,8 +6,10 @@ mod load;
mod lora;
mod metrics;
pub(crate) mod openai;
mod pause;
mod server_info;
mod sleep;
mod tokenize;
mod version;
use std::sync::Arc;
@@ -71,7 +73,9 @@ fn build_router_with_options(
.route("/v1/models", get(openai::list_models))
.route("/v1/completions", post(openai::completions))
.route("/v1/chat/completions", post(openai::chat_completions))
// vLLM specific inference endpoints
// vLLM specific endpoints
.route("/tokenize", post(tokenize::tokenize))
.route("/detokenize", post(tokenize::detokenize))
.route("/inference/v1/generate", post(inference::generate));
if runtime_lora_updating_enabled {
@@ -90,15 +94,32 @@ fn build_router_with_options(
.route("/sleep", post(sleep::sleep))
.route("/wake_up", post(sleep::wake_up))
.route("/is_sleeping", get(sleep::is_sleeping))
.route("/pause", post(pause::pause))
.route("/resume", post(pause::resume))
.route("/is_paused", get(pause::is_paused))
.route("/server_info", get(server_info::server_info))
}
let enable_request_id_headers = state.enable_request_id_headers;
let enable_api_key_auth = state.has_api_keys();
let mut router = router
.with_state(state.clone())
.layer(from_fn_with_state(state, middleware::track_server_load))
.layer(from_fn(middleware::track_http_metrics))
.layer(TraceLayer::new_for_http());
.layer(from_fn_with_state(
state.clone(),
middleware::track_server_load,
))
.layer(from_fn(middleware::track_http_metrics));
if enable_api_key_auth {
router = router.layer(from_fn_with_state(
state.clone(),
middleware::authenticate_api_key,
));
}
// Later layers wrap earlier ones. Keep tracing outside auth so rejected
// requests are visible, while metrics/load only see authenticated traffic.
router = router.layer(TraceLayer::new_for_http());
if enable_request_id_headers {
router = router.layer(from_fn(middleware::set_request_id_header));
@@ -156,8 +156,7 @@ pub(crate) fn prepare_chat_request(
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
})
}
fn normalize_generation_prompt_mode(
pub(crate) fn normalize_generation_prompt_mode(
add_generation_prompt: Option<bool>,
continue_final_message: bool,
messages: &[VllmChatMessage],
@@ -204,7 +203,7 @@ fn extract_last_assistant_content(messages: &[ChatMessage]) -> Option<String> {
}
/// Lower one OpenAI chat message into the `vllm-chat` message shape.
fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
pub(crate) fn convert_message(message: ChatMessage) -> Result<VllmChatMessage, ApiError> {
match message {
ChatMessage::System { content, .. } => {
Ok(VllmChatMessage::system(convert_content(content)?))
@@ -316,7 +315,7 @@ fn convert_assistant_tool_calls(
.collect()
}
fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
pub(crate) fn convert_tools(tools: Option<Vec<Tool>>) -> Result<Vec<ChatTool>, ApiError> {
tools
.unwrap_or_default()
.into_iter()
@@ -9,9 +9,9 @@ use vllm_chat::ReasoningEffort;
use crate::routes::openai::utils::structured_outputs::ResponseFormat;
use crate::routes::openai::utils::types::{
ChatLogProbs, ChatMessage, MessageContent, Normalizable, StreamOptions, StringOrArray, Tool,
ToolCall, ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage,
default_true, validate_stop, validate_top_p_value,
ChatLogProbs, ChatMessage, Normalizable, StreamOptions, StringOrArray, Tool, ToolCall,
ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage,
default_true, validate_messages, validate_stop, validate_top_p_value,
};
/// vLLM-compatible request type for the Chat Completions API.
@@ -430,32 +430,6 @@ fn default_model() -> String {
UNKNOWN_MODEL_ID.to_string()
}
/// Validates messages array is not empty and has valid content
fn validate_messages(messages: &[ChatMessage]) -> Result<(), validator::ValidationError> {
if messages.is_empty() {
return Err(validator::ValidationError::new("messages cannot be empty"));
}
for msg in messages {
if let ChatMessage::User { content, .. } = msg {
match content {
MessageContent::Text(text) if text.is_empty() => {
return Err(validator::ValidationError::new(
"message content cannot be empty",
));
}
MessageContent::Parts(parts) if parts.is_empty() => {
return Err(validator::ValidationError::new(
"message content parts cannot be empty",
));
}
_ => {}
}
}
}
Ok(())
}
/// Schema-level validation for cross-field dependencies
fn validate_chat_cross_parameters(
req: &ChatCompletionRequest,
@@ -412,6 +412,41 @@ pub struct ListModelsResponse {
pub data: Vec<ModelObject>,
}
// ============================================================================
// Shared validation helpers
// ============================================================================
/// Validates a messages array is non-empty and has valid user-message content.
///
/// Used by both `POST /v1/chat/completions` and `POST /tokenize` (chat form)
/// so validation behaviour stays in lockstep.
pub(crate) fn validate_messages(
messages: &[ChatMessage],
) -> Result<(), validator::ValidationError> {
if messages.is_empty() {
return Err(validator::ValidationError::new("messages cannot be empty"));
}
for msg in messages {
if let ChatMessage::User { content, .. } = msg {
match content {
MessageContent::Text(text) if text.is_empty() => {
return Err(validator::ValidationError::new(
"message content cannot be empty",
));
}
MessageContent::Parts(parts) if parts.is_empty() => {
return Err(validator::ValidationError::new(
"message content parts cannot be empty",
));
}
_ => {}
}
}
}
Ok(())
}
// ============================================================================
// Normalizable trait
// ============================================================================
+81
View File
@@ -0,0 +1,81 @@
use std::sync::Arc;
use axum::Json;
use axum::extract::rejection::QueryRejection;
use axum::extract::{Query, State};
use serde::{Deserialize, Serialize};
use vllm_engine_core_client::protocol::utility::PauseMode;
use crate::error::ApiError;
use crate::state::AppState;
use crate::utils::utility_call_error;
#[derive(Debug, Deserialize)]
pub(crate) struct PauseParams {
#[serde(default)]
mode: PauseMode,
#[serde(default = "default_clear_cache")]
clear_cache: bool,
}
#[derive(Serialize)]
pub(crate) struct StatusResponse {
status: &'static str,
}
#[derive(Serialize)]
pub(crate) struct IsPausedResponse {
is_paused: bool,
}
const fn default_clear_cache() -> bool {
true
}
fn invalid_query(error: QueryRejection) -> ApiError {
ApiError::invalid_request(error.body_text(), Some("mode"))
}
// TODO: the Python frontend also accepts the deprecated
// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is
// intentionally omitted here in favor of the `mode` parameter.
/// Pause the scheduler so generation can be halted (e.g. for weight updates).
pub async fn pause(
State(state): State<Arc<AppState>>,
params: Result<Query<PauseParams>, QueryRejection>,
) -> Result<Json<StatusResponse>, ApiError> {
let Query(params) = params.map_err(invalid_query)?;
state
.engine_core_client()
.pause_scheduler(params.mode, params.clear_cache)
.await
.map_err(|error| utility_call_error("pause", error))?;
Ok(Json(StatusResponse { status: "paused" }))
}
/// Resume the scheduler after a pause.
pub async fn resume(State(state): State<Arc<AppState>>) -> Result<Json<StatusResponse>, ApiError> {
state
.engine_core_client()
.resume_scheduler()
.await
.map_err(|error| utility_call_error("resume", error))?;
Ok(Json(StatusResponse { status: "resumed" }))
}
/// Return whether the scheduler is currently paused.
pub async fn is_paused(
State(state): State<Arc<AppState>>,
) -> Result<Json<IsPausedResponse>, ApiError> {
let is_paused = state
.engine_core_client()
.is_scheduler_paused()
.await
.map_err(|error| utility_call_error("is_paused", error))?;
Ok(Json(IsPausedResponse { is_paused }))
}
+10 -6
View File
@@ -1,9 +1,11 @@
use std::sync::Arc;
use axum::Json;
use axum::extract::rejection::QueryRejection;
use axum::extract::{Query, State};
use axum::http::StatusCode;
use serde::{Deserialize, Serialize};
use vllm_engine_core_client::protocol::utility::PauseMode;
use crate::error::ApiError;
use crate::state::AppState;
@@ -18,8 +20,8 @@ pub(crate) struct IsSleepingResponse {
pub(crate) struct SleepParams {
#[serde(default = "default_sleep_level")]
level: u32,
#[serde(default = "default_sleep_mode")]
mode: String,
#[serde(default)]
mode: PauseMode,
}
#[derive(Debug, Default, Deserialize)]
@@ -32,18 +34,20 @@ const fn default_sleep_level() -> u32 {
1
}
fn default_sleep_mode() -> String {
"abort".to_string()
fn invalid_query(error: QueryRejection) -> ApiError {
ApiError::invalid_request(error.body_text(), Some("mode"))
}
/// Put the engine to sleep.
pub async fn sleep(
State(state): State<Arc<AppState>>,
Query(params): Query<SleepParams>,
params: Result<Query<SleepParams>, QueryRejection>,
) -> Result<StatusCode, ApiError> {
let Query(params) = params.map_err(invalid_query)?;
state
.engine_core_client()
.sleep(params.level, &params.mode)
.sleep(params.level, params.mode)
.await
.map_err(|error| utility_call_error("sleep", error))?;
+670 -4
View File
@@ -407,6 +407,9 @@ struct FakeChatBackend {
multimodal_model_info: Option<vllm_chat::multimodal::MultimodalModelInfo>,
}
/// Synthetic BOS id used when `add_special_tokens` is true in tests.
const FAKE_BOS_TOKEN_ID: u32 = 1;
#[derive(Debug)]
struct FakeChatTokenizer;
@@ -414,9 +417,12 @@ impl Tokenizer for FakeChatTokenizer {
fn encode(
&self,
text: &str,
_add_special_tokens: bool,
add_special_tokens: bool,
) -> vllm_text::tokenizer::Result<Vec<u32>> {
let mut token_ids = Vec::new();
if add_special_tokens {
token_ids.push(FAKE_BOS_TOKEN_ID);
}
let mut rest = text;
while !rest.is_empty() {
if let Some(stripped) = rest.strip_prefix("<image>") {
@@ -460,6 +466,7 @@ impl Tokenizer for FakeChatTokenizer {
fn id_to_token(&self, id: u32) -> Option<String> {
match id {
FAKE_BOS_TOKEN_ID => Some("<bos>".to_string()),
999 => Some("<image>".to_string()),
151655 => Some("<|image_pad|>".to_string()),
0xF001 => Some("<think>".to_string()),
@@ -468,6 +475,7 @@ impl Tokenizer for FakeChatTokenizer {
0xF004 => Some("<|END_THINKING|>".to_string()),
0xF005 => Some("◁think▷".to_string()),
0xF006 => Some("◁/think▷".to_string()),
id if id < 128 => char::from_u32(id).map(|ch| ch.to_string()),
_ => None,
}
}
@@ -775,6 +783,19 @@ async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
(app, engine_task)
}
async fn test_app_with_api_keys(api_keys: Vec<String>) -> (axum::Router, MockEngineTask) {
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
b"engine-openai-api-key",
default_stream_output_specs(),
Arc::new(FakeChatBackend::new()),
)
.await;
let app = build_router(Arc::new(
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_keys(api_keys),
));
(app, engine_task)
}
async fn test_health_app_with_engine_script<F>(
script: F,
) -> (axum::Router, Arc<AppState>, MockEngineTask)
@@ -1073,6 +1094,101 @@ async fn request_id_header_echoes_incoming_header_when_enabled() {
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn api_key_auth_rejects_missing_token_on_guarded_route() {
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
let response = app
.call(
Request::builder()
.method("GET")
.uri("/v1/models")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
let json: serde_json::Value = serde_json::from_slice(&body).expect("json body");
assert_eq!(json, json!({ "error": "Unauthorized" }));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn api_key_auth_rejects_wrong_token_on_guarded_route() {
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
let response = app
.call(
Request::builder()
.method("GET")
.uri("/v1/models")
.header("authorization", "Bearer wrong")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn api_key_auth_accepts_matching_bearer_token_on_guarded_route() {
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
let response = app
.call(
Request::builder()
.method("GET")
.uri("/v1/models")
.header("authorization", "Bearer secret")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn api_key_auth_allows_options_without_token() {
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
let response = app
.call(
Request::builder()
.method("OPTIONS")
.uri("/v1/models")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_ne!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn api_key_auth_allows_unguarded_route_without_token() {
let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
let response = app
.call(
Request::builder()
.method("GET")
.uri("/health")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn version_returns_engine_vllm_version() {
@@ -2334,7 +2450,8 @@ async fn non_stream_completions_echo_prepends_prompt_text() {
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hello",
"echo": true,
"stream": false
"stream": false,
"add_special_tokens": false
})
.to_string(),
))
@@ -2537,7 +2654,8 @@ async fn non_stream_completions_include_prompt_logprobs() {
"prompt": "hello",
"stream": false,
"echo": true,
"logprobs": 1
"logprobs": 1,
"add_special_tokens": false
})
.to_string(),
))
@@ -3318,7 +3436,8 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {
"prompt": "hello",
"echo": true,
"stream": true,
"stream_options": {"include_usage": true}
"stream_options": {"include_usage": true},
"add_special_tokens": false
})
.to_string(),
))
@@ -4228,6 +4347,161 @@ async fn is_sleeping_route_returns_json_payload() {
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn pause_route_uses_python_compatible_default_query_values() {
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
boxed_test_future(async move {
let utility = recv_engine_message(dealer).await;
assert_eq!(utility[0].as_ref(), &[0x03]);
let payload = decode_value(&utility[1]).expect("decode utility payload");
let array = payload.as_array().expect("utility payload array");
let call_id = array[1].as_u64().expect("call id");
assert_eq!(array[2], Value::from("pause_scheduler"));
assert_eq!(
array[3],
Value::Array(vec![Value::from("abort"), Value::from(true)])
);
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
})
})
.await;
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/pause")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
json!({ "status": "paused" })
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn pause_route_rejects_invalid_mode() {
let (app, engine_task) =
test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/pause?mode=banana")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
assert_eq!(json["error"]["type"], "invalid_request_error");
assert_eq!(json["error"]["param"], "mode");
engine_task.abort_and_join().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn resume_route_sends_no_args() {
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
boxed_test_future(async move {
let utility = recv_engine_message(dealer).await;
assert_eq!(utility[0].as_ref(), &[0x03]);
let payload = decode_value(&utility[1]).expect("decode utility payload");
let array = payload.as_array().expect("utility payload array");
let call_id = array[1].as_u64().expect("call id");
assert_eq!(array[2], Value::from("resume_scheduler"));
assert_eq!(array[3], Value::Array(Vec::new()));
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
})
})
.await;
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/resume")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
json!({ "status": "resumed" })
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn is_paused_route_returns_json_payload() {
let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
boxed_test_future(async move {
let utility = recv_engine_message(dealer).await;
assert_eq!(utility[0].as_ref(), &[0x03]);
let payload = decode_value(&utility[1]).expect("decode utility payload");
let array = payload.as_array().expect("utility payload array");
let call_id = array[1].as_u64().expect("call id");
assert_eq!(array[2], Value::from("is_scheduler_paused"));
assert_eq!(array[3], Value::Array(Vec::new()));
send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
})
})
.await;
let response = app
.clone()
.call(
Request::builder()
.method("GET")
.uri("/is_paused")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
json!({ "is_paused": true })
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
@@ -4244,6 +4518,9 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
("GET", "/is_sleeping"),
("POST", "/sleep"),
("POST", "/wake_up"),
("GET", "/is_paused"),
("POST", "/pause"),
("POST", "/resume"),
("POST", "/collective_rpc"),
("POST", "/reset_prefix_cache"),
("POST", "/reset_mm_cache"),
@@ -4598,3 +4875,392 @@ async fn completions_empty_stop_string_returns_validation_error() {
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
async fn post_json(
app: &mut axum::Router,
uri: &str,
body: serde_json::Value,
) -> (StatusCode, serde_json::Value) {
let response = app
.call(
Request::builder()
.method("POST")
.uri(uri)
.header("content-type", "application/json")
.body(Body::from(body.to_string()))
.expect("build request"),
)
.await
.expect("call app");
let status = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
let json: serde_json::Value = serde_json::from_slice(&bytes)
.unwrap_or_else(|_| json!({ "raw": String::from_utf8_lossy(&bytes) }));
(status, json)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_completion_round_trips_through_detokenize() {
let mut app = test_app().await;
let prompt = "Hello world";
let (_, tokenize_json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": prompt,
"add_special_tokens": false,
}),
)
.await;
let tokens = tokenize_json["tokens"]
.as_array()
.expect("tokens array")
.iter()
.map(|v| v.as_u64().expect("token id") as u32)
.collect::<Vec<_>>();
let (status, detokenize_json) = post_json(
&mut app,
"/detokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"tokens": tokens,
}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(detokenize_json["prompt"], prompt);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_completion_add_special_tokens_changes_ids() {
let mut app = test_app().await;
let (_, with_special) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hi",
"add_special_tokens": true,
}),
)
.await;
let (_, without_special) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hi",
"add_special_tokens": false,
}),
)
.await;
let with_ids: Vec<u32> = with_special["tokens"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_u64().unwrap() as u32)
.collect();
let without_ids: Vec<u32> = without_special["tokens"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_u64().unwrap() as u32)
.collect();
assert_ne!(with_ids, without_ids);
assert_eq!(with_ids.first().copied(), Some(FAKE_BOS_TOKEN_ID));
assert_eq!(without_ids.first().copied(), Some(b'h' as u32));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_completion_return_token_strs_matches_tokens() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hi",
"add_special_tokens": false,
"return_token_strs": true,
}),
)
.await;
assert_eq!(status, StatusCode::OK);
let tokens = json["tokens"].as_array().expect("tokens");
let token_strs = json["token_strs"].as_array().expect("token_strs");
assert_eq!(tokens.len(), token_strs.len());
assert_eq!(token_strs.len(), json["count"].as_u64().unwrap() as usize);
assert!(!token_strs[0].as_str().unwrap().is_empty());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_completion_count_and_max_model_len() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"prompt": "hello",
"add_special_tokens": false,
}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(
json["count"].as_u64().unwrap() as usize,
json["tokens"].as_array().unwrap().len()
);
assert!(json["max_model_len"].as_u64().unwrap() > 0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_includes_generation_prompt_in_token_count() {
let mut app = test_app().await;
let messages = json!([{"role": "user", "content": "hi"}]);
let (_, with_prompt) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": messages,
"add_generation_prompt": true,
"add_special_tokens": false,
}),
)
.await;
let (_, without_prompt) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": messages,
"add_generation_prompt": false,
"add_special_tokens": false,
}),
)
.await;
let with_len = with_prompt["tokens"].as_array().unwrap().len();
let without_len = without_prompt["tokens"].as_array().unwrap().len();
assert!(with_len > without_len);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_conflicting_generation_flags_returns_400() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": [{"role": "user", "content": "hi"}],
"add_generation_prompt": true,
"continue_final_message": true,
}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(json["error"]["type"], "invalid_request_error");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_empty_messages_returns_400() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": [],
}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(json["error"]["type"], "invalid_request_error");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_empty_message_content_returns_400() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": [{"role": "user", "content": ""}],
}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(json["error"]["type"], "invalid_request_error");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_unknown_model_returns_404() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "does-not-exist",
"prompt": "hello",
}),
)
.await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(json["error"]["type"], "invalid_request_error");
assert_eq!(json["error"]["code"], "model_not_found");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn detokenize_unknown_model_returns_404() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/detokenize",
json!({
"model": "does-not-exist",
"tokens": [72, 101, 108, 108, 111],
}),
)
.await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(json["error"]["type"], "invalid_request_error");
assert_eq!(json["error"]["code"], "model_not_found");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn detokenize_empty_tokens_returns_empty_prompt() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/detokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"tokens": [],
}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(json["prompt"], "");
}
/// Decode an explicit token sequence — pins `/detokenize` independently of
/// `/tokenize` (the round-trip test alone would pass even if encode and decode
/// were both wrong in mirrored ways).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn detokenize_decodes_known_token_ids() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/detokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"tokens": [72, 101, 108, 108, 111],
}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(json["prompt"], "Hello");
}
/// `continue_final_message` without a trailing assistant message must 400.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_continue_without_assistant_returns_400() {
let mut app = test_app().await;
let (status, json) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": [{"role": "user", "content": "hi"}],
"add_generation_prompt": false,
"continue_final_message": true,
}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(json["error"]["type"], "invalid_request_error");
}
/// `continue_final_message` must not append a new generation suffix vs `add_generation_prompt`.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn tokenize_chat_continue_final_vs_new_assistant_differs() {
let mut app = test_app().await;
let messages = json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "partial,"}
]);
let (_, continue_final) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": messages,
"add_generation_prompt": false,
"continue_final_message": true,
"add_special_tokens": false,
}),
)
.await;
let (_, new_assistant) = post_json(
&mut app,
"/tokenize",
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"messages": messages,
"add_generation_prompt": true,
"continue_final_message": false,
"add_special_tokens": false,
}),
)
.await;
let continue_len = continue_final["tokens"].as_array().unwrap().len();
let new_len = new_assistant["tokens"].as_array().unwrap().len();
assert!(new_len > continue_len);
}
+148
View File
@@ -0,0 +1,148 @@
//! `POST /tokenize` and `POST /detokenize` (root paths, matching Python).
//!
//! Encode/decode runs entirely in-process via [`DynTokenizer`]; the inference
//! engine is not involved.
mod types;
use std::sync::Arc;
use axum::Json;
use axum::extract::State;
use axum::http::HeaderMap;
use axum::response::{IntoResponse, Response};
use thiserror_ext::AsReport as _;
use crate::error::{ApiError, server_error};
use crate::routes::openai::utils::validated_json::ValidatedJson;
use crate::routes::tokenize::types::{
DetokenizeRequest, DetokenizeResponse, TokenizeChatRequest, TokenizeCompletionRequest,
TokenizeRequest, TokenizeResponse,
};
use crate::state::AppState;
use crate::utils::resolve_base_request_id;
/// Match Python `tokenize-{base}` where base is `X-Request-Id` or a new UUID.
fn tokenize_request_id(headers: &HeaderMap) -> String {
let base = resolve_base_request_id(
headers.get("X-Request-Id").and_then(|value| value.to_str().ok()),
None,
);
format!("tokenize-{base}")
}
/// Reject an unknown model name, matching the other handlers.
fn check_model(state: &AppState, model: Option<&str>) -> Result<(), ApiError> {
if let Some(model) = model
&& !state.served_model_names().iter().any(|n| n == model)
{
return Err(ApiError::model_not_found(model.to_string()));
}
Ok(())
}
/// Build the `token_strs` vector when requested, via the tokenizer vocab.
fn token_strs(tokenizer: &vllm_text::tokenizer::DynTokenizer, ids: &[u32]) -> Vec<String> {
// Unknown IDs yield "" — intentional; matches Python's convert_ids_to_tokens behaviour.
ids.iter().map(|&id| tokenizer.id_to_token(id).unwrap_or_default()).collect()
}
pub async fn tokenize(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
ValidatedJson(body): ValidatedJson<TokenizeRequest>,
) -> Response {
let request_id = tokenize_request_id(&headers);
let tokenizer = state.chat.text().tokenizer();
let max_model_len = state.chat.engine_core_client().max_model_len();
let result = match body {
// Completion form: encode the raw `prompt` string (no chat template).
TokenizeRequest::Completion(req) => tokenize_completion(&state, &tokenizer, req),
// Chat form: render `messages` through the template, then encode (see `tokenize_chat`).
TokenizeRequest::Chat(req) => tokenize_chat(&state, &request_id, req).await,
};
match result {
Ok((tokens, want_strs)) => {
let token_strs = want_strs.then(|| token_strs(&tokenizer, &tokens));
Json(TokenizeResponse {
count: tokens.len(),
max_model_len,
tokens,
token_strs,
})
.into_response()
}
Err(error) => error.into_response(),
}
}
fn tokenize_completion(
state: &AppState,
tokenizer: &vllm_text::tokenizer::DynTokenizer,
req: TokenizeCompletionRequest,
) -> Result<(Vec<u32>, bool), ApiError> {
check_model(state, req.model.as_deref())?;
let tokens = tokenizer
.encode(&req.prompt, req.add_special_tokens)
.map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
Ok((tokens, req.return_token_strs))
}
/// HTTP adapter for the chat-shaped `/tokenize` body.
///
/// Not [`vllm_chat::ChatLlm::tokenize_chat`]: this checks the model name and maps
/// errors to [`ApiError`]; the chat-crate method does render → finalize → encode.
async fn tokenize_chat(
state: &AppState,
request_id: &str,
req: TokenizeChatRequest,
) -> Result<(Vec<u32>, bool), ApiError> {
check_model(state, req.model.as_deref())?;
let return_token_strs = req.return_token_strs;
// `continue_final_message` / `add_generation_prompt` mutual exclusion is
// enforced in `normalize_generation_prompt_mode` inside `into_chat_request`.
let tokens = state
.chat
.tokenize_chat(req.into_chat_request(request_id.to_string())?)
.await
.map_err(|e| server_error!("tokenize failed: {}", e.to_report_string()))?;
Ok((tokens, return_token_strs))
}
pub async fn detokenize(
State(state): State<Arc<AppState>>,
ValidatedJson(body): ValidatedJson<DetokenizeRequest>,
) -> Response {
if let Err(error) = check_model(&state, body.model.as_deref()) {
return error.into_response();
}
let tokenizer = state.chat.text().tokenizer();
match tokenizer.decode(&body.tokens, /* skip_special_tokens = */ false) {
Ok(prompt) => Json(DetokenizeResponse { prompt }).into_response(),
Err(e) => server_error!("detokenize failed: {}", e.to_report_string()).into_response(),
}
}
#[cfg(test)]
mod tests {
use axum::http::{HeaderMap, HeaderValue};
use super::tokenize_request_id;
#[test]
fn tokenize_request_id_prefers_x_request_id_header() {
let mut headers = HeaderMap::new();
headers.insert("X-Request-Id", HeaderValue::from_static("client-req-1"));
assert_eq!(tokenize_request_id(&headers), "tokenize-client-req-1");
}
#[test]
fn tokenize_request_id_generates_uuid_when_header_missing() {
let headers = HeaderMap::new();
let id = tokenize_request_id(&headers);
assert!(id.starts_with("tokenize-"));
assert_ne!(id, "tokenize-");
}
}
@@ -0,0 +1,209 @@
use std::collections::HashMap;
use itertools::Itertools as _;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use validator::{Validate, ValidationErrors};
use vllm_chat::{ChatOptions, ChatRequest, ChatToolChoice, SamplingParams};
use vllm_text::output::TextDecodeOptions;
use crate::error::ApiError;
use crate::routes::openai::chat_completions::convert::{
convert_message, convert_tools, normalize_generation_prompt_mode,
};
use crate::routes::openai::utils::types::{
ChatMessage, Normalizable, Tool, default_true, validate_messages,
};
/// `POST /tokenize` body. Untagged: a JSON object with `messages` parses as the
/// chat variant; one with `prompt` parses as the completion variant.
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum TokenizeRequest {
Chat(TokenizeChatRequest),
Completion(TokenizeCompletionRequest),
}
#[derive(Debug, Clone, Deserialize)]
pub struct TokenizeCompletionRequest {
pub model: Option<String>,
pub prompt: String,
#[serde(default = "default_true")]
pub add_special_tokens: bool,
#[serde(default)]
pub return_token_strs: bool,
}
#[derive(Debug, Clone, Deserialize, Validate)]
pub struct TokenizeChatRequest {
pub model: Option<String>,
#[validate(custom(function = "validate_messages"))]
pub messages: Vec<ChatMessage>,
#[serde(default = "default_true")]
pub add_generation_prompt: bool,
#[serde(default)]
pub continue_final_message: bool,
#[serde(default)] // chat default is FALSE (template adds specials)
pub add_special_tokens: bool,
#[serde(default)]
pub return_token_strs: bool,
#[serde(default)]
pub chat_template: Option<String>,
#[serde(default)]
pub chat_template_kwargs: Option<HashMap<String, Value>>,
#[serde(default)]
pub tools: Option<Vec<Tool>>,
}
impl TokenizeChatRequest {
/// Lower this tokenize body into a [`ChatRequest`] for template rendering.
///
/// Reuses [`convert_message`] and [`normalize_generation_prompt_mode`] from
/// `chat_completions/convert` so message lowering and generation-prompt
/// rules match chat completions. Only fields that affect rendering are set;
/// `sampling_params`, `decode_options`, etc. stay at default because
/// tokenize never generates.
pub fn into_chat_request(self, request_id: String) -> Result<ChatRequest, ApiError> {
let messages: Vec<_> = self.messages.into_iter().map(convert_message).try_collect()?;
let generation_prompt_mode = normalize_generation_prompt_mode(
Some(self.add_generation_prompt),
self.continue_final_message,
&messages,
)?;
Ok(ChatRequest {
request_id,
messages,
sampling_params: SamplingParams::default(),
chat_options: ChatOptions {
generation_prompt_mode,
chat_template: self.chat_template,
reasoning_effort: None,
template_kwargs: self.chat_template_kwargs.unwrap_or_default(),
},
tools: convert_tools(self.tools)?,
tool_choice: ChatToolChoice::Auto,
decode_options: TextDecodeOptions::default(),
intermediate: false,
priority: 0,
documents: None,
cache_salt: None,
add_special_tokens: self.add_special_tokens,
data_parallel_rank: None,
lora_request: None,
})
}
}
#[derive(Debug, Clone, Deserialize)]
pub struct DetokenizeRequest {
pub model: Option<String>,
pub tokens: Vec<u32>,
}
#[derive(Debug, Clone, Serialize)]
pub struct TokenizeResponse {
pub count: usize,
pub max_model_len: u32,
pub tokens: Vec<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub token_strs: Option<Vec<String>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct DetokenizeResponse {
pub prompt: String,
}
// ---- trait impls required by ValidatedJson ----
impl Validate for TokenizeRequest {
fn validate(&self) -> Result<(), ValidationErrors> {
if let Self::Chat(req) = self {
req.validate()?;
}
Ok(())
}
}
impl Validate for DetokenizeRequest {
fn validate(&self) -> Result<(), ValidationErrors> {
Ok(())
}
}
impl Normalizable for TokenizeRequest {} // default no-op normalize()
impl Normalizable for DetokenizeRequest {}
#[cfg(test)]
mod tests {
use super::*;
use crate::routes::openai::utils::types::{ChatMessage, MessageContent};
use serde_json::json;
use vllm_chat::ChatTool;
#[test]
fn tokenize_request_converts_openai_tools() {
// The untagged `TokenizeRequest` must resolve a messages+tools body to
// the chat variant and accept standard OpenAI tool objects
// (`{"type":"function",...}`), then convert them to `ChatTool`.
let request: TokenizeRequest = serde_json::from_value(json!({
"messages": [{"role": "user", "content": "hi"}],
"tools": [{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
},
},
}],
}))
.expect("OpenAI tool JSON deserializes to the chat variant");
let TokenizeRequest::Chat(req) = request else {
panic!("messages+tools body should parse as the chat variant");
};
let chat_request =
req.into_chat_request("tokenize-test".to_string()).expect("request is valid");
assert_eq!(
chat_request.tools,
vec![ChatTool {
name: "get_weather".to_string(),
description: Some("Get weather".to_string()),
parameters: json!({
"type": "object",
"properties": {"city": {"type": "string"}},
}),
strict: None,
}]
);
}
#[test]
fn into_chat_request_rejects_conflicting_generation_flags() {
let req = TokenizeChatRequest {
model: None,
messages: vec![ChatMessage::User {
content: MessageContent::Text("hi".to_string()),
name: None,
}],
add_generation_prompt: true,
continue_final_message: true,
add_special_tokens: false,
return_token_strs: false,
chat_template: None,
chat_template_kwargs: None,
tools: None,
};
let error = req
.into_chat_request("tokenize-test".to_string())
.expect_err("conflicting flags");
assert_eq!(
error.to_error_response().error.message,
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
);
}
}
+28
View File
@@ -2,6 +2,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::time::{Duration, Instant, sleep_until};
use tracing::warn;
use vllm_chat::ChatLlm;
@@ -14,6 +15,12 @@ use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
pub(crate) type ApiKeyHash = [u8; 32];
pub(crate) fn hash_api_key(api_key: &str) -> ApiKeyHash {
Sha256::digest(api_key.as_bytes()).into()
}
/// Shared router state for the minimal single-model OpenAI server.
pub struct AppState {
/// All public model IDs served by this frontend. The first entry is the
@@ -27,6 +34,8 @@ pub struct AppState {
pub enable_request_id_headers: bool,
/// Runtime server information returned by `/server_info`, when available.
server_info: Option<ServerInfoSnapshot>,
/// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
api_key_hashes: Vec<ApiKeyHash>,
/// Number of in-flight inference requests currently owned by this frontend.
server_load: AtomicU64,
/// Dynamic LoRA adapter registry.
@@ -53,6 +62,7 @@ impl AppState {
enable_log_requests: false,
enable_request_id_headers: false,
server_info: None,
api_key_hashes: Vec::new(),
server_load: AtomicU64::new(0),
lora_manager: LoraManager::new(),
}
@@ -84,6 +94,24 @@ impl AppState {
self.server_info.as_ref().map(|server_info| server_info.response(config_format))
}
/// Configure API keys accepted by guarded HTTP routes.
pub fn with_api_keys(mut self, api_keys: Vec<String>) -> Self {
self.api_key_hashes = api_keys
.into_iter()
.filter(|key| !key.is_empty())
.map(|key| hash_api_key(&key))
.collect();
self
}
pub(crate) fn has_api_keys(&self) -> bool {
!self.api_key_hashes.is_empty()
}
pub(crate) fn api_key_hashes(&self) -> &[ApiKeyHash] {
&self.api_key_hashes
}
/// The primary model name echoed back in API responses (the first served
/// name).
pub fn primary_model_name(&self) -> &str {
+37 -3
View File
@@ -1,3 +1,5 @@
use std::collections::BTreeMap;
use winnow::ascii::{digit1, multispace0 as ws0};
use winnow::combinator::{alt, eof, repeat, seq};
use winnow::prelude::*;
@@ -33,6 +35,7 @@ enum KimiK2Event {
ToolCallsStart,
ToolCallStart,
ToolCallHeader {
tool_call_id: String,
function_name: String,
function_index: usize,
},
@@ -60,6 +63,7 @@ pub struct KimiK2ToolParser {
buffer: String,
mode: KimiK2Mode,
active_tool_index: Option<usize>,
call_ids: BTreeMap<usize, String>,
}
impl KimiK2ToolParser {
@@ -69,6 +73,7 @@ impl KimiK2ToolParser {
buffer: String::new(),
mode: KimiK2Mode::Text,
active_tool_index: None,
call_ids: BTreeMap::new(),
}
}
@@ -81,6 +86,7 @@ impl KimiK2ToolParser {
KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,
KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,
KimiK2Event::ToolCallHeader {
tool_call_id,
function_name,
function_index,
} => {
@@ -89,6 +95,7 @@ impl KimiK2ToolParser {
self.mode = KimiK2Mode::Arguments {
json_scan: JsonObjectScanState::default(),
};
self.call_ids.insert(tool_index, tool_call_id);
output.calls.push(ToolCallDelta {
tool_index,
name: Some(function_name),
@@ -123,6 +130,7 @@ impl KimiK2ToolParser {
fn reset(&mut self) -> String {
self.mode = KimiK2Mode::Text;
self.active_tool_index = None;
self.call_ids.clear();
std::mem::take(&mut self.buffer)
}
}
@@ -139,6 +147,10 @@ impl ToolParser for KimiK2ToolParser {
true
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
self.call_ids.get(&tool_index).map(String::as_str)
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
self.buffer.push_str(chunk);
@@ -232,16 +244,18 @@ fn tool_call_end_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event>
/// Parse a Kimi K2 tool-call header before the argument marker.
fn tool_call_header_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {
let (header, _) = (
let (raw_header, _) = (
take_until(1.., TOOL_CALL_ARGUMENT_START),
literal(TOOL_CALL_ARGUMENT_START),
)
.parse_next(input)?;
let mut header_input = header;
let tool_call_id = raw_header.trim().to_string();
let mut header_input = raw_header;
let (header, _, _) = (tool_header, ws0, eof).parse_next(&mut header_input)?;
Ok(KimiK2Event::ToolCallHeader {
tool_call_id,
function_name: header.function_name,
function_index: header.function_index,
})
@@ -502,6 +516,25 @@ mod tests {
.assert_debug_eq(&output);
}
#[test]
fn kimi_k2_preserves_model_generated_tool_call_ids() {
let mut parser = KimiK2ToolParser::new(&test_tools());
let input = build_tool_section(&[
build_tool_call("get_weather", 0, r#"{"location":"Shanghai"}"#),
build_tool_call("add", 1, r#"{"x":1,"y":2}"#),
]);
for chunk in split_by_chars(&input, 7) {
parser.parse_chunk(chunk).unwrap();
}
// IDs are available after parsing but before finish(), which calls reset().
assert_eq!(parser.tool_call_id(0), Some("functions.get_weather:0"));
assert_eq!(parser.tool_call_id(1), Some("functions.add:1"));
parser.finish().unwrap();
assert_eq!(parser.tool_call_id(0), None);
}
#[test]
fn kimi_k2_accepts_non_functions_header_prefix() {
let mut parser = KimiK2ToolParser::new(&test_tools());
@@ -509,9 +542,10 @@ mod tests {
"{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"
);
let output = parser.parse_complete(&input).unwrap();
let output = parser.parse_chunk(&input).unwrap().coalesce_calls();
assert_eq!(output.calls[0].tool_index, 42);
assert_eq!(parser.tool_call_id(42), Some("api.tools.search:42"));
assert_eq!(output.calls[0].name.as_deref(), Some("search"));
assert_eq!(output.calls[0].arguments, "{}");
}
+6
View File
@@ -121,6 +121,12 @@ pub trait ToolParser: Send {
false
}
/// Return the parser-provided ID for a tool call by index, if the model
/// emitted one.
fn tool_call_id(&self, _tool_index: usize) -> Option<&str> {
None
}
/// Feed one decoded text delta into the parser, appending committed output
/// into `output`.
///
+26
View File
@@ -248,6 +248,32 @@ def test_deep_sleep_async():
asyncio.run(test())
@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
def test_cumem_without_sleep_mode():
"""Verify cumem allocator works independently of sleep mode."""
llm = LLM("hmellor/tiny-random-LlamaForCausalLM", enable_cumem_allocator=True)
prompt = "How are you?"
sampling_params = SamplingParams(temperature=0, max_tokens=10)
output = llm.generate(prompt, sampling_params)
assert output[0].outputs[0].text
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="Sleep mode requires CUDA or ROCm",
)
def test_sleep_mode_auto_enables_cumem():
"""Verify sleep mode automatically enables cumem allocator."""
from vllm.config.model import ModelConfig
cfg = ModelConfig(
"hmellor/tiny-random-LlamaForCausalLM",
enable_sleep_mode=True,
enable_cumem_allocator=False,
)
assert cfg.enable_cumem_allocator is True
@requires_fp8
def test_deep_sleep_fp8_kvcache():
model = "Qwen/Qwen2-0.5B"
@@ -189,6 +189,118 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
]
class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
"""Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
Four ``rms_norm`` sites that together hit every pattern registered by
``RocmAiterAllReduceFusionPass`` for the per-group FP8 quant path:
* ``norm[0]``: ``all_reduce -> rms_norm -> group_fp8_quant`` (no residual)
-> ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern``
* ``norm[1]``: ``all_reduce -> fused_add_rms_norm -> group_fp8_quant``
(single ``rms`` consumer)
-> ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern``
* ``norm[2..3]``: ``all_reduce -> fused_add_rms_norm
-> (group_fp8_quant + rocm_unquantized_gemm)`` (two ``rms`` consumers,
modeling the DSv3.2 indexer fan-out)
-> ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern``
The chain feeds the next AllReduce by dequantizing the FP8 output (FP8
cast back to bf16 multiplied by the per-group scale), which is enough to
keep the matmul chain bf16 without depending on a real FP8 block-scaled
GEMM kernel.
"""
quant_group_size = 128
indexer_out_dim = 8
def __init__(
self,
hidden_size=128,
token_num=16,
eps=1e-6,
dtype: torch.dtype = torch.bfloat16,
use_triton_quant: bool = False,
):
super().__init__()
self.hidden_size = hidden_size
self.eps = eps
self.use_triton_quant = use_triton_quant
assert hidden_size % self.quant_group_size == 0, (
f"hidden_size ({hidden_size}) must be a multiple of "
f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant"
)
self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
self.w = [torch.rand(hidden_size, hidden_size, dtype=dtype) for _ in range(3)]
self.indexer_w = [
torch.rand(self.indexer_out_dim, hidden_size, dtype=dtype) for _ in range(2)
]
def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
if self.use_triton_quant:
return torch.ops.vllm.triton_per_token_group_quant_fp8(
rms, self.quant_group_size
)
return torch.ops.vllm.rocm_aiter_group_fp8_quant.default(
rms, self.quant_group_size
)
def _dequantize_to_bf16(
self, q: torch.Tensor, s: torch.Tensor, ref: torch.Tensor
) -> torch.Tensor:
# Broadcast the per-group scale across each group of `quant_group_size`
# so we can chain the FP8 output back into a bf16 matmul. This avoids
# depending on a real FP8 block-scaled GEMM kernel in the test.
s_full = s.repeat_interleave(self.quant_group_size, dim=-1).to(ref.dtype)
return q.to(ref.dtype) * s_full
def forward(self, hidden_states):
z = torch.relu(hidden_states)
x = resid = tensor_model_parallel_all_reduce(z)
rms = self.norm[0](x)
q0, s0 = self._group_quant(rms)
y = self._dequantize_to_bf16(q0, s0, rms)
z2 = torch.mm(y, self.w[0])
x2 = tensor_model_parallel_all_reduce(z2)
rms2, resid = self.norm[1](x2, resid)
q1, s1 = self._group_quant(rms2)
y2 = self._dequantize_to_bf16(q1, s1, rms2)
z3 = torch.mm(y2, self.w[1])
x3 = tensor_model_parallel_all_reduce(z3)
rms3, resid = self.norm[2](x3, resid)
q2, s2 = self._group_quant(rms3)
# Second consumer of ``rms3``: forces the with-indexer pattern.
idx2 = torch.ops.vllm.rocm_unquantized_gemm(rms3, self.indexer_w[0], None)
y3 = self._dequantize_to_bf16(q2, s2, rms3)
z4 = torch.mm(y3, self.w[2])
x4 = tensor_model_parallel_all_reduce(z4)
rms4, resid = self.norm[3](x4, resid)
q3, s3 = self._group_quant(rms4)
# Second consumer of ``rms4``: forces the with-indexer pattern.
idx3 = torch.ops.vllm.rocm_unquantized_gemm(rms4, self.indexer_w[1], None)
y4 = self._dequantize_to_bf16(q3, s3, rms4)
return y4, idx2, idx3
def ops_in_model_before(self):
return [
torch.ops.vllm.all_reduce.default,
(
torch.ops.vllm.triton_per_token_group_quant_fp8.default
if self.use_triton_quant
else torch.ops.vllm.rocm_aiter_group_fp8_quant.default
),
]
def ops_in_model_after(self):
return [
rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op(),
rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op(), # noqa: E501
]
class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
def __init__(
self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
@@ -463,3 +575,167 @@ def all_reduce_fusion_pass_on_test_model(
assert fused_nodes
assert all(n.kwargs.get("weight_bias") == 1.0 for n in fused_nodes)
del all_reduce_fusion_pass
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("use_triton_quant", [True, False])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize("seq_len", [8])
@pytest.mark.parametrize("hidden_size", [128])
@pytest.mark.parametrize("dtype", [torch.bfloat16])
@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only",
)
@pytest.mark.skipif(not IS_AITER_FOUND, reason="aiter is not found")
def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
batch_size: int,
seq_len: int,
hidden_size: int,
dtype: torch.dtype,
enable_rms_norm_custom_op: bool,
use_triton_quant: bool,
monkeypatch: pytest.MonkeyPatch,
):
"""Sibling of ``test_all_reduce_fusion_pass_replace`` for the new
ROCm AITER AR+RMS+per-group-FP8-quant fusion patterns.
Validates the three new ``VllmPatternReplacement`` patterns added to
``RocmAiterAllReduceFusionPass``:
* ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern`` (no-residual)
* ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual,
single ``rms`` consumer)
* ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with-
residual, DSv3.2 indexer fan-out; parametrized over both
``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant``
producers).
"""
with monkeypatch.context() as m:
m.setenv("VLLM_ROCM_USE_AITER", "1")
rocm_aiter_ops.refresh_env_variables()
if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group():
pytest.skip(
"aiter build is missing 'fused_ar_rms_per_group_quant' (needs "
"ROCm/aiter PR #2823); the new patterns aren't registered."
)
num_processes = 2
def run_torch_spawn(fn, nprocs):
torch.multiprocessing.spawn(
fn,
args=(
num_processes,
TestAiterAllReduceRMSNormGroupQuantFP8Model,
batch_size,
seq_len,
hidden_size,
dtype,
enable_rms_norm_custom_op,
use_triton_quant,
monkeypatch,
),
nprocs=nprocs,
)
run_torch_spawn(rocm_aiter_group_quant_fusion_pass_on_test_model, num_processes)
def rocm_aiter_group_quant_fusion_pass_on_test_model(
local_rank: int,
world_size: int,
test_model_cls: torch.nn.Module,
batch_size: int,
seq_len: int,
hidden_size: int,
dtype: torch.dtype,
enable_rms_norm_custom_op: bool,
use_triton_quant: bool,
monkeypatch: pytest.MonkeyPatch,
):
set_random_seed(0)
device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
torch.accelerator.set_device_index(device)
torch.set_default_device(device)
torch.set_default_dtype(dtype)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12345",
"VLLM_ROCM_USE_AITER": "1",
}
)
rocm_aiter_ops.refresh_env_variables()
init_distributed_environment()
custom_ops = []
if enable_rms_norm_custom_op:
custom_ops.append("+rms_norm")
# ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip``
# only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at
# the call site). The patterns in this PR are robust to both Triton and
# rocm_aiter forms; we always enable +quant_fp8 so the matcher's example
# trace finds the same form the test model uses.
custom_ops.append("+quant_fp8")
vllm_config = VllmConfig(
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE, custom_ops=custom_ops
)
)
vllm_config.compilation_config.pass_config = PassConfig(
fuse_allreduce_rms=True, eliminate_noops=True
)
vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
vllm_config.parallel_config.rank = local_rank
model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8"
vllm_config.model_config = ModelConfig(
model=model_name, trust_remote_code=True, dtype=dtype, seed=42
)
with set_current_vllm_config(vllm_config):
initialize_model_parallel(tensor_model_parallel_size=world_size)
all_reduce_fusion_pass = RocmAiterAllReduceFusionPass(vllm_config)
noop_pass = NoOpEliminationPass(vllm_config)
func_pass = FixFunctionalizationPass(vllm_config)
cleanup_pass = PostCleanupPass(vllm_config)
backend = TestBackend(
noop_pass, all_reduce_fusion_pass, func_pass, cleanup_pass
)
token_num = batch_size * seq_len
model = test_model_cls(
hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant
)
hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
compiled_model = torch.compile(model, backend=backend)
compiled_model(hidden_states)
results_unfused = model(hidden_states)
results_fused = compiled_model(hidden_states)
# The fused per-group AR+RMS+QUANT op is bit-equivalent to the unfused
# chain modulo the small AllReduce + RMSNorm reordering inside aiter.
# Per-group FP8 quant introduces step noise <=1 per group; use the
# same tolerance as the sibling FP8 static test.
torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
# Four pattern firings: norm[0] (no-add quant), norm[1] (add quant,
# single ``rms`` consumer), norm[2..3] (add quant + indexer fan-out).
assert all_reduce_fusion_pass.matched_count == 4, (
f"{all_reduce_fusion_pass.matched_count=}"
)
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
backend.check_after_ops(model.ops_in_model_after())
del all_reduce_fusion_pass
+18 -17
View File
@@ -854,8 +854,13 @@ class HfRunner:
return self
def __exit__(self, exc_type, exc_value, traceback):
from tests.utils import wait_for_rocm_memory_to_settle
del self.model
cleanup_dist_env_and_memory()
# ROCm frees VRAM lazily; wait so a runner started right after this HF
# model exits does not OOM on its startup memory guard.
wait_for_rocm_memory_to_settle()
@pytest.fixture(scope="session")
@@ -1248,25 +1253,13 @@ class VllmRunner:
return self
def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None:
from tests.utils import wait_for_gpu_memory_to_clear
from vllm.platforms import current_platform
if not current_platform.is_rocm():
return
num_gpus = torch.accelerator.device_count()
if num_gpus == 0:
return
from tests.utils import wait_for_rocm_memory_to_settle
# V1 startup requires free_memory >= total * gpu_memory_utilization.
# Wait for the complementary used-memory ratio so the next runner does
# not fail the startup guard immediately after this runner exits. Bound
# the wait so cleanup failures fail this test instead of hanging.
wait_for_gpu_memory_to_clear(
devices=list(range(num_gpus)),
threshold_ratio=1.0 - gpu_memory_utilization,
timeout_s=120,
)
# not fail the startup guard immediately after this runner exits. The
# wait is bounded so cleanup failures fail this test instead of hanging.
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
def __exit__(self, exc_type, exc_value, traceback):
# Explicitly shutdown the engine core to release GPU resources
@@ -1276,8 +1269,16 @@ class VllmRunner:
gpu_memory_utilization = (
self.llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization
)
from vllm.platforms import current_platform
try:
self.llm.llm_engine.engine_core.shutdown()
# Give the engine core time to run its own graceful shutdown
# (model_executor teardown + empty_cache + process-group destroy)
# before the process manager SIGKILLs it at the default 5s. On ROCm
# a hard kill leaves the whole allocation for the driver's slow async
# VRAM reclamation, which starves the next test's startup.
shutdown_timeout = 60.0 if current_platform.is_rocm() else None
self.llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout)
except Exception:
# Ignore shutdown errors as cleanup will still proceed
pass
+1 -1
View File
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from tests.entrypoints.openai.chat_completion.test_oot_registration import (
from tests.plugins_tests.test_oot_registration_online import (
run_and_test_dummy_opt_api_server,
)
@@ -77,9 +77,9 @@ def make_fused_moe_layer(
intermediate_size=test_config.intermediate_size,
prefix=f"dummy_layer_{layer_idx}",
activation="silu",
is_act_and_mul=True,
params_dtype=test_config.weight_dtype,
)
re = fml.routed_experts
device = torch.device(f"cuda:{rank}")
@@ -92,12 +92,12 @@ def make_fused_moe_layer(
tensor_device=device,
)
assert isinstance(fml.w13_weight.data, torch.Tensor)
assert isinstance(fml.w2_weight.data, torch.Tensor)
fml.w13_weight.data = fml.w13_weight.data.to(device=device)
fml.w2_weight.data = fml.w2_weight.data.to(device=device)
w13_weight = fml.w13_weight.data
w2_weight = fml.w2_weight.data
assert isinstance(re.w13_weight.data, torch.Tensor)
assert isinstance(re.w2_weight.data, torch.Tensor)
re.w13_weight.data = re.w13_weight.data.to(device=device)
re.w2_weight.data = re.w2_weight.data.to(device=device)
w13_weight = re.w13_weight.data
w2_weight = re.w2_weight.data
assert w13_weight.size(0) == test_config.num_local_experts
for i in range(test_config.num_local_experts):
g_i = rank * test_config.num_local_experts + i
@@ -172,10 +172,10 @@ def make_fused_moe_layer(
assert not w2_weight_scale_inv.is_contiguous()
# Add scales to the parameter list
fml.w13_weight_scale_inv = torch.nn.Parameter(
re.w13_weight_scale_inv = torch.nn.Parameter(
w13_weight_scale_inv, requires_grad=False
)
fml.w2_weight_scale_inv = torch.nn.Parameter(
re.w2_weight_scale_inv = torch.nn.Parameter(
w2_weight_scale_inv, requires_grad=False
)
@@ -61,7 +61,6 @@ def make_fused_moe_layer(
intermediate_size=test_config.intermediate_size,
prefix=f"dummy_layer_{layer_idx}",
activation="silu",
is_act_and_mul=True,
params_dtype=torch.bfloat16,
quant_config=quant_config,
)
@@ -77,6 +76,7 @@ def make_fused_moe_layer(
)
fml = fml.to(device)
re = fml.routed_experts
w1_q, w2_q, quant_config = make_test_quant_config(
test_config.num_local_experts,
test_config.intermediate_size,
@@ -87,21 +87,21 @@ def make_fused_moe_layer(
per_act_token_quant=False,
)
fml.w13_weight.data = w1_q
fml.w2_weight.data = w2_q
re.w13_weight.data = w1_q
re.w2_weight.data = w2_q
fml.w2_input_scale.data = torch.randn_like(fml.w2_input_scale.data) / 5
fml.w13_input_scale.data = torch.randn_like(fml.w13_input_scale.data) / 5
fml.w2_weight_scale_2.data = torch.randn_like(fml.w2_weight_scale_2.data) / 5
fml.w13_weight_scale_2.data = torch.randn_like(fml.w13_weight_scale_2.data) / 5
fml.w2_weight_scale.data = (
torch.randn(fml.w2_weight_scale.data.shape, device=device) / 5
).to(fml.w2_weight_scale.data.dtype)
fml.w13_weight_scale.data = (
torch.randn(fml.w13_weight_scale.data.shape, device=device) / 5
).to(fml.w13_weight_scale.data.dtype)
re.w2_input_scale.data = torch.randn_like(re.w2_input_scale.data) / 5
re.w13_input_scale.data = torch.randn_like(re.w13_input_scale.data) / 5
re.w2_weight_scale_2.data = torch.randn_like(re.w2_weight_scale_2.data) / 5
re.w13_weight_scale_2.data = torch.randn_like(re.w13_weight_scale_2.data) / 5
re.w2_weight_scale.data = (
torch.randn(re.w2_weight_scale.data.shape, device=device) / 5
).to(re.w2_weight_scale.data.dtype)
re.w13_weight_scale.data = (
torch.randn(re.w13_weight_scale.data.shape, device=device) / 5
).to(re.w13_weight_scale.data.dtype)
nvfp4_fused_moe.process_weights_after_loading(fml)
nvfp4_fused_moe.process_weights_after_loading(re)
fml.maybe_init_modular_kernel()
+78
View File
@@ -19,6 +19,7 @@ from vllm.utils.flashinfer import (
has_flashinfer_nvlink_one_sided,
has_flashinfer_nvlink_two_sided,
)
from vllm.utils.import_utils import has_deep_ep_v2
from vllm.utils.network_utils import get_open_port
from ..utils import init_test_distributed_environment
@@ -194,6 +195,10 @@ requires_ptrace = pytest.mark.skipif(
not _has_sys_ptrace(),
reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
)
requires_deep_ep_v2 = pytest.mark.skipif(
not has_deep_ep_v2(),
reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
)
# NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
# their own @requires_two_sided / @requires_one_sided decorators, and
@@ -856,3 +861,76 @@ def _one_sided_data_worker(rank, world_size):
def test_one_sided_dispatch_combine(world_size):
"""Test FlashInfer one-sided dispatch/combine with actual data flow."""
_spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
# ---------------------------------------------------------------------------
# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
# ---------------------------------------------------------------------------
#
# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
#
# Uses EP group because the DeepEP v2 manager is constructed with an
# EP-scoped communicator in production. With tp=world_size the EP group
# spans all ranks.
# ---------------------------------------------------------------------------
def _deepep_v2_lifecycle_worker(rank, world_size):
from vllm.distributed.device_communicators.all2all import (
DeepEPV2All2AllManager,
)
cpu_group = get_ep_group().cpu_group
manager = DeepEPV2All2AllManager(cpu_group)
assert manager.rank == rank
assert manager.world_size == world_size
assert manager._num_sms is None
hidden_size = 7168
num_experts = world_size * 32
num_topk = 8
max_tokens = 256
handle_kwargs = dict(
num_max_tokens_per_rank=max_tokens,
hidden=hidden_size,
num_topk=num_topk,
num_experts=num_experts,
use_fp8_dispatch=False,
)
handle = manager.get_handle(handle_kwargs)
assert handle is not None
assert manager._num_sms is not None
assert manager._num_sms > 0
torch.distributed.barrier()
# get_handle again with same args should return cached handle
handle2 = manager.get_handle(dict(handle_kwargs))
assert handle2 is handle
torch.distributed.barrier()
# Destroy clears the cache
manager.destroy()
assert len(manager.handle_cache._cache) == 0
torch.distributed.barrier()
# Re-create after destroy
handle3 = manager.get_handle(dict(handle_kwargs))
assert handle3 is not None
torch.distributed.barrier()
manager.destroy()
@requires_multi_gpu
@requires_deep_ep_v2
@pytest.mark.parametrize("world_size", [2])
def test_deepep_v2_manager_lifecycle(world_size):
"""Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
_spawn_workers(_deepep_v2_lifecycle_worker, world_size)
-42
View File
@@ -4,7 +4,6 @@ import weakref
import pytest
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.sampling_params import SamplingParams
@@ -76,47 +75,6 @@ def test_multi_chat(text_llm):
assert len(outputs) == 2
@pytest.fixture(scope="function")
def vision_llm():
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model="microsoft/Phi-3.5-vision-instruct",
max_model_len=4096,
max_num_seqs=5,
enforce_eager=True,
trust_remote_code=True,
limit_mm_per_prompt={"image": 2},
seed=0,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.parametrize(
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
)
def test_chat_multi_image(vision_llm, image_urls: list[str]):
messages = [
{
"role": "user",
"content": [
*(
{"type": "image_url", "image_url": {"url": image_url}}
for image_url in image_urls
),
{"type": "text", "text": "What's in this image?"},
],
}
]
outputs = vision_llm.chat(messages)
assert len(outputs) >= 0
def test_llm_chat_tokenization_no_double_bos(text_llm):
"""
LLM.chat() should not add special tokens when using chat templates.
+10
View File
@@ -0,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
]
@@ -0,0 +1,50 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
@pytest.fixture(scope="function")
def vision_llm():
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model="microsoft/Phi-3.5-vision-instruct",
max_model_len=4096,
max_num_seqs=5,
enforce_eager=True,
trust_remote_code=True,
limit_mm_per_prompt={"image": 2},
seed=0,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.parametrize(
"image_urls", [[TEST_IMAGE_ASSETS[0], TEST_IMAGE_ASSETS[1]]], indirect=True
)
def test_chat_multi_image(vision_llm, image_urls: list[str]):
messages = [
{
"role": "user",
"content": [
*(
{"type": "image_url", "image_url": {"url": image_url}}
for image_url in image_urls
),
{"type": "text", "text": "What's in this image?"},
],
}
]
outputs = vision_llm.chat(messages)
assert len(outputs) >= 0
@@ -15,7 +15,7 @@ import logging
import pytest
import regex as re
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
from vllm import LLM, SamplingParams
from vllm.renderers.params import ChatParams
from vllm.v1.metrics import loggers as stat_loggers
@@ -6,7 +6,7 @@ import logging
import pytest
import regex as re
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
from vllm import LLM
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
from vllm.v1.metrics import loggers as stat_loggers
@@ -52,7 +52,7 @@ async def client_with_image_embeds(server_with_image_embeds):
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("dtype", [torch.half, torch.float16, torch.float32])
async def test_completions_with_image_embeds(
async def test_chat_completions_with_image_embeds(
client_with_image_embeds: openai.AsyncOpenAI,
model_name: str,
image_assets: ImageTestAssets,
@@ -8,6 +8,7 @@ import pytest
import pytest_asyncio
from transformers import AutoProcessor
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.multimodal.media import MediaWithBytes
from vllm.multimodal.utils import encode_image_url, fetch_image
@@ -16,14 +17,6 @@ from vllm.platforms import current_platform
MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
MAXIMUM_IMAGES = 2
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
]
# Required terms for beam search validation
# Each entry is a list of term groups - ALL groups must match
# Each group is a list of alternatives - at least ONE term in the group must appear
@@ -7,19 +7,13 @@ import openai
import pytest
import pytest_asyncio
from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.utils import encode_image_url
# Use a small vision model for testing
MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct"
MAXIMUM_IMAGES = 2
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
]
@pytest.fixture(scope="module")
@@ -14,6 +14,7 @@ import torch
from openai import BadRequestError
from tests.utils import VLLM_PATH, RemoteOpenAIServer
from vllm.platforms import current_platform
MODEL_NAME = "facebook/opt-125m"
CHAT_TEMPLATE = VLLM_PATH / "examples/template_chatml.jinja"
@@ -41,7 +42,11 @@ def server_args() -> list[str]:
@pytest.fixture(scope="module")
def server(server_args):
def server(server_args, request):
if current_platform.is_rocm():
# Materialize HF embeddings before the server reserves ROCm VRAM.
request.getfixturevalue("prompt_embeds_b64")
request.getfixturevalue("aligned_content_and_embeds_b64")
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
yield remote_server
@@ -5,7 +5,7 @@ import json
import openai
import pytest
from ....utils import RemoteOpenAIServer
from tests.utils import RemoteOpenAIServer
MODEL = "ibm-granite/granite-4.0-h-tiny"
@@ -9,12 +9,11 @@ import pytest_asyncio
from huggingface_hub import snapshot_download
from typing_extensions import TypedDict
from tests.utils import RemoteOpenAIServer
from vllm.tool_parsers.abstract_tool_parser import ToolParser
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
from ....utils import RemoteOpenAIServer
LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
TOOLS = [
@@ -9,7 +9,7 @@ import pytest
import pytest_asyncio
from rapidfuzz import fuzz
from ....utils import RemoteOpenAIServer
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "openai/gpt-oss-20b"

Some files were not shown because too many files have changed in this diff Show More