Compare commits

...
Author SHA1 Message Date
Bugen Zhao efa494e397 refactor setup.py to extract stuff into build_rust.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:39:17 +08:00
Bugen Zhao 74b77a593a update dockerfile image
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:38:24 +08:00
Bugen Zhao ec53889a3b simplify to only match .so, as .dylib will also be renamed to .so under macos
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:13:09 +08:00
Bugen Zhao ede3d4ddf6 skip tests when _rust_tool_parser is absent
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:12:48 +08:00
Bugen Zhao 7301a834fa fix glob issue
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 22:06:14 +08:00
Bugen Zhao 7104ac6b5d simplify setup.py
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 11:40:33 +00:00
Bugen Zhao 9a98eccb49 cover two calls case
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2ce1e2ba71 do not infer finished but just leave it unhandled
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 2496f66f7f Run tool parser tests before tokenizer tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao d18b7f2723 Build Rust parser extension in Docker rust stage
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 80b633f076 ci: configure PyO3 Python for Rust tests
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 28f589ebb1 do not enable pyo3 on --all-features
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:02 +00:00
Bugen Zhao 98998279d4 fix fmt & try import
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 71d208089e add tests with deepseek v4
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen Zhao 0ac6fc8352 Add Rust tool parser Python bridge
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 09:50:01 +00:00
Bugen ZhaoandGitHub fdfb2566c0 [Rust Frontend] [CI] Unify Rust artifact builds with setuptools-rust (#44981)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-10 17:48:34 +08:00
Juan Pérez de AlgabaandGitHub 8a5cf1ccd6 [Security] Fix remote DoS via invalid recovered token reinjection (#44744)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-10 02:31:43 -07:00
Kunshang JiandGitHub fe1d923afc [BUGFIX][XPU] fix xpu flash_attn_varlen_func interface (#45110)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-10 17:07:40 +08:00
JartXGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
32daf56b42 [Refactor] Rename rocm_moe.py to rocm_moe_rdna.py (#45011)
Signed-off-by: JartX <sagformas@epdcenter.es>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-10 17:02:09 +08:00
Andreas KaratzasandGitHub 82a42234be [ROCm][CI] Defer AITER sampler import and isolate server test PYTHONPATH (#44823)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-10 08:56:11 +00:00
Harry MellorandGitHub af9f583344 Revert "[Bugfix][CI] Gemma3 Transformers multimodal encoder profiling and build prompt-embedding fixtures" (#45029)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-10 01:37:03 -07:00
yihengandGitHub bd2d83ff31 [SpecDecode] Reduce TP communication for large-vocab draft models speculative decoding (#39419)
Signed-off-by: EanWang211123 <wangyiheng@sangfor.com.cn>
2026-06-10 07:59:24 +00:00
xiaohuguo2023andGitHub bb78168b21 [ROCm][gpt-oss] Hybrid CDNA4 swizzle gate for A8W4 MoE (#44804)
Signed-off-by: Xiaohu Guo <Xiaohu.Guo@amd.com>
2026-06-09 23:59:44 -07:00
89c6a41001 [Bench] Add BFCL dataset for vllm bench serve tool-calling workloads (#42457)
Signed-off-by: Li Zhang <lzhanga@amazon.com>
Co-authored-by: Li Zhang <lzhanga@amazon.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-06-09 23:59:18 -07:00
7fdfa6441d Model/colbert autoweightsloader (#44999)
Signed-off-by: Furkan Fidan <dev@yufufi.com>
Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-09 23:58:50 -07:00
Harry MellorandGitHub e9b728de8a Change from owning configs to owning config utils (#45058)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-10 06:40:25 +00:00
5828a205ef Fix Harmony tool descriptions for optional fields (#44686)
Signed-off-by: Varun Shenoy <varun.vinayak.shenoy@oracle.com>
Co-authored-by: Codex <codex@openai.com>
2026-06-09 23:29:22 -07:00
Yaoming ZhanandGitHub 7a74f31d2e [Rust Frontend] Add seed_oss and step3p5 reasoning parsers (#44552)
Signed-off-by: yzhan1 <zhanyaoming2014@gmail.com>
2026-06-09 23:01:33 -07:00
47930b59ca [Bugfix] Handle HWC images in ImageProcessorItems.get_image_size (#45057)
Signed-off-by: YellowFoxH4XOR <yellowfoxh4xor@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-10 05:35:50 +00:00
Flora FengandGitHub 6aec99f030 [Refactor] Remove dead states from chat completion serving (#45081)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-09 22:20:15 -07:00
bnellnmandGitHub f4966f8b3d [Bugfix] Fix weight loading issues caused by #41184 (#45054)
Signed-off-by: Bill Nell <bnell@redhat.com>
2026-06-10 01:20:13 -04:00
Mohammad Miadh AngkadandGitHub 2c9c07c85e [Bugfix][CI/Build] Fix Rust frontend build after chat conversion refactor (#45085)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
2026-06-09 20:04:41 -07:00
Change72andGitHub 320c52b134 [Bench] benchmark_serving_multi_turn: make non-standard conversation_id payload opt-in (#43756)
Signed-off-by: Change72 <cguo51@asu.edu>
2026-06-09 19:41:56 -07:00
6deb05e0e4 [Core][Model] Gemma4: Unified FA4 for all layers + FlashAttention mm_prefix support (#42175)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-06-09 17:45:39 -07:00
Flora FengandGitHub d82ac00923 [Refactor][Mistral] Extract parsing logic into MistralParser (#44596)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-10 00:12:23 +00:00
Bugen ZhaoandGitHub dac9e9a640 [Rust Frontend] Extract shared options in route helper params (#44884)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-09 17:02:35 -07:00
Wentao YeandGitHub d7607ad273 [Bug] Fix deepseek v4 OOM issue (#44914)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-09 15:47:06 -07:00
d955745d58 [ROCm][CI] fix test_rope_kvcache_fusion.py (#44678)
Signed-off-by: charlifu <charlifu@amd.com>
Co-authored-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com>
2026-06-09 21:53:46 +00:00
Micah WilliamsonandGitHub e1ed89dbee Revert "[Kernel] Speed up silu_and_mul_per_block_quant with warp-shuf… (#45066)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-09 14:12:06 -07:00
JimmyGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>simon-mo
1c2ffc6f88 feat(multi-turn-bench): add api_key and custom headers for multi turn benchmark (#44516)
Signed-off-by: Jimmy <jinmingyi1998@sina.cn>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: simon-mo <simon.mo@hey.com>
2026-06-09 14:00:07 -07:00
Jiangyun ZhuandGitHub ca4cfd8731 [Bugfix] fix qwen3.5 ep weight loading (#45002)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-06-09 13:55:30 -07:00
Micah WilliamsonandGitHub c9c1540e61 [ROCm][V2] Fix failed assertion in Llama models when using EAGLE with ROCM_AITER_FA (#44936)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-09 13:30:52 -05:00
c1d754d681 [Mooncake] Use all HCAs on multi-NIC hosts instead of GPU-indexed RNIC selection (#43799)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-06-09 11:05:36 -07:00
01d8cd92dd [ROCm][Perf] Use fused softplus-sqrt-topk router under AITER fused-MoE (#44945)
Co-authored-by: vLLM Contributor <contributor@vllm.ai>
2026-06-09 17:53:05 +00:00
a4b14b98c6 [Kernel] Speed up silu_and_mul_per_block_quant with warp-shuffle reduction + vectorized I/O (#44173)
Signed-off-by: SII-yangdian <yangdian@sii.edu.cn>
Co-authored-by: SII-yangdian <yangdian@sii.edu.cn>
2026-06-09 10:41:26 -07:00
Juan Pérez de AlgabaandGitHub cf1c906724 [Security] Fix image EXIF orientation and tRNS transparency handling (#44974)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-09 09:34:44 -07:00
766ce2bb6b Fix MiDashengLM TP>1 crash in audio encoder attention (#44408)
Signed-off-by: Michał Ganczarenko <michal.ganczarenko@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-09 09:29:14 -07:00
Ronen SchafferGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
3d119f78f7 [Docs] Add KV offloading usage guide (single- and multi-tier) (#44415)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-09 19:20:23 +03:00
Juan Pérez de AlgabaandGitHub 1b1359c332 [Security] Fix DoS via audio decompression bomb in speech-to-text endpoint (#44970)
Signed-off-by: jperezde <jperezde@redhat.com>
2026-06-10 00:18:53 +08:00
Tyko NiemiandGitHub cad4ca12b8 [Bugfix] Add X-Session-ID from conversation_id in multi-turn benchmark (#44663)
Signed-off-by: Tyko Niemi <tyko.niemi@amd.com>
2026-06-09 08:57:00 -07: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
357 changed files with 13780 additions and 7656 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}" \
@@ -110,6 +110,36 @@ install_uv() {
| env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
}
setup_pyo3_python() {
local python_version="${PYO3_PYTHON_VERSION:-3.12}"
log_section "Installing Python ${python_version} for PyO3 tests"
uv python install "$python_version"
PYO3_PYTHON="$(uv python find \
--managed-python \
--no-project \
--resolve-links \
"$python_version")"
export PYO3_PYTHON
local python_libdir
python_libdir="$("$PYO3_PYTHON" - <<'PY'
import pathlib
import sysconfig
libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
ldlibrary = sysconfig.get_config_var("LDLIBRARY")
assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
assert ldlibrary
assert (libdir / ldlibrary).exists(), libdir / ldlibrary
print(libdir)
PY
)"
export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
}
run_style_clippy() {
install_cargo_sort
@@ -132,6 +162,7 @@ run_style_clippy() {
run_tests() {
install_uv
setup_pyo3_python
install_cargo_nextest
log_section "Running cargo nextest"
+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
+2 -2
View File
@@ -300,9 +300,9 @@ steps:
- tests/multimodal
- tests/renderers
- tests/standalone_tests/lazy_imports.py
- tests/tokenizers_
- tests/reasoning
- tests/tool_parsers
- tests/tokenizers_
- tests/parser
- tests/transformers_utils
- tests/config
@@ -315,9 +315,9 @@ steps:
- pytest -v -s test_ray_env.py
- pytest -v -s -m 'cpu_test' multimodal
- pytest -v -s renderers
- pytest -v -s tokenizers_
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
- pytest -v -s tool_parsers
- pytest -v -s tokenizers_
- pytest -v -s parser
- pytest -v -s transformers_utils
- pytest -v -s config
+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
+6 -1
View File
@@ -23,9 +23,14 @@
# Any change to the VllmConfig changes can have a large user-facing impact,
# so spam a lot of people
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
/vllm/config/cache.py @heheda12345
# Config utils
/vllm/config/utils.py @hmellor
/vllm/engine/arg_utils.py @hmellor
/vllm/utils/argparse_utils.py
# Entrypoints
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
+1
View File
@@ -4,6 +4,7 @@ include requirements/cuda.txt
include requirements/rocm.txt
include requirements/cpu.txt
include CMakeLists.txt
include tools/build_rust.py
recursive-include cmake *
recursive-include csrc *
@@ -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,
@@ -65,6 +65,32 @@ class RequestArgs(NamedTuple):
limit_min_tokens: int # Use negative value for no limit
limit_max_tokens: int # Use negative value for no limit
timeout_sec: int
send_conversation_id: bool
headers: dict[str, str]
def parse_custom_header(header: str) -> tuple[str, str]:
separators = (":", "=")
for separator in separators:
if separator in header:
key, value = header.split(separator, 1)
key = key.strip()
value = value.strip()
if key:
return key, value
break
raise argparse.ArgumentTypeError(
"Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
)
def build_request_headers(
api_key: str | None, custom_headers: list[tuple[str, str]] | None
) -> dict[str, str]:
headers = dict(custom_headers or [])
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
class BenchmarkArgs(NamedTuple):
@@ -218,12 +244,11 @@ async def send_request(
max_tokens: int | None = None,
timeout_sec: int = 120,
conversation_id: str | None = None,
headers: dict[str, str] | None = None,
) -> ServerResponse:
payload = {
"model": model,
"messages": messages,
"seed": 0,
"temperature": 0.0,
}
if conversation_id is not None:
@@ -233,13 +258,17 @@ async def send_request(
payload["stream"] = True
payload["stream_options"] = {"include_usage": False}
if min_tokens is not None:
payload["min_tokens"] = min_tokens
# if min_tokens is not None:
# payload["min_tokens"] = min_tokens
if max_tokens is not None:
payload["max_tokens"] = max_tokens
headers = {"Content-Type": "application/json"}
request_headers = {"Content-Type": "application/json"}
if conversation_id is not None:
request_headers["X-Session-ID"] = str(conversation_id)
if headers is not None:
request_headers.update(headers)
# Calculate the timeout for the request
if max_tokens is not None:
@@ -265,7 +294,7 @@ async def send_request(
most_recent_timestamp: int = start_time
async with session.post(
url=chat_url, json=payload, headers=headers, timeout=timeout
url=chat_url, json=payload, headers=request_headers, timeout=timeout
) as response:
http_status = HTTPStatus(response.status)
if http_status == HTTPStatus.OK:
@@ -317,6 +346,8 @@ async def send_request(
latency = time.perf_counter_ns() - start_time
if ttft is None:
if stream:
valid_response = False
# The response was a single chunk
ttft = latency
@@ -423,7 +454,8 @@ async def send_turn(
min_tokens,
max_tokens,
req_args.timeout_sec,
conversation_id=conv_id,
conversation_id=conv_id if req_args.send_conversation_id else None,
headers=req_args.headers,
)
if response.valid is False:
@@ -872,6 +904,7 @@ def get_client_config(
# Arguments for API requests
chat_url = f"{args.url}/v1/chat/completions"
model_name = args.served_model_name if args.served_model_name else args.model
headers = build_request_headers(args.api_key, args.header)
req_args = RequestArgs(
chat_url=chat_url,
@@ -880,6 +913,8 @@ def get_client_config(
limit_min_tokens=args.limit_min_tokens,
limit_max_tokens=args.limit_max_tokens,
timeout_sec=args.request_timeout_sec,
send_conversation_id=args.send_conversation_id,
headers=headers,
)
return client_args, req_args
@@ -1245,19 +1280,19 @@ def process_statistics(
)
async def get_server_info(url: str) -> None:
async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
async with aiohttp.ClientSession() as session:
# Get server version (not mandatory, "version" endpoint may not exist)
url_version = f"{url}/version"
async with session.get(url_version) as response:
async with session.get(url_version, headers=headers) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
# Get available models
url_models = f"{url}/v1/models"
async with session.get(url_models) as response:
async with session.get(url_models, headers=headers) as response:
if HTTPStatus(response.status) == HTTPStatus.OK:
text = await response.text()
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1323,6 +1358,22 @@ async def main() -> None:
help="Base URL for the LLM API server",
)
parser.add_argument(
"--api-key",
type=str,
default=None,
help="API key to send as an Authorization bearer token",
)
parser.add_argument(
"--header",
action="append",
type=parse_custom_header,
default=None,
metavar="KEY=VALUE",
help="Custom request header. Can be specified multiple times. "
"Accepts 'Header-Name: value' or 'Header-Name=value'.",
)
parser.add_argument(
"-p",
"--num-clients",
@@ -1437,6 +1488,22 @@ async def main() -> None:
help="Disable stream/streaming mode (set 'stream' to False in the API request)",
)
parser.add_argument(
"--send-conversation-id",
default=False,
action="store_true",
help=(
"Inject a `conversation_id` field into each Chat Completions "
"payload. This is a non-standard OpenAI extension consumed by "
"vLLM's disaggregated multi-turn proxy "
"(examples/disaggregated/disaggregated_serving/"
"disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
"Leave disabled (default) when targeting strict "
"OpenAI-compatible endpoints; enable when benchmarking the "
"disaggregated proxy."
),
)
parser.add_argument(
"-e",
"--excel-output",
@@ -1525,7 +1592,8 @@ async def main() -> None:
args.model, trust_remote_code=args.trust_remote_code
)
await get_server_info(args.url)
headers = build_request_headers(args.api_key, args.header)
await get_server_info(args.url, headers=headers)
# Load the input file (either conversations of configuration file)
logger.info(f"Reading input file: {args.input_file}")
+4 -15
View File
@@ -1,5 +1,5 @@
#!/bin/bash
# Build the vllm-rs Rust frontend binary and install it into the vllm package.
# Build vLLM Rust artifacts and install them into the vllm package.
# Usage: ./build_rust.sh [--debug]
#
# By default builds in release mode. Pass --debug for faster compile times
@@ -8,8 +8,6 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
RUST_DIR="$REPO_ROOT/rust"
TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -27,18 +25,9 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
fi
if [[ "${1:-}" == "--debug" ]]; then
PROFILE_ARGS=()
PROFILE_DIR="debug"
PROFILE_ARG="--debug"
else
PROFILE_ARGS=(--release)
PROFILE_DIR="release"
PROFILE_ARG="--release"
fi
cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
--manifest-path "$RUST_DIR/Cargo.toml" \
--bin vllm-rs \
--features native-tls-vendored
mkdir -p "$(dirname "$TARGET_PATH")"
cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
echo "Installed vllm-rs to $TARGET_PATH"
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
+25 -1
View File
@@ -9,6 +9,7 @@
static const char* PYARGS_PARSE = "KKKK";
#else
#include <cstdlib>
#include <cstdint>
#include <cerrno>
#include <climits>
@@ -46,6 +47,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
@@ -325,7 +349,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;
}
+22 -22
View File
@@ -255,51 +255,49 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
#################### RUST BUILD IMAGE ####################
# Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
# build stage doesn't need the rust toolchain, protoc, or the rust source.
# This stage runs in parallel with csrc-build/extensions-build.
FROM ${BUILD_BASE_IMAGE} AS rust-build
# This stage reuses the Python environment from base and runs in parallel with
# csrc-build/extensions-build.
FROM base AS rust-build
ARG BUILD_OS
ENV DEBIAN_FRONTEND=noninteractive
# Install a basic C toolchain (some rust crates compile C in their build.rs
# scripts) and unzip (used to extract the pinned protoc release below).
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
dnf install -y --setopt=install_weak_deps=False \
ca-certificates curl git gcc gcc-c++ make unzip \
make unzip \
&& dnf clean all && rm -rf /var/cache/dnf; \
else \
apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
make unzip \
&& rm -rf /var/lib/apt/lists/*; \
fi
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
# Copy only the rust workspace — the binary is the sole artifact we need.
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
# 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.
# Build the release artifacts. Cache cargo registry/git, but not target/,
# because stale target metadata can outlive source updates across BuildKit
# cache reuse.
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
#################### RUST BUILD IMAGE ####################
#################### CSRC BUILD IMAGE ####################
@@ -342,6 +340,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
WORKDIR /workspace
COPY pyproject.toml setup.py CMakeLists.txt ./
COPY tools/build_rust.py tools/build_rust.py
COPY cmake cmake/
COPY csrc csrc/
COPY vllm/envs.py vllm/envs.py
@@ -506,9 +505,10 @@ WORKDIR /workspace
COPY --from=csrc-build /workspace/dist /precompiled-wheels
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
+17 -17
View File
@@ -93,35 +93,34 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
# Copy only the rust workspace — the binary is the sole artifact we need.
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
# 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 \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
# Build the release artifacts. Cache cargo registry/git, but not target/,
# because stale target metadata can outlive source updates across BuildKit
# cache reuse.
RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
--mount=type=cache,target=/root/.cargo/git,sharing=locked \
bash build_rust.sh
######################### BUILD IMAGE #########################
FROM base AS vllm-build
@@ -154,9 +153,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
+12 -11
View File
@@ -102,21 +102,22 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
@@ -125,8 +126,7 @@ ENV CARGO_BUILD_JOBS=4
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
#################### RUST BUILD IMAGE ####################
#################### WHEEL BUILD IMAGE ####################
@@ -139,9 +139,10 @@ ENV UV_HTTP_TIMEOUT=500
COPY . .
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
RUN python3 use_existing_torch.py
+17 -16
View File
@@ -138,27 +138,25 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
ENV CARGO_BUILD_JOBS=4
ENV CARGO_NET_RETRY=10
ENV RUSTUP_MAX_RETRIES=10
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
cd ${COMMON_WORKDIR}/vllm \
&& uv pip install --system -r requirements/build/rust.txt
# Build the release binary. Cargo's registry/git caches can be written by
# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
# keeping the cache benefit. Copy the binary out so it persists into the image
# layer for later COPY --from=rust-build.
# keeping the cache benefit. Do not cache target/, because stale target metadata
# can outlive source updates across BuildKit cache reuse.
RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
--mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
--mount=type=cache,id=vllm-rocm-cargo-target,target=${COMMON_WORKDIR}/vllm/rust/target,sharing=locked \
cd ${COMMON_WORKDIR}/vllm \
&& VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
&& test -x /tmp/vllm-rs
&& bash build_rust.sh \
&& test -x vllm/vllm-rs
# -----------------------
# vLLM native build stages
@@ -178,6 +176,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
# not invalidate the expensive native build layer.
COPY setup.py CMakeLists.txt ./
COPY tools/build_rust.py tools/build_rust.py
COPY cmake cmake/
COPY csrc csrc/
COPY vllm/envs.py vllm/envs.py
@@ -209,9 +208,10 @@ ENV VLLM_TARGET_DEVICE=rocm
COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
cd vllm \
@@ -418,9 +418,10 @@ FROM fetch_vllm AS build_vllm_wheel_release
ARG COMMON_WORKDIR
# Drop the pre-built rust frontend binary into the source tree. setup.py
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /tmp/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
# Create /install directory for custom wheels
RUN mkdir -p /install
+104 -90
View File
@@ -6,21 +6,22 @@ ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential unzip \
ca-certificates curl git build-essential unzip python3 python3-pip \
&& rm -rf /var/lib/apt/lists/*
COPY tools/install_protoc.sh /tmp/install_protoc.sh
RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
sh -s -- -y --profile minimal --default-toolchain none
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /workspace
COPY requirements/build/rust.txt requirements/build/rust.txt
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
# by the wheel build stage.
COPY rust rust
COPY rust-toolchain.toml rust-toolchain.toml
COPY tools/build_rust.py tools/build_rust.py
COPY build_rust.sh build_rust.sh
# Cap cargo parallelism to avoid exhausting the CI host's open-file limit
@@ -29,8 +30,7 @@ ENV CARGO_BUILD_JOBS=4
RUN --mount=type=cache,target=/root/.cargo/registry \
--mount=type=cache,target=/root/.cargo/git \
--mount=type=cache,target=/workspace/rust/target \
VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
bash build_rust.sh
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
@@ -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,24 +119,104 @@ 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
# detects it and ships it as-is, skipping the local cargo build.
COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs
# Drop the pre-built Rust artifacts into the source tree. setup.py detects
# them and ships them as-is, skipping the local Rust build.
COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
@@ -147,75 +227,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"]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 388 KiB

After

Width:  |  Height:  |  Size: 373 KiB

+41
View File
@@ -405,6 +405,47 @@ vllm bench serve \
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
#### BFCL (Tool-Calling) Benchmark
The Berkeley Function Calling Leaderboard (BFCL) dataset measures serving
latency and throughput on realistic tool-calling traffic. Each request
carries a per-sample `tools` schema and chat history, so the server must
expose `/v1/chat/completions` with an auto-tool-choice parser enabled.
The benchmark client always uses the `openai-chat` backend.
Start a tool-parser-enabled server, then run the bench. For example, with
`gpt-oss-20b`:
```bash
# Server
vllm serve openai/gpt-oss-20b \
--enable-auto-tool-choice \
--tool-call-parser openai \
--reasoning-parser openai_gptoss
# Client
vllm bench serve \
--backend openai-chat \
--endpoint /v1/chat/completions \
--model openai/gpt-oss-20b \
--dataset-name hf \
--dataset-path gorilla-llm/Berkeley-Function-Calling-Leaderboard \
--bfcl-categories simple,live_simple,multiple \
--num-prompts 200
```
`--bfcl-categories` is a comma-separated list of BFCL v3 category names
(without the `BFCL_v3_` prefix or `.json` suffix). Defaults to
`simple,live_simple,multiple`. Other supported non-multi-turn categories
include `parallel`, `live_parallel`, `parallel_multiple`,
`live_parallel_multiple`, `irrelevance`, `live_irrelevance`,
`live_relevance`, `java`, `javascript`, and `rest`. Multi-turn categories
are not yet supported.
The dataset class normalizes BFCL's loose schema dialect (`dict` →
`object`, `float` → `number`, `tuple` → `array`, `any` → `string`) so
modern grammar backends accept the translated tool definitions.
#### Other HuggingFaceDataset Examples
```bash
+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)
+2 -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:
@@ -44,16 +41,14 @@ Now supports 9 types of connectors:
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
```
For multi-tier offloading (e.g., CPU + filesystem tier) and the full configuration reference, see the [KV Offloading Usage Guide](kv_offloading_usage.md).
- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
```bash
--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.
+132
View File
@@ -0,0 +1,132 @@
# KV Offloading Usage Guide
This guide covers configuration of the [`OffloadingConnector`](disagg_prefill.md), which extends the prefix cache by offloading completed KV blocks to slower but larger tiers (CPU host memory, plus optional secondary tiers) as they are produced. Hits in the offload tiers are promoted back to GPU on demand. Transfers between GPU and CPU use DMA (`cudaMemcpyAsync`) and run asynchronously alongside model computation, so offloading adds minimal CPU- and GPU-core overhead.
!!! note
The `OffloadingConnector` currently supports CUDA, ROCm, and XPU only.
## Overview
Two specs are available, selected by the `spec_name` key in `kv_connector_extra_config`:
- `CPUOffloadingSpec` (default): single CPU tier. Completed GPU blocks are copied into pinned host memory.
- `TieringOffloadingSpec`: multi-tier. A CPU primary tier plus one or more secondary tiers.
Only the CPU primary tier has direct GPU access. Secondary tiers cannot read from or write to GPU memory; all GPU↔secondary transfers are staged through the CPU primary tier.
```mermaid
flowchart LR
GPU <--> CPU["CPU primary tier"]
CPU <--> S0["Secondary tier 0"]
CPU <--> S1["Secondary tier 1"]
CPU <--> SN["..."]
```
## Single-Tier Setup (CPU Only)
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"block_size": 64,
"cpu_bytes_to_use": 1000000000
}
}'
```
## Multi-Tier Setup
Set `spec_name` to `"TieringOffloadingSpec"` and supply a `secondary_tiers` list. Each entry is a dict with a required `type` key plus tier-specific fields. The list is ordered: tier 0 is consulted before tier 1, and so on. See [Secondary Tiers](#secondary-tiers) for tier-specific keys.
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "TieringOffloadingSpec",
"cpu_bytes_to_use": 10737418240,
"block_size": 16,
"eviction_policy": "lru",
"secondary_tiers": [
{
"type": "fs",
"root_dir": "/mnt/kv_cache",
"n_read_threads": 32,
"n_write_threads": 16
}
]
}
}'
```
## `kv_connector_extra_config` Reference
| Key | Required | Default | Scope | Notes |
| --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
Each entry in `secondary_tiers` is a dict with a required `type` field plus tier-specific fields.
### Filesystem (FS)
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
| Key | Required | Default | Notes |
| --- | --- | --- | --- |
| `type` | yes | — | Must be `fs`. |
| `root_dir` | yes | — | Base directory; vLLM creates subdirectories beneath it (see [On-Disk Layout](#on-disk-layout)). |
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
#### On-Disk Layout
Under `root_dir`, vLLM creates a subdirectory `<model>_<digest>`, where `<model>` is the model name with `/` replaced by `_` (so HuggingFace IDs like `meta-llama/Llama-3-8B` don't nest), and `<digest>` is a short SHA256 prefix derived from the run configuration (model, block size, parallelism, dtype, etc.). Runs with the same configuration share the same subdirectory; runs with different configurations live side-by-side under the same `root_dir` without colliding.
Inside that subdirectory, blocks are sharded across hash-prefix subdirectories to limit directory fan-out:
```text
<root_dir>/
<model>_<digest>/
config.json
<model>_<digest>_r<rank>/
<hhh>/ # first 3 hex chars of the block hash
<hh>_g<group_idx>/ # next 2 hex chars + KV cache group index
<hash_hex>.bin # full block hash (in hex)
```
`config.json` records the run (block size, number of KV groups, etc.) and is written on first start. Each rank writes blocks under its own `_r<rank>` sibling directory, so multiple ranks can safely share the same `root_dir`.
#### Cross-Process Sharing
To enable KV cache sharing between multiple vLLM instances using the same `root_dir` (e.g., via a shared PVC), the `PYTHONHASHSEED` environment variable must be set to the same fixed value (e.g., `"0"`) on every instance. Without this, each process initializes `NONE_HASH` (the chain-hash seed for block content hashes) with random bytes, producing different block filenames for identical token content.
```bash
PYTHONHASHSEED=0 vllm serve ...
```
## Tuning Tips
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
## Further Reading
- [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).
+15
View File
@@ -294,6 +294,21 @@ curl http://localhost:8000/v1/chat/completions \
!!! note
The `conversation_id` field is a non-standard extension to the OpenAI API. It is consumed by the proxy and not forwarded to the vLLM engine.
### Benchmarking the multi-turn proxy
[`benchmarks/multi_turn/benchmark_serving_multi_turn.py`](../../benchmarks/multi_turn/benchmark_serving_multi_turn.py) supports targeting the disaggregated multi-turn proxy with the `--send-conversation-id` flag, which injects a per-conversation `conversation_id` into every request payload so the proxy can key cross-turn KV cache reuse.
The flag is **off by default** so the benchmark is compatible with strict OpenAI-compatible frontends that reject unknown top-level fields. When benchmarking the multi-turn proxy you must pass it explicitly — otherwise every turn lands as a cache MISS and the bidirectional KV transfer path is never exercised.
```bash
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \
--model <MODEL> --served-model-name <NAME> \
--url http://<proxy_host>:8000 \
--input-file benchmarks/multi_turn/generate_multi_turn.json \
--num-clients 2 --max-active-conversations 6 \
--send-conversation-id
```
### Limitations
- Requires a stateful proxy (or equivalent router) to track and forward `kv_transfer_params` between turns.
+7
View File
@@ -504,6 +504,13 @@ Flags: `--tool-call-parser pythonic --chat-template {see_above}`
!!! warning
Llama's smaller models frequently fail to emit tool calls in the correct format. Results may vary depending on the model.
## Benchmarking Tool-Calling Performance
To measure serving latency and throughput on realistic tool-calling traffic,
use the BFCL (Berkeley Function Calling Leaderboard) dataset with
`vllm bench serve`. See the [BFCL benchmark example](../benchmarking/cli.md#bfcl-tool-calling-benchmark)
for the full server + client commands.
## How to Write a Tool Parser Plugin
A tool parser plugin is a Python file containing one or more ToolParser implementations. You can write a ToolParser similar to the `Hermes2ProToolParser` in [vllm/tool_parsers/hermes_tool_parser.py](../../vllm/tool_parsers/hermes_tool_parser.py).
@@ -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 ""
@@ -35,12 +35,36 @@ Conversation isolation:
the JSON body) to scope the KV cache across turns. Without it, the
proxy cannot link turns and falls back to no-cache behavior.
``conversation_id`` is a non-standard extension to the OpenAI Chat
Completions schema, consumed by this proxy and not forwarded to the
vLLM engine. Strict OpenAI-compatible frontends reject unknown
fields, so clients must opt in only when targeting this proxy.
Usage:
python disagg_proxy_multiturn.py \\
--host 0.0.0.0 --port 8000 \\
--prefiller-host 10.0.0.1 --prefiller-port 8100 \\
--decoder-host 10.0.0.2 --decoder-port 8200
Benchmarking:
Use ``benchmarks/multi_turn/benchmark_serving_multi_turn.py`` with
the ``--send-conversation-id`` flag to inject a per-conversation
``conversation_id`` into every request so this proxy can key
cross-turn KV cache reuse. The flag is *off by default*: without
it the benchmark sends OpenAI-schema-compliant payloads and every
turn lands as a cache MISS in this proxy.
Example:
python benchmarks/multi_turn/benchmark_serving_multi_turn.py \\
--model <MODEL> --served-model-name <NAME> \\
--url http://<proxy_host>:8000 \\
--input-file generate_multi_turn.json \\
--num-clients 2 --max-active-conversations 6 \\
--send-conversation-id
See ``docs/features/nixl_connector_usage.md`` for the broader
bidirectional-KV-transfer setup these benchmarks exercise.
Dependencies:
pip install fastapi uvicorn httpx
"""
@@ -373,7 +397,9 @@ async def _handle_request(api_path: str, request: Request):
logger.warning(
"[%s] No conversation_id provided — KV cache reuse disabled "
"for this request. Add a 'conversation_id' field to enable "
"cross-turn KV sharing.",
"cross-turn KV sharing. When using "
"benchmarks/multi_turn/benchmark_serving_multi_turn.py, pass "
"--send-conversation-id (off by default).",
request_id,
)
@@ -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",
]
+4
View File
@@ -0,0 +1,4 @@
# Dependencies for building Rust artifacts through setuptools-rust.
setuptools>=77.0.3,<81.0.0
setuptools-rust>=1.9.0
wheel
+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
+89
View File
@@ -3458,6 +3458,75 @@ version = "0.1.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
[[package]]
name = "pyo3"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
dependencies = [
"libc",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
]
[[package]]
name = "pyo3-build-config"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
dependencies = [
"libc",
"pyo3-build-config",
]
[[package]]
name = "pyo3-macros"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pyo3-macros-backend"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
dependencies = [
"heck",
"proc-macro2",
"pyo3-build-config",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pythonize"
version = "0.28.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
dependencies = [
"pyo3",
"serde",
"serde_json",
]
[[package]]
name = "qoi"
version = "0.4.1"
@@ -4669,6 +4738,12 @@ dependencies = [
"libc",
]
[[package]]
name = "target-lexicon"
version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "task-local"
version = "0.1.1"
@@ -5799,6 +5874,7 @@ dependencies = [
"axum",
"bytes",
"clap",
"educe",
"expect-test",
"futures",
"http-body",
@@ -5813,7 +5889,9 @@ dependencies = [
"serde_json",
"serde_with",
"serial_test",
"sha2",
"socket2",
"subtle",
"thiserror-ext",
"tokio",
"tokio-stream",
@@ -5901,6 +5979,17 @@ dependencies = [
"winnow",
]
[[package]]
name = "vllm-tool-parser-py"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"serde_json",
"thiserror-ext",
"vllm-tool-parser",
]
[[package]]
name = "walkdir"
version = "2.5.0"
+5
View File
@@ -12,6 +12,7 @@ members = [
"src/text",
"src/tokenizer",
"src/tool-parser",
"src/tool-parser/python",
]
resolver = "3"
@@ -60,6 +61,8 @@ prometheus-client = "0.24.0"
prometheus-client-derive-encode = "0.5.0"
prost = "0.14.3"
prost-types = "0.14.3"
pyo3 = "0.28.3"
pythonize = "0.28.0"
rand = "0.9.2"
reasoning-parser = "1.2.2"
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
@@ -75,8 +78,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"
+1 -1
View File
@@ -71,7 +71,7 @@ To build the `vllm-rs` in isolation:
```bash
# from the local checkout
cargo install --path src/cmd --bin vllm-rs
./build_rust.sh
```
### Example Request
+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>;
+29 -2
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?;
@@ -245,6 +272,6 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, step3)"].assert_eq(&error.to_report_string());
expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
}
}
+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])
}
+14 -2
View File
@@ -6,7 +6,8 @@ pub use vllm_reasoning_parser::{
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
Step3p5ReasoningParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -25,7 +26,9 @@ pub mod names {
pub const MINIMAX_M2: &str = "minimax_m2";
pub const NEMOTRON_V3: &str = "nemotron_v3";
pub const QWEN3: &str = "qwen3";
pub const SEED_OSS: &str = "seed_oss";
pub const STEP3: &str = "step3";
pub const STEP3P5: &str = "step3p5";
}
/// Constructor signature for one registered reasoning parser implementation.
@@ -61,7 +64,9 @@ impl ReasoningParserFactory {
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
.register_parser::<Step3ReasoningParser>(names::STEP3);
.register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
.register_parser::<Step3ReasoningParser>(names::STEP3)
.register_parser::<Step3p5ReasoningParser>(names::STEP3P5);
factory
.register_pattern("deepseek-r1", names::DEEPSEEK_R1)
@@ -77,7 +82,14 @@ impl ReasoningParserFactory {
.register_pattern("glm-4.5", names::GLM45)
.register_pattern("kimi-k2", names::KIMI_K2)
.register_pattern("kimi", names::KIMI)
// step3p5 patterns must precede `step3`: substring matching would
// otherwise route step3p5 IDs to step3.
.register_pattern("step-3p5", names::STEP3P5)
.register_pattern("step3p5", names::STEP3P5)
.register_pattern("step-3.5", names::STEP3P5)
.register_pattern("step3", names::STEP3)
.register_pattern("seed-oss", names::SEED_OSS)
.register_pattern("seedoss", names::SEED_OSS)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2)
.register_pattern("cohere", names::COHERE_CMD)
@@ -32,8 +32,12 @@ fn factory_contains_and_lists_registered_parsers() {
let factory = ReasoningParserFactory::new();
assert!(factory.contains(names::QWEN3));
assert!(factory.contains(names::DEEPSEEK_V4));
assert!(factory.contains(names::SEED_OSS));
assert!(factory.contains(names::STEP3P5));
assert!(factory.list().contains(&names::QWEN3.to_string()));
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
assert!(factory.list().contains(&names::SEED_OSS.to_string()));
assert!(factory.list().contains(&names::STEP3P5.to_string()));
}
#[test]
@@ -49,6 +53,41 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
);
}
#[test]
fn factory_routes_step3p5_models_to_dedicated_parser() {
let factory = ReasoningParserFactory::new();
// step3p5 patterns must beat the bare `step3` substring.
assert_eq!(
factory.resolve_name_for_model("step-3p5-instruct"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step3p5"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step-3.5-base"),
Some(names::STEP3P5)
);
assert_eq!(
factory.resolve_name_for_model("step3-base"),
Some(names::STEP3)
);
}
#[test]
fn factory_routes_seed_oss_models() {
let factory = ReasoningParserFactory::new();
assert_eq!(
factory.resolve_name_for_model("ByteDance-Seed/Seed-OSS-36B-Instruct"),
Some(names::SEED_OSS)
);
assert_eq!(
factory.resolve_name_for_model("seedoss-7b"),
Some(names::SEED_OSS)
);
}
#[test]
fn factory_rejects_unknown_parser_names() {
let tokenizer = Arc::new(FakeTokenizer);
+27 -2
View File
@@ -144,6 +144,30 @@ impl RoundtripCase {
json_fmt: spaced_json_fmt(),
}
}
/// SeedOSS with `<seed:think>` / `</seed:think>` reasoning tags.
fn seed_oss() -> Self {
Self {
model_id: "ByteDance-Seed/Seed-OSS-36B-Instruct",
assistant_stop_suffix: "<seed:eos>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
/// Step-3.5 with `<think>` / `</think>` reasoning tags and newline trimming.
fn step3p5() -> Self {
Self {
model_id: "stepfun-ai/Step-3.5-Flash",
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
}
macro_rules! roundtrip_tests {
@@ -168,10 +192,11 @@ roundtrip_tests! {
minimax_m25 => [reasoning_and_content, tool_call_mix],
deepseek_v4 => [reasoning_and_content, tool_call_mix],
glm47 => [reasoning_and_content, tool_call_mix],
seed_oss => [reasoning_and_content],
step3p5 => [reasoning_and_content],
// 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 {
@@ -68,6 +68,11 @@ impl DelimitedReasoningParser {
.unwrap_or(self.default_in_reasoning);
}
/// Return whether the parser is currently inside a reasoning section.
pub(crate) fn in_reasoning(&self) -> bool {
self.current_in_reasoning
}
/// Parse one decoded text delta and return its reasoning/content split.
pub(crate) fn push(&mut self, delta: &str) -> ReasoningDelta {
self.buffer.push_str(delta);
+4
View File
@@ -20,6 +20,8 @@ mod delimited;
mod gemma4;
mod kimi;
mod qwen3;
mod seed_oss;
mod step3p5;
use thiserror::Error;
use vllm_tokenizer::DynTokenizer;
@@ -30,6 +32,8 @@ pub(crate) use self::delimited::DelimitedReasoningParser;
pub use self::gemma4::Gemma4ReasoningParser;
pub use self::kimi::KimiReasoningParser;
pub use self::qwen3::Qwen3ReasoningParser;
pub use self::seed_oss::SeedOssReasoningParser;
pub use self::step3p5::Step3p5ReasoningParser;
/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
pub type DeepSeekV3ReasoningParser = Qwen3ReasoningParser;
+147
View File
@@ -0,0 +1,147 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
/// Reasoning parser for SeedOSS models using `<seed:think>`/`</seed:think>`
/// delimiters.
pub struct SeedOssReasoningParser {
inner: DelimitedReasoningParser,
}
impl SeedOssReasoningParser {
/// Create a SeedOSS parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(
tokenizer,
"<seed:think>",
"</seed:think>",
false,
)?,
})
}
}
impl ReasoningParser for SeedOssReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
Ok(self.inner.push(delta))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
Ok(self.inner.finish())
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::SeedOssReasoningParser;
use crate::{ReasoningParser, tests::FakeTokenizer};
#[test]
fn without_prompt_markers_expects_start_token() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("implicit reasoning</seed:think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(
delta.content.as_deref(),
Some("implicit reasoning</seed:think>answer")
);
}
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<seed:think>` (id 10), opening reasoning before the stream.
parser.initialize(&[10]).unwrap();
let delta = parser.push("reason</seed:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn respects_prompt_end_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
// Prompt already closed reasoning with `</seed:think>` (id 11).
parser.initialize(&[11]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn handles_explicit_start_token() {
// An explicit start delimiter must not leak into reasoning text.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<seed:think>reason</seed:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn streams_explicit_start_token_across_pushes() {
// Start token, reasoning body, end token, and content arrive in separate
// streaming deltas.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
let mut reasoning = String::new();
let mut content = String::new();
for delta_str in [
"<seed:think>",
"Some ",
"reasoning ",
"content",
"</seed:think>",
"Final ",
"answer",
] {
let delta = parser.push(delta_str).unwrap();
if let Some(r) = delta.reasoning {
reasoning.push_str(&r);
}
if let Some(c) = delta.content {
content.push_str(&c);
}
}
assert_eq!(reasoning, "Some reasoning content");
assert_eq!(content, "Final answer");
}
#[test]
fn handles_partial_delimiters_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = SeedOssReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[10]).unwrap();
// Closing delimiter `</seed:think>` arrives in two halves.
let first = parser.push("reason</seed:").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
let second = parser.push("think>answer").unwrap();
assert_eq!(second.reasoning, None);
assert_eq!(second.content.as_deref(), Some("answer"));
}
}
+308
View File
@@ -0,0 +1,308 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
/// Reasoning parser for Step3p5 outputs.
///
/// Step3p5 uses standard `<think>`/`</think>` delimiters but emits a `\n`
/// immediately before and/or after `</think>`. The parser drops these framing
/// newlines on both sides of the boundary, holding a trailing `\n` from
/// reasoning across pushes until either more reasoning text or `</think>`
/// arrives, and dropping a leading `\n` from the first content delta after
/// the boundary.
pub struct Step3p5ReasoningParser {
inner: DelimitedReasoningParser,
/// `\n` at end of last reasoning delta, held in case `</think>` follows.
pending_reasoning_newline: bool,
/// Last push ended on `</think>` without emitting content; the next
/// content delta's leading `\n` should be dropped.
just_ended_reasoning: bool,
}
impl Step3p5ReasoningParser {
/// Create a Step3p5 parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(tokenizer, "<think>", "</think>", false)?,
pending_reasoning_newline: false,
just_ended_reasoning: false,
})
}
/// Drop framing newlines around `</think>` and track held-newline state.
fn process(
&mut self,
mut inner_delta: ReasoningDelta,
was_in_reasoning: bool,
now_in_reasoning: bool,
) -> ReasoningDelta {
// A `<think>...</think>` round-trip in one push still counts as a
// transition: the inner emits reasoning while ending in content mode.
let transitioned =
!now_in_reasoning && (was_in_reasoning || inner_delta.reasoning.is_some());
// Replay or drop a previously-held trailing reasoning newline.
if self.pending_reasoning_newline {
if let Some(reasoning) = inner_delta.reasoning.as_mut() {
reasoning.insert(0, '\n');
self.pending_reasoning_newline = false;
} else if transitioned {
// The held `\n` was the one right before `</think>`: drop it.
self.pending_reasoning_newline = false;
}
}
// Hold back a trailing reasoning `\n` until we know if `</think>` follows.
if let Some(reasoning) = inner_delta.reasoning.as_mut()
&& reasoning.ends_with('\n')
{
reasoning.pop();
if !transitioned {
self.pending_reasoning_newline = true;
}
}
// Drop a leading `\n` of content emitted right after `</think>`.
if let Some(content) = inner_delta.content.as_mut()
&& (transitioned || self.just_ended_reasoning)
&& content.starts_with('\n')
{
content.remove(0);
}
self.just_ended_reasoning = transitioned && inner_delta.content.is_none();
if inner_delta.reasoning.as_deref() == Some("") {
inner_delta.reasoning = None;
}
if inner_delta.content.as_deref() == Some("") {
inner_delta.content = None;
}
inner_delta
}
}
impl ReasoningParser for Step3p5ReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
let was = self.inner.in_reasoning();
let inner_delta = self.inner.push(delta);
let now = self.inner.in_reasoning();
Ok(self.process(inner_delta, was, now))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
let was = self.inner.in_reasoning();
let inner_delta = self.inner.finish();
let now = self.inner.in_reasoning();
let mut delta = self.process(inner_delta, was, now);
// Emit a still-held newline rather than silently dropping it.
if self.pending_reasoning_newline {
match delta.reasoning.as_mut() {
Some(existing) => existing.push('\n'),
None => delta.reasoning = Some("\n".to_string()),
}
self.pending_reasoning_newline = false;
}
Ok(delta)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::Step3p5ReasoningParser;
use crate::{ReasoningParser, tests::FakeTokenizer};
#[test]
fn picks_up_prompt_start_boundary() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// Prompt prefills `<think>` (id 1), opening reasoning before the stream.
parser.initialize(&[1]).unwrap();
let delta = parser.push("This is a reasoning section</think>This is the rest").unwrap();
assert_eq!(
delta.reasoning.as_deref(),
Some("This is a reasoning section")
);
assert_eq!(delta.content.as_deref(), Some("This is the rest"));
}
#[test]
fn handles_unterminated_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("<think>reason without end").unwrap();
assert_eq!(pushed.reasoning.as_deref(), Some("reason without end"));
assert_eq!(pushed.content, None);
let flushed = parser.finish().unwrap();
assert!(flushed.is_empty());
}
#[test]
fn handles_empty_input() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let pushed = parser.push("").unwrap();
assert!(pushed.is_empty());
let flushed = parser.finish().unwrap();
assert!(flushed.is_empty());
}
#[test]
fn complex_newline_pattern_trims_only_single_framing_newline_each_side() {
// Only the immediately-adjacent framing `\n` is dropped on each side of
// `</think>`; surrounding newlines remain part of reasoning/content.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[1]).unwrap();
let delta = parser
.push("\n This is a \n reasoning section\n\n\n</think>\n\nThis is the rest")
.unwrap();
assert_eq!(
delta.reasoning.as_deref(),
Some("\n This is a \n reasoning section\n\n")
);
assert_eq!(delta.content.as_deref(), Some("\nThis is the rest"));
}
#[test]
fn drops_framing_newlines_in_single_push() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n</think>\nanswer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn drops_framing_newlines_across_pushes() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
// The trailing `\n` from the first push is held until we know whether
// `</think>` follows.
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
// `</think>` arrives standalone; the held newline should be dropped.
let second = parser.push("</think>").unwrap();
assert!(second.is_empty());
// The leading newline of the first content delta is dropped.
let third = parser.push("\nanswer").unwrap();
assert_eq!(third.reasoning, None);
assert_eq!(third.content.as_deref(), Some("answer"));
}
#[test]
fn replays_held_newline_when_more_reasoning_follows() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
let second = parser.push("more reason").unwrap();
assert_eq!(second.reasoning.as_deref(), Some("\nmore reason"));
assert_eq!(second.content, None);
}
#[test]
fn finish_flushes_held_newline_in_unterminated_stream() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason\n").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
let flushed = parser.finish().unwrap();
assert_eq!(flushed.reasoning.as_deref(), Some("\n"));
assert_eq!(flushed.content, None);
}
#[test]
fn preserves_inner_newlines_in_reasoning() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>line1\nline2</think>tail").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("line1\nline2"));
assert_eq!(delta.content.as_deref(), Some("tail"));
}
#[test]
fn trims_only_one_trailing_reasoning_newline() {
// Only the single framing newline immediately before `</think>` is
// dropped; earlier newlines in the reasoning body are preserved.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason\n\n</think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason\n"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn drops_only_first_content_newline_after_transition() {
// The leading-`\n` drop applies only to the first content delta after
// `</think>`; later deltas pass through untouched.
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let first = parser.push("<think>reason</think>").unwrap();
assert_eq!(first.reasoning.as_deref(), Some("reason"));
assert_eq!(first.content, None);
let second = parser.push("\nfirst").unwrap();
assert_eq!(second.reasoning, None);
assert_eq!(second.content.as_deref(), Some("first"));
// A `\n` arriving in a later content delta must NOT be dropped.
let third = parser.push("\nsecond").unwrap();
assert_eq!(third.reasoning, None);
assert_eq!(third.content.as_deref(), Some("\nsecond"));
}
#[test]
fn passes_through_clean_boundary_without_framing_newlines() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think>reason</think>tail").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("tail"));
}
#[test]
fn handles_empty_reasoning_section() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = Step3p5ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<think></think>answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
}
+3 -1
View File
@@ -6,7 +6,7 @@ use super::{
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
};
struct FakeTokenizer;
pub(crate) struct FakeTokenizer;
impl Tokenizer for FakeTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
@@ -32,6 +32,8 @@ impl Tokenizer for FakeTokenizer {
"<|END_THINKING|>" => Some(4),
"◁think▷" => Some(5),
"◁/think▷" => Some(6),
"<seed:think>" => Some(10),
"</seed:think>" => Some(11),
_ => None,
}
}
+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)
}
+4 -2
View File
@@ -14,7 +14,8 @@ mod utils;
use std::sync::{Arc, OnceLock};
use anyhow::{Context as _, Result};
use axum::{Router, serve::ListenerExt as _};
use axum::Router;
use axum::serve::ListenerExt as _;
pub use config::{Config, CoordinatorMode, HttpListenerMode};
use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until};
@@ -92,7 +93,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));
@@ -22,7 +22,7 @@ use vllm_llm::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
};
use self::convert::prepare_generate_request;
use self::convert::{ResponseOptions, prepare_generate_request};
use self::types::{
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
GenerateResponseStreamChoice, GenerateStreamResponse,
@@ -54,10 +54,7 @@ pub async fn generate(
);
let log_request = state.enable_log_requests;
let include_logprobs = prepared.include_logprobs;
let include_prompt_logprobs = prepared.include_prompt_logprobs;
let stream = prepared.stream;
let raw_stream = match state
.chat
.text()
@@ -80,9 +77,7 @@ pub async fn generate(
raw_stream,
prepared.request_id,
log_request,
prepared.include_usage,
prepared.include_continuous_usage,
include_logprobs,
prepared.options,
);
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
@@ -100,21 +95,11 @@ pub async fn generate(
}
};
if log_request {
info!(
parent: &request_span,
prompt_tokens = collected.prompt_token_ids.len(),
output_tokens = collected.token_ids.len(),
finish_reason = collected.finish_reason.as_str(),
"generate finished"
);
}
let response = match collect_generate(
collected,
prepared.request_id,
include_logprobs,
include_prompt_logprobs,
log_request,
prepared.options,
) {
Ok(response) => response,
Err(error) => return error.into_response(),
@@ -128,9 +113,13 @@ async fn generate_chunk_stream(
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
request_id: String,
log_request: bool,
include_usage: bool,
include_continuous_usage: bool,
include_logprobs: bool,
ResponseOptions {
include_usage,
include_continuous_usage,
include_logprobs,
// Ignored: raw generate streaming has no prompt-logprobs wire shape.
include_prompt_logprobs: _,
}: ResponseOptions,
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
@@ -222,8 +211,15 @@ async fn generate_chunk_stream(
fn collect_generate(
collected: CollectedGenerateOutput,
request_id: String,
include_logprobs: bool,
include_prompt_logprobs: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming raw generate responses do not include usage.
include_usage: _,
// Ignored: continuous usage is a streaming-only option.
include_continuous_usage: _,
include_logprobs,
include_prompt_logprobs,
}: ResponseOptions,
) -> Result<GenerateResponse, ApiError> {
let logprobs = if include_logprobs {
let logprobs = collected.logprobs.as_ref().ok_or_else(|| {
@@ -246,13 +242,23 @@ fn collect_generate(
} else {
None
};
let finish_reason = collected.finish_reason.as_str().to_string();
if log_request {
info!(
prompt_tokens = collected.prompt_token_ids.len(),
output_tokens = collected.token_ids.len(),
%finish_reason,
"generate finished"
);
}
Ok(GenerateResponse {
request_id,
choices: vec![GenerateResponseChoice {
index: 0,
logprobs,
finish_reason: Some(collected.finish_reason.as_str().to_string()),
finish_reason: Some(finish_reason),
token_ids: collected.token_ids,
}],
prompt_logprobs,
@@ -408,11 +414,19 @@ mod tests {
}),
]);
let chunks: Vec<_> =
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
.try_collect()
.await
.expect("collect chunks");
let chunks: Vec<_> = generate_chunk_stream(
stream,
"raw-stream".to_string(),
false,
ResponseOptions {
include_usage: true,
include_continuous_usage: true,
..Default::default()
},
)
.try_collect()
.await
.expect("collect chunks");
assert_eq!(chunks.len(), 2);
assert_eq!(
@@ -8,19 +8,29 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
/// Lowered generate request plus the response request ID.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
pub request_id: String,
pub text_request: TextRequest,
pub stream: bool,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
}
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Whether the caller asked for usage on every streamed chunk.
pub include_continuous_usage: bool,
/// Whether the caller requested output logprobs on generate choices.
pub include_logprobs: bool,
/// Whether the caller requested top-level prompt logprobs.
pub include_prompt_logprobs: bool,
}
/// Validate and lower one raw generate request into the internal
/// text-generation format.
pub fn prepare_generate_request(
pub(super) fn prepare_generate_request(
request: GenerateRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -65,10 +75,12 @@ pub fn prepare_generate_request(
request_id: ctx.request_id,
text_request,
stream,
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
options: ResponseOptions {
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
},
})
}
@@ -158,7 +170,7 @@ mod tests {
)
.expect("prepare");
assert!(!prepared.include_usage);
assert!(!prepared.include_continuous_usage);
assert!(!prepared.options.include_usage);
assert!(!prepared.options.include_continuous_usage);
}
}
@@ -1,4 +1,4 @@
pub mod convert;
pub(crate) mod convert;
mod types;
mod validate;
@@ -23,8 +23,8 @@ use vllm_chat::{
};
use vllm_engine_core_client::protocol::StopReason;
use self::convert::{ResponseOptions, prepare_chat_request};
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
use crate::routes::openai::chat_completions::types::{
AssistantRole, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest,
ChatCompletionResponse, ChatCompletionStreamChoice, ChatCompletionStreamResponse,
@@ -83,12 +83,7 @@ pub async fn chat_completions(
prepared.response_model,
created,
log_request,
prepared.include_usage,
prepared.requested_logprobs,
prepared.include_reasoning,
prepared.echo,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
prepared.options,
);
let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
@@ -99,12 +94,8 @@ pub async fn chat_completions(
prepared.request_id,
prepared.response_model,
created,
prepared.requested_logprobs,
prepared.include_prompt_logprobs,
prepared.include_reasoning,
prepared.echo,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
log_request,
prepared.options,
)
.instrument(request_span.clone())
.await
@@ -113,18 +104,6 @@ pub async fn chat_completions(
Err(error) => return error.into_response(),
};
if log_request {
let usage = response.usage.as_ref();
info!(
parent: &request_span,
model = %response.model,
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
"chat completion finished"
);
}
Json(response).into_response()
}
}
@@ -134,12 +113,17 @@ async fn collect_chat_completion(
request_id: String,
response_model: String,
created: u64,
requested_logprobs: bool,
include_prompt_logprobs: bool,
include_reasoning: bool,
echo: Option<String>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
echo,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
) -> Result<ChatCompletionResponse, ApiError> {
let collected = stream.collect_message().await.map_err(|error| {
server_error!(
@@ -201,6 +185,16 @@ async fn collect_chat_completion(
};
let usage = Usage::from_counts(prompt_token_count as u32, output_token_count as u32);
if log_request {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
output_tokens = usage.completion_tokens.unwrap_or(0),
finish_reason = %finish_reason,
"chat completion finished"
);
}
Ok(ChatCompletionResponse {
id: request_id,
object: "chat.completion".to_string(),
@@ -238,12 +232,16 @@ async fn chat_completion_chunk_stream(
response_model: String,
created: u64,
log_request: bool,
include_usage: bool,
requested_logprobs: bool,
include_reasoning: bool,
echo: Option<String>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
ResponseOptions {
include_usage,
requested_logprobs,
// Ignored: chat streaming prompt logprobs are rejected for Python parity.
include_prompt_logprobs: _,
include_reasoning,
echo,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
) -> Result<(), ApiError> {
let mut saw_tool_calls = false;
@@ -806,7 +804,7 @@ mod tests {
use vllm_engine_core_client::protocol::StopReason;
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::{block_delta_chunk, chat_completion_chunk_stream, final_chunk};
use super::{ResponseOptions, block_delta_chunk, chat_completion_chunk_stream, final_chunk};
#[test]
fn text_chunk_uses_content_only_delta() {
@@ -932,12 +930,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
true,
None,
false,
false,
ResponseOptions {
requested_logprobs: true,
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -996,12 +993,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
true,
None,
false,
false,
ResponseOptions {
requested_logprobs: true,
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1049,12 +1045,7 @@ mod tests {
"model".to_string(),
1,
false,
false,
false,
false,
None,
false,
false,
ResponseOptions::default(),
)
.collect::<Vec<_>>()
.await
@@ -1132,12 +1123,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
false,
None,
true,
false,
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1263,12 +1253,11 @@ mod tests {
"model".to_string(),
1,
false,
false,
true,
false,
None,
true,
false,
ResponseOptions {
requested_logprobs: true,
return_token_ids: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -1342,12 +1331,10 @@ mod tests {
"model".to_string(),
1,
false,
false,
false,
true,
None,
false,
false,
ResponseOptions {
include_reasoning: true,
..Default::default()
},
)
.collect::<Vec<_>>()
.await
@@ -18,11 +18,19 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
/// Lowered chat request plus the public response metadata carried by every SSE
/// chunk.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
/// Stable OpenAI-style request ID, reused as the external chat request ID.
pub request_id: String,
/// Public model ID echoed back to the client.
pub response_model: String,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
/// Lowered chat request for `vllm-chat`.
pub chat_request: ChatRequest,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Whether the caller requested output logprobs on chat choices.
@@ -31,8 +39,6 @@ pub struct PreparedRequest {
pub include_prompt_logprobs: bool,
/// Whether to include reasoning content in OpenAI responses.
pub include_reasoning: bool,
/// Lowered chat request for `vllm-chat`.
pub chat_request: ChatRequest,
/// Last assistant-role message content to echo back when `echo=true`.
pub echo: Option<String>,
/// Whether to include token IDs alongside generated text.
@@ -46,7 +52,7 @@ pub struct PreparedRequest {
///
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
/// the base `model` field in responses when no LoRA adapter is selected.
pub(crate) fn prepare_chat_request(
pub(super) fn prepare_chat_request(
request: ChatCompletionRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -146,18 +152,19 @@ pub(crate) fn prepare_chat_request(
Ok(PreparedRequest {
request_id,
response_model,
include_usage,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
options: ResponseOptions {
include_usage,
requested_logprobs,
include_prompt_logprobs,
include_reasoning,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
},
chat_request,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
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 +211,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 +323,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()
@@ -498,7 +505,7 @@ mod tests {
)
.expect("request is valid");
assert!(!prepared.include_reasoning);
assert!(!prepared.options.include_reasoning);
}
#[test]
@@ -868,8 +875,8 @@ mod tests {
)
.expect("request is valid");
assert!(prepared.requested_logprobs);
assert!(prepared.include_prompt_logprobs);
assert!(prepared.options.requested_logprobs);
assert!(prepared.options.include_prompt_logprobs);
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
assert_eq!(
prepared.chat_request.sampling_params.prompt_logprobs,
@@ -895,7 +902,7 @@ mod tests {
assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
assert!(!prepared.include_prompt_logprobs);
assert!(!prepared.options.include_prompt_logprobs);
}
#[test]
@@ -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,
@@ -18,13 +18,13 @@ use tracing::{debug, error, info, trace};
use tracing_futures::Instrument as _;
use vllm_text::{DecodedTextEvent, FinishReason, TextOutputStream, TextOutputStreamExt as _};
use self::convert::{ResponseOptions, prepare_completion_request};
use super::utils::logprobs::{
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
text_len,
};
use super::utils::types::Usage;
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::completions::convert::prepare_completion_request;
use crate::routes::openai::completions::types::{
CompletionChoice, CompletionRequest, CompletionResponse, CompletionSseChunk,
CompletionStreamChoice, CompletionStreamResponse,
@@ -42,7 +42,6 @@ pub async fn completions(
ValidatedJson(body): ValidatedJson<CompletionRequest>,
) -> Response {
let stream = body.stream;
let logprobs = body.logprobs;
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
@@ -57,9 +56,7 @@ pub async fn completions(
);
let created = unix_timestamp();
let include_prompt_logprobs = prepared.text_request.sampling_params.prompt_logprobs.is_some();
let log_request = state.enable_log_requests;
let text_stream = match state
.chat
.text()
@@ -84,11 +81,7 @@ pub async fn completions(
prepared.response_model,
created,
log_request,
prepared.include_usage,
prepared.echo,
logprobs,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
prepared.options,
);
let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
@@ -99,11 +92,8 @@ pub async fn completions(
prepared.request_id,
prepared.response_model,
created,
prepared.echo,
logprobs,
include_prompt_logprobs,
prepared.return_token_ids,
prepared.return_tokens_as_token_ids,
log_request,
prepared.options,
)
.instrument(request_span.clone())
.await
@@ -112,18 +102,6 @@ pub async fn completions(
Err(error) => return error.into_response(),
};
if log_request {
let usage = response.usage.as_ref();
info!(
parent: &request_span,
model = %response.model,
prompt_tokens = usage.map_or(0, |u| u.prompt_tokens),
output_tokens = usage.and_then(|u| u.completion_tokens).unwrap_or(0),
finish_reason = response.choices.first().and_then(|c| c.finish_reason.as_deref()).unwrap_or("unknown"),
"completion finished"
);
}
Json(response).into_response()
}
}
@@ -133,11 +111,16 @@ async fn collect_completion(
request_id: String,
response_model: String,
created: u64,
echo: Option<String>,
requested_logprobs: Option<u32>,
include_prompt_logprobs: bool,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
log_request: bool,
ResponseOptions {
// Ignored: non-streaming responses always include usage.
include_usage: _,
echo,
requested_logprobs,
include_prompt_logprobs,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
) -> Result<CompletionResponse, ApiError> {
let collected = stream
.collect_output()
@@ -175,6 +158,21 @@ async fn collect_completion(
None => collected.text,
Some(prompt) => format!("{prompt}{}", collected.text),
};
let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();
let usage = Usage::from_counts(
collected.prompt_token_ids.len() as u32,
collected.token_ids.len() as u32,
);
if log_request {
info!(
model = %response_model,
prompt_tokens = usage.prompt_tokens,
output_tokens = usage.completion_tokens.unwrap_or(0),
%finish_reason,
"completion finished"
);
}
Ok(CompletionResponse {
id: request_id,
@@ -185,16 +183,13 @@ async fn collect_completion(
index: 0,
text,
logprobs,
finish_reason: Some(completion_finish_reason_to_openai(finish_reason)?.into()),
finish_reason: Some(finish_reason),
stop_reason,
prompt_logprobs,
token_ids: return_token_ids.then(|| collected.token_ids.clone()),
prompt_token_ids: return_token_ids.then(|| collected.prompt_token_ids.to_vec()),
}],
usage: Some(Usage::from_counts(
collected.prompt_token_ids.len() as u32,
collected.token_ids.len() as u32,
)),
usage: Some(usage),
system_fingerprint: None,
kv_transfer_params: collected.kv_transfer_params,
})
@@ -208,11 +203,15 @@ async fn completion_chunk_stream(
response_model: String,
created: u64,
log_request: bool,
include_usage: bool,
echo: Option<String>,
requested_logprobs: Option<u32>,
return_token_ids: bool,
return_tokens_as_token_ids: bool,
ResponseOptions {
include_usage,
echo,
requested_logprobs,
// Ignored: streaming prompt logprobs are rejected for Python parity.
include_prompt_logprobs: _,
return_token_ids,
return_tokens_as_token_ids,
}: ResponseOptions,
mut y: TryYielder<CompletionSseChunk, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
@@ -432,7 +431,7 @@ mod tests {
FinishReason, Finished,
};
use super::{CompletionSseChunk, completion_chunk_stream, final_chunk};
use super::{CompletionSseChunk, ResponseOptions, completion_chunk_stream, final_chunk};
#[test]
fn final_chunk_maps_stop_finish_reason() {
@@ -527,11 +526,10 @@ mod tests {
"model".to_string(),
1,
false,
false,
None,
Some(1),
false,
false,
ResponseOptions {
requested_logprobs: Some(1),
..Default::default()
},
)
.collect::<Vec<_>>()
.await;
@@ -10,18 +10,28 @@ use crate::utils::{ResolvedRequestContext, convert_logit_bias, merge_kv_transfer
/// Lowered completion request plus the public response metadata carried by
/// every SSE chunk.
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedRequest {
pub(super) struct PreparedRequest {
/// Stable OpenAI-style request ID, reused as the external text request ID.
pub request_id: String,
/// Public model ID echoed back to the client.
pub response_model: String,
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Public response rendering options for route-layer helpers.
pub options: ResponseOptions,
/// Lowered text request for the shared `vllm-text` facade.
pub text_request: TextRequest,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(super) struct ResponseOptions {
/// Whether the caller asked for the final streamed usage chunk.
pub include_usage: bool,
/// Original text prompt that should be echoed back northbound when
/// `echo=true`.
pub echo: Option<String>,
/// Whether the caller requested output logprobs on completion choices.
pub requested_logprobs: Option<u32>,
/// Whether the caller requested choice-level prompt logprobs.
pub include_prompt_logprobs: bool,
/// Whether to include token IDs alongside generated text.
pub return_token_ids: bool,
/// Whether to format logprob tokens as `token_id:{id}`.
@@ -33,7 +43,7 @@ pub struct PreparedRequest {
///
/// `lora_resolution.model_names` must be non-empty; the first entry is used as
/// the base `model` field in responses when no LoRA adapter is selected.
pub(crate) fn prepare_completion_request(
pub(super) fn prepare_completion_request(
request: CompletionRequest,
lora_resolution: &LoraModelResolution,
ctx: ResolvedRequestContext,
@@ -64,6 +74,7 @@ pub(crate) fn prepare_completion_request(
let include_usage = (request.stream_options.as_ref())
.and_then(|options| options.include_usage)
.unwrap_or(false);
let include_prompt_logprobs = prompt_logprobs.is_some();
let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();
let structured_outputs =
@@ -116,11 +127,15 @@ pub(crate) fn prepare_completion_request(
Ok(PreparedRequest {
request_id,
response_model,
include_usage,
options: ResponseOptions {
include_usage,
echo,
requested_logprobs: request.logprobs,
include_prompt_logprobs,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
},
text_request,
echo,
return_token_ids: request.return_token_ids.unwrap_or(false),
return_tokens_as_token_ids: request.return_tokens_as_token_ids.unwrap_or(false),
})
}
@@ -206,7 +221,7 @@ mod tests {
)
.expect("prepare");
assert!(prepared.include_usage);
assert!(prepared.options.include_usage);
assert_eq!(
prepared.text_request.prompt,
Prompt::TokenIds(vec![11, 22, 33])
@@ -250,7 +265,7 @@ mod tests {
)
.expect("prepare");
assert_eq!(prepared.echo, Some("hello".to_string()));
assert_eq!(prepared.options.echo, Some("hello".to_string()));
assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
}
@@ -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))?;
+673 -96
View File
@@ -14,15 +14,14 @@ use std::{fmt, fs};
use axum::body::{Body, to_bytes};
use axum::http::{Request, StatusCode};
use bytes::Bytes;
use futures::StreamExt as _;
use rmpv::Value;
use serde_json::json;
use serial_test::serial;
use tower::{Service as _, ServiceExt as _};
use vllm_chat::{
ChatBackend, ChatContent, ChatContentPart, ChatEvent, ChatLlm, ChatMessage, ChatRenderer,
ChatRequest, ChatRole, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor,
DynChatRenderer, NewChatOutputProcessorOptions, SamplingParams,
ChatBackend, ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRenderer, ChatRequest,
ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
NewChatOutputProcessorOptions,
};
use vllm_engine_core_client::protocol::logprobs::{
Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
@@ -44,8 +43,6 @@ use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::{build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora};
use crate::lora::LoraModelResolution;
use crate::routes::openai::chat_completions::convert::prepare_chat_request;
use crate::state::AppState;
fn request_output(
@@ -407,6 +404,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 +414,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 +463,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 +472,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 +780,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 +1091,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 +2447,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 +2651,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 +3433,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(),
))
@@ -3359,92 +3475,6 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {
assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn chat_harness_streams_text_events() {
let (chat, engine_task) = test_chat_with_engine_handle().await;
let mut stream = chat
.chat(ChatRequest {
messages: vec![ChatMessage::text(ChatRole::User, "hello")],
sampling_params: SamplingParams {
max_tokens: Some(8),
..Default::default()
},
request_id: "chat-harness".to_string(),
..ChatRequest::for_test()
})
.await
.expect("submit chat request");
let mut saw_text = false;
let mut saw_done = false;
while let Some(event) = stream.next().await {
match event.expect("chat event") {
ChatEvent::BlockDelta { .. } => saw_text = true,
ChatEvent::Done { .. } => {
saw_done = true;
break;
}
ChatEvent::Start { .. }
| ChatEvent::LogprobsDelta { .. }
| ChatEvent::BlockStart { .. }
| ChatEvent::BlockEnd { .. }
| ChatEvent::ToolCallStart { .. }
| ChatEvent::ToolCallArgumentsDelta { .. }
| ChatEvent::ToolCallEnd { .. } => {}
}
}
engine_task.await.expect("mock engine task");
assert!(saw_text);
assert!(saw_done);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn prepared_openai_request_streams_text_events() {
let (chat, engine_task) = test_chat_with_engine_handle().await;
let prepared = prepare_chat_request(
serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
}))
.expect("decode request"),
&LoraModelResolution {
model_names: vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
lora_request: None,
},
crate::utils::ResolvedRequestContext::default(),
)
.expect("prepare request");
let mut stream = chat.chat(prepared.chat_request).await.expect("submit chat request");
let mut saw_text = false;
let mut saw_done = false;
while let Some(event) = stream.next().await {
match event.expect("chat event") {
ChatEvent::BlockDelta { .. } => saw_text = true,
ChatEvent::Done { .. } => {
saw_done = true;
break;
}
ChatEvent::Start { .. }
| ChatEvent::LogprobsDelta { .. }
| ChatEvent::BlockStart { .. }
| ChatEvent::BlockEnd { .. }
| ChatEvent::ToolCallStart { .. }
| ChatEvent::ToolCallArgumentsDelta { .. }
| ChatEvent::ToolCallEnd { .. } => {}
}
}
engine_task.await.expect("mock engine task");
assert!(saw_text);
assert!(saw_done);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {
@@ -4228,6 +4258,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 +4429,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 +4786,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 -1
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;
@@ -9,11 +10,16 @@ use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::lora::LoraRequest;
use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
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 +33,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 +61,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 +93,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 {
+19
View File
@@ -0,0 +1,19 @@
[package]
name = "vllm-tool-parser-py"
version.workspace = true
edition.workspace = true
license.workspace = true
[lib]
name = "_rust_tool_parser"
crate-type = ["cdylib", "rlib"]
[dependencies]
pyo3.workspace = true
pythonize = { workspace = true, features = ["serde_json"] }
serde_json.workspace = true
thiserror-ext.workspace = true
vllm-tool-parser.workspace = true
[lints]
workspace = true
+366
View File
@@ -0,0 +1,366 @@
//! Thin PyO3 bindings for `vllm_tool_parser`.
//!
//! This crate exposes the Rust tool parser trait and data shapes to Python
//! while keeping parser state, grammar, and schema-aware argument conversion in
//! Rust. Python callers should use this module as a typed bridge and keep any
//! vLLM protocol adaptation outside the binding.
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyModule};
use pythonize::{depythonize, pythonize};
use serde_json::Value;
use thiserror_ext::AsReport as _;
use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
macro_rules! tool_parser_factory {
($($parser:ident),+ $(,)?) => {
fn create_tool_parser(
name: &str,
tools: &[Tool],
) -> PyResult<Box<dyn ToolParser>> {
match name {
$(
stringify!($parser) => {
<vllm_tool_parser::$parser as ToolParser>::create(tools)
}
)+
_ => {
return Err(PyValueError::new_err(format!(
"unsupported tool parser `{name}`"
)));
}
}
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
};
}
// Export a tool parser to Python by registering it here.
tool_parser_factory! {
DeepSeekV4ToolParser, // for testing on Python side
}
#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
#[derive(Clone)]
struct PyTool(Tool);
#[pymethods]
impl PyTool {
#[new]
#[pyo3(signature = (name, description, parameters, strict=None))]
fn new(
name: String,
description: Option<String>,
parameters: &Bound<'_, PyAny>,
strict: Option<bool>,
) -> PyResult<Self> {
let parameters = depythonize::<Value>(parameters).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from Python to JSON: {error}"
))
})?;
Ok(Self(Tool {
name,
description,
parameters,
strict,
}))
}
#[getter]
fn name(&self) -> &str {
&self.0.name
}
#[getter]
fn description(&self) -> Option<&str> {
self.0.description.as_deref()
}
#[getter]
fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from JSON to Python: {error}"
))
})
}
#[getter]
fn strict(&self) -> Option<bool> {
self.0.strict
}
}
#[pyclass(
name = "ToolCallDelta",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolCallDelta(ToolCallDelta);
#[pymethods]
impl PyToolCallDelta {
#[new]
#[pyo3(signature = (tool_index, name, arguments))]
fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
Self(ToolCallDelta {
tool_index,
name,
arguments,
})
}
#[getter]
fn tool_index(&self) -> usize {
self.0.tool_index
}
#[getter]
fn name(&self) -> Option<&str> {
self.0.name.as_deref()
}
#[getter]
fn arguments(&self) -> &str {
&self.0.arguments
}
}
#[pyclass(
name = "ToolParserOutput",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolParserOutput(ToolParserOutput);
#[pymethods]
impl PyToolParserOutput {
#[new]
#[pyo3(signature = (normal_text="", calls=None))]
fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
let calls =
calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
Self(ToolParserOutput {
normal_text: normal_text.to_owned(),
calls,
})
}
#[getter]
fn normal_text(&self) -> &str {
&self.0.normal_text
}
#[getter]
fn calls(&self) -> Vec<PyToolCallDelta> {
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
}
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
self.0.append(other.0.clone());
}
fn coalesce_calls(&self) -> Self {
Self(self.0.clone().coalesce_calls())
}
}
#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
struct PyToolParser(Box<dyn ToolParser>);
impl PyToolParser {
fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
self.0
.parse_into(chunk, &mut output.0)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
}
#[pymethods]
impl PyToolParser {
#[new]
fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
create_tool_parser(parser_name, &tools).map(Self)
}
fn parse_into(
&mut self,
chunk: &str,
mut output: PyRefMut<'_, PyToolParserOutput>,
) -> PyResult<()> {
self.parse_into_output(chunk, &mut output)
}
fn finish(&mut self) -> PyResult<PyToolParserOutput> {
self.0
.finish()
.map(PyToolParserOutput)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
fn reset(&mut self) -> String {
self.0.reset()
}
fn preserve_special_tokens(&self) -> bool {
self.0.preserve_special_tokens()
}
}
#[pymodule]
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyTool>()?;
m.add_class::<PyToolCallDelta>()?;
m.add_class::<PyToolParserOutput>()?;
m.add_class::<PyToolParser>()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
Python::initialize();
Python::attach(f)
}
fn tool_schema() -> Value {
json!({
"type": "object",
"properties": {
"user_id": {"type": "integer"},
"shipping": {
"type": "object",
"properties": {
"city": {"type": "string"},
"zip": {"type": "integer"}
}
}
}
})
}
fn build_call() -> String {
r#"<DSMLtool_calls>
<DSMLinvoke name="create_order">
<DSMLparameter name="user_id" string="false">42</DSMLparameter>
<DSMLparameter name="shipping" string="false">{"city":"Singapore","zip":18956}</DSMLparameter>
</DSMLinvoke>
</DSMLtool_calls>"#
.to_owned()
}
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
let parameters = pythonize(py, &tool_schema()).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert test schema from JSON to Python: {error}"
))
})?;
Py::new(
py,
PyTool::new(
"create_order".to_owned(),
Some("Create an order".to_owned()),
&parameters,
None,
)?,
)
}
#[test]
fn tool_round_trips_typed_fields() {
with_python(|py| {
let tool = make_py_tool(py)?;
let borrowed = tool.borrow(py);
assert_eq!(borrowed.name(), "create_order");
assert_eq!(borrowed.description(), Some("Create an order"));
assert_eq!(borrowed.strict(), None);
let parameters = borrowed.parameters(py)?;
let parameters = depythonize::<Value>(parameters.bind(py))?;
assert_eq!(parameters, tool_schema());
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn output_append_and_coalesce_calls() {
with_python(|py| {
let first = Py::new(
py,
PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
)?;
let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
output.append(other.borrow(py));
let coalesced = output.coalesce_calls();
assert_eq!(coalesced.normal_text(), "text");
let calls = coalesced.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].tool_index(), 0);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(calls[0].arguments(), "{\"a\":1}");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_parse_finish_and_preserve_special_tokens() {
with_python(|py| {
let tool = make_py_tool(py)?;
let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
assert!(parser.preserve_special_tokens());
let mut output = PyToolParserOutput::new(py, "", None);
parser.parse_into_output(&build_call(), &mut output)?;
let finish = Py::new(py, parser.finish()?)?;
output.append(finish.borrow(py));
let output = output.coalesce_calls();
assert_eq!(output.normal_text(), "");
let calls = output.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(
serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
json!({
"user_id": 42,
"shipping": {
"city": "Singapore",
"zip": 18956
}
})
);
assert_eq!(parser.reset(), "");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_errors_for_unknown_name() {
with_python(|py| {
let tool = make_py_tool(py)?;
let error = match PyToolParser::new(py, "missing", vec![tool]) {
Ok(_) => panic!("missing parser name unexpectedly succeeded"),
Err(error) => error,
};
let message = format!("{error}");
assert!(message.contains("unsupported tool parser `missing`"));
PyResult::Ok(())
})
.unwrap();
}
}
+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 -52
View File
@@ -18,8 +18,6 @@ import torch
from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools_rust import Binding, RustExtension
from setuptools_rust.build import build_rust
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
@@ -35,11 +33,12 @@ def load_module_from_path(module_name, path):
ROOT_DIR = Path(__file__).parent
logger = logging.getLogger(__name__)
PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
# cannot import envs directly because it depends on vllm,
# which is not installed yet
envs = load_module_from_path("envs", os.path.join(ROOT_DIR, "vllm", "envs.py"))
rust_build = load_module_from_path(
"rust_build", os.path.join(ROOT_DIR, "tools", "build_rust.py")
)
VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
USE_PRECOMPILED_EXTENSIONS = envs.VLLM_USE_PRECOMPILED
@@ -49,11 +48,6 @@ USE_PRECOMPILED_RUST_FRONTEND = (
)
def should_require_rust_frontend() -> bool:
value = os.getenv("VLLM_REQUIRE_RUST_FRONTEND", "")
return value.lower() not in ("", "0", "false", "no")
if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
VLLM_TARGET_DEVICE = "cpu"
@@ -420,24 +414,6 @@ class precompiled_build_ext(build_ext):
return
class precompiled_build_rust(build_rust):
"""Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
def run(self) -> None:
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
logger.info(
"Skipping local Rust build: using precompiled %s",
PRECOMPILED_RUST_FRONTEND_PATH,
)
return
logger.warning(
"Precompiled wheel did not provide %s; falling back to local Rust build.",
PRECOMPILED_RUST_FRONTEND_PATH,
)
super().run()
class precompiled_wheel_utils:
"""Extracts libraries and other files from an existing wheel."""
@@ -731,9 +707,6 @@ class precompiled_wheel_utils:
"vllm/_rocm_C.abi3.so",
}
)
if extract_rust_frontend:
exact_members.add("vllm/vllm-rs")
flash_attn_regex = re.compile(
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
)
@@ -756,6 +729,12 @@ class precompiled_wheel_utils:
if member.filename in exact_members:
file_members.append(member)
continue
if (
extract_rust_frontend
and rust_build.is_precompiled_artifact_member(member.filename)
):
file_members.append(member)
continue
if not extract_extensions:
continue
@@ -1109,6 +1088,12 @@ package_data = {
}
def add_vllm_package_data(filename: str) -> None:
vllm_files = package_data.setdefault("vllm", [])
if filename not in vllm_files:
vllm_files.append(filename)
# If using precompiled artifacts, extract and patch package_data in advance.
if USE_PRECOMPILED_RUST_FRONTEND:
wheel_url, download_filename = precompiled_wheel_utils.determine_wheel_url()
@@ -1121,12 +1106,10 @@ if USE_PRECOMPILED_RUST_FRONTEND:
for pkg, files in patch.items():
package_data.setdefault(pkg, []).extend(files)
# If the rust frontend binary is already present in the source tree (e.g.,
# pre-built in a separate Docker build stage), ship it as-is.
if PRECOMPILED_RUST_FRONTEND_PATH.exists():
vllm_files = package_data.setdefault("vllm", [])
if "vllm-rs" not in vllm_files:
vllm_files.append("vllm-rs")
# Rust artifacts already present in the source tree (e.g., pre-built in a
# separate Docker build stage) are shipped as-is.
for rust_artifact in rust_build.find_precompiled_artifacts():
add_vllm_package_data(rust_artifact.name)
if _no_device():
ext_modules = []
@@ -1139,23 +1122,14 @@ else:
if USE_PRECOMPILED_EXTENSIONS
else cmake_build_ext,
}
if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():
cmdclass["build_rust"] = precompiled_build_rust
if USE_PRECOMPILED_RUST_FRONTEND or rust_build.find_precompiled_artifacts():
cmdclass["build_rust"] = rust_build.precompiled_build_rust
# Rust frontend binary, built via setuptools-rust and installed into the
# package directory alongside the Python modules.
# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
# requires extra work on using precompiled binaries.
rust_extensions = [
RustExtension(
target="vllm.vllm-rs",
path="rust/src/cmd/Cargo.toml",
args=["--bin", "vllm-rs"],
features=["native-tls-vendored"],
binding=Binding.Exec,
optional=not should_require_rust_frontend(),
),
]
# Rust artifacts, built via setuptools-rust and installed into the package
# directory alongside the Python modules.
rust_extensions = rust_build.rust_extensions(
optional=not rust_build.should_require_rust_frontend()
)
setup(
# static metadata should rather go in pyproject.toml
+341
View File
@@ -0,0 +1,341 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from transformers import AutoTokenizer, PreTrainedTokenizerBase
from vllm.benchmarks.datasets import BFCLDataset, get_samples
def _patch_hf_api(side_effect):
"""Return a patch context that swaps `hf_api()` to a stub whose
`.hf_hub_download` attribute uses `side_effect`."""
fake_api = MagicMock()
fake_api.hf_hub_download.side_effect = side_effect
return patch("vllm.benchmarks.datasets.datasets.hf_api", return_value=fake_api)
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
return AutoTokenizer.from_pretrained("gpt2")
_FAKE_ROWS = {
"simple": [
{
"id": "simple_0",
"question": [
[
{
"role": "user",
"content": "What is 2+2?",
}
]
],
"function": [
{
"name": "add",
"description": "Add two numbers.",
"parameters": {
"type": "dict",
"properties": {
"a": {"type": "integer", "description": "first"},
"b": {"type": "float", "description": "second"},
},
"required": ["a", "b"],
},
}
],
},
],
"live_simple": [
{
"id": "live_simple_0",
"question": [[{"role": "user", "content": "Tell me the weather."}]],
"function": [
{
"name": "get_weather",
"description": "Get weather.",
"parameters": {
"type": "dict",
"properties": {
"city": {"type": "any", "description": "city"},
"coords": {"type": "tuple", "description": "coords"},
},
"required": ["city"],
},
}
],
},
],
}
def _write_fake_files(tmp_path: Path) -> dict[str, Path]:
"""Write fake BFCL JSONL files mimicking the HF repo layout."""
paths = {}
for category, rows in _FAKE_ROWS.items():
p = tmp_path / f"BFCL_v3_{category}.json"
with p.open("w") as f:
for row in rows:
f.write(json.dumps(row) + "\n")
paths[category] = p
return paths
def _args_for_bfcl(categories: list[str] | None) -> argparse.Namespace:
return argparse.Namespace(
dataset_name="hf",
dataset_path="gorilla-llm/Berkeley-Function-Calling-Leaderboard",
hf_name=None,
hf_subset=None,
hf_split=None,
hf_output_len=64,
disable_shuffle=True,
num_prompts=2,
no_oversample=False,
no_stream=True,
seed=0,
request_id_prefix="",
trust_remote_code=False,
skip_chat_template=False,
enable_multimodal_chat=False,
backend="openai-chat",
bfcl_categories=categories,
)
@pytest.mark.benchmark
def test_bfcl_dataset_translates_schema_and_attaches_tools(
hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
) -> None:
"""BFCLDataset should translate schemas to OpenAI tool format, set
`messages` directly on SampleRequest, and attach tools/tool_choice via
request_overrides."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
args = _args_for_bfcl(categories=["simple", "live_simple"])
with _patch_hf_api(fake_download):
samples = get_samples(args, hf_tokenizer)
assert len(samples) == 2
for s in samples:
assert s.chat_messages is not None
assert isinstance(s.chat_messages, list)
assert s.chat_messages[0]["role"] == "user"
assert s.request_overrides is not None
assert "tools" in s.request_overrides
assert s.request_overrides["tool_choice"] == "auto"
# messages must NOT leak into request_overrides — it has its own
# typed field on SampleRequest.
assert "messages" not in s.request_overrides
tools = s.request_overrides["tools"]
assert len(tools) == 1
tool = tools[0]
assert tool["type"] == "function"
# Translated schema: dict -> object, float -> number,
# any -> string, tuple -> array.
params = tool["function"]["parameters"]
assert params["type"] == "object"
for prop in params["properties"].values():
assert prop["type"] in {"integer", "number", "string", "array"}
@pytest.mark.benchmark
def test_bfcl_dataset_requires_openai_chat_backend(
hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
args = _args_for_bfcl(categories=["simple"])
args.backend = "openai"
with pytest.raises(ValueError, match="openai-chat"):
get_samples(args, hf_tokenizer)
@pytest.mark.benchmark
def test_bfcl_dataset_missing_category_raises_clear_error(
hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
"""A typo'd category should produce an actionable ValueError, not an
opaque huggingface_hub exception."""
from huggingface_hub.errors import EntryNotFoundError
args = _args_for_bfcl(categories=["simpl"]) # typo
def raise_missing(_repo, filename, **_kwargs):
raise EntryNotFoundError(f"404 Not Found: {filename}")
with (
_patch_hf_api(raise_missing),
pytest.raises(ValueError, match=r"BFCL category 'simpl' not found"),
):
get_samples(args, hf_tokenizer)
@pytest.mark.benchmark
def test_chat_backend_uses_messages_field_when_set() -> None:
"""When RequestFuncInput.chat_messages is set, the chat backend must use
it verbatim and skip default content construction from `prompt`."""
import asyncio
from vllm.benchmarks.lib.endpoint_request_func import (
RequestFuncInput,
async_request_openai_chat_completions,
)
captured: dict = {}
class _FakeResp:
status = 500
reason = "stop-after-capture"
content = None
async def __aenter__(self):
return self
async def __aexit__(self, *exc):
return False
class _FakeSession:
def post(self, url, json, headers): # noqa: A002
captured["url"] = url
captured["payload"] = json
return _FakeResp()
messages = [
{"role": "system", "content": "you are helpful"},
{"role": "user", "content": "call add(3, 4)"},
]
req = RequestFuncInput(
prompt="IGNORED",
api_url="http://localhost:0/v1/chat/completions",
prompt_len=10,
output_len=16,
model="test-model",
chat_messages=messages,
extra_body={"tools": [{"type": "function", "function": {"name": "add"}}]},
)
asyncio.run(
async_request_openai_chat_completions(
request_func_input=req, session=_FakeSession()
)
)
payload = captured["payload"]
assert payload["messages"] is messages, (
"chat backend must forward RequestFuncInput.chat_messages verbatim "
"instead of constructing a default user message from `prompt`"
)
# extra_body still merges in as before (shallow, per-request wins).
assert payload["tools"][0]["function"]["name"] == "add"
@pytest.mark.benchmark
def test_bfcl_prompt_len_includes_tools(tmp_path: Path) -> None:
"""prompt_len must reflect tokens from both messages *and* tool schemas,
so percentile buckets and input-distribution summaries aren't biased
low for tool-heavy traffic."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
captured: dict = {}
class _FakeTokenizer:
def apply_chat_template(
self, messages, tools=None, tokenize=False, add_generation_prompt=True
):
captured["tools"] = tools
base = " ".join(m.get("content", "") for m in messages)
tool_text = json.dumps(tools) if tools else ""
return base + " " + tool_text
def __call__(self, text):
# 1 "token" per whitespace-separated word.
return type("Enc", (), {"input_ids": text.split()})()
fake = _FakeTokenizer()
args = _args_for_bfcl(categories=["simple"])
args.num_prompts = 1
with _patch_hf_api(fake_download):
samples = get_samples(args, fake)
assert len(samples) == 1
assert captured["tools"] is not None, (
"apply_chat_template must be called with tools= so the schema "
"contributes to the prompt-length estimate"
)
assert len(captured["tools"]) == 1
assert captured["tools"][0]["function"]["name"] == "add"
# Sanity: prompt_len exceeds a messages-only estimate. The fake row's
# user message is "What is 2+2?" (3 whitespace-separated tokens).
assert samples[0].prompt_len > 3
@pytest.mark.benchmark
def test_bfcl_prompt_len_falls_back_when_tokenizer_rejects_tools(
tmp_path: Path,
) -> None:
"""Older tokenizers don't accept tools=; fallback must still produce a
non-zero prompt_len without crashing."""
paths = _write_fake_files(tmp_path)
def fake_download(_repo, filename, **_kwargs):
category = filename.removeprefix("BFCL_v3_").removesuffix(".json")
return str(paths[category])
class _LegacyTokenizer:
def apply_chat_template(self, messages, **kwargs):
if "tools" in kwargs:
raise TypeError("unexpected keyword argument 'tools'")
return " ".join(m.get("content", "") for m in messages)
def __call__(self, text):
return type("Enc", (), {"input_ids": text.split()})()
args = _args_for_bfcl(categories=["simple"])
args.num_prompts = 1
with _patch_hf_api(fake_download):
samples = get_samples(args, _LegacyTokenizer())
assert len(samples) == 1
assert samples[0].prompt_len > 0
@pytest.mark.benchmark
def test_bfcl_schema_translation_is_recursive() -> None:
"""_translate_schema must recurse into nested properties."""
input_schema = {
"type": "dict",
"properties": {
"nested": {
"type": "dict",
"properties": {
"value": {"type": "float"},
"tags": {"type": "tuple", "items": {"type": "any"}},
},
}
},
}
out = BFCLDataset._translate_schema(input_schema)
assert out["type"] == "object"
assert out["properties"]["nested"]["type"] == "object"
assert out["properties"]["nested"]["properties"]["value"]["type"] == "number"
assert out["properties"]["nested"]["properties"]["tags"]["type"] == "array"
nested_props = out["properties"]["nested"]["properties"]
assert nested_props["tags"]["items"]["type"] == "string"
@@ -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
@@ -382,11 +382,12 @@ def test_rope_kvcache_fusion(
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
# Cannot compare fp8_* directly here, cast to model dtype instead
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
kv_cache_unfused.view(dtype),
kv_cache_fused.view(dtype),
atol=ATOL,
rtol=RTOL,
kv_cache_unfused.to(dtype),
kv_cache_fused.to(dtype),
atol=1e-1,
rtol=1e-1,
)
@@ -569,17 +570,19 @@ def test_rope_static_qquant_kvcache_fusion(
else:
ATOL, RTOL = (1e-2, 1e-2)
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
q_unfused.to(torch.float32),
q_fused.to(torch.float32),
atol=ATOL,
rtol=RTOL,
atol=1e-1,
rtol=1e-1,
)
torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
# TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged.
torch.testing.assert_close(
kv_cache_unfused.view(dtype),
kv_cache_fused.view(dtype),
atol=ATOL,
rtol=RTOL,
kv_cache_unfused.to(dtype),
kv_cache_fused.to(dtype),
atol=1e-1,
rtol=1e-1,
)
+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

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