Compare commits

...
Author SHA1 Message Date
Bugen ZhaoandOpenAI Codex e1a763558c [CI] Discover Rust coverage artifacts from build metadata
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 06:44:16 +00:00
Bugen ZhaoandOpenAI Codex 84aeec9f22 [CI] Simplify Rust coverage reporting
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 08:18:14 +00:00
Bugen ZhaoandOpenAI Codex 82a770ddbd [CI] Simplify Rust coverage aggregation
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:58:45 +00:00
Bugen Zhao a09a9bace1 [CI] Disable redundant Codecov file fixes 2026-07-21 13:57:38 +00:00
Bugen Zhao 6c20d467a2 [CI] Run Codecov from repository root 2026-07-21 13:34:22 +00:00
Bugen ZhaoandOpenAI Codex cb59d0a351 [CI] Collect Rust coverage in Buildkite
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 13:01:56 +00:00
Bugen ZhaoandOpenAI Codex 0ab1bded36 [CI] Instrument Rust artifacts for coverage
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 12:22:25 +00:00
Umut PolatandGitHub 040cbf95cc [Misc] Use VLLMValidationError in chat completion tool and batch validators (#49214)
Signed-off-by: Umut Polat <52835619+umut-polat@users.noreply.github.com>
2026-07-21 11:38:20 +00:00
SyaGitHubLi, Jiang <jiang1.li@intel.com>
5b3762a7f0 [Bugfix][CPU] Fix Clang OpenMP build on macOS (#49021)
Signed-off-by: markyangcc <mmdou3@163.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-21 09:58:52 +00:00
bastefaniakandGitHub 4d30c510ce [bugfix] Fix Cosmos3 Edge checkpoint weights filtering, video loading, prompt expansion (#49190)
Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>
2026-07-21 17:18:36 +08:00
6700813f86 [3/N][KV-Cache Layout Refactor] Standardize Mamba cache; drop get_transfer_cache_regions (#44456)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 09:16:15 +00:00
Bugen ZhaoandGitHub eb44b3aaa4 [Rust][Benchmark] Use async HTTP clients (#49295)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 16:53:57 +08:00
Nicolò LucchesiandGitHub 7a98c7a392 [Misc] Remove old now unsupported max_num_partial_prefills and max_long_partial_prefills (#49244)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-07-21 08:52:52 +00:00
Lena OnyshchenkoandGitHub 0d9e60619b [Misc][Docs] Fix XPU compute-runtime driver link version mismatch (#49299)
Signed-off-by: oonyshch <xonyshch@gmail.com>
2026-07-21 08:45:41 +00:00
1134545b6f Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) (#49033)
Co-authored-by: vllm-agent <vllm-agent@users.noreply.github.com>
2026-07-21 09:36:45 +01:00
Miłosz GrunwaldGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Harry Mellor
3e0c887511 [Bugfix] Fix Ovis2_5 special tokens for transformers v5 (#47298)
Signed-off-by: mgrunwal <milosz.grunwald@intel.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-21 08:09:20 +00:00
Stefan KaestleandGitHub adfbbc1005 Propagate Flash Attention cache configuration to Ray workers (#49177)
Signed-off-by: Stefan Kaestle <skaestle@nvidia.com>
2026-07-21 07:47:53 +00:00
Roy WangandGitHub adc98f04d0 [Misc] Add @esmeetu to codeowners for rust/src/bench (#49298)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-07-21 07:44:24 +00:00
8def3cdde2 [Bugfix] Propagate quant_config to LFM2 ShortConv projections (#48917)
Signed-off-by: Alex Yuan <alex.yuan@liquid.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-21 06:51:55 +00:00
616c9bd0f4 [Frontend] Support additional sampling parameters for translation API (#45839)
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-21 05:53:56 +00:00
Bugen ZhaoandGitHub 8688a06d67 [Rust][Benchmark] Use tracing for logs (#48937)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-21 05:18:13 +00:00
Salt SatoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f25953cc59 [Bugfix][Rust Frontend] Handle zero-column logprobs payloads without panicking (#49113)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
2026-07-21 04:30:29 +00:00
d9aa35161d Update BGE-M3 token expectations for leading spaces (#49269)
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-07-21 03:49:17 +00:00
6bcda970fd [CI][NIXL] Isolate concurrent engine internal ports (#49129)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 22:28:11 -05:00
Isotr0pyandGitHub ea0e9c8f2e [MRV2] Add encoder cache profiling implementation (#47985)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-20 20:18:14 -07:00
ChaunceyandGitHub 94ed0bf4e0 [Bugfix][KV Offloading] Handle queued request aborts without allocated KV blocks (#49146)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-07-21 11:16:26 +08:00
Connor CarpenterGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
1940c8441e [Rust Frontend][gRPC] Add engine-aware health reporting (#48992)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
2026-07-21 10:54:39 +08:00
Simon MoandGitHub 72d16aee15 [CI] Exercise FA3 FP8 attention on SM90 (#49231)
Signed-off-by: Simon Mo <simon@inferact.ai>
2026-07-21 10:26:55 +08:00
Kunshang JiandGitHub e78a0c8e59 [XPU][Doc] Update XPU docker image documents (#49148)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-21 10:11:27 +08:00
Chris LeonardandGitHub 97a98006b0 Update qutlass cmake for stable abi (#47879)
Signed-off-by: Chris Leonard <chleonar@redhat.com>
2026-07-20 18:31:10 -07:00
0a684ab0c0 [Bugfix] Fix WSL circular import from pin_memory warning_once (#48444)
Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:55 -07:00
0d9210a502 Fixes non-coalesced HBM access in marlin_int4_fp8_preprocess_kernel_awq (#47268)
Signed-off-by: xjx <493337577@qq.com>
Signed-off-by: flutist-alibaba <30485581+flutist@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
2026-07-20 18:30:38 -07:00
1d874867ea [Misc][Docs] Fix broken protocol link in speech_to_text doc (#47212)
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-21 01:05:05 +00:00
2e2e626b40 [Bugfix] Count per-group blocks in get_max_concurrency_for_kv_cache_config (#48317)
Signed-off-by: David Orman <ormandj@corenode.com>
Co-authored-by: Luke Alonso <lalonso@gmail.com>
Co-authored-by: Martin Vit <martin@voipmonitor.org>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-07-21 00:29:10 +00:00
Nick HillandGitHub af91f4b3e4 [Cleanup] Remove unused StructuredOutputRequest.status field (#49235)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-21 00:06:16 +00:00
2396a61108 [Attention][MLA][DCP] Query replication for MLA decode (DeepSeek-V2/R1 + Kimi-K2.5) (#45964)
Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 23:51:27 +00:00
97a668152b [RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel (#44214)
Signed-off-by: Xuanyu Zhang <xuanyu.zhang@mistral.ai>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
2026-07-20 16:45:10 -07:00
58b2012aa2 [copy of #45208] CuMem slept-L1 fragmentation accounting (#49208)
Signed-off-by: haosdent <haosdent@gmail.com>
Signed-off-by: Justin Wood <justin.m.wood@me.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: haosdent <haosdent@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Justin Wood <jwood@me.com>
2026-07-20 23:08:11 +00:00
Ning XieandGitHub b7c20d0cfa [chore] adjust logo be more friendly to white background terminal (#48938)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
2026-07-20 15:15:28 -07:00
TJianandGitHub a2b1f9fc3b [ROCm] [Release] [Bugfix] Fix the per commit wheel release pipeline. (#49245)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-07-20 22:12:38 +00:00
642076d26c Support loading sample_from_anchor flag from speculators config (#48639)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-07-20 14:52:43 -07:00
Charlie FuandGitHub 5feb3950e5 [ROCm][CI] fix test_rocm_quick_reduce.py (#49234)
Signed-off-by: charlifu <charlifu@amd.com>
2026-07-20 16:39:33 -05:00
4ec199b66a [Bugfix][Spec-Decode] Populate draft seq_lens_cpu_upper_bound for spec-decode attention metadata (#44492)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
2026-07-20 20:50:58 +00:00
Roberto L. CastroGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
7ca017778f [Feat][Perf] Add new warmup infrastructure for JITs (#47451)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 13:21:55 -07:00
fbfe58133d [Bugfix][KV Offload] Preserve reachable tails for hybrid SWA groups (#48911)
Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Colton Ottley <colton@ottleyengineering.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-07-20 22:12:36 +03:00
9dd62d80ab Cosmos3 FP8 ModelOpt/Diffusers remapping (#48952)
Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
Signed-off-by: wkutak <wkutak@nvidia.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-07-20 11:31:48 -07:00
f878367898 [Revert][Bugfix] Restore MiniCPM-V 4.6 ViT QKV weight loader (#49193)
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: wjinxu <1299461899@qq.com>
2026-07-20 18:17:46 +00:00
Matthew BonanniGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
bd091079cb [Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 10:53:27 -07:00
b23bd73f54 [XPU]add sycl path for Mhc (#47245)
Signed-off-by: root <xiaolong.guo@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:32:54 +00:00
Bugen ZhaoandGitHub e2d7adeb64 [Rust Frontend] Bump xgrammar-structural-tag and enable local extension (#49161)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-20 16:22:24 +01:00
Isotr0pyandGitHub 15cb8e140d [Multimodal] Allow keeping original image mode for ImageIO (#49159)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
2026-07-20 13:42:45 +00:00
f007cceb42 [KV Offload] Support self-describing KV events with TieringOffloadingSpec (#48679)
Signed-off-by: Change72 <changg@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 16:41:58 +03:00
0a5069e4e3 [Bugfix][Gemma4] Fix ModelOpt mixed-precision MoE config mapping (#48563)
Signed-off-by: wangqian <601731555@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 06:39:28 -07:00
8ce53a616e [Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches (#47574)
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
2026-07-20 13:18:17 +00:00
Lena OnyshchenkoandGitHub ae10e855ab [Misc][Docs] Remove duplicate CodeGeex4 row in XPU model table (#47210)
Signed-off-by: oonyshch <xonyshch@gmail.com>
2026-07-20 10:05:36 +00:00
hclandGitHub 530ee36a0d fix(openai): reject non-numeric logprobs with 400 instead of 500 (#49144)
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
2026-07-20 10:04:50 +00:00
Salt SatoandGitHub d835ad572c [Bugfix][Rust Frontend] Map missing prompt logprobs for single-token prompts in chat and raw generate (#49111)
Signed-off-by: Feathbow <feathbow@gmail.com>
2026-07-20 10:00:06 +00:00
47d0597ca2 [Misc][Docs] Fix broken csrc kernel links in fusions doc (#47211)
Signed-off-by: oonyshch <xonyshch@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-20 09:44:25 +00:00
ReidandGitHub 818cf61e91 [Rust Frontend] Fix macro-based content format detection (#49042)
Signed-off-by: reidliu41 <reid201711@gmail.com>
2026-07-20 09:39:13 +00:00
Bugen ZhaoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
c01618fdc8 [Rust][Benchmark] Integrate vllm-bench to vllm-rs & vllm CLI (#48930)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-20 09:31:25 +00:00
Xiaochang WuGitHubKunshang Jimergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
823eaf667d [XPU] FP8 o_proj with fp8_bmm and load-time scale transpose (#48334)
Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-20 16:32:03 +08:00
SageGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f1f1259692 [Rust Frontend] Use zero-copy slicing for multimodal tensors (#48781)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-07-20 16:28:25 +08:00
zofiaGitHubmayuyuacemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Kunshang Ji
df13b5aef5 [XPU] [MoE] add quant input when prepare for fusedmoe (#47122)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com>
Co-authored-by: mayuyuace <qiming1.zhang@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:47:26 +08:00
Sihan ChenGitHubLi, Jiang <jiang1.li@intel.com>
4938d44a3b [CPU] fixes heterogeneous NIXL KV transfer into CPU_ATTN decode workers (#47871)
Signed-off-by: Spycsh <sihan.chen@intel.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-20 07:33:13 +00:00
37bf988c2f [XPU][Bugfix] Fix GroupCoordinator device_index (#47295)
Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-20 15:25:56 +08:00
aoshen02andGitHub 9459fc6471 [Bugfix][RL] Set vLLM config during weight reload (#45989)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
2026-07-20 15:02:56 +08:00
5245c80564 [Doc] Document blocks_per_chunk in the KV offloading guide (#49100)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
2026-07-20 09:48:43 +03:00
9bc266d923 [Bugfix][KV Offload] Propagate EAGLE mode to SimpleCPU coordinator (#49071)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-20 06:39:11 +00:00
5c9f6557d7 [Hardware][CPU] Enable granite-4 model on cpu (#47641)
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan>
Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-07-20 06:15:16 +00:00
237 changed files with 8666 additions and 2441 deletions
+5 -1
View File
@@ -18,6 +18,8 @@ steps:
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
- tests/kernels/mamba/test_cpu_short_conv.py - tests/kernels/mamba/test_cpu_short_conv.py
- tests/kernels/mamba/test_causal_conv1d.py
- tests/kernels/mamba/test_mamba_ssm.py
commands: commands:
- | - |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# Note: SDE can't be downloaded from CI host because of AWS WAF # Note: SDE can't be downloaded from CI host because of AWS WAF
# - label: CPU-Compatibility Tests # - label: CPU-Compatibility Tests
+2 -2
View File
@@ -813,8 +813,8 @@ steps:
# Download artifacts from current build # Download artifacts from current build
echo "Downloading artifacts from current build" echo "Downloading artifacts from current build"
# buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" . buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
# buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" . buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# # Run upload script # # Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh bash .buildkite/scripts/upload-rocm-wheels.sh
@@ -40,7 +40,9 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# skip tests requiring model downloads if HF_TOKEN is not set # skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits # due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
# All of CPU tests are expected to be finished less than 40 mins. # All of CPU tests are expected to be finished less than 40 mins.
export -f cpu_tests export -f cpu_tests
timeout 2h bash -c cpu_tests timeout 2h bash -c cpu_tests
@@ -29,6 +29,7 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}" CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}" CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}" CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
log_section() { log_section() {
echo "--- $*" echo "--- $*"
@@ -106,6 +107,18 @@ install_cargo_nextest() {
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}" "cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
} }
install_cargo_llvm_cov() {
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
local toolchain
toolchain="$(rust_toolchain)"
rustup component add --toolchain "$toolchain" llvm-tools-preview
cargo binstall \
--no-confirm \
--force \
--secure \
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
}
install_uv() { install_uv() {
log_section "Installing uv ${UV_VERSION}" log_section "Installing uv ${UV_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf \ curl -L --proto '=https' --tlsv1.2 -sSf \
@@ -176,14 +189,41 @@ run_tests() {
setup_pyo3_python setup_pyo3_python
install_cargo_binstall install_cargo_binstall
install_cargo_nextest install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest" log_section "Running cargo nextest with Rust coverage"
cargo nextest run \ mkdir -p artifacts
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
cargo llvm-cov clean \
--manifest-path rust/Cargo.toml \
--profraw-only
set +e
cargo llvm-cov nextest \
--manifest-path rust/Cargo.toml \ --manifest-path rust/Cargo.toml \
--workspace \ --workspace \
--all-features \ --all-features \
--locked \ --locked \
--no-fail-fast --no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
} }
install_protoc install_protoc
+182
View File
@@ -0,0 +1,182 @@
#!/bin/sh
RUST_CODECOV_VERSION="v11.3.1"
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
rust_coverage_repo_root() {
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
printf '%s\n' /vllm-workspace
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
else
git rev-parse --show-toplevel
fi
}
rust_coverage_start() {
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
mkdir -p "$RUST_COVERAGE_DIR"
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
export LLVM_PROFILE_FILE
trap rust_coverage_finalize 0
}
rust_coverage_objects() {
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
python3 - "$rust_cov_objects_manifest" <<'PY'
from pathlib import Path
import sys
for relative in Path(sys.argv[1]).read_text().splitlines():
for entry in sys.path:
path = Path(entry or ".").resolve() / relative
if path.is_file():
print(path)
break
else:
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
PY
}
rust_coverage_collect() {
rust_cov_collect_flag=${1:?coverage flag is required}
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
rust_cov_collect_primary=
set --
while IFS= read -r rust_cov_collect_object; do
if [ -z "$rust_cov_collect_primary" ]; then
rust_cov_collect_primary=$rust_cov_collect_object
else
set -- "$@" "--object=$rust_cov_collect_object"
fi
done <<EOF
$rust_cov_collect_objects
EOF
llvm-profdata merge \
-sparse \
"$RUST_COVERAGE_DIR"/*.profraw \
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
llvm-cov export \
"$rust_cov_collect_primary" \
"$@" \
--format=lcov \
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
> "$rust_cov_collect_lcov" || return 1
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
export RUST_COVERAGE_LCOV
}
rust_coverage_upload() {
rust_cov_upload_lcov=${1:?LCOV path is required}
rust_cov_upload_flag=${2:?coverage flag is required}
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
if [ "$(uname -m)" != "x86_64" ]; then
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
return 1
fi
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|| return 1
curl -fsSL \
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
| sha256sum -c - || return 1
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
rust_cov_upload_slug="vllm-project/vllm"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
case "$rust_cov_upload_slug" in
*/*) ;;
*) rust_cov_upload_slug="vllm-project/vllm" ;;
esac
fi
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
if [ -z "${CODECOV_TOKEN:-}" ]; then
# Codecov accepts tokenless public uploads on unprotected branch names.
# A colon-separated prefix keeps feature-branch and fork uploads from
# requiring a repository secret.
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
else
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
fi
fi
set --
set -- "$@" upload-process
set -- "$@" --file "$rust_cov_upload_lcov"
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
# source-line fix scanning, which is unrelated to path mapping.
set -- "$@" --disable-search --disable-file-fixes
set -- "$@" --fail-on-error --git-service github
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
set -- "$@" --branch "$rust_cov_upload_branch"
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
set -- "$@" --slug "$rust_cov_upload_slug"
set -- "$@" --flag "$rust_cov_upload_flag"
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
set -- "$@" --dir "$rust_cov_upload_repo_root"
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
fi
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
# E2E steps run from tests/, so execute from the repository root to resolve
# codecov.yml and repository paths consistently.
(
cd "$rust_cov_upload_repo_root" || exit 1
"$rust_cov_upload_codecov_dir/codecov" "$@"
) >"$rust_cov_upload_log" 2>&1
rust_cov_upload_rc=$?
cat "$rust_cov_upload_log"
# v11.3.1 can log API failures while returning zero even with
# --fail-on-error. Preserve the strict CI contract explicitly.
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
echo "Codecov CLI reported an upload error" >&2
rust_cov_upload_rc=1
fi
rm -rf "$rust_cov_upload_codecov_dir"
return "$rust_cov_upload_rc"
}
rust_coverage_finalize() {
rust_cov_finalize_test_rc=$?
trap - 0
set +e
rust_coverage_collect "$RUST_COVERAGE_FLAG"
rust_cov_finalize_collect_rc=$?
rust_cov_finalize_upload_rc=0
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
rust_cov_finalize_upload_rc=$?
fi
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
exit "$rust_cov_finalize_test_rc"
fi
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
exit "$rust_cov_finalize_collect_rc"
fi
exit "$rust_cov_finalize_upload_rc"
}
@@ -10,7 +10,9 @@ steps:
- vllm/engine/arg_utils.py - vllm/engine/arg_utils.py
- vllm/config/model.py - vllm/config/model.py
- vllm/model_executor - vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor - tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands: commands:
- apt-get update && apt-get install -y curl libsodium23 - apt-get update && apt-get install -y curl libsodium23
@@ -34,7 +36,9 @@ steps:
- vllm/engine/arg_utils.py - vllm/engine/arg_utils.py
- vllm/config/model.py - vllm/config/model.py
- vllm/model_executor - vllm/model_executor
- vllm/model_executor/warmup
- tests/model_executor - tests/model_executor
- tests/model_executor/test_jit_warmup.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- vllm/_aiter_ops.py - vllm/_aiter_ops.py
- vllm/platforms/rocm.py - vllm/platforms/rocm.py
+30
View File
@@ -8,6 +8,11 @@ steps:
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
source_file_dependencies: source_file_dependencies:
- rust/ - rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/benchmarks/ - vllm/benchmarks/
- vllm/entrypoints/openai/ - vllm/entrypoints/openai/
- vllm/entrypoints/serve/ - vllm/entrypoints/serve/
@@ -23,6 +28,7 @@ steps:
- tests/entrypoints/openai/test_uds.py - tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py - tests/v1/sample/test_logprobs_e2e.py
commands: commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1 - export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)" - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
@@ -43,6 +49,11 @@ steps:
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
source_file_dependencies: source_file_dependencies:
- rust/ - rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/ - vllm/entrypoints/openai/
- vllm/entrypoints/serve/ - vllm/entrypoints/serve/
- vllm/v1/engine/ - vllm/v1/engine/
@@ -54,6 +65,7 @@ steps:
# - tests/entrypoints/serve/dev/test_sleep.py # - tests/entrypoints/serve/dev/test_sleep.py
- tests/entrypoints/serve/tokenize/test_tokenization.py - tests/entrypoints/serve/tokenize/test_tokenization.py
commands: commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1 - export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
@@ -72,10 +84,16 @@ steps:
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
source_file_dependencies: source_file_dependencies:
- rust/ - rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/ - vllm/entrypoints/openai/
- tests/utils.py - tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py - tests/entrypoints/openai/correctness/test_lmeval.py
commands: commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1 - export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
@@ -86,11 +104,17 @@ steps:
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
source_file_dependencies: source_file_dependencies:
- rust/ - rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/ - vllm/entrypoints/openai/
- vllm/tool_parsers/ - vllm/tool_parsers/
- tests/utils.py - tests/utils.py
- tests/tool_use/ - tests/tool_use/
commands: commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1 - export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
@@ -101,6 +125,11 @@ steps:
working_dir: "/vllm-workspace/tests" working_dir: "/vllm-workspace/tests"
source_file_dependencies: source_file_dependencies:
- rust/ - rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/distributed/ - vllm/distributed/
- vllm/engine/ - vllm/engine/
- vllm/executor/ - vllm/executor/
@@ -111,6 +140,7 @@ steps:
- tests/v1/distributed/test_hybrid_lb_dp.py - tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py - tests/v1/distributed/test_internal_lb_dp.py
commands: commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1 - export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0 - export NCCL_CUMEM_HOST_ENABLE=0
@@ -26,5 +26,7 @@ steps:
- rust-toolchain.toml - rust-toolchain.toml
- .buildkite/test_areas/rust_frontend_cargo.yaml - .buildkite/test_areas/rust_frontend_cargo.yaml
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh - .buildkite/scripts/run-rust-frontend-cargo-ci.sh
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
commands: commands:
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test - .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
+1
View File
@@ -47,6 +47,7 @@
# Rust Frontend # Rust Frontend
/rust/ @BugenZhao @njhill /rust/ @BugenZhao @njhill
/rust/src/bench @esmeetu
/build_rust.sh @BugenZhao @njhill /build_rust.sh @BugenZhao @njhill
/rust-toolchain.toml @BugenZhao @njhill /rust-toolchain.toml @BugenZhao @njhill
/.buildkite/test_areas/rust* @BugenZhao @njhill /.buildkite/test_areas/rust* @BugenZhao @njhill
+1
View File
@@ -257,3 +257,4 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers # Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
+1 -1
View File
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
- Tool calling and reasoning parsers - Tool calling and reasoning parsers
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support - OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
- Efficient multi-LoRA support for dense and MoE layers - Efficient multi-LoRA support for dense and MoE layers
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. - Support for NVIDIA GPUs, AMD GPUs, Intel GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
vLLM seamlessly supports 200+ model architectures on Hugging Face, including: vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
+37
View File
@@ -8,6 +8,8 @@
set -euo pipefail set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)" REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
CARGO_LLVM_COV_VERSION="0.8.7"
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
# Read the required toolchain from rust-toolchain.toml. # Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/') TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -30,4 +32,39 @@ else
PROFILE_ARG="--release" PROFILE_ARG="--release"
fi fi
rm -rf "$COVERAGE_TOOLS_DIR"
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
# Keep instrumented intermediates isolated when local builds switch modes.
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
cargo +"$TOOLCHAIN" install \
--locked \
--version "$CARGO_LLVM_COV_VERSION" \
cargo-llvm-cov
eval "$(
cargo +"$TOOLCHAIN" llvm-cov show-env \
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
--sh
)"
# Build scripts and proc macros can run during compilation. Their profiles
# are unrelated to runtime coverage and would otherwise pollute the tree.
export LLVM_PROFILE_FILE=/dev/null
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG" python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
--print target-libdir)")/bin"
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
fi
+3
View File
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
"csrc/cpu/layernorm.cpp" "csrc/cpu/layernorm.cpp"
"csrc/cpu/mla_decode.cpp" "csrc/cpu/mla_decode.cpp"
"csrc/cpu/pos_encoding.cpp" "csrc/cpu/pos_encoding.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp" "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
"csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp") "csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_attn.cpp"
"csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/torch_bindings.cpp" "csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files # TODO: Remove these files
"csrc/cpu/activation.cpp" "csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/utils.cpp" "csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp" "csrc/cpu/cpu_attn.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/torch_bindings.cpp" "csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files # TODO: Remove these files
+5 -5
View File
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused") set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
else() else()
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git") set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65") set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}") set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
if(NOT _qutlass_fc_root) if(NOT _qutlass_fc_root)
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}" CUDA_ARCHS "${QUTLASS_ARCHS}"
) )
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target( define_extension_target(
_qutlass_C _qutlass_C
DESTINATION vllm DESTINATION vllm
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
WITH_SOABI) WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1 QUTLASS_MINIMAL_BUILD=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC} TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1) CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
TORCH_TARGET_VERSION=0x020B000000000000ULL
USE_CUDA)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3> $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -39,7 +39,7 @@ else()
FetchContent_Declare( FetchContent_Declare(
vllm-flash-attn vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8 GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types # Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+13
View File
@@ -10,3 +10,16 @@ fixes:
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/" - "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/" - "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
- "/usr/lib/python3.*/site-packages/vllm/::vllm/" - "/usr/lib/python3.*/site-packages/vllm/::vllm/"
# Map Rust sources built in the E2E image and on Buildkite agents.
- "/workspace/rust/::rust/"
- "/var/lib/buildkite-agent/.*/rust/::rust/"
flags:
rust-unit:
paths:
- rust/
carryforward: false
rust-e2e:
paths:
- rust/
carryforward: false
+3 -1
View File
@@ -102,7 +102,9 @@ class TileGemm82 {
kv_cache_t* __restrict__ curr_b = b_tile; kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) { for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b); auto fp32_b_regs = load_b_pair_vec(curr_b);
auto fp32_b_0_reg = fp32_b_regs.first;
auto fp32_b_1_reg = fp32_b_regs.second;
float* __restrict__ curr_m_a = curr_a; float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) { vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
+8 -7
View File
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
reg.val[1] = fp16_to_fp32_bits(raw_lo); reg.val[1] = fp16_to_fp32_bits(raw_lo);
} }
float reduce_sum() const { float reduce_sum() const {
AliasReg ar; // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
ar.reg = reg; // Step 1: pairwise sum of the two 4-wide halves
float result = 0; __vector float s = vec_add(reg.val[0], reg.val[1]);
unroll_loop<int, VEC_ELEM_NUM>( // Step 2: rotate by 8 bytes (2 floats) and add
[&result, &ar](int i) { result += ar.values[i]; }); s = vec_add(s, vec_sld(s, s, 8));
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
return result; s = vec_add(s, vec_sld(s, s, 4));
return vec_extract(s, 0);
} }
FP32Vec8 exp() const { FP32Vec8 exp() const {
f32x4x2_t out; f32x4x2_t out;
+285
View File
@@ -0,0 +1,285 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
// mamba_kernels.hpp.
#include "cpu/mamba_kernels.hpp"
#include <ATen/ATen.h>
#include <torch/library.h>
#include <c10/util/Optional.h>
#include "cpu_types.hpp"
// ---------------------------------------------------------------------------
// causal_conv1d_update
// ---------------------------------------------------------------------------
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
bool do_silu = false;
if (activation.has_value()) {
const std::string& act = activation.value();
do_silu = (act == "silu" || act == "swish");
}
at::ScalarType dtype = x.scalar_type();
// Input x: contiguous in native dtype.
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
// conv_state: NEVER copy the full paged tensor just for layout reasons.
// If the dtype matches we work directly on conv_state (contiguous or not)
// by extracting strides and passing them to the kernel.
// Only a dtype-conversion copy is made when types differ (rare for BF16).
bool state_type_ok = (conv_state.scalar_type() == dtype);
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
// state_c and conv_state may be non-contiguous — that is intentional.
// Weight: coerce to same dtype if needed (should match in practice)
at::Tensor w_c =
(weight.scalar_type() != dtype)
? weight.to(dtype).contiguous()
: (weight.is_contiguous() ? weight : weight.contiguous());
// Bias stays float32 (small scalar, used only for fp32 accumulation)
at::Tensor bias_f32;
if (bias.has_value() && bias.value().defined())
bias_f32 = bias.value().to(at::kFloat).contiguous();
int64_t batch = x_c.size(0);
int64_t dim = x_c.size(1);
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
int64_t width = w_c.size(1);
int64_t state_len = state_c.size(2);
// Extract strides — works for contiguous AND non-contiguous (transposed)
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
// contiguous) stride(1): between conv channels (dim stride) stride(2):
// between state elements (=1 when contiguous, =dim when transposed)
int64_t stride_s_slot = state_c.stride(0);
int64_t stride_s_dim = state_c.stride(1);
int64_t stride_s_state = state_c.stride(2);
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
const int32_t* cache_idx_ptr = nullptr;
at::Tensor cache_idx_int;
if (conv_state_indices.has_value()) {
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
}
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<scalar_t>(), cache_idx_ptr,
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
do_silu);
});
// Write back only when a type-conversion copy was made.
// Layout-only non-contiguity is handled via strides above — no copy needed.
if (!state_type_ok) conv_state.copy_(state_c);
return out;
}
// ---------------------------------------------------------------------------
// selective_state_update
// ---------------------------------------------------------------------------
void selective_state_update_cpu_impl(
at::Tensor& state, // (nstates, nheads, dim, dstate)
const at::Tensor& x, // (N, nheads, dim)
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C, const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens) {
at::ScalarType state_type = state.scalar_type();
at::ScalarType input_type = x.scalar_type();
// x, B, C must be contiguous and match input_type
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_input(x);
at::Tensor B_in = ensure_input(B);
at::Tensor C_in = ensure_input(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
// A, D, dt_bias are float32 model parameters that arrive here as expanded
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
// We need just the scalar value per head as a (nheads,) 1-D array so that
// A_ptr[h] in the kernel correctly reads head h's value.
//
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
// → .select(2,0) → (nheads, head_dim) strides (1,0)
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
// No allocation, no type conversion (A is already float32).
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
at::Tensor D_f32, dt_bias_f32;
if (D.has_value() && D.value().defined())
D_f32 = to_per_head_1d_f32(D.value());
if (dt_bias.has_value() && dt_bias.value().defined())
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
// the type conversion so we convert head_dim x fewer elements.
at::Tensor dt_f32;
{
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
// take a zero-copy view of index 0 along that dim first.
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
}
int64_t nheads = state.size(1);
int64_t dim = state.size(2);
int64_t dstate = state.size(3);
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
? cu_seqlens.value().size(0) - 1
: x_in.size(0);
int64_t ngroups = B_in.size(1);
// Strides
int64_t stride_state_n = state.stride(0);
int64_t stride_state_h = state.stride(1);
int64_t stride_state_d = state.stride(2);
int64_t stride_x_n = x_in.stride(0);
int64_t stride_x_h = x_in.stride(1);
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
int64_t stride_BC_n = B_in.stride(0);
int64_t stride_BC_g = B_in.stride(1);
int64_t stride_out_n = out.stride(0);
int64_t stride_out_h = out.stride(1);
// Optional index pointers
auto get_int32_ptr =
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
return (opt.has_value() && opt.value().defined())
? opt.value().data_ptr<int32_t>()
: nullptr;
};
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
// Dispatch on (state_t, input_t, out_t): write directly into `out`
// without any intermediate float32 buffer.
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
using state_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
using input_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
using out_t = scalar_t;
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
nheads, ngroups, dim, dstate, dt_softplus);
});
});
});
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd_cpu
// ---------------------------------------------------------------------------
void mamba_chunk_scan_fwd_cpu_impl(
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
at::Tensor&
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
const at::Tensor& x, // [seqlen, nheads, headdim]
const at::Tensor&
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
const at::Tensor& A, // [nheads] float32
const at::Tensor& B, // [seqlen, ngroups, dstate]
const at::Tensor& C, // [seqlen, ngroups, dstate]
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
const at::Tensor& cu_seqlens // [batch+1] int32
) {
const at::ScalarType input_type = x.scalar_type();
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_contig(x);
at::Tensor B_in = ensure_contig(B);
at::Tensor C_in = ensure_contig(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
// A and D are float32 model parameters, potentially broadcast-expanded.
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_f32(A);
at::Tensor D_f32;
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
// Python.
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
const int64_t batch = final_states.size(0);
const int64_t nheads = final_states.size(1);
const int64_t headdim = final_states.size(2);
const int64_t dstate = final_states.size(3);
const int64_t ngroups = B_in.size(1);
TORCH_CHECK(final_states.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
TORCH_CHECK(out.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
"raw data_ptr)");
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
ngroups, headdim, dstate);
});
}
+382
View File
@@ -0,0 +1,382 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Fused CPU vector kernels for Mamba decode-step hotspots:
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
// - selective_state_update (SSM recurrence, single-step)
#pragma once
#include "cpu_types.hpp"
#include <cmath>
#include <cstring>
#include <cstdint>
#include <algorithm>
namespace mamba_cpu {
// ---------------------------------------------------------------------------
// causal_conv1d_update — templated for native BF16/FP32
//
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
// Explicit strides are passed so the kernel writes directly into the
// correct memory locations without making a contiguous copy of the full
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
//
// stride_s_slot = state.stride(0) — between cache slots
// stride_s_dim = state.stride(1) — between conv_dim channels
// stride_s_state = state.stride(2) — between state elements
//
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
// ---------------------------------------------------------------------------
template <typename scalar_t>
inline void causal_conv1d_update_kernel(
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
int64_t width, int64_t state_len, bool do_silu) {
#pragma omp parallel for
for (int64_t b = 0; b < batch; ++b) {
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
if (cache_idx == pad_slot_id) continue;
for (int64_t t = 0; t < seqlen; ++t) {
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
// Base of this slot in the (possibly non-contiguous) paged state
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_b[d * seqlen]);
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
const scalar_t* w = weight_ptr + d * width;
// Accumulate in float32 for precision
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
for (int64_t k = 0; k < state_len; ++k) {
acc += static_cast<float>(w[k]) *
static_cast<float>(sd[k * stride_s_state]);
}
acc += static_cast<float>(w[state_len]) * x_val;
// Shift state left and append new input.
// Use memmove when contiguous (stride==1); element loop otherwise.
if (stride_s_state == 1) {
if (state_len > 1)
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
} else {
for (int64_t k = 0; k < state_len - 1; ++k)
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
if (state_len > 0)
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
}
if (do_silu) {
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
: std::exp(acc) / (1.0f + std::exp(acc));
acc *= sigmoid;
}
out_b[d * seqlen] = static_cast<scalar_t>(acc);
}
}
}
}
// ---------------------------------------------------------------------------
// selective_state_update
//
// Template parameters:
// state_t - dtype of ssm_state cache (typically BFloat16)
// input_t - dtype of x, B, C (typically BFloat16)
// out_t - dtype of output tensor (typically BFloat16)
// Write directly — no float32 intermediate buffer needed.
//
// A, D, dt_bias are accepted as const float* (they are always float32
// model parameters in Mamba2). This eliminates the per-call float32→BF16
// conversion and the .contiguous() materialisation of the broadcast-expand.
//
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
// ---------------------------------------------------------------------------
template <typename state_t, typename input_t, typename out_t = float>
inline void selective_state_update_kernel(
state_t* __restrict__ state_ptr, int64_t stride_state_n,
int64_t stride_state_h, int64_t stride_state_d,
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
// A: (nheads,) float32 — scalar per head
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
// D: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ D_ptr,
// z: same shape as x (optional)
const input_t* __restrict__ z_ptr,
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
int64_t stride_out_n, int64_t stride_out_h,
const int32_t* __restrict__ state_batch_indices,
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
const int32_t* __restrict__ num_accepted_tokens,
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
using state_vec_t = vec_op::vec_t<state_t>;
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
int64_t nheads_per_group = nheads / ngroups;
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
int64_t bos, seq_len;
if (cu_seqlens != nullptr) {
bos = cu_seqlens[seq_idx];
seq_len = cu_seqlens[seq_idx + 1] - bos;
} else {
bos = seq_idx;
seq_len = 1;
}
int64_t state_read_idx = (state_batch_indices != nullptr)
? state_batch_indices[seq_idx]
: seq_idx;
if (state_read_idx == null_block_id) continue;
int64_t state_write_idx = (num_accepted_tokens == nullptr)
? ((dst_state_batch_indices != nullptr)
? dst_state_batch_indices[seq_idx]
: state_read_idx)
: -1;
state_t* s = state_ptr + state_read_idx * stride_state_n;
for (int64_t t = 0; t < seq_len; ++t) {
int64_t token_idx = bos + t;
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
// dt: (N, nheads) — one float per head per token
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
out_t* out_tok = out_ptr + token_idx * stride_out_n;
#pragma omp parallel for
for (int64_t h = 0; h < nheads; ++h) {
int64_t g = h / nheads_per_group;
const input_t* x_h = x_tok + h * stride_x_h;
const input_t* B_g = B_tok + g * stride_BC_g;
const input_t* C_g = C_tok + g * stride_BC_g;
out_t* out_h = out_tok + h * stride_out_h;
state_t* s_h = s + h * stride_state_h;
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
float dt_val = dt_tok[h];
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
if (dt_softplus) {
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
}
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
const input_t* z_h =
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
: nullptr;
vec_op::FP32Vec8 dt_vec(dt_val);
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
// and broadcast. This saves 7 redundant std::exp() calls that
// FP32Vec8::exp() would otherwise make on the broadcast vector.
const float dA_scalar = std::exp(A_val * dt_val);
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_h[d]);
vec_op::FP32Vec8 out_vec(0.0f);
state_t* s_hd = s_h + d * stride_state_d;
const input_t* B_g_base = B_g;
const input_t* C_g_base = C_g;
vec_op::FP32Vec8 x_vec(x_val);
// dBx = B * x * dt — same dA for all dstate (A is scalar)
// s_new = s * dA + B * x * dt
int64_t n = 0;
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
state_vec_t(s_new).save(s_hd + n);
out_vec = out_vec + s_new * C_v;
}
float out_val = out_vec.reduce_sum();
for (; n < dstate; ++n) {
// Reuse dA_scalar computed once per head — no exp() re-call
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
s_hd[n] = static_cast<state_t>(s_new);
out_val += s_new * static_cast<float>(C_g[n]);
}
if (D_ptr != nullptr) out_val += x_val * D_val;
if (z_h != nullptr) {
float z_val = static_cast<float>(z_h[d]);
float sigmoid = (z_val >= 0)
? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
out_val *= z_val * sigmoid;
}
out_h[d] = static_cast<out_t>(out_val);
}
}
if (num_accepted_tokens != nullptr &&
dst_state_batch_indices != nullptr) {
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
state_write_idx != state_read_idx) {
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd
//
// Prefill SSM recurrence for Mamba2 / SSD models.
//
// Key difference from selective_state_update_kernel (decode path):
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
// Each thread owns a (batch, head) slice and runs the entire token
// sequence without any per-token OpenMP synchronisation overhead.
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
//
// `dt` arrives already processed (float32, after bias + softplus + clamp)
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
//
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
// tensor, pre-initialised by the caller (zero or from initial_states).
// Each (b, h) slice is private to exactly one thread via collapse(2), so
// there are no write conflicts.
//
// D is treated as a scalar per head ([nheads] float32).
// ---------------------------------------------------------------------------
template <typename input_t>
inline void mamba_chunk_scan_fwd_kernel(
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
const float* __restrict__ A_ptr, // [nheads] f32
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
int64_t dstate) {
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
const int64_t nheads_per_group = nheads / ngroups;
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
// guarantee)
const int64_t stride_s_b = nheads * headdim * dstate;
const int64_t stride_s_h = headdim * dstate;
// stride_s_d = dstate, stride_s_n = 1
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t b = 0; b < batch; ++b) {
for (int64_t h = 0; h < nheads; ++h) {
const int64_t seq_start = cu_seqlens[b];
const int64_t seq_end = cu_seqlens[b + 1];
const int64_t g = h / nheads_per_group;
const float A_val = A_ptr[h];
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
for (int64_t t = seq_start; t < seq_end; ++t) {
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
const float* dt_h = dt_ptr + t * nheads + h;
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
const input_t* z_h = (z_ptr != nullptr)
? z_ptr + t * nheads * headdim + h * headdim
: nullptr;
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
const float dt_val = *dt_h;
const float dA_val = std::exp(A_val * dt_val);
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
const vec_op::FP32Vec8 dt_vec(dt_val);
for (int64_t d = 0; d < headdim; ++d) {
const float x_val = static_cast<float>(x_h[d]);
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
// Vectorised SSM update + readout over dstate:
// s_new = s * dA + x * dt * B
// y += s_new * C
int64_t n = 0;
vec_op::FP32Vec8 y_vec(0.0f);
const vec_op::FP32Vec8 x_vec(x_val);
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
const vec_op::FP32Vec8 s_v(s_bhd + n);
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
s_new.save(s_bhd + n);
y_vec = y_vec + s_new * C_v;
}
float y_val = y_vec.reduce_sum();
// Scalar tail for remaining dstate elements
for (; n < dstate; ++n) {
const float B_n = static_cast<float>(B_g[n]);
const float C_n = static_cast<float>(C_g[n]);
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
s_bhd[n] = s_new;
y_val += s_new * C_n;
}
// D skip connection (scalar per head)
if (D_ptr != nullptr) y_val += x_val * D_val;
// z gating: out = y * z * sigmoid(z) (SiLU)
if (z_h != nullptr) {
const float z_val = static_cast<float>(z_h[d]);
const float sigmoid =
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
y_val *= z_val * sigmoid;
}
out_h[d] = static_cast<input_t>(y_val);
}
}
}
}
}
} // namespace mamba_cpu
+50
View File
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping, torch::Tensor slot_mapping,
const int64_t block_size); const int64_t block_size);
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
void selective_state_update_cpu_impl(
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens);
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C,
const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const at::Tensor& cu_seqlens);
void init_cpu_memory_env(std::vector<int64_t> node_ids); void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils { namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()", "block_size) -> ()",
&compute_slot_mapping_kernel_impl); &compute_slot_mapping_kernel_impl);
// Mamba CPU kernels
ops.def(
"causal_conv1d_update_cpu_vec("
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
&causal_conv1d_update_cpu_impl);
ops.def(
"selective_state_update_cpu("
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
"SymInt null_block_id, Tensor(a13!) out, "
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
&selective_state_update_cpu_impl);
ops.def(
"mamba_chunk_scan_fwd_cpu("
"Tensor(a0!) out, Tensor(a1!) final_states, "
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
&mamba_chunk_scan_fwd_cpu_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels // Speculative decoding kernels
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
// AWQ zeros: (size_k // group_size, size_n // 8) // AWQ zeros: (size_k // group_size, size_n // 8)
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
int32_t group_size) { int32_t group_size) {
int32_t val = // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
int32_t zero = // same row (stride 1) instead of striding across rows (stride size_n/8).
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + int col = blockIdx.y * 32 + threadIdx.x;
blockIdx.y]; if (col >= size_n / 8) return;
(void)size_k;
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
int32_t new_val = 0; int32_t new_val = 0;
#pragma unroll #pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
zero >>= 4; zero >>= 4;
} }
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; output[blockIdx.x * (size_n / 8) + col] = new_val;
} }
torch::stable::Tensor marlin_int4_fp8_preprocess( torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
"qweight.size(0) % qzeros.size(0) != 0"); "qweight.size(0) % qzeros.size(0) != 0");
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
dim3 blocks(size_k / 32, size_n / 8); dim3 blocks(size_k, (size_n / 8 + 31) / 32);
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>( marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()), reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
reinterpret_cast<int32_t*>(output.mutable_data_ptr()), reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
+11
View File
@@ -294,6 +294,9 @@ FROM base AS rust-build
ARG BUILD_OS ARG BUILD_OS
ARG USE_SCCACHE ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT ARG SCCACHE_ENDPOINT
# Temporary default for the initial CI validation. Set this back to 0 when
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
ARG VLLM_RUST_COVERAGE=1
# Install native tools needed only for Rust/protoc builds. # Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \ RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -902,6 +905,14 @@ COPY ./vllm/collect_env.py .
# note that this uses vllm installed by `pip` # note that this uses vllm installed by `pip`
FROM vllm-base AS test FROM vllm-base AS test
COPY --from=rust-build \
/workspace/rust-coverage-tools/ \
/opt/vllm-rust-coverage/
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
ENV LLVM_PROFILE_FILE=/dev/null
ADD . /vllm-workspace/ ADD . /vllm-workspace/
ARG PYTHON_VERSION ARG PYTHON_VERSION
+3
View File
@@ -46,6 +46,9 @@
"TORCH_CUDA_ARCH_LIST": { "TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0" "default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
}, },
"VLLM_RUST_COVERAGE": {
"default": "1"
},
"MAX_JOBS": { "MAX_JOBS": {
"default": "2" "default": "2"
}, },
+2 -2
View File
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 | | `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x | | `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 | | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x | | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any | | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 | | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
+2 -2
View File
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py) - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu) - CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
### SiLU+Mul + Quantization (`fuse_act_quant`) ### SiLU+Mul + Quantization (`fuse_act_quant`)
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) - Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
### RMSNorm + Padding (`fuse_act_padding`) ### RMSNorm + Padding (`fuse_act_padding`)
+4 -3
View File
@@ -68,13 +68,14 @@ vllm serve <model> \
| --- | --- | --- | --- | --- | | --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. | | `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). | | `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. | | `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. | | `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`. | | `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. | | `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). | | `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. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `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). | | `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 ## Secondary Tiers
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
- `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. - `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. - 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. - `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
- 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. - 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. - 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.
@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers). - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)): - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue. - Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
```bash ```bash
git clone https://github.com/vllm-project/vllm.git git clone https://github.com/vllm-project/vllm.git
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
--8<-- [end:build-wheel-from-source] --8<-- [end:build-wheel-from-source]
--8<-- [start:pre-built-images] --8<-- [start:pre-built-images]
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm). vLLM offers official Docker images for deployment.
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
```bash
docker run --rm \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
vllm/vllm-openai-xpu:<tag> \
--model Qwen/Qwen3-0.6B
```
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
???+ console "Commands"
```bash
docker run --rm -it \
--network=host \
--device /dev/dri:/dev/dri \
-v /dev/dri/by-path:/dev/dri/by-path \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--ipc=host \
--privileged \
--entrypoint /bin/bash \
vllm/vllm-openai-xpu:<tag>
```
--8<-- [end:pre-built-images] --8<-- [end:pre-built-images]
--8<-- [start:build-image-from-source] --8<-- [start:build-image-from-source]
+9
View File
@@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform:
!!! tip !!! tip
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds. A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
=== "Intel GPU"
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
!!! tip
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
=== "Google TPU" === "Google TPU"
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package. To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
@@ -31,10 +31,8 @@
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | | | THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | | | chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | | | 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | | | deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | | | meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | | | Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | | | Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | | | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
- `response_format`: Format of the response ("json", "text") (optional) - `response_format`: Format of the response ("json", "text") (optional)
- `temperature`: Sampling temperature between 0 and 1 (optional) - `temperature`: Sampling temperature between 0 and 1 (optional)
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182). For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
**Response Format:** **Response Format:**
+30 -12
View File
@@ -3446,7 +3446,6 @@ dependencies = [
"base64 0.22.1", "base64 0.22.1",
"bytes", "bytes",
"encoding_rs", "encoding_rs",
"futures-channel",
"futures-core", "futures-core",
"futures-util", "futures-util",
"h2", "h2",
@@ -4876,6 +4875,7 @@ dependencies = [
"futures-core", "futures-core",
"pin-project-lite", "pin-project-lite",
"tokio", "tokio",
"tokio-util",
] ]
[[package]] [[package]]
@@ -4937,9 +4937,9 @@ dependencies = [
[[package]] [[package]]
name = "tonic" name = "tonic"
version = "0.14.5" version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
dependencies = [ dependencies = [
"async-trait", "async-trait",
"axum", "axum",
@@ -4966,9 +4966,9 @@ dependencies = [
[[package]] [[package]]
name = "tonic-build" name = "tonic-build"
version = "0.14.5" version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734" checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
dependencies = [ dependencies = [
"prettyplease", "prettyplease",
"proc-macro2", "proc-macro2",
@@ -4977,10 +4977,23 @@ dependencies = [
] ]
[[package]] [[package]]
name = "tonic-prost" name = "tonic-health"
version = "0.14.5" version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
dependencies = [
"prost",
"tokio",
"tokio-stream",
"tonic",
"tonic-prost",
]
[[package]]
name = "tonic-prost"
version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
dependencies = [ dependencies = [
"bytes", "bytes",
"prost", "prost",
@@ -4989,9 +5002,9 @@ dependencies = [
[[package]] [[package]]
name = "tonic-prost-build" name = "tonic-prost-build"
version = "0.14.5" version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a" checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
dependencies = [ dependencies = [
"prettyplease", "prettyplease",
"proc-macro2", "proc-macro2",
@@ -5484,10 +5497,13 @@ dependencies = [
"serde", "serde",
"serde_json", "serde_json",
"thiserror 2.0.18", "thiserror 2.0.18",
"thiserror-ext",
"tiktoken-rs 0.9.1", "tiktoken-rs 0.9.1",
"tokenizers", "tokenizers",
"tokio", "tokio",
"tokio-stream", "tokio-stream",
"tracing",
"tracing-subscriber",
"url", "url",
"uuid", "uuid",
] ]
@@ -5560,6 +5576,7 @@ dependencies = [
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
"uuid", "uuid",
"vllm-bench",
"vllm-chat", "vllm-chat",
"vllm-engine-core-client", "vllm-engine-core-client",
"vllm-managed-engine", "vllm-managed-engine",
@@ -5729,6 +5746,7 @@ dependencies = [
"tokio-stream", "tokio-stream",
"tokio-util", "tokio-util",
"tonic", "tonic",
"tonic-health",
"tonic-prost", "tonic-prost",
"tonic-prost-build", "tonic-prost-build",
"tower", "tower",
@@ -6320,9 +6338,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
[[package]] [[package]]
name = "xgrammar-structural-tag" name = "xgrammar-structural-tag"
version = "0.1.0+xgrammar.0.2.2.4d145cc" version = "0.2.0+xgrammar.0.2.4.dd729e7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6" checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
dependencies = [ dependencies = [
"auto_impl", "auto_impl",
"serde", "serde",
+7 -5
View File
@@ -118,10 +118,11 @@ tokio = { version = "1.47.1", features = [
tokio-openssl = "0.6" tokio-openssl = "0.6"
tokio-stream = "0.1" tokio-stream = "0.1"
tokio-util = { version = "0.7.18", features = ["rt"] } tokio-util = { version = "0.7.18", features = ["rt"] }
tonic = "0.14.5" tonic = "0.14.6"
tonic-build = "0.14.5" tonic-build = "0.14.6"
tonic-prost = "0.14.5" tonic-health = "0.14.6"
tonic-prost-build = "0.14.5" tonic-prost = "0.14.6"
tonic-prost-build = "0.14.6"
tool-parser = "1.2.0" tool-parser = "1.2.0"
tower = { version = "0.5.3", features = ["util"] } tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.8", features = ["cors", "trace"] } tower-http = { version = "0.6.8", features = ["cors", "trace"] }
@@ -132,6 +133,7 @@ trait-set = "0.3.0"
url = "2.5.7" url = "2.5.7"
uuid = { version = "1.22.0", features = ["v4"] } uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] } validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" } vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" } vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" } vllm-llm = { path = "src/llm" }
@@ -142,7 +144,7 @@ vllm-server = { path = "src/server" }
vllm-text = { path = "src/text" } vllm-text = { path = "src/text" }
vllm-tokenizer = { path = "src/tokenizer" } vllm-tokenizer = { path = "src/tokenizer" }
winnow = { version = "1.0.2", features = ["simd"] } winnow = { version = "1.0.2", features = ["simd"] }
xgrammar-structural-tag = "0.1.0" xgrammar-structural-tag = "0.2.0"
zeromq = { version = "0.6.0", default-features = false, features = [ zeromq = { version = "0.6.0", default-features = false, features = [
"tokio-runtime", "tokio-runtime",
"all-transport", "all-transport",
+4 -1
View File
@@ -20,16 +20,19 @@ mimalloc.workspace = true
rand.workspace = true rand.workspace = true
rand_distr.workspace = true rand_distr.workspace = true
rayon.workspace = true rayon.workspace = true
reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] } reqwest = { workspace = true, features = ["json", "stream", "http2"] }
rlimit.workspace = true rlimit.workspace = true
rustc-hash.workspace = true rustc-hash.workspace = true
serde = { workspace = true, features = ["rc"] } serde = { workspace = true, features = ["rc"] }
serde_json = { workspace = true, features = ["raw_value"] } serde_json = { workspace = true, features = ["raw_value"] }
thiserror.workspace = true thiserror.workspace = true
thiserror-ext.workspace = true
tiktoken-rs.workspace = true tiktoken-rs.workspace = true
tokenizers.workspace = true tokenizers.workspace = true
tokio.workspace = true tokio.workspace = true
tokio-stream.workspace = true tokio-stream.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true url.workspace = true
uuid.workspace = true uuid.workspace = true
+8 -7
View File
@@ -158,9 +158,10 @@ impl PoolingBackend {
// (mirrors Python async_request_vllm_rerank). // (mirrors Python async_request_vllm_rerank).
if let Some(ref list) = input.prompt_list { if let Some(ref list) = input.prompt_list {
if list.len() < 2 { if list.len() < 2 {
eprintln!( tracing::warn!(
"WARNING: vllm-rerank request has no documents \ backend = "vllm-rerank",
(prompt_list needs [query, doc, ...])" inputs = list.len(),
"rerank request has no documents"
); );
} }
let query = list.first().map(|s| s.as_ref()).unwrap_or(""); let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -175,10 +176,10 @@ impl PoolingBackend {
// Legacy path: text prompt as query, documents via --extra-body. // Legacy path: text prompt as query, documents via --extra-body.
let query = input.prompt.as_ref(); let query = input.prompt.as_ref();
if query.is_empty() && input.prompt_token_ids.is_some() { if query.is_empty() && input.prompt_token_ids.is_some() {
eprintln!( tracing::warn!(
"WARNING: vllm-rerank received empty query (random dataset uses \ backend = "vllm-rerank",
token IDs only). Use --dataset-name random-rerank for meaningful \ dataset = "random",
rerank benchmarks." "rerank request has an empty query; use the random-rerank dataset"
); );
} }
serde_json::json!({ serde_json::json!({
+114 -84
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Instant; use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle}; use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend}; use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
@@ -72,12 +73,12 @@ pub fn pre_resolve_dns(
v4.extend(v6); v4.extend(v6);
if !v4.is_empty() { if !v4.is_empty() {
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect(); let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
println!("Pre-resolved {host} -> {ips:?}"); tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
builder = builder.resolve_to_addrs(host, &v4); builder = builder.resolve_to_addrs(host, &v4);
} }
} }
Err(e) => { Err(e) => {
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}"); tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
} }
} }
@@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let (model_id, model_name) = if let Some(ref m) = config.model { let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone()) (m.clone(), config.model_name.clone())
} else { } else {
println!("Model not specified, fetching first model from server..."); tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) = let (name, id) =
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?; get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}"); tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name)) (id, Some(name))
}; };
@@ -358,10 +363,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None None
} else { } else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
println!("Loading tokenizer: {tid}"); tracing::info!(tokenizer = tid, "loading tokenizer");
let server_info = Some((config.base_url.as_str(), model_id.as_str())); let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?; let t =
println!("Tokenizer loaded successfully."); crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
Some(t) Some(t)
}; };
let has_tokenizer = tokenizer.is_some(); let has_tokenizer = tokenizer.is_some();
@@ -421,7 +426,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.num_prompts, config.random_batch_size, config.is_reranker, config.num_prompts, config.random_batch_size, config.is_reranker,
), ),
}; };
println!("Generating {dataset_label}..."); tracing::info!(
dataset = ?config.dataset_name,
prompts = config.num_prompts,
description = %dataset_label,
"generating benchmark dataset"
);
let gen_start = Instant::now(); let gen_start = Instant::now();
let mut input_requests = match config.dataset_name { let mut input_requests = match config.dataset_name {
@@ -472,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let path = match config.dataset_path.as_deref() { let path = match config.dataset_path.as_deref() {
Some(p) => p, Some(p) => p,
None => { None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded.as_str() downloaded.as_str()
} }
}; };
@@ -512,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None => { None => {
downloaded = crate::datasets::speed_bench::download_speed_bench( downloaded = crate::datasets::speed_bench::download_speed_bench(
config.speed_bench_config, config.speed_bench_config,
)?; )
.await?;
downloaded.as_str() downloaded.as_str()
} }
}; };
@@ -543,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.hf_subset.as_deref(), config.hf_subset.as_deref(),
config.hf_split.as_deref(), config.hf_split.as_deref(),
config.num_prompts, config.num_prompts,
)?; )
.await?;
crate::datasets::hf_dataset::load_hf_dataset( crate::datasets::hf_dataset::load_hf_dataset(
tok, tok,
&downloaded_path, &downloaded_path,
@@ -608,18 +620,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
}; };
let gen_elapsed = gen_start.elapsed(); let gen_elapsed = gen_start.elapsed();
println!( tracing::info!(
"Generated {} prompts in {:.2}s", prompts = input_requests.len(),
input_requests.len(), elapsed_seconds = gen_elapsed.as_secs_f64(),
gen_elapsed.as_secs_f64() "generated benchmark dataset"
); );
let filtered_count = let filtered_count =
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len); filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
if filtered_count > 0 { if filtered_count > 0 {
println!( tracing::info!(
"Filtered {filtered_count} prompt(s) above --max-model-len {}.", filtered_prompts = filtered_count,
config.max_model_len.unwrap() max_model_len = config.max_model_len.unwrap(),
"filtered prompts above maximum model length"
); );
} }
if input_requests.is_empty() { if input_requests.is_empty() {
@@ -670,7 +683,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Ready check // Ready check
if config.ready_check_timeout_sec > 0 { if config.ready_check_timeout_sec > 0 {
println!("Starting initial single prompt test run..."); tracing::info!("starting initial single-prompt test run");
let test_output = wait_for_endpoint( let test_output = wait_for_endpoint(
config.backend, config.backend,
&client, &client,
@@ -685,7 +698,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
test_output.error test_output.error
))); )));
} }
println!("Initial test run completed."); tracing::info!("initial single-prompt test run completed");
} }
// Verify and fix prompt token lengths against the server's /tokenize endpoint. // Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -703,12 +716,15 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
DatasetName::Random | DatasetName::PrefixRepetition DatasetName::Random | DatasetName::PrefixRepetition
); );
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() { if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
println!("Using prompt_token_ids, skipping server-side tokenizer verification."); tracing::info!(
reason = "prompt_token_ids",
"skipping server tokenizer verification"
);
} }
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() { if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id); let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
if is_tokenizer_verified(&cache_key) { if is_tokenizer_verified(&cache_key) {
println!("Tokenizer verified in previous run (cached), skipping verification."); tracing::info!(reason = "cached", "skipping server tokenizer verification");
} else { } else {
let num_special = let num_special =
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0); tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -723,14 +739,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await? .await?
{ {
SampleVerifyOutcome::Passed => { SampleVerifyOutcome::Passed => {
println!("Sample verification passed, skipping full verification."); tracing::info!("tokenizer sample verification passed");
mark_tokenizer_verified(&cache_key); mark_tokenizer_verified(&cache_key);
} }
SampleVerifyOutcome::Skipped(reason) => { SampleVerifyOutcome::Skipped(reason) => {
println!("Server /tokenize unavailable ({reason}), skipping verification."); tracing::warn!(
reason = %reason,
"server tokenizer unavailable; skipping prompt verification"
);
} }
SampleVerifyOutcome::Mismatch => { SampleVerifyOutcome::Mismatch => {
println!("Sample verification found mismatch, running full verify+fix..."); tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
match verify_and_fix_prompt_lengths( match verify_and_fix_prompt_lengths(
&client, &client,
&config.base_url, &config.base_url,
@@ -742,16 +761,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await .await
{ {
Ok(()) => { Ok(()) => {
println!( tracing::info!(
"All {} prompts verified: exact token length match.", prompts = input_requests.len(),
input_requests.len() "verified exact prompt token lengths"
); );
mark_tokenizer_verified(&cache_key); mark_tokenizer_verified(&cache_key);
} }
Err(BenchError::TokenizeUnavailable(reason)) => { Err(BenchError::TokenizeUnavailable(reason)) => {
println!( tracing::warn!(
"Server /tokenize became unavailable during verification \ reason = %reason,
({reason}); proceeding with client-side token counts." "server tokenizer became unavailable; using client token counts"
); );
} }
Err(e) => return Err(e), Err(e) => return Err(e),
@@ -763,7 +782,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Warmup // Warmup
if config.num_warmups > 0 { if config.num_warmups > 0 {
println!("Warming up with {} requests...", config.num_warmups); tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
run_warmup( run_warmup(
config.backend, config.backend,
&client, &client,
@@ -776,7 +795,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.disable_tqdm, config.disable_tqdm,
) )
.await; .await;
println!("Warmup run completed."); tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
} }
// Start profiler if requested (immediate mode — no batch threshold) // Start profiler if requested (immediate mode — no batch threshold)
@@ -814,28 +833,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let spec_decode_before = let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() { if spec_decode_before.is_some() {
println!("Speculative decoding detected, will collect metrics."); tracing::info!("detected speculative decoding; collecting metrics");
} }
// Main benchmark // Main benchmark
println!("Starting main benchmark run...");
let distribution = if config.burstiness == 1.0 { let distribution = if config.burstiness == 1.0 {
"Poisson process" "Poisson process"
} else { } else {
"Gamma distribution" "Gamma distribution"
}; };
println!( tracing::info!(
"Traffic request rate: {}", request_rate = config.request_rate,
if config.request_rate.is_infinite() { burstiness = config.burstiness,
"inf".to_string() distribution,
} else { max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
format!("{}", config.request_rate) prompts = config.num_prompts,
} "starting main benchmark run"
);
println!("Burstiness factor: {} ({distribution})", config.burstiness);
println!(
"Maximum request concurrency: {}",
config.max_concurrency.unwrap_or(config.num_prompts)
); );
// Pre-assign LoRA adapters to each request (None when --lora-modules not set). // Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -847,11 +860,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
); );
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) { if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect(); let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!( tracing::info!(
"LoRA adapters ({}): {:?} [assignment={:?}]", adapters = modules.len(),
modules.len(), names = ?names,
names, assignment = ?config.lora_assignment,
config.lora_assignment "assigned LoRA adapters"
); );
} }
@@ -1125,7 +1138,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
if let Some((cancel_tx, task)) = profile_task { if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(()); let _ = cancel_tx.send(());
if let Err(e) = task.await { if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}"); tracing::error!(error = %e.as_report(), "profiler background task failed");
} }
} }
@@ -1289,12 +1302,14 @@ pub(crate) async fn start_profiler_immediate(
base_url: &str, base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>, extra_headers: &Option<std::collections::HashMap<String, String>>,
) { ) {
println!("Starting profiler...");
let profile_url = format!("{base_url}/start_profile"); let profile_url = format!("{base_url}/start_profile");
tracing::info!(url = %profile_url, "starting profiler");
match send_profile_request(client, &profile_url, extra_headers).await { match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => println!("Profiler started"), Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"), Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"), Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
}
} }
} }
@@ -1304,12 +1319,14 @@ pub(crate) async fn stop_profiler_immediate(
base_url: &str, base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>, extra_headers: &Option<std::collections::HashMap<String, String>>,
) { ) {
println!("Stopping profiler...");
let profile_url = format!("{base_url}/stop_profile"); let profile_url = format!("{base_url}/stop_profile");
tracing::info!(url = %profile_url, "stopping profiler");
match send_profile_request(client, &profile_url, extra_headers).await { match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => println!("Profiler stopped"), Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
}
} }
} }
@@ -1371,25 +1388,30 @@ pub(crate) async fn profile_on_batch_threshold(
duration_secs: f64, duration_secs: f64,
mut cancel_rx: tokio::sync::oneshot::Receiver<()>, mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
) { ) {
println!( tracing::info!(
"Waiting for batch size >= {threshold} before starting profiler \ threshold,
(will capture {duration_secs}s)..." duration_seconds = duration_secs,
"waiting for profiler batch threshold"
); );
loop { loop {
if let Some(running) = fetch_num_requests_running(client, base_url).await if let Some(running) = fetch_num_requests_running(client, base_url).await
&& running >= threshold && running >= threshold
{ {
println!("Batch size {running} >= {threshold}, starting profiler..."); tracing::info!(
running_requests = running,
threshold,
"profiler batch threshold reached"
);
break; break;
} }
// Wait 500ms or until the benchmark signals cancellation // Wait 500ms or until the benchmark signals cancellation
tokio::select! { tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = &mut cancel_rx => { _ = &mut cancel_rx => {
eprintln!( tracing::warn!(
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \ threshold,
profiling skipped." "benchmark finished before profiler batch threshold; skipping profiling"
); );
return; return;
} }
@@ -1398,13 +1420,13 @@ pub(crate) async fn profile_on_batch_threshold(
let start_url = format!("{base_url}/start_profile"); let start_url = format!("{base_url}/start_profile");
match send_profile_request(client, &start_url, extra_headers).await { match send_profile_request(client, &start_url, extra_headers).await {
Ok(true) => println!("Profiler started"), Ok(true) => tracing::info!(url = %start_url, "profiler started"),
Ok(false) => { Ok(false) => {
eprintln!("WARNING: Profiler start request returned non-success"); tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
return; return;
} }
Err(e) => { Err(e) => {
eprintln!("WARNING: Failed to start profiler: {e}"); tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
return; return;
} }
} }
@@ -1413,15 +1435,17 @@ pub(crate) async fn profile_on_batch_threshold(
tokio::select! { tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {} _ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
_ = &mut cancel_rx => { _ = &mut cancel_rx => {
println!("Benchmark finished, stopping profiler early..."); tracing::info!("benchmark finished; stopping profiler early");
} }
} }
let stop_url = format!("{base_url}/stop_profile"); let stop_url = format!("{base_url}/stop_profile");
match send_profile_request(client, &stop_url, extra_headers).await { match send_profile_request(client, &stop_url, extra_headers).await {
Ok(true) => println!("Profiler stopped after capturing"), Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), Err(e) => {
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
}
} }
} }
@@ -1502,10 +1526,11 @@ async fn verify_and_fix_prompt_lengths(
let excess = tokens.len().saturating_sub(expected_input_len); let excess = tokens.len().saturating_sub(expected_input_len);
let compensate = if excess > 0 && last_excess == Some(excess) { let compensate = if excess > 0 && last_excess == Some(excess) {
if _iter == 1 { if _iter == 1 {
eprintln!( tracing::warn!(
"Prompt {i}: server consistently adds {excess} extra token(s) \ prompt_index = i,
(likely BOS), compensating target to {}.", extra_tokens = excess,
expected_input_len.saturating_sub(excess), adjusted_target = expected_input_len.saturating_sub(excess),
"server consistently adds prompt tokens; compensating verification target"
); );
} }
excess excess
@@ -1563,7 +1588,10 @@ async fn verify_and_fix_prompt_lengths(
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed); let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
if fc > 0 { if fc > 0 {
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence."); tracing::info!(
fixed_prompts = fc,
"fixed prompt lengths using server tokenizer"
);
} }
Ok(()) Ok(())
@@ -1818,7 +1846,7 @@ async fn sample_verify_prompts(
let tokenize_url = format!("{base_url}/tokenize"); let tokenize_url = format!("{base_url}/tokenize");
let api_key = std::env::var("OPENAI_API_KEY").ok(); let api_key = std::env::var("OPENAI_API_KEY").ok();
println!("Sampling {sample_size} prompts for verification..."); tracing::info!(sample_size, "sampling prompts for tokenizer verification");
for (i, request) in requests.iter().enumerate().take(sample_size) { for (i, request) in requests.iter().enumerate().take(sample_size) {
let tokens = match server_tokenize( let tokens = match server_tokenize(
@@ -1841,9 +1869,11 @@ async fn sample_verify_prompts(
let expected = request.prompt_len + num_special; let expected = request.prompt_len + num_special;
if tokens.len() != expected { if tokens.len() != expected {
println!( tracing::warn!(
"Prompt {i}: expected {expected} tokens, server returned {}", prompt_index = i,
tokens.len() expected_tokens = expected,
actual_tokens = tokens.len(),
"tokenizer verification sample mismatch"
); );
return Ok(SampleVerifyOutcome::Mismatch); return Ok(SampleVerifyOutcome::Mismatch);
} }
+4 -11
View File
@@ -3,8 +3,6 @@
use std::fmt; use std::fmt;
use clap::Parser;
/// Backend type for the benchmark endpoint. /// Backend type for the benchmark endpoint.
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendKind { pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
ShareGpt, ShareGpt,
#[value(name = "sonnet")] #[value(name = "sonnet")]
Sonnet, Sonnet,
#[value(name = "speed-bench")] #[value(name = "speed-bench", alias = "speed_bench")]
SpeedBench, SpeedBench,
#[value(name = "hf")] #[value(name = "hf")]
Hf, Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
} }
/// High-performance benchmark client for vLLM serving endpoints. /// High-performance benchmark client for vLLM serving endpoints.
#[derive(Parser, Debug, Clone)] #[derive(clap::Args, Debug, Clone)]
#[command( pub struct BenchServeArgs {
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
pub struct Cli {
/// The type of backend or endpoint to use for the benchmark. /// The type of backend or endpoint to use for the benchmark.
#[arg(long, default_value = "openai")] #[arg(long, default_value = "openai")]
pub backend: BackendKind, pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
pub lora_assignment: LoraAssignment, pub lora_assignment: LoraAssignment,
} }
impl Cli { impl BenchServeArgs {
/// Resolve the base URL from explicit --base-url or from --host/--port. /// Resolve the base URL from explicit --base-url or from --host/--port.
pub fn resolve_base_url(&self) -> String { pub fn resolve_base_url(&self) -> String {
if let Some(ref base) = self.base_url { if let Some(ref base) = self.base_url {
+230 -197
View File
@@ -4,7 +4,9 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig}; use crate::cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt}; use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
use crate::error::{BenchError, Result}; use crate::error::{BenchError, Result};
@@ -215,63 +217,63 @@ pub struct BenchConfig {
} }
impl BenchConfig { impl BenchConfig {
pub fn from_cli(cli: &Cli) -> Result<Self> { pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
if cli.burstiness <= 0.0 { if args.burstiness <= 0.0 {
return Err(BenchError::Config("Burstiness must be positive".into())); return Err(BenchError::Config("Burstiness must be positive".into()));
} }
if cli.num_prompts == 0 { if args.num_prompts == 0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--num-prompts must be at least 1".into(), "--num-prompts must be at least 1".into(),
)); ));
} }
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() { if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
return Err(BenchError::Config( return Err(BenchError::Config(
"--request-rate must be positive (or inf)".into(), "--request-rate must be positive (or inf)".into(),
)); ));
} }
if cli.max_model_len == Some(0) { if args.max_model_len == Some(0) {
return Err(BenchError::Config( return Err(BenchError::Config(
"--max-model-len must be at least 1".into(), "--max-model-len must be at least 1".into(),
)); ));
} }
let base_url = cli.resolve_base_url(); let base_url = args.resolve_base_url();
let api_url = cli.resolve_api_url(); let api_url = args.resolve_api_url();
let extra_headers = cli.parse_headers()?; let extra_headers = args.parse_headers()?;
let mut extra_body = cli.parse_extra_body()?; let mut extra_body = args.parse_extra_body()?;
// Merge sampling parameters into extra_body (matches Python behavior). // Merge sampling parameters into extra_body (matches Python behavior).
// Python collects non-None sampling params and merges them UNDER extra_body, // Python collects non-None sampling params and merges them UNDER extra_body,
// meaning extra_body keys take precedence over sampling params. // meaning extra_body keys take precedence over sampling params.
{ {
let mut sampling_params = serde_json::Map::new(); let mut sampling_params = serde_json::Map::new();
if let Some(v) = cli.top_p { if let Some(v) = args.top_p {
sampling_params.insert("top_p".into(), serde_json::json!(v)); sampling_params.insert("top_p".into(), serde_json::json!(v));
} }
if let Some(v) = cli.top_k { if let Some(v) = args.top_k {
sampling_params.insert("top_k".into(), serde_json::json!(v)); sampling_params.insert("top_k".into(), serde_json::json!(v));
} }
if let Some(v) = cli.min_p { if let Some(v) = args.min_p {
sampling_params.insert("min_p".into(), serde_json::json!(v)); sampling_params.insert("min_p".into(), serde_json::json!(v));
} }
if let Some(v) = cli.temperature { if let Some(v) = args.temperature {
sampling_params.insert("temperature".into(), serde_json::json!(v)); sampling_params.insert("temperature".into(), serde_json::json!(v));
} }
if let Some(v) = cli.frequency_penalty { if let Some(v) = args.frequency_penalty {
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v)); sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
} }
if let Some(v) = cli.presence_penalty { if let Some(v) = args.presence_penalty {
sampling_params.insert("presence_penalty".into(), serde_json::json!(v)); sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
} }
if let Some(v) = cli.repetition_penalty { if let Some(v) = args.repetition_penalty {
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v)); sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
} }
if !sampling_params.is_empty() { if !sampling_params.is_empty() {
if !cli.backend.is_openai_compatible() { if !args.backend.is_openai_compatible() {
return Err(BenchError::Config( return Err(BenchError::Config(
"Sampling parameters are only supported by openai-compatible backends." "Sampling parameters are only supported by openai-compatible backends."
.into(), .into(),
@@ -286,10 +288,18 @@ impl BenchConfig {
} }
Some(other) => { Some(other) => {
// extra_body was not an object — just use sampling params // extra_body was not an object — just use sampling params
eprintln!( let value_type = match &other {
"Warning: --extra-body is not a JSON object, sampling params may be lost" serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => unreachable!(),
};
tracing::warn!(
value_type,
"sampling parameters may be lost because --extra-body is not a JSON object"
); );
let _ = other;
sampling_params sampling_params
} }
None => sampling_params, None => sampling_params,
@@ -299,7 +309,7 @@ impl BenchConfig {
} }
// Parse metadata // Parse metadata
let metadata = match &cli.metadata { let metadata = match &args.metadata {
None => None, None => None,
Some(items) => { Some(items) => {
let mut pairs = Vec::new(); let mut pairs = Vec::new();
@@ -314,24 +324,24 @@ impl BenchConfig {
}; };
// Parse goodput SLOs // Parse goodput SLOs
let goodput = parse_goodput(&cli.goodput)?; let goodput = parse_goodput(&args.goodput)?;
// Parse ramp-up config // Parse ramp-up config
let ramp_up = parse_ramp_up(cli)?; let ramp_up = parse_ramp_up(args)?;
// Default percentile metrics based on backend type // Default percentile metrics based on backend type
let default_percentile_metrics = if cli.backend.is_pooling() { let default_percentile_metrics = if args.backend.is_pooling() {
"e2el" "e2el"
} else { } else {
"ttft,tpot,itl,e2el" "ttft,tpot,itl,e2el"
}; };
let percentile_metrics_str = let percentile_metrics_str =
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics); args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
let selected_percentile_metrics: Vec<String> = let selected_percentile_metrics: Vec<String> =
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect(); percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?; let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
let sweep_summary_percentiles = cli let sweep_summary_percentiles = args
.sweep_summary_percentiles .sweep_summary_percentiles
.as_deref() .as_deref()
.map(|raw| parse_percentiles(raw, true)) .map(|raw| parse_percentiles(raw, true))
@@ -344,38 +354,38 @@ impl BenchConfig {
selected_percentiles.push(90.0); selected_percentiles.push(90.0);
} }
let tokenizer_id = if cli.skip_tokenizer_init { let tokenizer_id = if args.skip_tokenizer_init {
None None
} else { } else {
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default()) args.tokenizer.clone().or_else(|| args.model.clone())
}; };
// Resolve input/output lengths // Resolve input/output lengths
let random_input_len = cli.resolved_random_input_len(); let random_input_len = args.resolved_random_input_len();
let random_output_len = cli.resolved_random_output_len(); let random_output_len = args.resolved_random_output_len();
let per_turn_input_len = cli.resolved_per_turn_input_len(); let per_turn_input_len = args.resolved_per_turn_input_len();
// Normalized multi-turn turn counts (computed in validation block below, defaults // Normalized multi-turn turn counts (computed in validation block below, defaults
// to num_turns if multi-turn mode is not active) // to num_turns if multi-turn mode is not active)
let mut multi_turn_min_turns = cli.multi_turn_num_turns; let mut multi_turn_min_turns = args.multi_turn_num_turns;
let mut multi_turn_max_turns = cli.multi_turn_num_turns; let mut multi_turn_max_turns = args.multi_turn_num_turns;
// For random datasets with openai-compatible backends, default to ignore_eos. // For random datasets with openai-compatible backends, default to ignore_eos.
// Exception: multi-turn mode, where ignore_eos causes unbounded context growth // Exception: multi-turn mode, where ignore_eos causes unbounded context growth
// across turns. Multi-turn uses min_tokens instead for output length control. // across turns. Multi-turn uses min_tokens instead for output length control.
// Pooling backends don't generate tokens, so ignore_eos is irrelevant. // Pooling backends don't generate tokens, so ignore_eos is irrelevant.
let ignore_eos = if cli.backend.is_pooling() { let ignore_eos = if args.backend.is_pooling() {
false false
} else { } else {
cli.ignore_eos args.ignore_eos
|| ((cli.dataset_name == DatasetName::Random || ((args.dataset_name == DatasetName::Random
|| cli.dataset_name == DatasetName::RandomMm) || args.dataset_name == DatasetName::RandomMm)
&& cli.backend.is_openai_compatible() && args.backend.is_openai_compatible()
&& !cli.multi_turn) && !args.multi_turn)
}; };
// Pooling backends don't support multi-turn // Pooling backends don't support multi-turn
if cli.backend.is_pooling() && cli.multi_turn { if args.backend.is_pooling() && args.multi_turn {
return Err(BenchError::Config( return Err(BenchError::Config(
"Pooling/embedding backends do not support --multi-turn".into(), "Pooling/embedding backends do not support --multi-turn".into(),
)); ));
@@ -383,7 +393,7 @@ impl BenchConfig {
// LoRA validation. Adapter names must be non-empty after trim; pooling // LoRA validation. Adapter names must be non-empty after trim; pooling
// backends are out of scope (vLLM LoRA routing is for generative paths). // backends are out of scope (vLLM LoRA routing is for generative paths).
let lora_modules = match cli.lora_modules.as_ref() { let lora_modules = match args.lora_modules.as_ref() {
None => None, None => None,
Some(names) => { Some(names) => {
if names.is_empty() { if names.is_empty() {
@@ -391,7 +401,7 @@ impl BenchConfig {
"--lora-modules requires at least one adapter name".into(), "--lora-modules requires at least one adapter name".into(),
)); ));
} }
if cli.backend.is_pooling() { if args.backend.is_pooling() {
return Err(BenchError::Config( return Err(BenchError::Config(
"--lora-modules is not supported for pooling/embedding backends".into(), "--lora-modules is not supported for pooling/embedding backends".into(),
)); ));
@@ -411,18 +421,18 @@ impl BenchConfig {
}; };
// Random-MM validation and config parsing // Random-MM validation and config parsing
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm { let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
if cli.backend != BackendKind::OpenaiChat { if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config( return Err(BenchError::Config(
"Multi-modal content (images) is only supported on 'openai-chat' backend." "Multi-modal content (images) is only supported on 'openai-chat' backend."
.into(), .into(),
)); ));
} }
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt( let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
&cli.random_mm_limit_mm_per_prompt, &args.random_mm_limit_mm_per_prompt,
)?; )?;
let buckets = let buckets =
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?; crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
(limit, buckets) (limit, buckets)
} else { } else {
(MmLimitPerPrompt::default(), Vec::new()) (MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +442,18 @@ impl BenchConfig {
// sonnet (uses built-in Shakespeare's sonnets). // sonnet (uses built-in Shakespeare's sonnets).
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1)) // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?; let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
// Batched inputs only make sense for pooling backends (the generation // Batched inputs only make sense for pooling backends (the generation
// backends send one prompt per request). // backends send one prompt per request).
if cli.random_batch_size == 0 { if args.random_batch_size == 0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--random-batch-size must be at least 1".into(), "--random-batch-size must be at least 1".into(),
)); ));
} }
if cli.random_batch_size > 1 if args.random_batch_size > 1
&& !cli.backend.is_pooling() && !args.backend.is_pooling()
&& cli.dataset_name != DatasetName::RandomRerank && args.dataset_name != DatasetName::RandomRerank
{ {
return Err(BenchError::Config( return Err(BenchError::Config(
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(), "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +461,16 @@ impl BenchConfig {
} }
// random-rerank validation (mirrors Python RandomDatasetForReranking) // random-rerank validation (mirrors Python RandomDatasetForReranking)
let is_reranker = !cli.no_reranker; let is_reranker = !args.no_reranker;
if cli.dataset_name == DatasetName::RandomRerank { if args.dataset_name == DatasetName::RandomRerank {
if !cli.backend.is_pooling() { if !args.backend.is_pooling() {
return Err(BenchError::Config( return Err(BenchError::Config(
"--dataset-name random-rerank requires an embeddings/pooling backend \ "--dataset-name random-rerank requires an embeddings/pooling backend \
(e.g. --backend vllm-rerank)" (e.g. --backend vllm-rerank)"
.into(), .into(),
)); ));
} }
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) { if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
return Err(BenchError::Config( return Err(BenchError::Config(
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \ "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
(the query is folded into the first batch slot)" (the query is folded into the first batch slot)"
@@ -470,8 +480,8 @@ impl BenchConfig {
} }
// Custom dataset validation // Custom dataset validation
if cli.dataset_name == DatasetName::Custom { if args.dataset_name == DatasetName::Custom {
match cli.dataset_path.as_deref() { match args.dataset_path.as_deref() {
None => { None => {
return Err(BenchError::Config( return Err(BenchError::Config(
"--dataset-path is required for --dataset-name custom \ "--dataset-path is required for --dataset-name custom \
@@ -486,38 +496,38 @@ impl BenchConfig {
} }
_ => {} _ => {}
} }
if !cli.skip_chat_template { if !args.skip_chat_template {
eprintln!( tracing::warn!(
"NOTE: client-side chat template rendering is not supported; custom \ dataset = "custom",
dataset prompts are sent raw (equivalent to --skip-chat-template)." "client-side chat template rendering is unsupported; sending prompts raw"
); );
} }
} }
// Prefix repetition validation // Prefix repetition validation
if cli.dataset_name == DatasetName::PrefixRepetition { if args.dataset_name == DatasetName::PrefixRepetition {
if cli.prefix_repetition_num_prefixes == 0 { if args.prefix_repetition_num_prefixes == 0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--prefix-repetition-num-prefixes must be at least 1".into(), "--prefix-repetition-num-prefixes must be at least 1".into(),
)); ));
} }
if cli.num_prompts < cli.prefix_repetition_num_prefixes { if args.num_prompts < args.prefix_repetition_num_prefixes {
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})", "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
cli.num_prompts, cli.prefix_repetition_num_prefixes args.num_prompts, args.prefix_repetition_num_prefixes
))); )));
} }
} }
// HF dataset validation // HF dataset validation
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() { if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
return Err(BenchError::Config( return Err(BenchError::Config(
"--dataset-path is required for --dataset-name hf \ "--dataset-path is required for --dataset-name hf \
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')" (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
.into(), .into(),
)); ));
} }
if let Some(len) = cli.hf_output_len if let Some(len) = args.hf_output_len
&& len == 0 && len == 0
{ {
return Err(BenchError::Config( return Err(BenchError::Config(
@@ -526,13 +536,13 @@ impl BenchConfig {
} }
// Multi-turn validation // Multi-turn validation
if cli.multi_turn { if args.multi_turn {
if cli.backend != BackendKind::OpenaiChat { if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config( return Err(BenchError::Config(
"--multi-turn requires --backend openai-chat".into(), "--multi-turn requires --backend openai-chat".into(),
)); ));
} }
if cli.multi_turn_num_turns == 0 { if args.multi_turn_num_turns == 0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--multi-turn-num-turns must be at least 1".into(), "--multi-turn-num-turns must be at least 1".into(),
)); ));
@@ -541,18 +551,18 @@ impl BenchConfig {
// Normalize and validate min/max turns. ShareGPT only consumes max_turns // Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the // (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply. // min/num/max coupling used for synthetic generation does not apply.
if cli.dataset_name == DatasetName::ShareGpt { if args.dataset_name == DatasetName::ShareGpt {
if cli.multi_turn_max_turns == 1 { if args.multi_turn_max_turns == 1 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
)); ));
} }
} else { } else {
(multi_turn_min_turns, multi_turn_max_turns) = (multi_turn_min_turns, multi_turn_max_turns) =
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) { match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns), (0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
(m, 0) => (m, cli.multi_turn_num_turns), (m, 0) => (m, args.multi_turn_num_turns),
(0, x) => (cli.multi_turn_num_turns, x), (0, x) => (args.multi_turn_num_turns, x),
(m, x) => (m, x), (m, x) => (m, x),
}; };
if multi_turn_min_turns < 1 { if multi_turn_min_turns < 1 {
@@ -568,15 +578,16 @@ impl BenchConfig {
} }
if ignore_eos { if ignore_eos {
eprintln!( tracing::warn!(
"WARNING: --ignore-eos is set with --multi-turn. The server may not \ ignore_eos,
respect output length limits, causing unbounded context growth." multi_turn = true,
"output length limits may be ignored, causing unbounded context growth"
); );
} }
// Validate prefix sharing ratios // Validate prefix sharing ratios
let pg = cli.multi_turn_prefix_global_ratio; let pg = args.multi_turn_prefix_global_ratio;
let pc = cli.multi_turn_prefix_conversation_ratio; let pc = args.multi_turn_prefix_conversation_ratio;
if !(0.0..=1.0).contains(&pg) { if !(0.0..=1.0).contains(&pg) {
return Err(BenchError::Config( return Err(BenchError::Config(
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(), "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +603,20 @@ impl BenchConfig {
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(), "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
)); ));
} }
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random { if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
return Err(BenchError::Config( return Err(BenchError::Config(
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(), "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
)); ));
} }
} }
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) { if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
"--steady-state-threshold must be in (0.0, 1.0], got {}", "--steady-state-threshold must be in (0.0, 1.0], got {}",
cli.steady_state_threshold args.steady_state_threshold
))); )));
} }
if let Some(mw) = cli.steady_state_min_window if let Some(mw) = args.steady_state_min_window
&& mw < 0.0 && mw < 0.0
{ {
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
@@ -613,122 +624,122 @@ impl BenchConfig {
))); )));
} }
if cli.profile_batch_threshold.is_some() && !cli.profile { if args.profile_batch_threshold.is_some() && !args.profile {
return Err(BenchError::Config( return Err(BenchError::Config(
"--profile-batch-threshold requires --profile".into(), "--profile-batch-threshold requires --profile".into(),
)); ));
} }
if cli.profile_duration <= 0.0 { if args.profile_duration <= 0.0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--profile-duration must be positive".into(), "--profile-duration must be positive".into(),
)); ));
} }
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 { if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
return Err(BenchError::Config( return Err(BenchError::Config(
"--profile-duration requires --profile-batch-threshold".into(), "--profile-duration requires --profile-batch-threshold".into(),
)); ));
} }
Ok(BenchConfig { Ok(BenchConfig {
backend: cli.backend, backend: args.backend,
base_url, base_url,
api_url, api_url,
model: cli.model.clone(), model: args.model.clone(),
model_name: cli.served_model_name.clone(), model_name: args.served_model_name.clone(),
tokenizer_id, tokenizer_id,
tokenizer_mode: cli.tokenizer_mode.clone(), tokenizer_mode: args.tokenizer_mode.clone(),
trust_remote_code: cli.trust_remote_code, trust_remote_code: args.trust_remote_code,
skip_tokenizer_init: cli.skip_tokenizer_init, skip_tokenizer_init: args.skip_tokenizer_init,
dataset_name: cli.dataset_name, dataset_name: args.dataset_name,
dataset_path: cli.dataset_path.clone(), dataset_path: args.dataset_path.clone(),
max_model_len: cli.max_model_len, max_model_len: args.max_model_len,
random_input_len, random_input_len,
random_output_len, random_output_len,
random_prefix_len: cli.random_prefix_len, random_prefix_len: args.random_prefix_len,
random_range_ratio, random_range_ratio,
random_batch_size: cli.random_batch_size, random_batch_size: args.random_batch_size,
is_reranker, is_reranker,
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len), custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len, prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len, prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes, prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
prefix_repetition_output_len: cli prefix_repetition_output_len: args
.output_len .output_len
.unwrap_or(cli.prefix_repetition_output_len), .unwrap_or(args.prefix_repetition_output_len),
random_cache_hit_fraction: cli.random_cache_hit_fraction, random_cache_hit_fraction: args.random_cache_hit_fraction,
random_cache_ratio: cli.random_cache_ratio, random_cache_ratio: args.random_cache_ratio,
sharegpt_output_len: cli.sharegpt_output_len, sharegpt_output_len: args.sharegpt_output_len,
sonnet_input_len: cli.sonnet_input_len, sonnet_input_len: args.sonnet_input_len,
sonnet_output_len: cli.sonnet_output_len, sonnet_output_len: args.sonnet_output_len,
sonnet_prefix_len: cli.sonnet_prefix_len, sonnet_prefix_len: args.sonnet_prefix_len,
no_oversample: cli.no_oversample, no_oversample: args.no_oversample,
disable_shuffle: cli.disable_shuffle, disable_shuffle: args.disable_shuffle,
num_prompts: cli.num_prompts, num_prompts: args.num_prompts,
request_rate: cli.request_rate, request_rate: args.request_rate,
burstiness: cli.burstiness, burstiness: args.burstiness,
max_concurrency: cli.max_concurrency, max_concurrency: args.max_concurrency,
steady_state_threshold: cli.steady_state_threshold, steady_state_threshold: args.steady_state_threshold,
steady_state_min_window: cli.steady_state_min_window, steady_state_min_window: args.steady_state_min_window,
no_steady_state: cli.no_steady_state, no_steady_state: args.no_steady_state,
disable_tqdm: cli.disable_tqdm, disable_tqdm: args.disable_tqdm,
num_warmups: cli.num_warmups, num_warmups: args.num_warmups,
profile: cli.profile, profile: args.profile,
profile_batch_threshold: cli.profile_batch_threshold, profile_batch_threshold: args.profile_batch_threshold,
profile_duration: cli.profile_duration, profile_duration: args.profile_duration,
save_result: cli.save_result, save_result: args.save_result,
save_detailed: cli.save_detailed, save_detailed: args.save_detailed,
append_result: cli.append_result, append_result: args.append_result,
result_dir: cli.result_dir.clone(), result_dir: args.result_dir.clone(),
result_filename: cli.result_filename.clone(), result_filename: args.result_filename.clone(),
seed: cli.seed, seed: args.seed,
ignore_eos, ignore_eos,
insecure: cli.insecure, insecure: args.insecure,
selected_percentile_metrics, selected_percentile_metrics,
selected_percentiles, selected_percentiles,
sweep_summary_percentiles, sweep_summary_percentiles,
label: cli.label.clone(), label: args.label.clone(),
logprobs: cli.logprobs, logprobs: args.logprobs,
request_id_prefix: cli.get_request_id_prefix(), request_id_prefix: args.get_request_id_prefix(),
ready_check_timeout_sec: cli.ready_check_timeout_sec, ready_check_timeout_sec: args.ready_check_timeout_sec,
extra_headers, extra_headers,
extra_body, extra_body,
metadata, metadata,
dry_run: cli.dry_run, dry_run: args.dry_run,
goodput, goodput,
ramp_up, ramp_up,
multi_turn: cli.multi_turn, multi_turn: args.multi_turn,
multi_turn_num_turns: cli.multi_turn_num_turns, multi_turn_num_turns: args.multi_turn_num_turns,
multi_turn_min_turns, multi_turn_min_turns,
multi_turn_max_turns, multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if cli.multi_turn sharegpt_multi_turn_max_turns: if args.multi_turn
&& cli.dataset_name == DatasetName::ShareGpt && args.dataset_name == DatasetName::ShareGpt
&& cli.multi_turn_max_turns != 0 && args.multi_turn_max_turns != 0
{ {
Some(cli.multi_turn_max_turns) Some(args.multi_turn_max_turns)
} else { } else {
None None
}, },
per_turn_input_len, per_turn_input_len,
multi_turn_concurrency: cli.multi_turn_concurrency, multi_turn_concurrency: args.multi_turn_concurrency,
multi_turn_delay_ms: cli.multi_turn_delay_ms, multi_turn_delay_ms: args.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio, multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio, multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
speed_bench_config: cli.speed_bench_config, speed_bench_config: args.speed_bench_config,
speed_bench_category: cli.speed_bench_category.clone(), speed_bench_category: args.speed_bench_category.clone(),
speed_bench_max_input_len: cli.speed_bench_max_input_len, speed_bench_max_input_len: args.speed_bench_max_input_len,
hf_split: cli.hf_split.clone(), hf_split: args.hf_split.clone(),
hf_subset: cli.hf_subset.clone(), hf_subset: args.hf_subset.clone(),
hf_output_len: cli.hf_output_len, hf_output_len: args.hf_output_len,
hf_text_column: cli.hf_text_column.clone(), hf_text_column: args.hf_text_column.clone(),
reset_prefix_cache: cli.reset_prefix_cache, reset_prefix_cache: args.reset_prefix_cache,
prompt_token_ids: cli.prompt_token_ids, prompt_token_ids: args.prompt_token_ids,
random_mm_base_items_per_request: cli.random_mm_base_items_per_request, random_mm_base_items_per_request: args.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio, random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
random_mm_limit, random_mm_limit,
random_mm_buckets, random_mm_buckets,
enable_multimodal_chat: cli.enable_multimodal_chat, enable_multimodal_chat: args.enable_multimodal_chat,
lora_modules, lora_modules,
lora_assignment: cli.lora_assignment, lora_assignment: args.lora_assignment,
}) })
} }
} }
@@ -811,17 +822,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
Ok(config) Ok(config)
} }
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> { fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
let strategy = match cli.ramp_up_strategy { let strategy = match args.ramp_up_strategy {
None => return Ok(None), None => return Ok(None),
Some(s) => s, Some(s) => s,
}; };
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| { let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into()) BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
})?; })?;
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| { let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into()) BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
})?; })?;
@@ -843,7 +854,21 @@ mod tests {
use clap::Parser; use clap::Parser;
use super::*; use super::*;
use crate::cli::Cli; use crate::cli::BenchServeArgs;
#[derive(Parser)]
struct TestCli {
#[command(flatten)]
args: BenchServeArgs,
}
fn parse_args<I, T>(args: I) -> BenchServeArgs
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString> + Clone,
{
TestCli::parse_from(args).args
}
fn base_multi_turn_args() -> Vec<&'static str> { fn base_multi_turn_args() -> Vec<&'static str> {
vec![ vec![
@@ -859,8 +884,8 @@ mod tests {
#[test] #[test]
fn test_prefix_sharing_defaults_to_zero() { fn test_prefix_sharing_defaults_to_zero() {
let args = base_multi_turn_args(); let args = base_multi_turn_args();
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0); assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0); assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
} }
@@ -874,8 +899,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio", "--multi-turn-prefix-conversation-ratio",
"0.8", "0.8",
]); ]);
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10); assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10); assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
} }
@@ -889,8 +914,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio", "--multi-turn-prefix-conversation-ratio",
"0.6", "0.6",
]); ]);
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
@@ -902,16 +927,16 @@ mod tests {
"--multi-turn-prefix-conversation-ratio", "--multi-turn-prefix-conversation-ratio",
"0.5", "0.5",
]); ]);
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
fn test_prefix_sharing_out_of_range_fails() { fn test_prefix_sharing_out_of_range_fails() {
let mut args = base_multi_turn_args(); let mut args = base_multi_turn_args();
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]); args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
@@ -928,8 +953,8 @@ mod tests {
"--multi-turn-prefix-global-ratio", "--multi-turn-prefix-global-ratio",
"0.1", "0.1",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
@@ -944,8 +969,8 @@ mod tests {
"--dataset-name", "--dataset-name",
"sharegpt", "sharegpt",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_max_turns, 3); assert_eq!(config.multi_turn_max_turns, 3);
assert_eq!(config.sharegpt_multi_turn_max_turns, None); assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +993,8 @@ mod tests {
"--multi-turn-max-turns", "--multi-turn-max-turns",
"2", "2",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
} }
@@ -987,8 +1012,8 @@ mod tests {
"--multi-turn-max-turns", "--multi-turn-max-turns",
"1", "1",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string(); let err = BenchConfig::from_args(&args).unwrap_err().to_string();
assert!( assert!(
err.contains("at least 2 for ShareGPT"), err.contains("at least 2 for ShareGPT"),
"expected ShareGPT-specific error, got: {err}" "expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1034,8 @@ mod tests {
"--multi-turn-max-turns", "--multi-turn-max-turns",
"20", "20",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
} }
@@ -1018,8 +1043,8 @@ mod tests {
#[test] #[test]
fn test_sweep_summary_percentiles_default_empty() { fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args(); let args = base_multi_turn_args();
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert!(config.sweep_summary_percentiles.is_empty()); assert!(config.sweep_summary_percentiles.is_empty());
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]); assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1059,8 @@ mod tests {
"--sweep-summary-percentiles", "--sweep-summary-percentiles",
"90,95,90", "90,95,90",
]); ]);
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]); assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]); assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1070,8 @@ mod tests {
fn test_invalid_sweep_summary_percentile_fails() { fn test_invalid_sweep_summary_percentile_fails() {
let mut args = base_multi_turn_args(); let mut args = base_multi_turn_args();
args.extend(["--sweep-summary-percentiles", "101"]); args.extend(["--sweep-summary-percentiles", "101"]);
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
@@ -1058,12 +1083,20 @@ mod tests {
"--max-model-len", "--max-model-len",
"4096", "4096",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
let config = BenchConfig::from_cli(&cli).unwrap(); let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.max_model_len, Some(4096)); assert_eq!(config.max_model_len, Some(4096));
} }
#[test]
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
let args = parse_args(["vllm-bench"]);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.tokenizer_id, None);
}
#[test] #[test]
fn test_zero_max_model_len_fails() { fn test_zero_max_model_len_fails() {
let args = vec![ let args = vec![
@@ -1073,9 +1106,9 @@ mod tests {
"--max-model-len", "--max-model-len",
"0", "0",
]; ];
let cli = Cli::parse_from(args); let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err()); assert!(BenchConfig::from_args(&args).is_err());
} }
#[test] #[test]
fn test_range_ratio_parse_float() { fn test_range_ratio_parse_float() {
+4 -2
View File
@@ -128,8 +128,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access. /// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind { fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None) TokenizerKind::Tiktoken(
.expect("gpt2 built-in tiktoken should always load without network") crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
} }
#[test] #[test]
+74 -37
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng}; use rand::{Rng, SeedableRng};
use super::SampleRequest; use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::error::{BenchError, Result}; use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind; use crate::tokenizer::TokenizerKind;
@@ -50,18 +51,19 @@ enum ColumnFormat {
/// Make a GET request with retry logic (3 retries with exponential backoff). /// Make a GET request with retry logic (3 retries with exponential backoff).
/// Returns the parsed JSON response. /// Returns the parsed JSON response.
fn get_with_retry( async fn get_with_retry(
client: &reqwest::blocking::Client, client: &reqwest::Client,
url: &str, url: &str,
label: &str, label: &str,
) -> Result<serde_json::Value> { ) -> Result<serde_json::Value> {
let max_retries = 3; let max_retries = 3;
for attempt in 0..=max_retries { for attempt in 0..=max_retries {
let resp = match client.get(url).send() { let resp = match client.get(url).send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if attempt < max_retries { if attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
.await;
continue; continue;
} }
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
@@ -80,7 +82,7 @@ fn get_with_retry(
} }
if status.is_server_error() && attempt < max_retries { if status.is_server_error() && attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
continue; continue;
} }
@@ -92,6 +94,7 @@ fn get_with_retry(
let data: serde_json::Value = resp let data: serde_json::Value = resp
.json() .json()
.await
.map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?; .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
return Ok(data); return Ok(data);
} }
@@ -105,7 +108,7 @@ fn get_with_retry(
/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization. /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected /// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
/// or the dataset is exhausted. /// or the dataset is exhausted.
pub fn download_hf_dataset( pub async fn download_hf_dataset(
dataset: &str, dataset: &str,
subset: Option<&str>, subset: Option<&str>,
split: Option<&str>, split: Option<&str>,
@@ -115,7 +118,7 @@ pub fn download_hf_dataset(
url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
let mut client_builder = let mut client_builder =
reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120)); reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
// Add HF_TOKEN auth header if available // Add HF_TOKEN auth header if available
if let Ok(token) = std::env::var("HF_TOKEN") { if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -138,7 +141,7 @@ pub fn download_hf_dataset(
// Call /info to discover available configs and splits // Call /info to discover available configs and splits
let info_url = let info_url =
format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}"); format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
let info = get_with_retry(&client, &info_url, "HF dataset /info")?; let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
let dataset_info = let dataset_info =
info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| { info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -201,7 +204,12 @@ pub fn download_hf_dataset(
(resolved_config, resolved_split) (resolved_config, resolved_split)
}; };
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})"); tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"resolved Hugging Face dataset"
);
// Check cache // Check cache
let dir = cache_dir(); let dir = cache_dir();
@@ -215,11 +223,16 @@ pub fn download_hf_dataset(
if cache_path.exists() { if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string(); let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset cached: {path_str}"); tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
return Ok((path_str, resolved_config, resolved_split)); return Ok((path_str, resolved_config, resolved_split));
} }
println!("Downloading HF dataset '{dataset}' from datasets-server..."); tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"downloading Hugging Face dataset"
);
let encoded_config: String = let encoded_config: String =
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect(); url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -229,6 +242,7 @@ pub fn download_hf_dataset(
let mut all_rows: Vec<serde_json::Value> = Vec::new(); let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize; let mut offset = 0usize;
let page_size = 100usize; let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop { loop {
let url = format!( let url = format!(
@@ -240,7 +254,7 @@ pub fn download_hf_dataset(
&length={page_size}" &length={page_size}"
); );
let data = get_with_retry(&client, &url, "HF dataset /rows")?; let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
let rows = data["rows"] let rows = data["rows"]
.as_array() .as_array()
@@ -260,14 +274,14 @@ pub fn download_hf_dataset(
offset += fetched; offset += fetched;
let total = data["num_rows_total"].as_u64().unwrap_or(0); let total = data["num_rows_total"].as_u64().unwrap_or(0);
eprint!("\r Fetched {offset}/{total} rows..."); progress.update(offset, total);
// Stop if we have enough rows or reached end of dataset // Stop if we have enough rows or reached end of dataset
if all_rows.len() >= num_rows_needed || fetched < page_size { if all_rows.len() >= num_rows_needed || fetched < page_size {
break; break;
} }
} }
eprintln!(); // newline after progress progress.finish();
if all_rows.is_empty() { if all_rows.is_empty() {
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
@@ -280,7 +294,12 @@ pub fn download_hf_dataset(
std::fs::write(&cache_path, &json_str)?; std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string(); let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset: {} rows saved to {path_str}", all_rows.len()); tracing::info!(
dataset,
rows = all_rows.len(),
path = %path_str,
"saved Hugging Face dataset"
);
Ok((path_str, resolved_config, resolved_split)) Ok((path_str, resolved_config, resolved_split))
} }
@@ -483,21 +502,31 @@ pub fn load_hf_dataset(
// Detect column format from first row // Detect column format from first row
let format = detect_column_format(&entries[0], text_column_override)?; let format = detect_column_format(&entries[0], text_column_override)?;
// Print detected format
match &format { match &format {
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"), ColumnFormat::Chat(col) => {
tracing::info!(
format = "chat",
column = col,
"detected Hugging Face dataset format"
);
}
ColumnFormat::Text { ColumnFormat::Text {
prompt_col, prompt_col,
output_col, output_col,
} => { } => {
let out_msg = output_col.as_deref().unwrap_or("none"); tracing::info!(
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}"); format = "text",
prompt_column = prompt_col,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
} }
ColumnFormat::Combined { cols, output_col } => { ColumnFormat::Combined { cols, output_col } => {
let out_msg = output_col.as_deref().unwrap_or("none"); tracing::info!(
println!( format = "combined",
"HF dataset: detected combined columns {:?}, output column: {out_msg}", prompt_columns = ?cols,
cols output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
); );
} }
} }
@@ -608,9 +637,10 @@ pub fn load_hf_dataset(
if len == 0 { 128 } else { len } if len == 0 { 128 } else { len }
} else { } else {
if !warned_no_output { if !warned_no_output {
eprintln!( tracing::warn!(
"WARNING: No output column detected and --hf-output-len not set. \ path = dataset_path,
Using default output length of 128 tokens." default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
); );
warned_no_output = true; warned_no_output = true;
} }
@@ -618,9 +648,10 @@ pub fn load_hf_dataset(
} }
} else { } else {
if !warned_no_output { if !warned_no_output {
eprintln!( tracing::warn!(
"WARNING: No output column detected and --hf-output-len not set. \ path = dataset_path,
Using default output length of 128 tokens." default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
); );
warned_no_output = true; warned_no_output = true;
} }
@@ -640,9 +671,11 @@ pub fn load_hf_dataset(
// Oversample if needed // Oversample if needed
if samples.len() < num_requests { if samples.len() < num_requests {
if no_oversample { if no_oversample {
println!( tracing::info!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})", dataset = "hf",
samples.len() samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
); );
} else if !samples.is_empty() { } else if !samples.is_empty() {
let original_len = samples.len(); let original_len = samples.len();
@@ -652,9 +685,11 @@ pub fn load_hf_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req); samples.push(req);
} }
println!( tracing::info!(
"Oversampled HF dataset from {original_len} to {} total samples.", dataset = "hf",
samples.len() original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
); );
} }
} }
@@ -1002,8 +1037,10 @@ mod tests {
/// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required). /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind { fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None) crate::tokenizer::TokenizerKind::Tiktoken(
.expect("gpt2 built-in tiktoken should always load without network") crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
} }
/// Write JSON data to a unique temp file and return the path string. /// Write JSON data to a unique temp file and return the path string.
+9 -6
View File
@@ -5,6 +5,7 @@ pub mod custom;
pub mod hf_dataset; pub mod hf_dataset;
pub mod multi_turn; pub mod multi_turn;
pub mod prefix_repetition; pub mod prefix_repetition;
mod progress;
pub mod random; pub mod random;
pub mod random_mm; pub mod random_mm;
pub mod random_rerank; pub mod random_rerank;
@@ -90,9 +91,10 @@ pub fn oversample_requests(
return; return;
} }
if no_oversample { if no_oversample {
println!( tracing::info!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})", samples = requests.len(),
requests.len() requested = num_requests,
"skipping dataset oversampling"
); );
return; return;
} }
@@ -103,9 +105,10 @@ pub fn oversample_requests(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
requests.push(req); requests.push(req);
} }
println!( tracing::info!(
"Oversampled requests from {original_len} to {} total samples.", original_samples = original_len,
requests.len() samples = requests.len(),
"oversampled dataset"
); );
} }
+29 -18
View File
@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i); conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
conversations.push(conv); conversations.push(conv);
} }
println!( tracing::info!(
"Oversampled multi-turn conversations from {original_len} to {} total.", original_conversations = original_len,
conversations.len() conversations = conversations.len(),
"oversampled multi-turn conversations"
); );
} }
@@ -525,10 +526,12 @@ mod tests {
len len
} }
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_prefix_sharing_structure() { async fn test_prefix_sharing_structure() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig { let cfg = MultiTurnRandomConfig {
num_conversations: 5, num_conversations: 5,
@@ -610,10 +613,12 @@ mod tests {
println!("All prefix sharing checks passed!"); println!("All prefix sharing checks passed!");
} }
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_per_turn_input_len_default_mode() { async fn test_per_turn_input_len_default_mode() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig { let cfg = MultiTurnRandomConfig {
num_conversations: 4, num_conversations: 4,
@@ -650,10 +655,12 @@ mod tests {
println!("per_turn_input_len default-mode checks passed!"); println!("per_turn_input_len default-mode checks passed!");
} }
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_variable_turns_range() { async fn test_variable_turns_range() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig { let cfg = MultiTurnRandomConfig {
num_conversations: 50, num_conversations: 50,
@@ -684,10 +691,12 @@ mod tests {
println!("variable_turns_range checks passed! counts: {distinct_counts:?}"); println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
} }
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_variable_turns_fixed() { async fn test_variable_turns_fixed() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
let cfg = MultiTurnRandomConfig { let cfg = MultiTurnRandomConfig {
num_conversations: 10, num_conversations: 10,
@@ -709,10 +718,12 @@ mod tests {
println!("variable_turns_fixed checks passed!"); println!("variable_turns_fixed checks passed!");
} }
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_per_turn_input_len_prefix_sharing() { async fn test_per_turn_input_len_prefix_sharing() {
let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
.await
.unwrap();
// Turn 0 input_len=1000, turns 1+ per_turn_input_len=600 // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
// global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100 // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
@@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
} }
let total = prompts_per_prefix * num_prefixes; let total = prompts_per_prefix * num_prefixes;
if total != num_requests { if total != num_requests {
println!( tracing::info!(
"prefix_repetition: generating {total} requests \ requested = num_requests,
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \ generated = total,
{} dropped to divide evenly)", prefixes = num_prefixes,
num_requests - total prompts_per_prefix,
dropped = num_requests - total,
"adjusted prefix-repetition request count"
); );
} }
@@ -109,8 +111,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access. /// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind { fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None) TokenizerKind::Tiktoken(
.expect("gpt2 built-in tiktoken should always load without network") crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
} }
#[test] #[test]
+77
View File
@@ -0,0 +1,77 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
/// Reports row download progress to an interactive progress bar, or through
/// periodic tracing events when the progress bar is hidden on a non-TTY.
pub(super) struct RowDownloadReporter {
progress: ProgressBar,
next_report: Instant,
}
impl RowDownloadReporter {
/// Creates a reporter that emits non-TTY updates every 10 seconds.
pub fn new() -> Self {
let progress = ProgressBar::new(0);
progress.set_style(
ProgressStyle::with_template(
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
)
.unwrap()
.progress_chars("#>-"),
);
Self {
progress,
next_report: Instant::now() + REPORT_INTERVAL,
}
}
/// Updates the current row count and reports progress when due.
pub fn update(&mut self, rows: usize, total: u64) {
let rows = rows as u64;
let total = total.max(rows);
self.progress.set_length(total);
self.progress.set_position(rows);
if self.should_report(Instant::now()) {
tracing::info!(rows, total, "fetching dataset rows");
}
}
/// Clears the interactive progress bar after the download completes.
pub fn finish(self) {
self.progress.finish_and_clear();
}
fn should_report(&mut self, now: Instant) -> bool {
if !self.progress.is_hidden() || now < self.next_report {
return false;
}
self.next_report = now + REPORT_INTERVAL;
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hidden_reporter_uses_ten_second_deadline() {
let start = Instant::now();
let mut reporter = RowDownloadReporter {
progress: ProgressBar::hidden(),
next_report: start + REPORT_INTERVAL,
};
assert!(!reporter.should_report(start + Duration::from_secs(9)));
assert!(reporter.should_report(start + Duration::from_secs(10)));
assert!(!reporter.should_report(start + Duration::from_secs(19)));
assert!(reporter.should_report(start + Duration::from_secs(20)));
}
}
+18 -12
View File
@@ -49,9 +49,12 @@ pub fn generate_random_dataset(
let (input_low, input_high) = range_ratio.input_bounds(real_input_len); let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
let (output_low, output_high) = range_ratio.output_bounds(output_len); let (output_low, output_high) = range_ratio.output_bounds(output_len);
if !range_ratio.is_fixed() { if !range_ratio.is_fixed() {
println!( tracing::info!(
"Sampling input_len from [{input_low}, {input_high}] and \ input_low,
output_len from [{output_low}, {output_high}]" input_high,
output_low,
output_high,
"sampling random request lengths"
); );
} }
@@ -305,7 +308,8 @@ mod tests {
#[test] #[test]
#[ignore] #[ignore]
fn test_generate_random_dataset_token_ids() { fn test_generate_random_dataset_token_ids() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let requests = generate_random_dataset( let requests = generate_random_dataset(
&tokenizer, &tokenizer,
10, // num_requests 10, // num_requests
@@ -337,7 +341,8 @@ mod tests {
#[test] #[test]
#[ignore] #[ignore]
fn test_generate_random_dataset_text() { fn test_generate_random_dataset_text() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let requests = generate_random_dataset( let requests = generate_random_dataset(
&tokenizer, &tokenizer,
10, // num_requests 10, // num_requests
@@ -371,7 +376,8 @@ mod tests {
#[test] #[test]
#[ignore] #[ignore]
fn test_token_length_exact_local() { fn test_token_length_exact_local() {
let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); let tokenizer =
TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
let target_len = 512; let target_len = 512;
let requests = generate_random_dataset( let requests = generate_random_dataset(
&tokenizer, &tokenizer,
@@ -405,11 +411,11 @@ mod tests {
} }
/// Test that tiktoken tokenizer produces exact target token lengths (token ID mode). /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_token_length_exact_tiktoken() { async fn test_token_length_exact_tiktoken() {
// Use Qwen2.5 which has a tiktoken-format tokenizer // Use Qwen2.5 which has a tiktoken-format tokenizer
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = match tokenizer { let tokenizer = match tokenizer {
Ok(t) => t, Ok(t) => t,
Err(e) => { Err(e) => {
@@ -453,10 +459,10 @@ mod tests {
/// Test encode/decode roundtrip stability for tiktoken. /// Test encode/decode roundtrip stability for tiktoken.
/// After one decode→encode cycle with UTF-8-safe tokens, length must not drift. /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
#[test] #[tokio::test]
#[ignore] #[ignore]
fn test_tiktoken_roundtrip_stability() { async fn test_tiktoken_roundtrip_stability() {
let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
let tokenizer = match tokenizer { let tokenizer = match tokenizer {
Ok(t) => t, Ok(t) => t,
Err(e) => { Err(e) => {
+4 -2
View File
@@ -139,8 +139,10 @@ mod tests {
/// gpt2 via built-in tiktoken encoding — loads without network access. /// gpt2 via built-in tiktoken encoding — loads without network access.
fn test_tokenizer() -> TokenizerKind { fn test_tokenizer() -> TokenizerKind {
crate::tokenizer::load_tokenizer("gpt2", false, None) TokenizerKind::Tiktoken(
.expect("gpt2 built-in tiktoken should always load without network") crate::tiktoken::load_builtin_tiktoken("gpt2")
.expect("gpt2 built-in tiktoken should always load without network"),
)
} }
fn fixed_ratio() -> RangeRatio { fn fixed_ratio() -> RangeRatio {
+19 -12
View File
@@ -22,18 +22,21 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
/// Download the default ShareGPT dataset from HuggingFace Hub. /// Download the default ShareGPT dataset from HuggingFace Hub.
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly. /// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
pub fn download_sharegpt_dataset() -> Result<String> { pub async fn download_sharegpt_dataset() -> Result<String> {
println!( tracing::info!(
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..." repository = DEFAULT_SHAREGPT_REPO,
file = DEFAULT_SHAREGPT_FILE,
"downloading ShareGPT dataset"
); );
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string()); let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| { .map_err(BenchError::Config)?;
let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
BenchError::Config(format!( BenchError::Config(format!(
"Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}" "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
)) ))
})?; })?;
let path_str = path.to_string_lossy().to_string(); let path_str = path.to_string_lossy().to_string();
println!("ShareGPT dataset ready: {path_str}"); tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
Ok(path_str) Ok(path_str)
} }
@@ -135,9 +138,11 @@ pub fn load_sharegpt_dataset(
// Oversample if dataset is smaller than requested // Oversample if dataset is smaller than requested
if samples.len() < num_requests { if samples.len() < num_requests {
if no_oversample { if no_oversample {
println!( tracing::info!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})", dataset = "sharegpt",
samples.len() samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
); );
} else if !samples.is_empty() { } else if !samples.is_empty() {
let needed = num_requests - samples.len(); let needed = num_requests - samples.len();
@@ -147,9 +152,11 @@ pub fn load_sharegpt_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req); samples.push(req);
} }
println!( tracing::info!(
"Oversampled requests from {original_len} to {} total samples.", dataset = "sharegpt",
samples.len() original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
); );
} }
} }
+30 -23
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng}; use rand::{Rng, SeedableRng};
use super::SampleRequest; use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::cli::SpeedBenchConfig; use crate::cli::SpeedBenchConfig;
use crate::error::{BenchError, Result}; use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind; use crate::tokenizer::TokenizerKind;
@@ -25,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf {
/// Download SPEED-Bench dataset from HuggingFace datasets-server API. /// Download SPEED-Bench dataset from HuggingFace datasets-server API.
/// Results are cached as JSON locally for subsequent runs. /// Results are cached as JSON locally for subsequent runs.
pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> { pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let config_name = config.as_str(); let config_name = config.as_str();
let dir = cache_dir(); let dir = cache_dir();
@@ -35,13 +36,13 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
// Return cached file if it exists // Return cached file if it exists
if cache_path.exists() { if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string(); let path_str = cache_path.to_string_lossy().to_string();
println!("SPEED-Bench ({config_name}) cached: {path_str}"); tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
return Ok(path_str); return Ok(path_str);
} }
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server..."); tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
let client = reqwest::blocking::Client::builder() let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120)) .timeout(std::time::Duration::from_secs(120))
.build() .build()
.map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?; .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let mut all_rows: Vec<serde_json::Value> = Vec::new(); let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize; let mut offset = 0usize;
let page_size = 100usize; let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop { loop {
let url = format!( let url = format!(
@@ -64,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let max_retries = 3; let max_retries = 3;
let mut data: Option<serde_json::Value> = None; let mut data: Option<serde_json::Value> = None;
for attempt in 0..=max_retries { for attempt in 0..=max_retries {
let resp = match client.get(&url).send() { let resp = match client.get(&url).send().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
if attempt < max_retries { if attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs( tokio::time::sleep(std::time::Duration::from_secs(
2 * (attempt as u64 + 1), 2 * (attempt as u64 + 1),
)); ))
.await;
continue; continue;
} }
return Err(BenchError::Config(format!( return Err(BenchError::Config(format!(
@@ -80,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
}; };
if resp.status().is_server_error() && attempt < max_retries { if resp.status().is_server_error() && attempt < max_retries {
std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
continue; continue;
} }
@@ -91,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
))); )));
} }
data = Some(resp.json().map_err(|e| { data = Some(resp.json().await.map_err(|e| {
BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}")) BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
})?); })?);
break; break;
@@ -116,15 +119,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let fetched = rows.len(); let fetched = rows.len();
offset += fetched; offset += fetched;
// Print progress
let total = data["num_rows_total"].as_u64().unwrap_or(0); let total = data["num_rows_total"].as_u64().unwrap_or(0);
eprint!("\r Fetched {offset}/{total} rows..."); progress.update(offset, total);
if fetched < page_size { if fetched < page_size {
break; break;
} }
} }
eprintln!(); // newline after progress progress.finish();
if all_rows.is_empty() { if all_rows.is_empty() {
return Err(BenchError::Config( return Err(BenchError::Config(
@@ -137,9 +139,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
std::fs::write(&cache_path, &json_str)?; std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string(); let path_str = cache_path.to_string_lossy().to_string();
println!( tracing::info!(
"SPEED-Bench ({config_name}): {} rows saved to {path_str}", config = config_name,
all_rows.len() rows = all_rows.len(),
path = %path_str,
"saved SPEED-Bench dataset"
); );
Ok(path_str) Ok(path_str)
} }
@@ -263,9 +267,11 @@ pub fn load_speed_bench_dataset(
// Oversample if needed // Oversample if needed
if samples.len() < num_requests { if samples.len() < num_requests {
if no_oversample { if no_oversample {
println!( tracing::info!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})", dataset = "speed-bench",
samples.len() samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
); );
} else if !samples.is_empty() { } else if !samples.is_empty() {
let original_len = samples.len(); let original_len = samples.len();
@@ -275,9 +281,11 @@ pub fn load_speed_bench_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req); samples.push(req);
} }
println!( tracing::info!(
"Oversampled SPEED-Bench from {original_len} to {} total samples.", dataset = "speed-bench",
samples.len() original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
); );
} }
} }
@@ -288,7 +296,6 @@ pub fn load_speed_bench_dataset(
)); ));
} }
// Print category distribution
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new(); let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
for entry in &filtered[..filtered.len().min(samples.len())] { for entry in &filtered[..filtered.len().min(samples.len())] {
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown"); let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -297,7 +304,7 @@ pub fn load_speed_bench_dataset(
let mut cats: Vec<_> = cat_counts.into_iter().collect(); let mut cats: Vec<_> = cat_counts.into_iter().collect();
cats.sort_by_key(|b| std::cmp::Reverse(b.1)); cats.sort_by_key(|b| std::cmp::Reverse(b.1));
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect(); let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
println!("SPEED-Bench categories: {}", cat_str.join(", ")); tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
Ok(samples) Ok(samples)
} }
+20 -32
View File
@@ -1,51 +1,39 @@
// SPDX-License-Identifier: Apache-2.0 // SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Sync facade over the async `hf_hub` API.
//!
//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
//! backend unconditionally pulls ureq's default rustls feature. So we use the
//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
//! loaders, tokenizer fallback in rayon threads) by running each download on a
//! dedicated thread with its own single-threaded runtime.
use std::path::PathBuf; use std::path::PathBuf;
use hf_hub::Repo;
use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache. /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
pub struct HubRepo { pub struct HubRepo {
repo: hf_hub::Repo, repo: ApiRepo,
} }
impl HubRepo { impl HubRepo {
pub fn model(model_id: String) -> Self { pub fn model(model_id: String) -> Result<Self, String> {
Self { Self::new(Repo::model(model_id))
repo: hf_hub::Repo::model(model_id),
}
} }
pub fn dataset(repo_id: String) -> Self { pub fn dataset(repo_id: String) -> Result<Self, String> {
Self { Self::new(Repo::dataset(repo_id))
repo: hf_hub::Repo::dataset(repo_id), }
fn new(repo: Repo) -> Result<Self, String> {
let mut builder = ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
} }
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
Ok(Self {
repo: api.repo(repo),
})
} }
/// Download (or fetch from cache) a single file from the repo. /// Download (or fetch from cache) a single file from the repo.
/// Auth is handled by hf-hub via HF_TOKEN / the cached login token. /// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
pub fn get(&self, filename: &str) -> Result<PathBuf, String> { pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
let repo = self.repo.clone(); self.repo.get(filename).await.map_err(|e| format!("{e}"))
let filename = filename.to_string();
std::thread::spawn(move || {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let api = hf_hub::api::tokio::Api::new()
.map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
.join()
.map_err(|_| "HF Hub download thread panicked".to_string())?
} }
} }
+86
View File
@@ -0,0 +1,86 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
use anyhow::Context;
pub use cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use config::BenchConfig;
/// Prepare process-wide resources for a benchmark run.
pub fn prepare_process() {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
tracing::info!(soft_limit = new, "raised open-file limit");
}
}
/// Run the online serving benchmark.
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = args.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_args(&args).context("Configuration error")?;
async {
if config.multi_turn {
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
args.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if args.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, args.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
}
.await
.context("Benchmark failed")
}
+25 -73
View File
@@ -1,92 +1,44 @@
// SPDX-License-Identifier: Apache-2.0 // SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
#[cfg(not(target_env = "msvc"))] #[cfg(not(target_env = "msvc"))]
#[global_allocator] #[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use anyhow::Context; use anyhow::Context;
use clap::Parser; use clap::Parser;
use cli::Cli;
use config::BenchConfig; #[derive(Parser)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
struct Cli {
#[command(flatten)]
args: vllm_bench::BenchServeArgs,
}
// TODO: unify the tracing subscriber used by different binaries.
fn init_tracing() {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
let _ = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_writer(std::io::stderr)
.try_init();
}
fn main() -> anyhow::Result<()> { fn main() -> anyhow::Result<()> {
// Raise the open-file soft limit to the hard limit. High-concurrency init_tracing();
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
let cli = Cli::parse(); let cli = Cli::parse();
vllm_bench::prepare_process();
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = cli.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
let runtime = tokio::runtime::Builder::new_multi_thread() let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all() .enable_all()
.build() .build()
.expect("Failed to build tokio runtime"); .context("Failed to build tokio runtime")?;
runtime runtime.block_on(vllm_bench::run(cli.args))
.block_on(async {
if config.multi_turn {
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
cli.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
.await?;
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if cli.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, cli.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
})
.context("Benchmark failed")
} }
+18 -16
View File
@@ -7,6 +7,22 @@ use crate::datasets::SampleRequest;
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
use crate::multi_turn::ConversationOutput; use crate::multi_turn::ConversationOutput;
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
if failed_outputs.is_empty() {
return;
}
tracing::warn!(
failed_requests = failed_outputs.len(),
displayed_errors = failed_outputs.len().min(10),
"benchmark requests failed"
);
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
tracing::warn!(index, error = %output.error, "benchmark request failed");
}
}
/// Calculate benchmark metrics from request outputs. /// Calculate benchmark metrics from request outputs.
/// ///
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599. /// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
let failed = outputs.len() - completed; let failed = outputs.len() - completed;
// Print failed request errors (capped to 10) log_failed_requests(outputs);
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Calculate max output tokens per second and max concurrent requests // Calculate max output tokens per second and max concurrent requests
let mut max_output_tokens_per_s = 0.0_f64; let mut max_output_tokens_per_s = 0.0_f64;
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
let failed = outputs.len() - completed; let failed = outputs.len() - completed;
// Print failed request errors (capped to 10) log_failed_requests(outputs);
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
// Compute peak concurrent requests from start_time + latency windows // Compute peak concurrent requests from start_time + latency windows
let successful_outputs: Vec<&RequestFuncOutput> = let successful_outputs: Vec<&RequestFuncOutput> =
+53 -39
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Instant; use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle}; use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend}; use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
@@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (model_id, model_name) = if let Some(ref m) = config.model { let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone()) (m.clone(), config.model_name.clone())
} else { } else {
println!("Model not specified, fetching first model from server..."); tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?; let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}"); tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name)) (id, Some(name))
}; };
@@ -83,15 +88,19 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
None None
} else { } else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
println!("Loading tokenizer: {tid}"); tracing::info!(tokenizer = tid, "loading tokenizer");
let server_info = Some((config.base_url.as_str(), model_id.as_str())); let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?; let t =
println!("Tokenizer loaded successfully."); crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
Some(t) Some(t)
}; };
// Generate/load conversations // Generate/load conversations
println!("Generating multi-turn conversations..."); tracing::info!(
dataset = ?config.dataset_name,
conversations = config.num_prompts,
"generating multi-turn conversations"
);
let gen_start = Instant::now(); let gen_start = Instant::now();
let mut conversations = match config.dataset_name { let mut conversations = match config.dataset_name {
@@ -131,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let path = match config.dataset_path.as_deref() { let path = match config.dataset_path.as_deref() {
Some(p) => p, Some(p) => p,
None => { None => {
downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
downloaded.as_str() downloaded.as_str()
} }
}; };
@@ -179,8 +188,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (filtered_conversations, filtered_turns) = let (filtered_conversations, filtered_turns) =
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history); filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
if filtered_turns > 0 || filtered_conversations > 0 { if filtered_turns > 0 || filtered_conversations > 0 {
println!( tracing::info!(
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}." filtered_turns,
filtered_conversations,
max_model_len,
"filtered conversations above maximum model length"
); );
} }
if conversations.is_empty() { if conversations.is_empty() {
@@ -192,11 +204,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let gen_elapsed = gen_start.elapsed(); let gen_elapsed = gen_start.elapsed();
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum(); let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
println!( tracing::info!(
"Generated {} conversations ({} total turns) in {:.2}s", conversations = conversations.len(),
conversations.len(),
total_turns, total_turns,
gen_elapsed.as_secs_f64() elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated multi-turn conversations"
); );
// Log prefix sharing info // Log prefix sharing info
@@ -208,18 +220,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let conv_tokens = let conv_tokens =
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize; (real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens); let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
println!( tracing::info!(
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)", global_ratio = config.multi_turn_prefix_global_ratio,
config.multi_turn_prefix_global_ratio * 100.0,
global_tokens, global_tokens,
config.multi_turn_prefix_conversation_ratio * 100.0, conversation_ratio = config.multi_turn_prefix_conversation_ratio,
conv_tokens, conversation_tokens = conv_tokens,
(1.0 - config.multi_turn_prefix_global_ratio
- config.multi_turn_prefix_conversation_ratio)
* 100.0,
unique_tokens, unique_tokens,
history_accumulation = false,
"configured multi-turn prefix sharing"
); );
println!("No history accumulation: each turn sends fixed-length prompt only.");
} }
if config.dry_run { if config.dry_run {
@@ -253,7 +262,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
..Default::default() ..Default::default()
}; };
println!("Starting initial single prompt test run..."); tracing::info!("starting initial single-prompt test run");
let test_output = crate::ready_checker::wait_for_endpoint( let test_output = crate::ready_checker::wait_for_endpoint(
config.backend, config.backend,
&client, &client,
@@ -268,7 +277,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
test_output.error test_output.error
))); )));
} }
println!("Initial test run completed."); tracing::info!("initial single-prompt test run completed");
} }
// For random datasets in multi-turn mode, auto-set min_tokens to enforce // For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -283,9 +292,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
"min_tokens".to_string(), "min_tokens".to_string(),
serde_json::json!(config.random_output_len), serde_json::json!(config.random_output_len),
); );
println!( tracing::info!(
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)", min_tokens = config.random_output_len,
config.random_output_len dataset = "random",
"set minimum output tokens for multi-turn dataset"
); );
} }
Some(body) Some(body)
@@ -297,7 +307,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let spec_decode_before = let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() { if spec_decode_before.is_some() {
println!("Speculative decoding detected, will collect metrics."); tracing::info!("detected speculative decoding; collecting metrics");
} }
// Start profiler if requested (immediate mode — no batch threshold) // Start profiler if requested (immediate mode — no batch threshold)
@@ -330,10 +340,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
}; };
// Main benchmark // Main benchmark
println!("Starting multi-turn benchmark..."); tracing::info!(
println!("Conversations: {}", conversations.len()); conversations = conversations.len(),
println!("Concurrency: {concurrency}"); total_turns,
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms); concurrency,
inter_turn_delay_ms = config.multi_turn_delay_ms,
"starting multi-turn benchmark"
);
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0); let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
@@ -364,11 +377,12 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
); );
if let Some(modules) = config.lora_modules.as_ref() { if let Some(modules) = config.lora_modules.as_ref() {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect(); let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!( tracing::info!(
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]", adapters = modules.len(),
modules.len(), names = ?names,
names, assignment = ?config.lora_assignment,
config.lora_assignment scope = "conversation",
"assigned LoRA adapters"
); );
} }
@@ -433,7 +447,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
match handle.await { match handle.await {
Ok(output) => all_outputs.push(output), Ok(output) => all_outputs.push(output),
Err(e) => { Err(e) => {
eprintln!("Conversation task panicked: {e}"); tracing::error!(error = %e.as_report(), "conversation task panicked");
} }
} }
} }
@@ -453,7 +467,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
if let Some((cancel_tx, task)) = profile_task { if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(()); let _ = cancel_tx.send(());
if let Err(e) = task.await { if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}"); tracing::error!(error = %e.as_report(), "profiler background task failed");
} }
} }
+2 -2
View File
@@ -603,7 +603,7 @@ fn add_metric_stats(
pub fn save_result(json: &Value, file_path: &str) -> Result<()> { pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
let content = serde_json::to_string(json)?; let content = serde_json::to_string(json)?;
std::fs::write(file_path, content)?; std::fs::write(file_path, content)?;
println!("Results saved to {file_path}"); tracing::info!(path = file_path, "saved benchmark results");
Ok(()) Ok(())
} }
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
file.write_all(b"\n")?; file.write_all(b"\n")?;
} }
file.write_all(content.as_bytes())?; file.write_all(content.as_bytes())?;
println!("Results appended to {file_path}"); tracing::info!(path = file_path, "appended benchmark results");
Ok(()) Ok(())
} }
+8 -2
View File
@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
let backend = get_backend(backend)?; let backend = get_backend(backend)?;
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds); let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
println!("Waiting for endpoint to become up in {timeout_seconds}s"); tracing::info!(
timeout_seconds,
retry_interval,
"waiting for endpoint readiness"
);
let pb = ProgressBar::new(timeout_seconds); let pb = ProgressBar::new(timeout_seconds);
pb.set_style( pb.set_style(
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
Ok(output) => { Ok(output) => {
let err = output.error.clone(); let err = output.error.clone();
let err_last_line = err.lines().last().unwrap_or(&err); let err_last_line = err.lines().last().unwrap_or(&err);
eprintln!("Endpoint is not ready. Error='{err_last_line}'"); pb.suspend(|| {
tracing::warn!(error = err_last_line, "endpoint is not ready");
});
last_error = err; last_error = err;
} }
Err(e) => { Err(e) => {
+1 -1
View File
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
.await .await
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?; .map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
if resp.status().is_success() { if resp.status().is_success() {
println!("Prefix cache reset successfully."); tracing::info!(url = %url, "reset prefix cache");
} else { } else {
let status = resp.status(); let status = resp.status();
let body = resp.text().await.unwrap_or_default(); let body = resp.text().await.unwrap_or_default();
+36 -24
View File
@@ -199,12 +199,17 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
} }
}; };
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?; let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})"); tracing::info!(
encoding,
kind = "built-in-tiktoken",
vocab_size,
"loaded tokenizer"
);
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size)) Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
} }
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo. /// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> { pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
// Phase 1: If model_id is a local directory, look for tiktoken files there // Phase 1: If model_id is a local directory, look for tiktoken files there
let local_dir = Path::new(model_id); let local_dir = Path::new(model_id);
if local_dir.is_dir() { if local_dir.is_dir() {
@@ -212,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
} }
// Phase 2: Fall back to HuggingFace Hub download // Phase 2: Fall back to HuggingFace Hub download
try_load_tiktoken_from_hf(model_id) try_load_tiktoken_from_hf(model_id).await
} }
/// Common tiktoken model filenames to search for. /// Common tiktoken model filenames to search for.
@@ -247,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
} }
/// Load a tiktoken tokenizer from a HuggingFace model repo. /// Load a tiktoken tokenizer from a HuggingFace model repo.
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> { async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string()); let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
let model_path = repo let mut model_path = None;
.get("tiktoken.model") for filename in TIKTOKEN_MODEL_FILENAMES {
.or_else(|_| repo.get("qwen.tiktoken")) if let Ok(path) = repo.get(filename).await {
.or_else(|_| repo.get("vocab.tiktoken")) model_path = Some(path);
.map_err(|_| { break;
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) }
})?; }
let model_path = model_path.ok_or_else(|| {
BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
})?;
let num_base_tokens = count_base_tokens(&model_path)?; let num_base_tokens = count_base_tokens(&model_path)?;
let config = match repo.get("tokenizer_config.json") { let config = match repo.get("tokenizer_config.json").await {
Ok(config_path) => read_tokenizer_config(&config_path), Ok(config_path) => read_tokenizer_config(&config_path),
Err(_) => None, Err(_) => None,
}; };
let pattern = extract_pat_str_from_repo(&repo); let pattern = extract_pat_str_from_repo(&repo).await;
build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens) build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
} }
@@ -309,15 +317,16 @@ fn build_tiktoken(
} }
} }
println!( tracing::info!(
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...", model = model_id,
num_base_tokens, base_tokens = num_base_tokens,
all_special_tokens.len(), special_tokens = all_special_tokens.len(),
if pattern.is_some() { pattern = if pattern.is_some() {
"custom" "custom"
} else { } else {
"default" "default"
}, },
"loading tiktoken model"
); );
TiktokenTokenizer::from_file( TiktokenTokenizer::from_file(
@@ -397,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
/// Try to download the Python tokenizer source file and extract pat_str via regex. /// Try to download the Python tokenizer source file and extract pat_str via regex.
/// Returns None if unavailable or unparsable. /// Returns None if unavailable or unparsable.
fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> { async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
// Try common Python tokenizer filenames // Try common Python tokenizer filenames
let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?; let py_path = match repo.get("tokenization_kimi.py").await {
Ok(path) => path,
Err(_) => repo.get("tokenizer.py").await.ok()?,
};
let source = std::fs::read_to_string(&py_path).ok()?; let source = std::fs::read_to_string(&py_path).ok()?;
@@ -438,9 +450,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
if !fragments.is_empty() { if !fragments.is_empty() {
let pattern = fragments.join("|"); let pattern = fragments.join("|");
println!( tracing::debug!(
"Extracted pat_str from Python source: {} fragments", fragments = fragments.len(),
fragments.len() "extracted tiktoken pattern from Python source"
); );
return Some(pattern); return Some(pattern);
} }
+98 -21
View File
@@ -2,8 +2,10 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::HashSet; use std::collections::HashSet;
use std::future::Future;
use std::path::Path; use std::path::Path;
use thiserror_ext::AsReport as _;
use tokenizers::Tokenizer; use tokenizers::Tokenizer;
use crate::error::{BenchError, Result}; use crate::error::{BenchError, Result};
@@ -18,7 +20,8 @@ pub enum TokenizerKind {
/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints. /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
pub struct ServerTokenizer { pub struct ServerTokenizer {
client: reqwest::blocking::Client, client: reqwest::Client,
runtime: tokio::runtime::Handle,
tokenize_url: String, tokenize_url: String,
detokenize_url: String, detokenize_url: String,
model: String, model: String,
@@ -27,8 +30,8 @@ pub struct ServerTokenizer {
impl ServerTokenizer { impl ServerTokenizer {
/// Create a new server tokenizer and verify connectivity. /// Create a new server tokenizer and verify connectivity.
pub fn new(base_url: &str, model: &str) -> Result<Self> { pub async fn new(base_url: &str, model: &str) -> Result<Self> {
let client = reqwest::blocking::Client::builder() let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30)) .timeout(std::time::Duration::from_secs(30))
.build() .build()
.map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?; .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -38,6 +41,7 @@ impl ServerTokenizer {
let st = Self { let st = Self {
client, client,
runtime: tokio::runtime::Handle::current(),
tokenize_url, tokenize_url,
detokenize_url, detokenize_url,
model: model.to_string(), model: model.to_string(),
@@ -45,7 +49,7 @@ impl ServerTokenizer {
}; };
// Probe the endpoint to verify it works and discover vocab size // Probe the endpoint to verify it works and discover vocab size
let test_tokens = st.encode_inner("test")?; let test_tokens = st.encode_async("test").await?;
let max_id = test_tokens.iter().copied().max().unwrap_or(0); let max_id = test_tokens.iter().copied().max().unwrap_or(0);
let estimated_vocab = (max_id * 2).max(131072); let estimated_vocab = (max_id * 2).max(131072);
@@ -56,6 +60,10 @@ impl ServerTokenizer {
} }
fn encode_inner(&self, text: &str) -> Result<Vec<u32>> { fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
self.block_on(self.encode_async(text))
}
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
let payload = serde_json::json!({ let payload = serde_json::json!({
"model": self.model, "model": self.model,
"prompt": text, "prompt": text,
@@ -66,6 +74,7 @@ impl ServerTokenizer {
.post(&self.tokenize_url) .post(&self.tokenize_url)
.json(&payload) .json(&payload)
.send() .send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?; .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
if !resp.status().is_success() { if !resp.status().is_success() {
@@ -75,7 +84,7 @@ impl ServerTokenizer {
))); )));
} }
let data: serde_json::Value = resp.json().map_err(|e| { let data: serde_json::Value = resp.json().await.map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}")) BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
})?; })?;
@@ -95,6 +104,10 @@ impl ServerTokenizer {
} }
fn decode_inner(&self, ids: &[u32]) -> Result<String> { fn decode_inner(&self, ids: &[u32]) -> Result<String> {
self.block_on(self.decode_async(ids))
}
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
let payload = serde_json::json!({ let payload = serde_json::json!({
"model": self.model, "model": self.model,
"tokens": ids, "tokens": ids,
@@ -105,6 +118,7 @@ impl ServerTokenizer {
.post(&self.detokenize_url) .post(&self.detokenize_url)
.json(&payload) .json(&payload)
.send() .send()
.await
.map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?; .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
if !resp.status().is_success() { if !resp.status().is_success() {
@@ -114,7 +128,7 @@ impl ServerTokenizer {
))); )));
} }
let data: serde_json::Value = resp.json().map_err(|e| { let data: serde_json::Value = resp.json().await.map_err(|e| {
BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}")) BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
})?; })?;
@@ -123,6 +137,26 @@ impl ServerTokenizer {
.map(|s| s.to_string()) .map(|s| s.to_string())
.ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into())) .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
} }
fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
if matches!(
self.runtime.runtime_flavor(),
tokio::runtime::RuntimeFlavor::CurrentThread
) {
return Err(BenchError::Tokenizer(
"Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
));
}
// Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
// Tokio workers must enter a blocking region before re-entering the runtime;
// Rayon workers can drive the future directly with the saved runtime handle.
if tokio::runtime::Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
} }
// --- TokenizerKind methods --- // --- TokenizerKind methods ---
@@ -192,7 +226,7 @@ impl TokenizerKind {
/// 3. Server-side /tokenize + /detokenize endpoints /// 3. Server-side /tokenize + /detokenize endpoints
/// ///
/// `server_info` is `Some((base_url, model))` to enable server-side fallback. /// `server_info` is `Some((base_url, model))` to enable server-side fallback.
pub fn load_tokenizer( pub async fn load_tokenizer(
model_id: &str, model_id: &str,
_trust_remote_code: bool, _trust_remote_code: bool,
server_info: Option<(&str, &str)>, server_info: Option<(&str, &str)>,
@@ -212,31 +246,48 @@ pub fn load_tokenizer(
} }
// 1. Try local HuggingFace tokenizer (tokenizer.json) // 1. Try local HuggingFace tokenizer (tokenizer.json)
match try_load_local(model_id) { match try_load_local(model_id).await {
Ok(tok) => { Ok(tok) => {
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true)); tracing::info!(
model = model_id,
kind = "local",
vocab_size = tok.get_vocab_size(true),
"loaded tokenizer"
);
Ok(TokenizerKind::Local(Box::new(tok))) Ok(TokenizerKind::Local(Box::new(tok)))
} }
Err(local_err) => { Err(local_err) => {
// 2. Try tiktoken format // 2. Try tiktoken format
println!("No tokenizer.json for '{model_id}', trying tiktoken format..."); tracing::info!(
match crate::tiktoken::try_load_tiktoken(model_id) { model = model_id,
error = %local_err.as_report(),
"local tokenizer unavailable; trying tiktoken"
);
match crate::tiktoken::try_load_tiktoken(model_id).await {
Ok(tok) => { Ok(tok) => {
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size()); tracing::info!(
model = model_id,
kind = "tiktoken",
vocab_size = tok.vocab_size(),
"loaded tokenizer"
);
Ok(TokenizerKind::Tiktoken(tok)) Ok(TokenizerKind::Tiktoken(tok))
} }
Err(tiktoken_err) => { Err(tiktoken_err) => {
// 3. Try server-side fallback // 3. Try server-side fallback
if let Some((base_url, model)) = server_info { if let Some((base_url, model)) = server_info {
println!( tracing::info!(
"Tiktoken also not available ({tiktoken_err}), \ model = model_id,
trying server-side tokenization..." error = %tiktoken_err.as_report(),
"tiktoken unavailable; trying server-side tokenization"
); );
match ServerTokenizer::new(base_url, model) { match ServerTokenizer::new(base_url, model).await {
Ok(srv) => { Ok(srv) => {
println!( tracing::info!(
"Tokenizer: Server (vocab_size≈{})", model = model_id,
srv.cached_vocab_size kind = "server",
vocab_size = srv.cached_vocab_size,
"loaded tokenizer"
); );
return Ok(TokenizerKind::Server(srv)); return Ok(TokenizerKind::Server(srv));
} }
@@ -264,7 +315,7 @@ pub fn load_tokenizer(
} }
/// Try loading tokenizer.json from local path or HuggingFace Hub. /// Try loading tokenizer.json from local path or HuggingFace Hub.
fn try_load_local(model_id: &str) -> Result<Tokenizer> { async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
// 1. Try local directory with tokenizer.json // 1. Try local directory with tokenizer.json
let local_path = Path::new(model_id).join("tokenizer.json"); let local_path = Path::new(model_id).join("tokenizer.json");
if local_path.exists() { if local_path.exists() {
@@ -290,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result<Tokenizer> {
} }
// 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token) // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
let repo = crate::hub::HubRepo::model(model_id.to_string()); let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
let tokenizer_path = repo let tokenizer_path = repo
.get("tokenizer.json") .get("tokenizer.json")
.await
.map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?; .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
Tokenizer::from_file(&tokenizer_path) Tokenizer::from_file(&tokenizer_path)
.map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}"))) .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
} }
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_tokenizer_sync_bridge() {
let tokenizer = std::sync::Arc::new(ServerTokenizer {
client: reqwest::Client::new(),
runtime: tokio::runtime::Handle::current(),
tokenize_url: String::new(),
detokenize_url: String::new(),
model: String::new(),
cached_vocab_size: 0,
});
assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
let (tx, rx) = tokio::sync::oneshot::channel();
rayon::spawn(move || {
let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
});
assert_eq!(rx.await.unwrap().unwrap(), 2);
}
}
+4 -4
View File
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key); let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
let (value, field) = match spec.field_layout_for(key) { let (value, field) = match spec.field_layout_for(key) {
Some(FieldLayout::Batched) => ( Some(FieldLayout::Batched) => (
tensor.batched_value_at(index)?, tensor.batched_wire_value_at(index)?,
MmField::Batched(MmBatchedField { keep_on_cpu }), MmField::Batched(MmBatchedField { keep_on_cpu }),
), ),
Some(FieldLayout::Flat { sizes_key }) => { Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
})?; })?;
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?; let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
( (
tensor.flat_value_range(start, end)?, tensor.flat_wire_value_range(start, end)?,
MmField::Flat(MmFlatField { MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec { slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0), start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
) )
} }
None => ( None => (
tensor.clone(), tensor.try_into()?,
MmField::Shared(MmSharedField { MmField::Shared(MmSharedField {
batch_size: len, batch_size: len,
keep_on_cpu, keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
data.insert( data.insert(
key.clone(), key.clone(),
MmFieldElem { MmFieldElem {
data: Some(value.try_into()?), data: Some(value),
field, field,
}, },
); );
+68 -81
View File
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
use crate::error::{Error, Result, bail_multimodal, multimodal}; use crate::error::{Error, Result, bail_multimodal, multimodal};
/// Representation for multimodal kwarg values for transformation. /// Representation for multimodal kwarg values for transformation.
#[derive(Debug, Clone)] #[derive(Debug)]
pub(super) enum KwargValue { pub(super) enum KwargValue {
/// Float tensor with row-major flat data and shape. /// Float tensor with row-major flat data and shape.
F32Tensor { data: Vec<f32>, shape: Vec<usize> }, F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
} }
} }
impl TryFrom<KwargValue> for ProtocolKwargValue { impl TryFrom<&KwargValue> for ProtocolKwargValue {
type Error = Error; type Error = Error;
fn try_from(value: KwargValue) -> Result<Self> { fn try_from(value: &KwargValue) -> Result<Self> {
match value { let tensor = match value {
KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor( KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?, KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
)), KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor( KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?, KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
)), KwargValue::Passthrough(value) => return Ok(value.clone()),
KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor( };
WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?, tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
)),
KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Passthrough(value) => Ok(value),
}
} }
} }
@@ -145,63 +136,55 @@ impl KwargValue {
} }
} }
/// Extract one media item from a batched tensor field. /// Convert one media item from a batched tensor field to wire bytes.
/// ///
/// Batched fields use their first axis as media-item index and drop that /// Batched fields use their first axis as media-item index and drop that
/// axis in the per-feature value, matching vLLM's batched-field semantics. /// axis in the per-feature value, matching vLLM's batched-field semantics.
pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> { pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
match self { self.wire_value_range(index, index + 1, true)
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F32Tensor { data, shape })
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F16Tensor { data, shape })
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::Bf16Tensor { data, shape })
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::I64Tensor { data, shape })
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::U32Tensor { data, shape })
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
} }
/// Extract one media item's variable-length range from a flat tensor field. /// Convert one media item's flat tensor range directly to wire bytes.
/// ///
/// Flat fields keep the first axis as the sliced length for this item. /// Flat fields keep the first axis as the sliced length for this item.
pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> { pub(super) fn flat_wire_value_range(
match self { &self,
start: usize,
end: usize,
) -> Result<ProtocolKwargValue> {
self.wire_value_range(start, end, false)
}
fn wire_value_range(
&self,
start: usize,
end: usize,
drop_axis: bool,
) -> Result<ProtocolKwargValue> {
let tensor = match self {
Self::F32Tensor { data, shape } => { Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
Ok(Self::F32Tensor { data, shape }) WireTensor::from_f32(shape, data)
} }
Self::F16Tensor { data, shape } => { Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
Ok(Self::F16Tensor { data, shape }) WireTensor::from_f16(shape, data)
} }
Self::Bf16Tensor { data, shape } => { Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
Ok(Self::Bf16Tensor { data, shape }) WireTensor::from_bf16(shape, data)
} }
Self::I64Tensor { data, shape } => { Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
Ok(Self::I64Tensor { data, shape }) WireTensor::from_i64(shape, data)
} }
Self::U32Tensor { data, shape } => { Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?; let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
Ok(Self::U32Tensor { data, shape }) WireTensor::from_u32(shape, data)
} }
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())), Self::Passthrough(value) => return Ok(value.clone()),
} };
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
} }
} }
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
} }
/// Slice a flat row-major tensor along its first axis. /// Slice a flat row-major tensor along its first axis.
fn slice_first_axis_range<T: Clone>( fn slice_first_axis_range<'a, T>(
shape: &[usize], shape: &[usize],
data: &[T], data: &'a [T],
start: usize, start: usize,
end: usize, end: usize,
drop_axis: bool, drop_axis: bool,
) -> Result<(Vec<usize>, Vec<T>)> { ) -> Result<(Vec<usize>, &'a [T])> {
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?; let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
if start > end || end > first_dim { if start > end || end > first_dim {
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}"); bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
shape[0] = end - start; shape[0] = end - start;
shape shape
}; };
Ok((out_shape, data[data_start..data_end].to_vec())) Ok((out_shape, &data[data_start..data_end]))
} }
#[cfg(test)] #[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
use super::*; use super::*;
#[test] #[test]
fn batched_value_at_drops_first_axis() { fn batched_wire_value_at_drops_first_axis() {
let value = KwargValue::F32Tensor { let value = KwargValue::F32Tensor {
data: vec![1.0, 2.0, 3.0, 4.0], data: vec![1.0, 2.0, 3.0, 4.0],
shape: vec![2, 2], shape: vec![2, 2],
}; };
let value = value.batched_value_at(1).unwrap(); let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
panic!("expected tensor");
};
assert!(matches!( assert_eq!(tensor.shape, vec![2]);
value, assert_eq!(
KwargValue::F32Tensor { data, shape } tensor.data.into_raw_view().unwrap(),
if shape == vec![2] && data == vec![3.0, 4.0] [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
)); );
} }
#[test] #[test]
fn flat_value_range_keeps_first_axis() { fn flat_wire_value_range_keeps_first_axis() {
let value = KwargValue::U32Tensor { let value = KwargValue::U32Tensor {
data: (0..10).collect(), data: (0..10).collect(),
shape: vec![5, 2], shape: vec![5, 2],
}; };
let value = value.flat_value_range(1, 3).unwrap(); let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
panic!("expected tensor");
};
assert!(matches!( assert_eq!(tensor.shape, vec![2, 2]);
value, assert_eq!(
KwargValue::U32Tensor { data, shape } tensor.data.into_raw_view().unwrap(),
if shape == vec![2, 2] && data == vec![2, 3, 4, 5] [2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
)); );
} }
#[test] #[test]
@@ -336,7 +323,7 @@ mod tests {
let value = let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap(); KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap() let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else { else {
panic!("expected tensor"); panic!("expected tensor");
}; };
@@ -351,7 +338,7 @@ mod tests {
let value = let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap(); KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap() let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else { else {
panic!("expected tensor"); panic!("expected tensor");
}; };
+4 -4
View File
@@ -130,7 +130,7 @@ fn build_video_item(
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key); let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
let (value, field) = match support.spec.field_layout_for(&key) { let (value, field) = match support.spec.field_layout_for(&key) {
Some(FieldLayout::Batched) => ( Some(FieldLayout::Batched) => (
tensor.batched_value_at(0)?, tensor.batched_wire_value_at(0)?,
MmField::Batched(MmBatchedField { keep_on_cpu }), MmField::Batched(MmBatchedField { keep_on_cpu }),
), ),
Some(FieldLayout::Flat { .. }) => { Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
.first_dim() .first_dim()
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?; .ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
( (
tensor, (&tensor).try_into()?,
MmField::Flat(MmFlatField { MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec { slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0), start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
) )
} }
None => ( None => (
tensor, (&tensor).try_into()?,
MmField::Shared(MmSharedField { MmField::Shared(MmSharedField {
batch_size: 1, batch_size: 1,
keep_on_cpu, keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
data.insert( data.insert(
key, key,
MmFieldElem { MmFieldElem {
data: Some(value.try_into()?), data: Some(value),
field, field,
}, },
); );
+1 -1
View File
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser> Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
}; };
apply_structural_tag_constraint(request, parser.structural_tag_model())?; apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
if parser.preserve_special_tokens() { if parser.preserve_special_tokens() {
request.decode_options.skip_special_tokens = false; request.decode_options.skip_special_tokens = false;
@@ -7,7 +7,8 @@ use thiserror_ext::AsReport;
use vllm_engine_core_client::protocol::structured_outputs::{ use vllm_engine_core_client::protocol::structured_outputs::{
StructuredOutputBackend, StructuredOutputsParams, StructuredOutputBackend, StructuredOutputsParams,
}; };
use vllm_parser::tool::StructuralTagModel; use vllm_parser::tool::StructuralTagBuilder;
use xgrammar_structural_tag::builders::StructuralTagOptions;
use xgrammar_structural_tag::{ use xgrammar_structural_tag::{
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam, FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
build_structural_tag, build_structural_tag,
@@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult};
/// support and the request's tool choice. /// support and the request's tool choice.
pub(super) fn apply_structural_tag_constraint( pub(super) fn apply_structural_tag_constraint(
request: &mut ChatRequest, request: &mut ChatRequest,
model: Option<StructuralTagModel>, builder: Option<&dyn StructuralTagBuilder>,
) -> ChatResult<()> { ) -> ChatResult<()> {
let Some(model) = model else { let Some(builder) = builder else {
return Ok(()); return Ok(());
}; };
let Some(tool_choice) = structural_tag_tool_choice(request) else { let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint(
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let structural_tag = build_structural_tag(model, &tools, tool_choice, false) let structural_tag = build_structural_tag(
.and_then(|tag| tag.to_json_string()) builder,
.map_err(|error| Error::StructuralTag { &tools,
message: error.to_report_string(), tool_choice,
})?; StructuralTagOptions::default().with_reasoning(false),
)
.and_then(|tag| tag.to_json_string())
.map_err(|error| Error::StructuralTag {
message: error.to_report_string(),
})?;
// Overwrite any existing structured output settings with the structural tag constraint. // Overwrite any existing structured output settings with the structural tag constraint.
request.sampling_params.structured_outputs = Some(StructuredOutputsParams { request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -141,7 +147,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build"); .expect("structural tag should build");
let tag = structural_tag_value(&request); let tag = structural_tag_value(&request);
@@ -154,7 +160,7 @@ mod tests {
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]); let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed"); .expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none()); assert!(request.sampling_params.structured_outputs.is_none());
@@ -169,7 +175,7 @@ mod tests {
}); });
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build"); .expect("structural tag should build");
let params = structured_outputs(&request); let params = structured_outputs(&request);
@@ -184,7 +190,7 @@ mod tests {
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build"); .expect("structural tag should build");
let tag = structural_tag_value(&request); let tag = structural_tag_value(&request);
@@ -201,7 +207,7 @@ mod tests {
}); });
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build"); .expect("structural tag should build");
let params = structured_outputs(&request); let params = structured_outputs(&request);
@@ -221,7 +227,7 @@ mod tests {
); );
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag should build"); .expect("structural tag should build");
let tag = structural_tag_value(&request).to_string(); let tag = structural_tag_value(&request).to_string();
@@ -234,7 +240,7 @@ mod tests {
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed"); .expect("structural tag decision should succeed");
assert!(request.sampling_params.structured_outputs.is_none()); assert!(request.sampling_params.structured_outputs.is_none());
@@ -249,7 +255,7 @@ mod tests {
}); });
let parser = qwen3_coder_parser(&request.tools); let parser = qwen3_coder_parser(&request.tools);
apply_structural_tag_constraint(&mut request, parser.structural_tag_model()) apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
.expect("structural tag decision should succeed"); .expect("structural tag decision should succeed");
let params = structured_outputs(&request); let params = structured_outputs(&request);
+14 -3
View File
@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
loops.into_iter().any(|loop_ast| { loops.into_iter().any(|loop_ast| {
matches!(loop_ast.target, Expr::Var(_)) matches!(loop_ast.target, Expr::Var(_))
&& message_varnames && (is_var_access(&loop_ast.iter, "content")
.iter() || message_varnames.iter().any(|varname| {
.any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))) is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
}))
}) })
} }
@@ -315,6 +316,16 @@ mod tests {
); );
} }
#[test]
fn detects_openai_template_with_content_parameter_loop() {
assert_eq!(
detect(
"{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
),
ChatTemplateContentFormat::OpenAi
);
}
#[test] #[test]
fn detects_openai_template_with_messages_alias() { fn detects_openai_template_with_messages_alias() {
assert_eq!( assert_eq!(
+20
View File
@@ -1309,6 +1309,26 @@ mod tests {
.assert_eq(&rendered); .assert_eq(&rendered);
} }
#[test]
fn qwen35_template_auto_detects_openai_multimodal_content() {
let mut request = image_request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_mm(
QWEN3_5_0_8B_TEMPLATE,
&request,
ChatTemplateContentFormatOption::Auto,
)
.unwrap();
expect![[r#"
Text(
"<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
)
"#]]
.assert_debug_eq(&rendered.prompt);
}
#[test] #[test]
fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() { fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
+1
View File
@@ -29,6 +29,7 @@ tokio-util.workspace = true
tracing.workspace = true tracing.workspace = true
tracing-subscriber.workspace = true tracing-subscriber.workspace = true
uuid.workspace = true uuid.workspace = true
vllm-bench.workspace = true
vllm-chat.workspace = true vllm-chat.workspace = true
vllm-engine-core-client.workspace = true vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true vllm-managed-engine.workspace = true
+11 -1
View File
@@ -79,13 +79,23 @@ impl Cli {
} }
/// Supported top-level CLI commands. /// Supported top-level CLI commands.
#[derive(Debug, Subcommand, PartialEq, Eq)] #[derive(Debug, Subcommand)]
pub enum Command { pub enum Command {
/// Run the Rust OpenAI frontend as a Python-supervised worker. /// Run the Rust OpenAI frontend as a Python-supervised worker.
Frontend(FrontendArgs), Frontend(FrontendArgs),
/// Launch a managed Python headless engine, then run the Rust OpenAI /// Launch a managed Python headless engine, then run the Rust OpenAI
/// frontend. /// frontend.
Serve(ServeArgs), Serve(ServeArgs),
/// Run vLLM benchmarks.
#[command(subcommand)]
Bench(BenchCommand),
}
/// Supported benchmark commands.
#[derive(Debug, Subcommand)]
pub enum BenchCommand {
/// Benchmark online serving throughput.
Serve(vllm_bench::BenchServeArgs),
} }
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for /// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
+21 -1
View File
@@ -5,7 +5,27 @@ use expect_test::expect;
use vllm_engine_core_client::TransportMode; use vllm_engine_core_client::TransportMode;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection}; use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{Cli, Command}; use super::{BenchCommand, Cli, Command};
#[test]
fn bench_serve_args_parse_without_managed_engine_repartition() {
let cli = Cli::try_parse_from([
"vllm-rs",
"bench",
"serve",
"--backend",
"openai-chat",
"--request-rate",
"inf",
])
.unwrap();
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
panic!("expected bench serve args");
};
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
assert!(args.request_rate.is_infinite());
}
#[test] #[test]
fn serve_args_forward_python_flags_with_separator() { fn serve_args_forward_python_flags_with_separator() {
+17 -7
View File
@@ -26,18 +26,21 @@ const RESET: &str = "\x1b[0m";
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] = const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
format_description!("[month]-[day] [hour]:[minute]:[second]"); format_description!("[month]-[day] [hour]:[minute]:[second]");
const PROCESS_LABEL: &str = "RustFrontend";
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary. /// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
pub(crate) fn init_tracing() { pub(crate) fn init_tracing(process_label: &str) {
let filter = build_targets_filter( let filter = build_targets_filter(
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(), env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
env::var("RUST_LOG").ok().as_deref(), env::var("RUST_LOG").ok().as_deref(),
); );
let formatter = VllmEventFormatter::new(); let formatter = VllmEventFormatter::new(process_label);
let _ = tracing_subscriber::registry() let _ = tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter)) .with(
tracing_subscriber::fmt::layer()
.event_format(formatter)
.with_writer(std::io::stderr)
.with_filter(filter),
)
.try_init(); .try_init();
} }
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
} }
impl VllmEventFormatter { impl VllmEventFormatter {
fn new() -> Self { fn new(process_label: &str) -> Self {
Self { Self {
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()), prefix: format!("({process_label} pid={})", process::id()),
timer: VllmLocalTimer::default(), timer: VllmLocalTimer::default(),
} }
} }
@@ -291,6 +294,13 @@ fn map_python_log_level(level: &str) -> LevelFilter {
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn formatter_prefix_uses_process_label() {
let formatter = VllmEventFormatter::new("Bench");
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
}
#[test] #[test]
fn rust_log_target_overrides_are_merged_with_vllm_default_level() { fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error")); let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
+14 -2
View File
@@ -5,6 +5,7 @@ mod cli;
mod logging; mod logging;
use std::env; use std::env;
use std::ffi::OsStr;
use std::process::ExitStatus; use std::process::ExitStatus;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -12,7 +13,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{info, warn}; use tracing::{info, warn};
use vllm_managed_engine::ManagedEngineHandle; use vllm_managed_engine::ManagedEngineHandle;
use crate::cli::{Cli, Command}; use crate::cli::{BenchCommand, Cli, Command};
#[global_allocator] #[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
} }
fn main() -> Result<()> { fn main() -> Result<()> {
logging::init_tracing(); let process_label =
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
"bench" => "Bench",
"serve" | "frontend" => "RustFrontend",
_ => "Rust",
};
logging::init_tracing(process_label);
let cli = Cli::parse(); let cli = Cli::parse();
let mut runtime = tokio::runtime::Builder::new_multi_thread(); let mut runtime = tokio::runtime::Builder::new_multi_thread();
@@ -100,6 +108,10 @@ fn main() -> Result<()> {
async fn async_main(cli: Cli) -> Result<()> { async fn async_main(cli: Cli) -> Result<()> {
match cli.command { match cli.command {
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await, Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
Command::Bench(BenchCommand::Serve(bench_args)) => {
vllm_bench::prepare_process();
vllm_bench::run(bench_args).await
}
Command::Serve(args) => { Command::Serve(args) => {
let handshake_port = args.managed_engine.resolve_handshake_port()?; let handshake_port = args.managed_engine.resolve_handshake_port()?;
@@ -460,6 +460,12 @@ impl EngineCoreClient {
self.inner.is_healthy() self.inner.is_healthy()
} }
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {
self.inner.subscribe_health()
}
/// Return the first persistent health error observed by the client, if any. /// Return the first persistent health error observed by the client, if any.
pub fn health_error(&self) -> Option<Arc<Error>> { pub fn health_error(&self) -> Option<Arc<Error>> {
self.inner.health_error() self.inner.health_error()
+20 -1
View File
@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;
use parking_lot::Mutex; use parking_lot::Mutex;
use thiserror_ext::AsReport as _; use thiserror_ext::AsReport as _;
use tokio::runtime::Handle; use tokio::runtime::Handle;
use tokio::sync::mpsc; use tokio::sync::{mpsc, watch};
use tracing::{debug, info, trace, warn}; use tracing::{debug, info, trace, warn};
use vllm_metrics::METRICS; use vllm_metrics::METRICS;
use zeromq::RouterSendHalf; use zeromq::RouterSendHalf;
@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {
request_reg: Mutex<RequestRegistry>, request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>, utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>, health_error: ArcSwapOption<Error>,
health_tx: watch::Sender<bool>,
} }
impl ClientInner { impl ClientInner {
@@ -57,6 +58,7 @@ impl ClientInner {
request_reg: Mutex::new(RequestRegistry::new(engines)), request_reg: Mutex::new(RequestRegistry::new(engines)),
utility_reg: Mutex::new(UtilityRegistry::default()), utility_reg: Mutex::new(UtilityRegistry::default()),
health_error: ArcSwapOption::empty(), health_error: ArcSwapOption::empty(),
health_tx: watch::Sender::new(true),
} }
} }
@@ -169,6 +171,7 @@ impl ClientInner {
/// persistent health error. /// persistent health error.
pub fn close_registries(&self, error: Arc<Error>) { pub fn close_registries(&self, error: Arc<Error>) {
let persistent_error = self.record_health_error(error); let persistent_error = self.record_health_error(error);
self.publish_unhealthy();
let request_senders = self.request_reg.lock().close(); let request_senders = self.request_reg.lock().close();
let utility_senders = self.utility_reg.lock().close(); let utility_senders = self.utility_reg.lock().close();
@@ -191,6 +194,12 @@ impl ClientInner {
self.health_error.load().is_none() self.health_error.load().is_none()
} }
/// Subscribe to engine health changes. The current value is `true` while
/// the client is healthy and changes permanently to `false` on failure.
pub fn subscribe_health(&self) -> watch::Receiver<bool> {
self.health_tx.subscribe()
}
/// Resolve one utility output to the waiting caller. Returns `true` if a /// Resolve one utility output to the waiting caller. Returns `true` if a
/// waiting caller existed. /// waiting caller existed.
pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool { pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {
@@ -280,6 +289,11 @@ impl ClientInner {
.expect("health error must be recorded before registries close") .expect("health error must be recorded before registries close")
} }
/// Publish the sticky healthy-to-unhealthy transition.
fn publish_unhealthy(&self) {
self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));
}
/// Assert there is a recorded health error and return a `Shared` variant /// Assert there is a recorded health error and return a `Shared` variant
/// wrapping it for error returns when the client is already closed. /// wrapping it for error returns when the client is already closed.
fn closed_error(&self) -> Error { fn closed_error(&self) -> Error {
@@ -461,13 +475,18 @@ mod tests {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn close_registries_records_first_health_error_only() { async fn close_registries_records_first_health_error_only() {
let inner = test_inner().await; let inner = test_inner().await;
let mut health = inner.subscribe_health();
assert!(*health.borrow());
inner.close_registries(Arc::new(Error::EngineCoreDead)); inner.close_registries(Arc::new(Error::EngineCoreDead));
health.changed().await.expect("health sender remains open");
assert!(!inner.is_healthy()); assert!(!inner.is_healthy());
assert!(!*health.borrow());
assert!(matches!( assert!(matches!(
inner.health_error().as_deref(), inner.health_error().as_deref(),
Some(Error::EngineCoreDead) Some(Error::EngineCoreDead)
)); ));
assert!(!*inner.subscribe_health().borrow());
inner.close_registries(Arc::new(client_closed!("shutdown"))); inner.close_registries(Arc::new(client_closed!("shutdown")));
assert!(matches!( assert!(matches!(
@@ -269,6 +269,19 @@ impl WireLogprobs {
); );
} }
// Empty position lists may be encoded as either [0, 0] or [0, k + 1].
if token_ids.rows == 0 {
return Ok(Logprobs {
positions: Vec::new(),
});
}
if token_ids.cols == 0 {
bail_ext_value_decode!(
"{field_prefix}: zero-column logprobs payload with {} rows",
token_ids.rows
);
}
let mut positions = Vec::with_capacity(token_ids.rows); let mut positions = Vec::with_capacity(token_ids.rows);
for ((token_ids_row, logprobs_row), sampled_rank) in token_ids for ((token_ids_row, logprobs_row), sampled_rank) in token_ids
.data .data
@@ -303,3 +303,49 @@ fn rejects_non_none_cu_num_generated_tokens() {
"messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]" "messagepack ext value decode failed: new_logprobs.cu_num_generated_tokens: expected None for per-request engine-core logprobs payload, got [0, 1]"
); );
} }
#[test]
fn decodes_zero_row_logprobs_as_empty() {
for shape in [[0usize, 0], [0, 3]] {
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
None,
Some(Value::Array(vec![
ndarray_value("<i8", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<f4", &shape, Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[0], Value::Ext(3, Vec::new())),
Value::Nil,
])),
)))];
let decoded = decode_engine_core_outputs(&frames).unwrap().into_request_batch().unwrap();
let logprobs = decoded.outputs[0]
.new_prompt_logprobs_tensors
.clone()
.unwrap()
.into_direct()
.unwrap();
assert!(logprobs.is_empty());
}
}
#[test]
fn rejects_zero_column_logprobs_with_rows() {
let ranks = Value::Ext(3, vec![1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
let frames = vec![Bytes::from(encode_value(&output_wire_with_custom_fields(
Some(Value::Array(vec![
ndarray_value("<i8", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<f4", &[2, 0], Value::Ext(3, Vec::new())),
ndarray_value("<i8", &[2], ranks),
Value::Nil,
])),
None,
)))];
let error = decode_engine_core_outputs(&frames).unwrap_err();
let crate::error::Error::ExtValueDecode { message } = &error else {
panic!("expected ExtValueDecode");
};
assert_eq!(
message,
"new_logprobs: zero-column logprobs payload with 2 rows"
);
}
@@ -55,52 +55,57 @@ pub struct WireNdArray {
impl WireNdArray { impl WireNdArray {
/// Build a float32 tensor/ndarray backed by native-endian raw-view bytes. /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> { pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?; validate_element_count(&shape, data.len())?;
Ok(Self { Ok(Self {
dtype: "float32".to_string(), dtype: "float32".to_string(),
shape, shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)), data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
}) })
} }
/// Build a float16 tensor/ndarray backed by native-endian raw-view bytes. /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> { pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?; validate_element_count(&shape, data.len())?;
Ok(Self { Ok(Self {
dtype: "float16".to_string(), dtype: "float16".to_string(),
shape, shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)), data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
}) })
} }
/// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes. /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> { pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?; validate_element_count(&shape, data.len())?;
Ok(Self { Ok(Self {
dtype: "bfloat16".to_string(), dtype: "bfloat16".to_string(),
shape, shape,
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)), data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
}) })
} }
/// Build an int64 tensor/ndarray backed by native-endian raw-view bytes. /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> { pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?; validate_element_count(&shape, data.len())?;
Ok(Self { Ok(Self {
dtype: "int64".to_string(), dtype: "int64".to_string(),
shape, shape,
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)), data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
}) })
} }
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes. /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> { pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?; validate_element_count(&shape, data.len())?;
Ok(Self { Ok(Self {
dtype: "uint32".to_string(), dtype: "uint32".to_string(),
shape, shape,
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)), data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
}) })
} }
+3 -3
View File
@@ -4,7 +4,7 @@
use std::sync::Arc; use std::sync::Arc;
use vllm_parser::tool::{ use vllm_parser::tool::{
Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput, Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
}; };
use vllm_parser::unified::{ use vllm_parser::unified::{
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput, UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
self.inner.preserve_special_tokens() self.inner.preserve_special_tokens()
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.inner.structural_tag_model() self.inner.structural_tag_builder()
} }
fn tool_call_id(&self, tool_index: usize) -> Option<&str> { fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens}; use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.2 models. /// Tool parser for DeepSeek V3.2 models.
/// ///
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
true true
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::DeepSeekV32) Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekDsmlToolParser, DsmlTokens}; use super::{DeepSeekDsmlToolParser, DsmlTokens};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V4 models. /// Tool parser for DeepSeek V4 models.
/// ///
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
true true
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::DeepSeekV4) Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -73,7 +73,7 @@ mod tests {
use super::DeepSeekV4ToolParser; use super::DeepSeekV4ToolParser;
use crate::tool::test_utils::{collect_stream, test_tools}; use crate::tool::test_utils::{collect_stream, test_tools};
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _}; use crate::tool::{ToolParser, ToolParserTestExt as _};
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
let params = params let params = params
@@ -91,13 +91,10 @@ mod tests {
} }
#[test] #[test]
fn deepseek_v4_exposes_structural_tag_model() { fn deepseek_v4_exposes_structural_tag_builder() {
let parser = DeepSeekV4ToolParser::new(&test_tools()); let parser = DeepSeekV4ToolParser::new(&test_tools());
assert_eq!( assert!(parser.structural_tag_builder().is_some());
parser.structural_tag_model(),
Some(StructuralTagModel::DeepSeekV4)
);
} }
#[test] #[test]
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3 JSON-fenced tool calls. /// Tool parser for DeepSeek V3 JSON-fenced tool calls.
/// ///
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::DeepSeekR1) Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for DeepSeek V3.1 raw JSON tool calls. /// Tool parser for DeepSeek V3.1 raw JSON tool calls.
/// ///
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::DeepSeekV31) Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{GlmXmlToolParser, Separator}; use super::{GlmXmlToolParser, Separator};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Tool parser for GLM-4.7 MoE XML-style tool calls. /// Tool parser for GLM-4.7 MoE XML-style tool calls.
/// ///
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Glm47) Some(xgrammar_structural_tag::Model::Glm47.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas; use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool}; use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<tool_calls>"; const TOOL_CALLS_START: &str = "<tool_calls>";
const TOOL_CALLS_END: &str = "</tool_calls>"; const TOOL_CALLS_END: &str = "</tool_calls>";
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::HyV3) Some(xgrammar_structural_tag::Model::HyV3.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig { const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Hermes", parser_name: "Hermes",
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Hermes) Some(xgrammar_structural_tag::Model::Hermes.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+5 -3
View File
@@ -12,7 +12,9 @@ use super::{
argument_delta_event, tool_call_header_event, argument_delta_event, tool_call_header_event,
}; };
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event}; use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; use crate::tool::{
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
};
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
enum LlamaJsonMode { enum LlamaJsonMode {
@@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Llama) Some(xgrammar_structural_tag::Model::Llama.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -2,7 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "Qwen XML", parser_name: "Qwen XML",
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Qwen3) Some(xgrammar_structural_tag::Model::Qwen3.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object}; use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool}; use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>"; const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>"; const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
true true
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Kimi) Some(xgrammar_structural_tag::Model::Kimi.builder())
} }
fn tool_call_id(&self, tool_index: usize) -> Option<&str> { fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
+3 -3
View File
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
use super::parameters::ToolSchemas; use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::{StructuralTagModel, Tool}; use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<minimax:tool_call>"; const TOOL_CALL_START: &str = "<minimax:tool_call>";
const TOOL_CALL_END: &str = "</minimax:tool_call>"; const TOOL_CALL_END: &str = "</minimax:tool_call>";
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Minimax) Some(xgrammar_structural_tag::Model::Minimax.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+3 -3
View File
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
pub use seed_oss::SeedOssToolParser; pub use seed_oss::SeedOssToolParser;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
pub use xgrammar_structural_tag::Model as StructuralTagModel; pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
use crate::utils; use crate::utils;
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
false false
} }
/// Return the xgrammar structural-tag model used for strict tool calling. /// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None None
} }
+7 -10
View File
@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
use super::parameters::ToolSchemas; use super::parameters::ToolSchemas;
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
use crate::tool::Tool; use crate::tool::{StructuralTagBuilder, Tool};
const TOOL_CALL_START: &str = "<tool_call>"; const TOOL_CALL_START: &str = "<tool_call>";
const TOOL_CALL_END: &str = "</tool_call>"; const TOOL_CALL_END: &str = "</tool_call>";
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
Ok(Box::new(Self::new(tools))) Ok(Box::new(Self::new(tools)))
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
Some(StructuralTagModel::Qwen3Coder) Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
} }
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -294,7 +294,7 @@ mod tests {
use serde_json::{Value, json}; use serde_json::{Value, json};
use thiserror_ext::AsReport; use thiserror_ext::AsReport;
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser}; use super::{Qwen3CoderToolParser, ToolParser};
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
@@ -308,13 +308,10 @@ mod tests {
} }
#[test] #[test]
fn qwen_coder_exposes_structural_tag_model() { fn qwen_coder_exposes_structural_tag_builder() {
let parser = Qwen3CoderToolParser::new(&test_tools()); let parser = Qwen3CoderToolParser::new(&test_tools());
assert_eq!( assert!(parser.structural_tag_builder().is_some());
parser.structural_tag_model(),
Some(StructuralTagModel::Qwen3Coder)
);
} }
#[test] #[test]
+4 -7
View File
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
use crate::reasoning::ReasoningParser; use crate::reasoning::ReasoningParser;
use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput}; use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
/// Unified parser that composes existing reasoning and tool parsers. /// Unified parser that composes existing reasoning and tool parsers.
pub struct CombinedParser { pub struct CombinedParser {
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
|| self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens()) || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
} }
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
self.tool.as_ref().and_then(|parser| parser.structural_tag_model()) self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
} }
fn tool_call_id(&self, tool_index: usize) -> Option<&str> { fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -269,10 +269,7 @@ mod tests {
fn combined_parser_emits_tool_calls_from_visible_content() { fn combined_parser_emits_tool_calls_from_visible_content() {
let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap(); let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
let mut parser = CombinedParser::new(None, Some(tool)); let mut parser = CombinedParser::new(None, Some(tool));
assert!(matches!( assert!(parser.structural_tag_builder().is_some());
parser.structural_tag_model(),
Some(crate::tool::StructuralTagModel::Qwen3)
));
let output = collect( let output = collect(
&mut parser, &mut parser,
+3 -3
View File
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
use crate::reasoning::ReasoningError; use crate::reasoning::ReasoningError;
use crate::tool::{ use crate::tool::{
StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput, StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
}; };
/// Result alias for unified parser operations. /// Result alias for unified parser operations.
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
false false
} }
/// Return the xgrammar structural-tag model used for strict tool calling. /// Return the xgrammar structural-tag builder used for strict tool calling.
fn structural_tag_model(&self) -> Option<StructuralTagModel> { fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
None None
} }
+1
View File
@@ -35,6 +35,7 @@ tokio-openssl.workspace = true
tokio-stream.workspace = true tokio-stream.workspace = true
tokio-util.workspace = true tokio-util.workspace = true
tonic.workspace = true tonic.workspace = true
tonic-health.workspace = true
tonic-prost.workspace = true tonic-prost.workspace = true
tower.workspace = true tower.workspace = true
tower-http.workspace = true tower-http.workspace = true
+70
View File
@@ -0,0 +1,70 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tonic::server::NamedService;
use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::GenerateGrpcService;
pub(crate) async fn monitor_health(
mut health_reporter: HealthReporter,
mut engine_health: watch::Receiver<bool>,
shutdown: CancellationToken,
) {
let generate_service = GenerateGrpcService::NAME;
let status = ServingStatus::NotServing;
let health_event_first = tokio::select! {
result = engine_health.wait_for(|healthy| !*healthy) => {
match result {
Ok(_) => warn!(
generate_service,
overall_service = true,
status = ?status,
reason = "engine_unhealthy",
"marking gRPC health services as not serving"
),
Err(error) => warn!(
%error,
generate_service,
overall_service = true,
status = ?status,
reason = "health_channel_closed",
"engine health channel closed; marking gRPC health services as not serving"
),
}
true
}
_ = shutdown.cancelled() => {
info!(
generate_service,
overall_service = true,
status = ?status,
reason = "server_shutdown",
"server shutting down; marking gRPC health services as not serving"
);
false
}
};
health_reporter.set_not_serving::<GenerateGrpcService>().await;
// Generate is currently the only engine-backed gRPC service, so overall
// server health intentionally mirrors it.
health_reporter.set_service_status("", status).await;
if health_event_first {
shutdown.cancelled().await;
info!(
generate_service,
overall_service = true,
reason = "server_shutdown",
"server shutting down; closing gRPC health watches"
);
}
health_reporter.clear_service_status(generate_service).await;
health_reporter.clear_service_status("").await;
}
+4
View File
@@ -4,6 +4,7 @@
//! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade. //! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.
mod convert; mod convert;
mod health;
use std::pin::Pin; use std::pin::Pin;
use std::sync::Arc; use std::sync::Arc;
@@ -24,8 +25,11 @@ pub mod pb {
tonic::include_proto!("vllm"); tonic::include_proto!("vllm");
} }
pub(crate) use health::monitor_health;
pub use pb::generate_server::GenerateServer; pub use pb::generate_server::GenerateServer;
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
+175 -12
View File
@@ -16,6 +16,10 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use tokio::net::TcpStream; use tokio::net::TcpStream;
use tokio_openssl::SslStream; use tokio_openssl::SslStream;
use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri}; use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};
use tonic_health::pb::HealthCheckRequest;
use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;
use tonic_health::pb::health_client::HealthClient;
use tonic_health::server::health_reporter;
use tower::service_fn; use tower::service_fn;
use vllm_chat::{ use vllm_chat::{
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor, ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
@@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend {
async fn setup_grpc_service( async fn setup_grpc_service(
engine_id: impl Into<EngineId>, engine_id: impl Into<EngineId>,
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>, output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) { ) -> (
GenerateServer<GenerateServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
let ipc = IpcNamespace::new().expect("create ipc namespace"); let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint(); let handshake_address = ipc.handshake_endpoint();
let engine_id = engine_id.into(); let engine_id = engine_id.into();
@@ -232,6 +240,7 @@ async fn setup_grpc_service(
) )
.await .await
.expect("connect client"); .expect("connect client");
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend( let chat = ChatLlm::from_shared_backend(
test_llm(client), test_llm(client),
@@ -240,6 +249,7 @@ async fn setup_grpc_service(
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
( (
GenerateServer::new(GenerateServiceImpl::new(state)), GenerateServer::new(GenerateServiceImpl::new(state)),
engine_health,
engine_task, engine_task,
) )
} }
@@ -254,25 +264,51 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>, tokio::task::JoinHandle<()>,
MockEngineTask, MockEngineTask,
) { ) {
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
svc,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
(GenerateClient::new(channel), server_task, engine_task)
}
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr"); let addr = listener.local_addr().expect("local addr");
let server_task = tokio::spawn(async move { let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
TonicServer::builder() let server = TonicServer::builder()
.add_service(svc) .add_service(health_service)
.serve_with_incoming(incoming) .add_service(generate_service)
.await .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
.expect("grpc server"); let health_monitor =
super::monitor_health(health_reporter, engine_health, shutdown.clone());
let server = async move {
let result = server.await;
shutdown.cancel();
result
};
let (server_result, ()) = tokio::join!(server, health_monitor);
server_result.expect("grpc server");
}); });
let grpc_client = GenerateClient::connect(format!("http://{addr}")) let channel = Endpoint::from_shared(format!("http://{addr}"))
.expect("grpc endpoint")
.connect()
.await .await
.expect("connect grpc client"); .expect("connect grpc channel");
(grpc_client, server_task, engine_task) (channel, server_task)
} }
/// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode). /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).
@@ -283,7 +319,7 @@ async fn grpc_tls_test_server(
certs: &TestCerts, certs: &TestCerts,
cert_reqs: i32, cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await; let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs)) let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config"); .expect("build grpc tls config");
@@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>, engine_id: impl Into<EngineId>,
keepalive: Option<Duration>, keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) { ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await; let (svc, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr").to_string(); let addr = listener.local_addr().expect("local addr").to_string();
@@ -1035,3 +1072,129 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort(); server_task.abort();
} }
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
service
};
let mut stream = health_client
.watch(HealthCheckRequest {
service: service.to_string(),
})
.await
.unwrap_or_else(|error| {
panic!("failed to start health watch for {service_label}: {error}")
})
.into_inner();
let initial = stream
.message()
.await
.unwrap_or_else(|error| {
panic!("failed to read initial health status for {service_label}: {error}")
})
.unwrap_or_else(|| {
panic!("health watch for {service_label} ended before its initial status")
});
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for {service_label}"
);
health_streams.push((service_label, stream));
}
engine_health_tx.send(false).expect("publish unhealthy engine state");
for (service_label, mut stream) in health_streams {
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))
.unwrap_or_else(|error| {
panic!("failed to read health update for {service_label}: {error}")
})
.unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected health status for {service_label}"
);
}
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
service: "vllm.Generate".to_string(),
})
.await
.expect("start health watch for vllm.Generate")
.into_inner();
let initial = stream
.message()
.await
.expect("read initial health status for vllm.Generate")
.expect("health watch ended before its initial status");
assert_eq!(
initial.status,
HealthServingStatus::Serving as i32,
"unexpected initial health status for vllm.Generate"
);
shutdown.cancel();
let update = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for shutdown health update for vllm.Generate")
.expect("failed to read shutdown health update for vllm.Generate")
.expect("health watch ended before its shutdown update");
assert_eq!(
update.status,
HealthServingStatus::NotServing as i32,
"unexpected shutdown health status for vllm.Generate"
);
let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for vllm.Generate health watch to close")
.expect("failed while closing vllm.Generate health watch");
assert!(
stream_end.is_none(),
"vllm.Generate health watch remained open"
);
tokio::time::timeout(Duration::from_secs(2), server_task)
.await
.expect("timed out waiting for gRPC server shutdown")
.expect("gRPC server task failed");
}
+28 -14
View File
@@ -39,6 +39,7 @@ use tokio::net::TcpListener;
use tokio::time::{Instant, sleep_until}; use tokio::time::{Instant, sleep_until};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use tonic::transport::Server as TonicServer; use tonic::transport::Server as TonicServer;
use tonic_health::server::health_reporter;
use tower::ServiceExt as _; use tower::ServiceExt as _;
use tracing::{info, trace, warn}; use tracing::{info, trace, warn};
use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};
@@ -203,14 +204,19 @@ where
.map(tls::build_grpc_server_config) .map(tls::build_grpc_server_config)
.transpose() .transpose()
.context("invalid gRPC TLS configuration")?; .context("invalid gRPC TLS configuration")?;
let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone())); let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder() let svc = TonicServer::builder()
.http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone())) .layer(middleware::request_runtime_layer(state.clone()))
.add_service(svc); .add_service(health_service)
.add_service(generate_service);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls)) Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
} else { } else {
None None
}; };
@@ -294,7 +300,8 @@ where
let server_shutdown = server_shutdown.clone(); let server_shutdown = server_shutdown.clone();
let force_shutdown = force_shutdown.clone(); let force_shutdown = force_shutdown.clone();
async move { async move {
let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else { let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup
else {
// No gRPC configured: just wait for shutdown so we do not race the // No gRPC configured: just wait for shutdown so we do not race the
// join! by resolving early and tripping the cancellation token. // join! by resolving early and tripping the cancellation token.
shutdown.cancelled().await; shutdown.cancelled().await;
@@ -304,19 +311,26 @@ where
Some(context) => MaybeTlsListener::tls(grpc_listener, context), Some(context) => MaybeTlsListener::tls(grpc_listener, context),
None => MaybeTlsListener::plain(grpc_listener), None => MaybeTlsListener::plain(grpc_listener),
}; };
let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()); let server =
svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);
let result = tokio::select! { let server = async move {
result = server => { let result = tokio::select! {
result.context("gRPC server failed") result = server => {
} result.context("gRPC server failed")
_ = force_shutdown.cancelled() => { }
warn!("gRPC graceful shutdown deadline elapsed; aborting server"); _ = force_shutdown.cancelled() => {
Ok(()) warn!("gRPC graceful shutdown deadline elapsed; aborting server");
} Ok(())
}
};
server_shutdown.cancel();
result
}; };
server_shutdown.cancel(); let (result, ()) = tokio::join!(server, health_monitor);
result result
} }
}; };
@@ -238,13 +238,18 @@ fn collect_generate(
None None
}; };
let prompt_logprobs = if include_prompt_logprobs { let prompt_logprobs = if include_prompt_logprobs {
let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| { match collected.prompt_logprobs.as_ref() {
ApiError::server_error( Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
"raw generate response requested prompt_logprobs but generation returned none" // A single-token prompt has no scored positions; same mapping
.to_string(), // as /v1/completions.
) None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
})?; None => {
Some(raw_prompt_logprobs_to_maps(prompt_logprobs)) return Err(ApiError::server_error(
"raw generate response requested prompt_logprobs but generation returned none"
.to_string(),
));
}
}
} else { } else {
None None
}; };
@@ -472,4 +477,48 @@ mod tests {
Some(2) Some(2)
); );
} }
#[test]
fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
request_id: "raw-1".to_string(),
prompt_logprobs: None,
token_ids: vec![3],
logprobs: None,
finish_reason: FinishReason::stop_eos(),
usage: vllm_llm::TokenUsage {
prompt_token_count: prompt_token_ids.len(),
output_token_count: 1,
cached_token_count: 0,
},
kv_transfer_params: None,
ec_transfer_params: None,
prompt_token_ids,
};
let response = collect_generate(
output_without_payload(vec![9707]),
"raw-1".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect("single-token prompt without payload maps to [None]");
let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
assert_eq!(prompt_logprobs.len(), 1);
assert!(prompt_logprobs[0].is_none());
collect_generate(
output_without_payload(vec![9707, 11]),
"raw-2".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect_err("multi-token prompt without payload is an engine failure");
}
} }
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
ChatMessageDelta, ChatMessageDelta,
}; };
use crate::routes::openai::utils::logprobs::{ use crate::routes::openai::utils::logprobs::{
decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps, decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
}; };
use crate::routes::openai::utils::types::{ use crate::routes::openai::utils::types::{
ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage, ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
None None
}; };
let prompt_logprobs = if include_prompt_logprobs { let prompt_logprobs = if include_prompt_logprobs {
Some(decoded_prompt_logprobs_to_maps( Some(prompt_logprobs_to_maps(
prompt_logprobs.as_ref().ok_or_else(|| { prompt_logprobs.as_ref(),
server_error!( &prompt_token_ids,
"chat response requested prompt_logprobs but generation returned none"
)
})?,
return_tokens_as_token_ids, return_tokens_as_token_ids,
)) )?)
} else { } else {
None None
}; };
@@ -5,7 +5,6 @@ mod convert;
mod types; mod types;
mod validate; mod validate;
use std::collections::HashMap;
use std::convert::Infallible; use std::convert::Infallible;
use std::result::Result; use std::result::Result;
use std::sync::Arc; use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{
use self::convert::{ResponseOptions, prepare_completion_request}; use self::convert::{ResponseOptions, prepare_completion_request};
use super::utils::logprobs::{ use super::utils::logprobs::{
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps, collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
decoded_prompt_logprobs_to_openai, text_len, prompt_logprobs_to_maps, text_len,
}; };
use super::utils::types::Usage; use super::utils::types::Usage;
use crate::config::ApiServerOptions; use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
)) ))
} }
fn prompt_logprobs_to_maps(
prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool,
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
if let Some(prompt_logprobs) = prompt_logprobs {
return Ok(decoded_prompt_logprobs_to_maps(
prompt_logprobs,
return_tokens_as_token_ids,
));
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"completion response requested prompt_logprobs but generation returned none"
))
}
fn usage_chunk( fn usage_chunk(
request_id: &str, request_id: &str,
response_model: &str, response_model: &str,
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
}) })
} }
/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response /// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
/// shape. /// missing single-token payload as `[None]`.
pub fn decoded_prompt_logprobs_to_maps( pub fn prompt_logprobs_to_maps(
prompt_logprobs: &DecodedPromptLogprobs, prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool, return_tokens_as_token_ids: bool,
) -> Vec<Option<HashMap<String, f32>>> { ) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
std::iter::once(None) if let Some(prompt_logprobs) = prompt_logprobs {
.chain(prompt_logprobs.scored_positions.iter().map(|position| { return Ok(std::iter::once(None)
Some(position_top_logprobs_map( .chain(prompt_logprobs.scored_positions.iter().map(|position| {
position, Some(position_top_logprobs_map(
return_tokens_as_token_ids, position,
)) return_tokens_as_token_ids,
})) ))
.collect() }))
.collect());
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"prompt_logprobs were requested but generation returned none"
))
} }
/// Convert decoded token-position logprobs into the OpenAI chat `logprobs` /// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
mod tests { mod tests {
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::decoded_logprobs_to_openai_chat; use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
#[test]
fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
prompt_logprobs_to_maps(None, &[9707, 11], false)
.expect_err("multi-token prompt without payload is an engine failure");
}
fn sample_logprobs() -> DecodedLogprobs { fn sample_logprobs() -> DecodedLogprobs {
DecodedLogprobs { DecodedLogprobs {

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