Compare commits

..
Author SHA1 Message Date
Tyler Michael SmithandGitHub 6257a61c3d Merge branch 'main' into worktree-fix-cudagraph-flaky 2026-07-06 13:59:00 -04:00
Xiaohong (Sean) ChenGitHubAndreas Karatzasmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
9fde043f54 [Kernel][Helion][1/N] Add Helion kernel for silu_and_mul_per_block_quant (#43994)
Signed-off-by: Sean Chen <seachen@redhat.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-07 00:19:01 +08:00
24dd2aec81 [Bugfix] Preserve FP8 indexer WK pairs across incremental load_weights (#46168)
Signed-off-by: lcheng <lcheng321@gatech.edu>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>
2026-07-06 09:16:46 -07:00
RanranandGitHub 3ee9eea928 [macOS][CPU][Installation] Fix the broken installation of vllm 0.24.0 in macos + cpu (#47457)
Signed-off-by: Ranran Haoran Zhang <ranranhaoranzhang@gmail.com>
2026-07-06 08:59:16 -07:00
Harry MellorGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
5bce653e09 Make the Transformers modeling backend as fast as native vLLM (#47187)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 16:59:14 +01:00
Kevin_XiongGitHubCodexIsotr0pymergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Isotr0pyIsotr0py
5ad11172b7 [perf]Add fused Kimi image preprocessing (#47416)
Signed-off-by: Kevin-XiongC <kevin_xiong1997@outlook.com>
Signed-off-by: Kevin_Xiong <kevin_xiong1997@outlook.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
2026-07-06 08:46:32 -07:00
Wentao YeandGitHub f70caef48b [Perf] Cache token_to_req_indices for dsv4, 5x~6x kernel performance improvement (#47474)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-06 11:17:46 -04:00
Laurent-ZhangGitHubClaudemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
8d8ec38361 [Bugfix][Spec Decode] Add missing draft_id_to_target_id to DSparkDeepseekV4ForCausalLM (#47429)
Signed-off-by: Laurent-Zhang <zhangdongsheng80@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 10:55:47 -04:00
Wentao YeandGitHub b1c6dba558 [Refactor] Remove multiple dead code (#47329)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-06 07:54:08 -07:00
jescoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
598d51153a [Bugfix][Distributed] Delegate MNNVL allreduce one-shot selection (#47589)
Signed-off-by: jesco-absolut <team@srswti.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 07:47:06 -07:00
Yifan QiaoandGitHub 095adf1fdc [Bugfix] Fix int32 overflow in triton_decode_attention page offsets (#47671)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-07-06 10:36:15 -04:00
Harry MellorandGitHub 51ee564e56 [CI] Skip test for checkpoint that was deleted (#47748)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-06 07:24:09 -07:00
Ting SUNGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
373eb314af [Bugfix][Core] Fix num_output_placeholders underflow with async scheduling + spec decode (#46066)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 13:50:38 +00:00
641cb59592 [Doc] Clarify fastokens availability (#45813)
Signed-off-by: LjjJzd <3542531707@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-06 13:33:05 +00:00
07f9baf756 Revert "[Platform] Replace torch.cuda.Event with torch.Event (#47140)" (#47668)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-06 14:18:33 +01:00
7a90eb98ab [Bugfix] [Gemma4] Fix Gemma4 MTP draft model layers ignoring quant_config (#47091)
Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com>
Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
2026-07-06 14:04:00 +01:00
Bugen ZhaoGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
8f4c69b222 [Rust Frontend] Cache metric handles for scheduler & request stats (#47444)
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-06 13:02:59 +00:00
8b79971bb9 attention: pass None for unused args in unified attention TD path (#43597)
Signed-off-by: Artur Fierka <artur.fierka@intel.com>
Co-authored-by: quinnlp <quinnlp@users.noreply.github.com>
2026-07-06 21:01:21 +08:00
Nick HillandGitHub f676808ba0 [CI] Use TTY for AMD CI tests for colored buildkite logs (#47730)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-07-06 20:50:29 +08:00
Qiming ZhangandGitHub 98e4726a14 [fix][run_batch]: respect proxy env vars when downloading media URLs (#47697)
Signed-off-by: mauyuyuace <qiming1.zhang@intel.com>
2026-07-06 12:45:48 +00:00
BadrBasowidandGitHub 740f379fae [ROCm][AITER] Directly Implement AITER Custom All-reduce in CudaCommunicator (#46065)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
2026-07-06 12:16:32 +00:00
Alexis K.andGitHub 40cc2e8327 [Bugfix] Return HTTP 422 for unprocessable image URLs instead of 500 (#47165)
Signed-off-by: Alexis Kinsella <alexis.kinsella@gmail.com>
2026-07-06 11:56:23 +00:00
Juan Pérez de AlgabaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
ba22152096 fix(security): block request-level GPU video backend selection withou… (#47259)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 02:36:49 -07:00
Yan MaandGitHub 90ce3a09be [bugfix] fix MOSS-Audio deepstack_input_embeds initialization in PP (#47607)
Signed-off-by: Yan Ma <yan.ma@intel.com>
2026-07-06 17:15:50 +08:00
26c754d847 [XPU][Bugfix] Do not transpose weight_scale_inv at load time (#47116)
Signed-off-by: Ma Jian <jian1.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-06 17:15:26 +08:00
Sungjae LeeandGitHub 3d7f357ebf [Doc] docs: fix note formatting for pooling models (#47701)
Signed-off-by: Sungjae Lee <33976427+llsj14@users.noreply.github.com>
Signed-off-by: Sungjae Lee <sung-jae.lee@navercorp.com>
2026-07-06 09:01:10 +00:00
liuzhenweiandGitHub 736f1a5907 [XPU] Route mm_prefix models to Triton attention backend (#47688)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
2026-07-06 16:52:44 +08:00
Li, JiangandGitHub 344609ab17 [CI/Build] Fix pre-commit check (#47695)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-07-06 08:24:24 +00:00
xiaozhoupyandGitHub d039c17114 [Bugfix] Recycle post-final-norm hidden in GLM MTP (single norm) (#47448) 2026-07-06 01:07:56 -07:00
xiangdongandGitHub cdab28319f [XPU][CI]Add agent tags for Basic Models Tests (Initialization) in Intel GPU CI (#47675)
Signed-off-by: zengxian <xiangdong.zeng@intel.com>
2026-07-06 15:15:45 +08:00
Qiming ZhangandGitHub 2fa10566e3 [Core][DP] Rotate load-balancer tie-break to avoid systematic engine bias (#47420)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
2026-07-06 07:09:16 +00:00
Andreas KaratzasandGitHub fb265fc8fb [ROCm][CI] Increasing parallelism in Basic Models Tests (Extra Initialization) (#47591)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-07-06 15:06:16 +08:00
Andreas KaratzasandGitHub 8f0e75e16b [ROCm][CI] Adding nixl multiconn (#47481)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-07-06 15:04:58 +08:00
98ba9b9583 [Frontend] Support OpenAI Responses API namespace tools (#47024)
Signed-off-by: zhongjing123 <jimzhong5193@gmail.com>
Co-authored-by: zhongjing123 <jimzhong5193@gmail.com>
2026-07-06 06:21:27 +00:00
velonica0andGitHub 990c2a0187 [RISC-V] Enable BF16 on VLEN=256 hardware (#45243)
Signed-off-by: velonica0 <like@mail.nankai.edu.cn>
2026-07-06 06:05:16 +00:00
e433634c78 [Performance][Hardware][RISC-V] Reduce LMUL pressure in INT4 LUT dequant (#47538)
Signed-off-by: liutong <liutong@iscas.ac.cn>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-06 05:58:56 +00:00
16f8110935 [Bugfix][CPU][RISC-V] Fix VLEN detection for RVV attention path (#47532)
Signed-off-by: liutong <liutong@iscas.ac.cn>
Co-authored-by: Claude <noreply@anthropic.com>
2026-07-06 05:58:03 +00:00
Zhenzhong XuGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d9c1767cd4 [INC][ARK] Direct Register Custom Op for ARK (#46361)
Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-06 13:45:50 +08:00
Li, JiangandGitHub e9cc1fd093 [CI/Build][CPU] Remove global extra index (#47687)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-07-06 13:42:01 +08:00
Fadi ArafehandGitHub f1073c050c [CPU][BugFix] Multiple fixes to w4a8_int8 CPU MoE path (#46739)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
2026-07-06 05:39:20 +00:00
Qiming ZhangandGitHub 394edc8108 [XPU] limit max-num-seqs in test_lmeval.py for XPU (#47682)
Signed-off-by: mauyuyuace <qiming1.zhang@intel.com>
2026-07-06 05:34:16 +00:00
69715823df [Test][XPU] Skip fork in kv_sharing_fast_prefill test on XPU (#47406)
Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-06 11:32:26 +08:00
Chaojun ZhangandGitHub 6569df6a3e [Test][LoRA] Use lightweight CPU reference and skip heavy cleanup in punica ops tests (#47534)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
2026-07-06 11:29:59 +08:00
f2aaf59151 [Feature] Support MTP speculative decoding for Bailing hybrid models (#44880)
Signed-off-by: zc02384840 <zc02384840@antgroup.com>
Co-authored-by: zc02384840 <zc02384840@antgroup.com>
2026-07-06 10:38:50 +08:00
95a248faed [Attention Backend] HPC_ATTN backend support mtp and dynamic scheduled attention (#47433)
Signed-off-by: chengvjiang <chengvjiang@tencent.com>
Co-authored-by: chengvjiang <chengvjiang@tencent.com>
2026-07-05 18:18:25 -07:00
Chaojun ZhangGitHubAndreas Karatzasmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Kunshang Ji
d2ec433e37 [XPU] Fix Eagle3 initialization on XPU (#43957)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-07-06 08:46:05 +08:00
Łukasz ŚlusarczykGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Kunshang Ji
78a04c208d [XPU] Fix CUDA API shims breaking Torch Dynamo during AOT compile (#43092)
Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@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-06 08:29:20 +08:00
Spandan TiwariandGitHub b71218107f [ROCm][Test] Fix test_per_token_group_quant_fp8 tolerance for 1-ULP FP8 rounding on gfx950 (#46944)
Signed-off-by: Spandan Tiwari <sptiwari@amd.com>
2026-07-05 18:02:30 -05:00
cc1d020d01 [MRV2] Enable mm prefix bidi attention support on MRV2 (#46942)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <2037008807@qq.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-05 14:45:29 +00:00
Ting SUNandGitHub 8974ed89cd [Bugfix][Voxtral Realtime] Fix token feedback timeout silent hang (#44461)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-07-05 05:42:36 -07:00
Ting SUNGitHubWentao Yemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
fb2faceacd [Bugfix][Model] Fix crash loading Mamba/Mamba2 checkpoints without an architectures field (#46037)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Signed-off-by: Ting SUN <suntcrick@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-05 05:42:32 -07:00
b6cc46ec3b [Feature] Support sequence parallel without the need for DP, 1.9%~5.0% E2E Throughput Improvement (#47070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-07-05 05:41:30 -07:00
Lucas WilkinsonandGitHub fa4321de3d [Bugfix][TurboQuant] Preserve KV cache dtype in backend shape (#47609) 2026-07-05 08:20:48 +00:00
Ting SUNandGitHub 9226613043 [Bugfix][Pooling] Forward instruction to Jina reranker scoring prompts (#47590)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-07-05 05:39:13 +00:00
34b560b725 [Bugfix][Gemma4] Fix FA4 mm_prefix mask: add sliding window and absolute q_idx (#47332)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
2026-07-04 17:46:40 -07:00
Ting SUNandGitHub 91b5647300 [Bugfix][Model] Allow Run:ai memory_limit sentinel values (#47337)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-07-05 00:08:34 +00:00
Carl PerssonandGitHub 4a6bf3c77f [ROCm][CI] Fix Kernels and Kernels attention test failures (#47519)
Signed-off-by: Carl Persson <carl.persson@amd.com>
2026-07-04 15:59:51 -05:00
Ting SUNandGitHub d2afe39647 [Bugfix][Frontend] Preserve default sampling params in batch chat (#47597)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
2026-07-04 19:06:39 +00:00
Wentao YeandGitHub 2a9113f998 [Perf] Remove redundant op for GLM 5.2 (#47198)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-07-04 13:25:02 -04:00
yzong-rhandGitHub 0cd6f767e3 [Bugfix][Frontend][gpt-oss] Recover raw tail when Harmony parser ends non-terminal (#47379) 2026-07-04 10:46:24 -04:00
Harry MellorandGitHub f1445f6dbd [CI] Bump huggingface-hub from v1.10.2 to v1.22.0 (#47551)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-07-04 07:45:45 -07:00
1d354c694e [Misc] Validate Pooling cache_salt Values (#46966)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-07-04 10:19:28 -04:00
Taneem IbrahimandGitHub 2f21224527 [Misc] Update request-extras parity for batch chat completion (#47333)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
2026-07-04 10:19:04 -04:00
Taneem IbrahimGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Wentao Ye
fa1fa968c4 [Misc] Forward request-level prompt extras for cross-encoder scoring (#46939)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
2026-07-04 10:18:36 -04:00
Taneem IbrahimGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
6eac8e0070 [Misc] Preserve cross-encoder pooling extra kwargs (#47082)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-04 08:14:13 -04:00
1a308c449c [XPU] Add W8A8 FP8 linear kernel with multi-granularity quant support (#43645)
Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com>
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-04 18:10:01 +08:00
e7c9df9449 [Bugfix][Structured Output][Spec Decode] Constrain bitmask and trim grammar advance at the reasoning boundary (#44297)
Signed-off-by: Allen.Yu <yuyue0225sc@163.com>
Signed-off-by: yue.yu <yuyue0225sc@163.com>
Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
2026-07-04 09:08:45 +00:00
gausah01andGitHub 26eb87204d [Bugfix] Fix CPU split-KV scratchpad sizing (#45844)
Signed-off-by: Gauri Sahnan <gauri.sahnan@arm.com>
2026-07-04 06:47:23 +00:00
4c3c17d43b [ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations (#47567)
Signed-off-by: Rohan Potdar <rohan.potdar@amd.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-04 01:21:43 -05:00
f329ce405b [ROCm][CI][Bugfix] Use VllmRunner for voxtral_realtime tests to avoid OOM on AMD GPU (#47536)
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-04 12:26:10 +08:00
07516fda67 [MRV2][SD] Make Dynamic SD comatible with Full Cuda Graphs (#45953)
Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com>
Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-07-03 23:58:27 -04:00
sychen52GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
67ff0ae30f Support nvfp4 kv with kv-cache-dtype-skip-layers sliding_window (#42890)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-04 02:29:13 +00:00
Bugen ZhaoandGitHub ab3b6d97aa [Frontend] Limit SO_REUSEPORT to multi-worker serving (#47529)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-04 01:26:24 +00:00
Ben BrowningandGitHub fb5291b35b [Frontend] [Parser] Port DeepSeek V4 to streaming parser engine framework (#45877)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-07-03 20:55:23 -04:00
labAxiaomingandGitHub d6d39c111e [GLM4V] Avoid GLM4V processor init during startup metadata reads (#47155)
Signed-off-by: xiaoming <1259730330@qq.com>
2026-07-03 15:03:16 -07:00
379950191f [Bugfix][Multimodal] Normalize direct PIL image inputs (#47566)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
2026-07-03 14:27:14 -07:00
576bf75d0e [AMD][EPLB] Enable EPLB for Quark OCP MXFP4 MoE (#47220)
Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-03 14:41:52 -05:00
TresandGitHub f006e5a24c [CI][AMD] Allow git operations on previously created work trees (#47554)
Signed-off-by: Tres Popp <tres.popp@amd.com>
2026-07-03 14:41:01 -05:00
f63dca6838 [ROCm] Fix encoder-decoder cross-attention KV layout aliasing (#47035)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
2026-07-03 13:53:29 -05:00
Bugen ZhaoandGitHub 8651f043b8 [Rust Frontend] Speed up chat roundtrip tests (#47523)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-07-03 19:25:06 +01:00
Andreas KaratzasandGitHub 3775d5fcab [ROCm][CI] Adding test groups for parity with upstream (#47479)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-07-03 19:15:01 +04:00
Roberto L. CastroGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d7192cfccf [CI Bugfix] Lazily import Qwen warmup dependencies (#47539)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-07-03 23:10:49 +08:00
AgenticSparkandGitHub 978de83353 [Bugfix][CPU] Ship examples/ in the CPU release image (#47447)
Signed-off-by: liejiang <jianglie2023@gmail.com>
2026-07-03 11:46:24 +00:00
Tyler Michael SmithandClaude Opus 4.6 cc17a5e0d1 [CI] Fix flaky cudagraph mode test by isolating each case in its own process
The test_cudagraph_compilation_combo and test_backend_and_cudagraph_mode_combo
tests ran all parametrized cases in a single pytest process, relying on
weakref + wait_for_gpu_memory_to_clear(120s) between cases. After many
LLM create/destroy cycles, accumulated CUDA driver state could delay
subprocess memory reclamation past the 120s timeout, causing flaky failures
(e.g. FA2-FULL_DECODE_ONLY-3-True which captures 51 full CUDA graphs).

Wrap both tests with @create_new_process_for_each_test("spawn") so each
parametrized case gets a clean process, eliminating the cross-test memory
accumulation and removing the need for manual teardown.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-12 17:57:10 -04:00
292 changed files with 16506 additions and 4324 deletions
+24
View File
@@ -76,6 +76,30 @@ steps:
pytest -v -s v1/sample/test_logprobs.py &&
pytest -v -s v1/sample/test_logprobs_e2e.py'
- label: Basic Models Tests (Initialization)
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
working_dir: "."
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/models/test_initialization.py
- tests/models/registry.py
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'export VLLM_XPU_FUSED_MOE_USE_REF=1 &&
cd tests &&
pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]'
- label: XPU CPU Offload
timeout_in_minutes: 60
device: intel_gpu
+48 -1
View File
@@ -49,6 +49,30 @@ steps:
VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
'
- label: "XPU W8A8 FP8 Linear Examples"
depends_on:
- image-build-xpu
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 24+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 &&
python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 &&
python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
'
- label: "XPU V1 test"
depends_on:
- image-build-xpu
@@ -120,4 +144,27 @@ steps:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_auto_round.py'
pytest -v -s quantization/test_auto_round.py'
- label: "XPU compressed tensors FP8 test"
depends_on:
- image-build-xpu
timeout_in_minutes: 60
device: intel_gpu
agent_tags:
label: production
gpu: 1+
mem: 16+
no_plugin: true
env:
REGISTRY: "public.ecr.aws/q9t5s3a7"
REPO: "vllm-ci-test-repo"
VLLM_TEST_DEVICE: "xpu"
source_file_dependencies:
- vllm/
- tests/quantization/test_compressed_tensors.py
- .buildkite/intel_jobs/test-intel.yaml
commands:
- >-
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
'cd tests &&
pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
@@ -448,9 +448,11 @@ checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}"
if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then
checkout="."
fi
if git -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
# Pass safe.directory per-command (-c) because buildkite runs will always fail
# the next check on git 2.35.2+ due to mixed uses of root and buildkite-agent/uids.
if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
vllm_standalone_merge_base="$(
git -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true
)"
fi
if [[ -z "${vllm_standalone_merge_base}" ]]; then
@@ -549,6 +551,7 @@ else
fi
docker run \
-t -i \
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
$RDMA_FLAGS \
--network=host \
@@ -38,7 +38,8 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/core/test_cpu_activation.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"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
+3 -1
View File
@@ -109,7 +109,9 @@ run_nodes() {
if [ "$node" -ne 0 ]; then
docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
else
docker exec "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
# Allocate a TTY (-t -i) for the foreground head node so its output
# keeps ANSI color in the Buildkite log (see run-amd-test.sh).
docker exec -t -i "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
fi
done
}
+251 -9
View File
@@ -472,7 +472,7 @@ steps:
commands:
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
@@ -1025,6 +1025,23 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
- label: MRCR Eval Small Models # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/models/
- vllm/model_executor/model_loader/
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- tests/evals/mrcr/
commands:
- pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
- label: LM Eval Small Models (MI300) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1283,6 +1300,21 @@ steps:
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
- label: vLLM IR Tests # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/"
source_file_dependencies:
- vllm/ir
- vllm/kernels
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s tests/ir
- pytest -v -s tests/kernels/ir
- label: Kernels Attention Test %N # TBD
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1317,6 +1349,21 @@ steps:
commands:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
- label: Kernels KDA Test # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/layers/fla/ops/kda.py
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
- vllm/model_executor/layers/fla/ops/l2norm.py
- tests/kernels/test_kda.py
- vllm/platforms/rocm.py
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels MoE Test %N # TBD
timeout_in_minutes: 95
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1429,13 +1476,135 @@ steps:
- pytest -v -s model_executor -m '(not slow_test)'
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
#---------------------------------------------------- mi300 · model_runner_v2 -------------------------------------------------------#
- label: Model Runner V2 Core Tests # TBD
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- vllm/v1/core/sched/
- vllm/v1/attention/
- tests/v1/engine/test_llm_engine.py
- tests/v1/e2e/
- tests/entrypoints/llm/test_struct_output_generate.py
- vllm/platforms/rocm.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
- pytest -v -s v1/e2e/general/test_context_length.py
- pytest -v -s v1/e2e/general/test_min_tokens.py
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
- label: Model Runner V2 Examples # TBD
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/core/sched/
- vllm/v1/worker/gpu_worker.py
- examples/basic/offline_inference/
- examples/generate/multimodal/
- examples/features/
- examples/pooling/embed/vision_embedding_offline.py
- examples/features/tensorize_vllm_model.py
- examples/deployment/llm_engine_example.py
- vllm/platforms/rocm.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pip install tensorizer
- python3 basic/offline_inference/chat.py
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
- python3 pooling/embed/vision_embedding_offline.py --seed 0
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 deployment/llm_engine_example.py
- python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- label: Model Runner V2 Distributed (2 GPUs) # TBD
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/basic_correctness/test_basic_correctness.py
- tests/v1/distributed/test_async_llm_dp.py
- tests/v1/distributed/test_eagle_dp.py
- vllm/platforms/rocm.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/distributed/test_pipeline_parallel.py
- tests/distributed/test_pp_cudagraph.py
- vllm/platforms/rocm.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- label: Model Runner V2 Spec Decode # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/spec_decode/test_rejection_sampler_utils.py
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- tests/v1/e2e/spec_decode/test_spec_decode.py
- vllm/platforms/rocm.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
#------------------------------------------------------ mi300 · models / basic -------------------------------------------------------#
- label: Basic Models Tests (Extra Initialization) %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
parallelism: 2
parallelism: 6
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/models/
@@ -1468,10 +1637,10 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/test_terratorch.py
- tests/models/test_transformers.py
- tests/models/transformers/test_backend.py
- tests/models/test_registry.py
commands:
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
#----------------------------------------------------- mi300 · models / language -----------------------------------------------------#
@@ -1691,7 +1860,7 @@ steps:
- examples/
commands:
- pip install --upgrade git+https://github.com/huggingface/transformers
- pytest -v -s tests/models/test_transformers.py
- pytest -v -s tests/models/transformers/test_backend.py
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/basic/offline_inference/chat.py
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
@@ -1745,6 +1914,22 @@ steps:
- pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
- label: GGUF Plugin # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
soft_fail: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/plugins_tests/test_gguf_plugin.py
- tests/plugins_tests/gguf
- vllm/platforms/rocm.py
commands:
- pip install "vllm-gguf-plugin >= 0.0.2"
- pytest -v -s plugins_tests/gguf
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
- label: Quantization # TBD
@@ -2063,7 +2248,7 @@ steps:
- pytest -v -s v1/worker
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
# - export HSA_NO_SCRATCH_RECLAIM=1
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
@@ -2254,6 +2439,59 @@ steps:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/core/sched/
- vllm/v1/core/kv_cache_coordinator.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- label: V1 e2e (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2544,15 +2782,19 @@ steps:
- vllm/envs.py
- examples/offline_inference/data_parallel.py
- tests/distributed/test_context_parallel.py
- tests/distributed/test_rocm_aiter_custom_ar.py
- tests/distributed/test_rocm_quick_reduce.py
- tests/distributed/test_quick_all_reduce.py
- tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
- tests/v1/distributed/test_dbo.py
- tests/utils.py
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
- pytest -v -s tests/v1/distributed/test_dbo.py
- pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
- pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
- pytest -v -s tests/distributed/test_rocm_quick_reduce.py
- pytest -v -s tests/distributed/test_quick_all_reduce.py
- pytest -v -s tests/v1/distributed/test_dbo.py
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
@@ -2770,7 +3012,7 @@ steps:
#--------------------------------------------------------- mi355 · examples ----------------------------------------------------------#
- label: Examples # TBD
timeout_in_minutes: 45
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
agent_pool: mi355_1
working_dir: "/vllm-workspace/examples"
@@ -3133,7 +3375,7 @@ steps:
- pytest -v -s v1/worker
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
- label: V1 Sample + Logits # TBD
+1 -1
View File
@@ -103,7 +103,7 @@ steps:
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
# Integration test for streaming correctness (requires special branch).
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
mirror:
amd:
+4 -3
View File
@@ -36,10 +36,10 @@ steps:
source_file_dependencies:
- vllm/
- tests/models/test_terratorch.py
- tests/models/test_transformers.py
- tests/models/transformers/test_backend.py
- tests/models/test_registry.py
commands:
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
mirror:
amd:
device: mi325_1
@@ -55,6 +55,7 @@ steps:
- vllm/
- tests/models/test_utils.py
- tests/models/test_vision.py
- tests/models/transformers/fusers/
device: cpu-small
commands:
- pytest -v -s models/test_utils.py models/test_vision.py
- pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
@@ -17,7 +17,7 @@ steps:
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
# Avoid importing model tests that cause CUDA reinitialization error
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
+2 -1
View File
@@ -27,6 +27,7 @@ steps:
- tests/models/multimodal
commands:
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
- pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
mirror:
amd:
@@ -58,7 +59,7 @@ steps:
- vllm/
- tests/models/multimodal
commands:
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
+1 -1
View File
@@ -119,7 +119,7 @@
# Transformers modeling backend
/vllm/model_executor/models/transformers @hmellor
/tests/models/test_transformers.py @hmellor
/tests/models/transformers @hmellor
# Docs
/docs/mkdocs @hmellor
-1
View File
@@ -400,7 +400,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/topk.cu"
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"csrc/libtorch_stable/cache_kernels_fused.cu"
"csrc/libtorch_stable/custom_all_reduce.cu"
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+4 -2
View File
@@ -132,8 +132,10 @@ def benchmark_function(
reset_memory_stats()
# Benchmark
start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
start_events = [
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
]
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
for i in range(benchmark_iters):
logits_copy = logits.clone()
+2 -2
View File
@@ -134,8 +134,8 @@ def benchmark_config(
torch.accelerator.synchronize()
# Benchmark
start = torch.Event(enable_timing=True)
end = torch.Event(enable_timing=True)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(num_iters):
with override_config(config):
@@ -170,8 +170,8 @@ def benchmark_config(
graph.replay()
torch.accelerator.synchronize()
start = torch.Event(enable_timing=True)
end = torch.Event(enable_timing=True)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
latencies: list[float] = []
for _ in range(num_iters):
start.record()
+23 -5
View File
@@ -15,6 +15,7 @@ endif()
#
set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
include_directories("${CMAKE_SOURCE_DIR}/csrc")
@@ -110,6 +111,13 @@ else()
set(ARM_BF16_FOUND ON)
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
endif()
# Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
# in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
# overrides the detection result.
if (ENABLE_RVV_BF16)
set(RVV_BF16_FOUND ON)
message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable")
endif()
endif()
if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
@@ -178,7 +186,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
# Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
if(NOT DEFINED VLLM_RVV_VLEN)
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
if(EXISTS /proc/cpuinfo)
# Skip when cross-compiling — /proc/cpuinfo describes the build host.
if(CMAKE_CROSSCOMPILING)
message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo")
elseif(EXISTS /proc/cpuinfo)
file(READ /proc/cpuinfo _cpuinfo)
set(_best 0)
foreach(_n IN ITEMS 128 256 512 1024)
@@ -186,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
set(_best ${_n})
endif()
endforeach()
# Only VLEN=128 and VLEN=256 are supported by the RVV kernels.
if(_best GREATER 256)
message(WARNING
"Detected VLEN=${_best} but only 128/256 are supported; "
"clamping to 256")
set(_best 256)
endif()
if(_best GREATER 0)
set(VLLM_RVV_VLEN ${_best})
endif()
@@ -195,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
message(FATAL_ERROR
"RISC-V RVV is available but VLEN could not be auto-detected. "
"Please specify VLEN explicitly:\n"
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)")
"Please specify VLEN explicitly via CMAKE_ARGS:\n"
" CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n"
" CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)")
endif()
endif()
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -209,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
message(STATUS "BF16 extension detected")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
elseif(RVV_FP16_FOUND)
message(WARNING "BF16 functionality is not available")
message(WARNING "BF16 functionality is not available.")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
else()
message(STATUS "compile riscv with scalar (no FP16/BF16)")
+2 -1
View File
@@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
bool cpu_attn_has_isa(const std::string& isa) {
if (isa == "rvv") {
#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
return true;
#else
return false;
+56 -14
View File
@@ -417,8 +417,10 @@ class AttentionScheduler {
has_decode_request = has_decode_request || (q_token_num == 1);
decode_only_batch = decode_only_batch && (q_token_num == 1);
}
int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
const int32_t original_q_head_per_kv =
input.num_heads_q / input.num_heads_kv;
int32_t q_head_per_kv = original_q_head_per_kv;
const bool supports_gqa = original_q_head_per_kv <= max_num_q_per_iter;
const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
if (!use_gqa_scratchpad) {
@@ -671,22 +673,62 @@ class AttentionScheduler {
metadata_ptr->effective_thread_num = effective_thread_num;
{
// when q_tile_size = max_num_q_per_iter, requires max
// attention_scratchpad_size
AttentionScratchPad sc(0, *metadata_ptr, 0x0);
int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
input.logits_buffer_elem_size, input.output_buffer_elem_size,
max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true);
sc.update(input.head_dim, input.q_buffer_elem_size,
input.logits_buffer_elem_size, input.output_buffer_elem_size,
max_num_q_per_iter, max_num_q_per_iter, n);
int64_t max_attention_scratchpad_size = 0;
for (const AttentionWorkItemGroup& item : workitems) {
const bool curr_use_gqa =
use_gqa_fast_path || (supports_gqa && item.q_token_num == 1);
const int32_t curr_q_heads_per_kv =
curr_use_gqa ? original_q_head_per_kv : 1;
const int32_t curr_default_q_tile_token_num =
default_tile_size / curr_q_heads_per_kv;
for (int32_t q_token_offset = 0; q_token_offset < item.q_token_num;
q_token_offset += curr_default_q_tile_token_num) {
const int32_t actual_q_token_num = std::min(
curr_default_q_tile_token_num, item.q_token_num - q_token_offset);
const int32_t q_head_tile_size =
actual_q_token_num * curr_q_heads_per_kv;
const int32_t rounded_q_head_tile_size =
((q_head_tile_size + max_num_q_per_iter - 1) /
max_num_q_per_iter) *
max_num_q_per_iter;
const int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
cache_size, input.head_dim, input.elem_size,
input.q_buffer_elem_size, input.logits_buffer_elem_size,
input.output_buffer_elem_size, max_num_q_per_iter,
kv_len_alignment, rounded_q_head_tile_size,
rounded_q_head_tile_size <= max_num_q_per_iter);
sc.update(input.head_dim, input.q_buffer_elem_size,
input.logits_buffer_elem_size,
input.output_buffer_elem_size, max_num_q_per_iter,
rounded_q_head_tile_size, n);
max_attention_scratchpad_size = std::max(
max_attention_scratchpad_size, sc.get_thread_scratchpad_size());
}
}
metadata_ptr->attention_scratchpad_size_per_thread =
((sc.get_thread_scratchpad_size() + 63) / 64) * 64;
((max_attention_scratchpad_size + 63) / 64) * 64;
int32_t max_reduction_q_head_tile_size = 0;
for (const ReductionWorkItemGroup& item : reduce_workitems) {
const bool curr_use_gqa =
use_gqa_fast_path || (supports_gqa && item.q_token_id_num == 1);
const int32_t curr_q_heads_per_kv =
curr_use_gqa ? original_q_head_per_kv : 1;
max_reduction_q_head_tile_size =
std::max(max_reduction_q_head_tile_size,
item.q_token_id_num * curr_q_heads_per_kv);
}
sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
q_head_per_kv * split_kv_q_token_num_threshold,
input.output_buffer_elem_size);
max_reduction_q_head_tile_size, input.output_buffer_elem_size);
metadata_ptr->reduction_scratchpad_size_per_kv_head =
((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
}
+35 -16
View File
@@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
explicit BF16Vec32(const BF16Vec8& v) {
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
fixed_u16x32_t u16_combined =
RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
u16_val, u16_val, u16_val, u16_val);
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
LMUL_512)(u16_combined);
// Widen LMUL_128 → LMUL_256 so vslideup operands share a type.
// At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2.
fixed_u16x16_t ext =
RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val);
// Build 16-element half: place the 8 elements at offsets 0 and 8.
fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16);
half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8);
half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16);
// Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128).
fixed_u16x32_t dst =
RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half);
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst);
};
void save(void* ptr) const {
@@ -623,17 +630,29 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
const uint64_t q_values = static_cast<uint64_t>(value);
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
auto shifts =
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
auto shifted =
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
auto idx64 =
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
// Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on
// VLEN=128 / m4 on VLEN=256), which causes heavy register spilling.
constexpr int HALF = VEC_ELEM_NUM / 2;
const auto q = static_cast<uint64_t>(value);
const uint32_t lo = static_cast<uint32_t>(q);
const uint32_t hi = static_cast<uint32_t>(q >> 32);
auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF);
auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF);
auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF);
auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)(
RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF,
HALF);
auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF);
auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)(
RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF,
HALF);
auto idx =
RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi);
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM);
}
explicit FP32Vec16(const FP16Vec16& v);
+2 -1
View File
@@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"dynamic_4bit_int_moe("
"Tensor x, Tensor topk_ids, Tensor topk_weights,"
"Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
"Tensor w13_packed, Tensor w2_packed,"
"int hidden_size, int intermediate_size,"
"int group_size, bool apply_router_weight_on_input, int activation_kind"
") -> Tensor");
+53 -28
View File
@@ -29,25 +29,37 @@ enum ActivationKind : int64_t {
torch::Tensor dynamic_4bit_int_moe_cpu(
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
int64_t activation_kind) {
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
int64_t intermediate_size, int64_t group_size,
bool apply_router_weight_on_input, int64_t activation_kind) {
TORCH_CHECK(x.dim() == 2, "x must be 2D");
TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
"topk tensors must be [T, K]");
TORCH_CHECK(
w13_packed.size(0) == w2_packed.size(0),
"w13_packed and w2_packed must have same number of experts in dim 0");
TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");
const int64_t T = x.size(0);
const int64_t K = topk_ids.size(1);
const int64_t E = w13_packed.size(0);
const int64_t N = T * K;
const int64_t w13_out_features = 2 * intermediate_size;
auto x_c = x.contiguous();
// _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
// dtypes:
// - fp32: with channelwise and groupwise
// - bf16: with channelwise -> upcast to fp32 for groupwise
// - fp16: not supported -> upcast to fp32 for groupwise & channelwise
const auto output_dtype = x_c.scalar_type();
const bool should_cast_input =
((group_size != -1) && output_dtype == at::kBFloat16) ||
output_dtype == at::kHalf;
if (should_cast_input) {
x_c = x_c.to(at::kFloat);
}
auto ids_c = topk_ids.contiguous();
auto gates_c = topk_weights.to(at::kFloat).contiguous();
auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
// bucketing tokens -> experts
c10::SmallVector<int64_t, 64> counts(
@@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
c10::SmallVector<int64_t, 65> offsets(E + 1, 0); // ( E +1 )
for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
// expert_tokens = [tokens indices for expert 0, ...]
// expert_gates = [router weights for tokens assigned to expert 0, ...]
auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
auto expert_gates = at::empty({offsets[E]}, gates_c.options());
{
c10::SmallVector<int64_t, 64> cursor(E, 0);
const auto* ids_ptr = ids_c.data_ptr<int64_t>();
const auto* gts_ptr = gates_c.data_ptr<float>();
auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
auto* gate_ptr = expert_gates.data_ptr<float>();
AT_DISPATCH_FLOATING_TYPES_AND2(
at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
"bucket_expert_tokens_and_gates", [&] {
const auto* ids_ptr = ids_c.data_ptr<int64_t>();
const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
for (int64_t t = 0; t < T; ++t) {
const int64_t base = t * K;
for (int64_t k = 0; k < K; ++k) {
const int64_t idx = base + k;
const int64_t e = ids_ptr[idx];
const int64_t p = offsets[e] + (cursor[e]++);
tok_ptr[p] = t;
gate_ptr[p] = gts_ptr[idx];
}
}
for (int64_t t = 0; t < T; ++t) {
const int64_t base = t * K;
for (int64_t k = 0; k < K; ++k) {
const int64_t idx = base + k;
const int64_t e = ids_ptr[idx];
const int64_t p = offsets[e] + (cursor[e]++);
tok_ptr[p] = t;
gate_ptr[p] = gts_ptr[idx];
}
}
});
}
const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
const int64_t g_eff_2 = (group_size != -1) ? group_size : I;
const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
// X_all [num_tokens * K, hidden_size]
auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
if (apply_router_weight_on_input) {
X_all = X_all.mul(expert_gates.unsqueeze(1));
}
auto Y_all = at::empty({offsets[E], H}, x_c.options());
auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
c10::InferenceMode guard;
@@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
auto w2_e = w2_packed.select(/*dim=*/0, e);
// W13
auto y13 =
mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
/*out_features=*/w13_out_features);
auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
auto g_part =
y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
/*length=*/intermediate_size);
torch::Tensor act;
if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI
@@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
}
// W2
auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
/*out_features=*/hidden_size);
// Store per-expert result
Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
@@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
if (!apply_router_weight_on_input) {
Y_all = Y_all.mul(expert_gates.unsqueeze(1));
}
if (Y_all.scalar_type() != output_dtype) {
Y_all = Y_all.to(output_dtype);
}
auto out = at::zeros({T, H}, x.options());
auto out = at::zeros({T, hidden_size}, x.options());
out =
at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
+3 -3
View File
@@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
torch::Tensor dynamic_4bit_int_moe_cpu(
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
int64_t activation_kind);
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
int64_t intermediate_size, int64_t group_size,
bool apply_router_weight_on_input, int64_t activation_kind);
using fptr_t = int64_t;
#ifdef USE_ROCM
+10 -7
View File
@@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common
WORKDIR /workspace
ARG PYTHON_VERSION=3.12
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
ARG max_jobs=32
ENV MAX_JOBS=${max_jobs}
@@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH"
ENV UV_HTTP_TIMEOUT=500
# Install Python dependencies
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
@@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --upgrade pip && \
uv pip install -r requirements/cpu.txt
uv pip install -r requirements/cpu.txt --torch-backend cpu
ARG TARGETARCH
ENV TARGETARCH=${TARGETARCH}
@@ -149,7 +146,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \
COPY requirements/build/cpu.txt requirements/build/cpu.txt
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/build/cpu.txt
uv pip install -r requirements/build/cpu.txt --torch-backend cpu
COPY . .
@@ -205,7 +202,7 @@ RUN case "$(uname -m)" in \
esac
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/test/cpu.txt
uv pip install -r requirements/test/cpu.txt --torch-backend cpu
######################### DEV IMAGE #########################
FROM vllm-build AS vllm-dev
@@ -231,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements/lint.txt && \
uv pip install -r requirements/test/cpu.txt && \
uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
pre-commit install --hook-type pre-commit --hook-type commit-msg
ENTRYPOINT ["bash"]
@@ -301,6 +298,12 @@ LABEL ai.vllm.build.cpu-x86="${VLLM_CPU_X86:-false}"
LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
# Copy the examples directory (including the chat/tool templates) so it is
# present in the released image, as the CUDA image ships it too. The vllm-test
# stage above adds examples/ for testing only, so without this the published
# vllm-openai-cpu image would not ship examples/*.jinja.
COPY examples examples
ENTRYPOINT ["vllm", "serve"]
+2 -1
View File
@@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
uv pip install -v \
uv pip install -v \
$ARROW_WHL_FILE \
$VISION_WHL_FILE \
$HF_XET_WHL_FILE \
@@ -257,6 +257,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
$NUMBA_WHL_FILE \
$OPENCV_WHL_FILE \
$GUIDANCE_WHL_FILE \
--torch-backend cpu \
--index-strategy unsafe-best-match \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt
+3 -2
View File
@@ -276,8 +276,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power
the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.)
you can switch to the [fastokens](https://github.com/crusoecloud/fastokens)
Rust backend, a drop-in replacement that's substantially faster on
encode/decode and on streaming detokenization. Enable it by setting
`VLLM_USE_FASTOKENS=1`:
encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is
available in vLLM v0.23.0 and later. If your installed vLLM version does not
recognize the environment variable, upgrade vLLM before enabling the override:
```console
VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B
+1 -1
View File
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
+1 -4
View File
@@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \
--max-model-len 4096
```
!!! note
To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now.
## Evaluating the Quantized Model with vLLM
```bash
lm_eval --model vllm \
--model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \
--model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \
--tasks gsm8k \
--num_fewshot 5 \
--batch_size 128
+2
View File
@@ -62,6 +62,8 @@ weight name. Unset fields fall back to the `--quantization` shorthand's
defaults, or for already-quantized checkpoints to whatever the checkpoint
declares.
On XPU, non-block FP8 scaled-mm linear layers default to W8A16; setting `--linear-backend xpu` forces W8A8. Use `--linear-backend xpu_woq` to explicitly select weight-only quantization (W8A16).
The CLI accepts the same shape as JSON or as dotted keys:
```bash
@@ -49,6 +49,32 @@ You can configure how the quantization scales are computed in vLLM using three d
- `kv_cache_dtype="fp8_e4m3"`: Supported on CUDA 11.8+ and ROCm (AMD GPUs)
- `kv_cache_dtype="fp8_e5m2"`: Supported on CUDA 11.8+
### Skipping Specific Layers from KV-Cache Quantization
Some attention layer types (e.g. sliding-window) are more sensitive to KV-cache quantization. The `--kv-cache-dtype-skip-layers` flag leaves the specified layers at the model's native dtype while keeping the rest of the layers under the chosen quantized dtype. The flag accepts either layer indices or layer-type names:
```bash
# Skip every sliding-window attention layer.
vllm serve <model> \
--kv-cache-dtype fp8 \
--kv-cache-dtype-skip-layers sliding_window
# Skip specific layer indices.
vllm serve <model> \
--kv-cache-dtype fp8 \
--kv-cache-dtype-skip-layers 0 1 23
```
Programmatic usage:
```python
llm = LLM(
model="meta-llama/Llama-3.1-8B-Instruct",
kv_cache_dtype="fp8",
kv_cache_dtype_skip_layers=["sliding_window"],
)
```
---
## Examples
@@ -71,8 +71,5 @@ VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
## Limitations
* only tested with Eagle and Eagle-3. Other SD methods may or may not work out of the box
* only usable with Model Runner V1
* not compatible with full cuda graph so we force piece-wise cuda graph with this feature
We are working on enabling it on MRv2 with full cuda graph support.
* Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box
* Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature
@@ -35,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
uv pip install -r requirements/cpu.txt
uv pip install -e .
```
!!! tip
The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts.
The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation.
!!! note
On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device.
@@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source.
```bash
uv pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
--torch-backend auto \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
--torch-backend cpu \
--index-strategy unsafe-best-match && \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
uv pip install dist/*.whl
```
+2 -3
View File
@@ -2,8 +2,7 @@
!!! note
We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance
improvements over using Hugging Face Transformers or Sentence Transformers directly.
improvements over using Hugging Face Transformers or Sentence Transformers directly.
We plan to optimize pooling models in vLLM. Please comment on <https://github.com/vllm-project/vllm/issues/21796> if you have any suggestions!
## What are pooling models?
@@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
!!! note
Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
### Pooling Types
+4 -2
View File
@@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang
### Transformers
vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend".
vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation.
Currently, the Transformers modeling backend works for the following:
@@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded:
That's it!
For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
<details class="code">
<summary>configuration_my_model.py</summary>
@@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
</details>
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported).
- vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated.
- `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s:
- You only need to do this for layers which are not present on all pipeline stages
- vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages
- When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage
- The `list` in the first element of the `tuple` contains the names of the input arguments
- The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
@@ -0,0 +1,73 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ruff: noqa: E501
"""
Example online usage of the Jina Reranker v3 score and rerank APIs with a task
instruction.
Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
server in vLLM.
"""
import argparse
import json
import requests
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
headers = {"User-Agent": "Test Client"}
response = requests.post(api_url, headers=headers, json=prompt)
return response
def print_response(name: str, prompt: dict, response: requests.Response) -> None:
print(f"\n{name} request:")
print(json.dumps(prompt, indent=2))
print(f"\n{name} response:")
print(json.dumps(response.json(), indent=2))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3")
return parser.parse_args()
def main(args):
score_url = f"http://{args.host}:{args.port}/score"
rerank_url = f"http://{args.host}:{args.port}/rerank"
model_name = args.model
query = "Which passage is about sports?"
documents = [
"Basketball is played by two teams on a court.",
"Green tea contains antioxidants and may support metabolism.",
]
instruction = "Rank passages about sports higher than passages about nutrition."
score_prompt = {
"model": model_name,
"queries": query,
"documents": documents,
"instruction": instruction,
}
score_response = post_http_request(prompt=score_prompt, api_url=score_url)
print_response("Score", score_prompt, score_response)
rerank_prompt = {
"model": model_name,
"query": query,
"documents": documents,
"instruction": instruction,
}
rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url)
print_response("Rerank", rerank_prompt, rerank_response)
if __name__ == "__main__":
args = parse_args()
main(args)
+2 -2
View File
@@ -129,9 +129,9 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
"docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
"tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*",
"rust/src/chat/src/renderer/deepseek_v32/fixtures/*",
"rust/src/parser/*",
"rust/src/parser/src/tool/gemma4.rs", "rust/src/parser/src/unified/gemma4.rs",
"rust/src/text/src/output/decoded.rs",
"rust/src/tokenizer/src/incremental.rs",]
"rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"]
ignore-hidden = false
[tool.typos.default]
-1
View File
@@ -1,4 +1,3 @@
--extra-index-url https://download.pytorch.org/whl/cpu
cmake>=3.26.1
ninja
packaging>=24.2
-1
View File
@@ -1,4 +1,3 @@
--extra-index-url https://download.pytorch.org/whl/cpu
# Common dependencies
-r common.txt
+1 -1
View File
@@ -25,7 +25,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2
quack-kernels>=0.3.3
# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.1.2
tokenspeed-mla==0.1.2; platform_system == "Linux"
# Humming kernels for quantization gemm
humming-kernels[cu13]==0.1.6
+4 -4
View File
@@ -112,9 +112,10 @@ charset-normalizer==3.4.0
# via requests
chz==0.3.0
# via gpt-oss
click==8.1.7
click==8.4.2
# via
# black
# huggingface-hub
# jiwer
# nltk
# ray
@@ -309,7 +310,7 @@ h2==4.3.0
# via httpx
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.5.1
# via huggingface-hub
hiredis==3.0.0
# via tensorizer
@@ -335,7 +336,7 @@ httpx==0.27.2
# schemathesis
httpx-sse==0.4.3
# via mcp
huggingface-hub==1.10.2
huggingface-hub==1.22.0
# via
# accelerate
# datasets
@@ -1182,7 +1183,6 @@ typer==0.26.8
# fastapi-cli
# fastapi-cloud-cli
# fastsafetensors
# huggingface-hub
# perceptron
# transformers
typing-extensions==4.15.0
+4 -4
View File
@@ -117,9 +117,10 @@ charset-normalizer==3.4.0
# via requests
chz==0.3.0
# via gpt-oss
click==8.1.7
click==8.4.2
# via
# black
# huggingface-hub
# jiwer
# nltk
# ray
@@ -330,7 +331,7 @@ h2==4.3.0
# via httpx
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.5.1
# via huggingface-hub
hiredis==3.0.0
# via tensorizer
@@ -356,7 +357,7 @@ httpx==0.27.2
# schemathesis
httpx-sse==0.4.3
# via mcp
huggingface-hub==1.10.2
huggingface-hub==1.22.0
# via
# accelerate
# datasets
@@ -1285,7 +1286,6 @@ typer==0.26.8
# fastapi-cli
# fastapi-cloud-cli
# fastsafetensors
# huggingface-hub
# perceptron
# transformers
typing-extensions==4.15.0
+4 -4
View File
@@ -116,9 +116,10 @@ choreographer==1.2.1
# via kaleido
chz==0.4.0
# via gpt-oss
click==8.3.1
click==8.4.2
# via
# black
# huggingface-hub
# jiwer
# nltk
# ray
@@ -323,7 +324,7 @@ h2==4.3.0
# via httpx
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.5.1
# via huggingface-hub
hiredis==3.3.1
# via tensorizer
@@ -349,7 +350,7 @@ httpx==0.27.2
# schemathesis
httpx-sse==0.4.3
# via mcp
huggingface-hub==1.10.2
huggingface-hub==1.22.0
# via
# accelerate
# datasets
@@ -1244,7 +1245,6 @@ typer==0.24.1
# fastapi-cli
# fastapi-cloud-cli
# fastsafetensors
# huggingface-hub
# perceptron
# transformers
typing-extensions==4.15.0
+4 -4
View File
@@ -83,8 +83,9 @@ charset-normalizer==3.4.6
# via requests
chz==0.4.0
# via gpt-oss
click==8.3.1
click==8.4.2
# via
# huggingface-hub
# jiwer
# nltk
# rich-toolkit
@@ -206,7 +207,7 @@ h11==0.16.0
# uvicorn
harfile==0.4.0
# via schemathesis
hf-xet==1.4.3
hf-xet==1.5.1
# via huggingface-hub
html2text==2025.4.15
# via gpt-oss
@@ -227,7 +228,7 @@ httpx==0.28.1
# schemathesis
httpx-sse==0.4.3
# via mcp
huggingface-hub==1.10.2
huggingface-hub==1.22.0
# via
# accelerate
# datasets
@@ -959,7 +960,6 @@ typer==0.24.1
# via
# fastapi-cli
# fastapi-cloud-cli
# huggingface-hub
# transformers
typing-extensions==4.15.0
# via
+1 -1
View File
@@ -16,5 +16,5 @@ torch==2.12.0
torchaudio
torchvision
auto_round_lib>=0.13.3
auto_round_lib>=0.14.0
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl
+7
View File
@@ -144,6 +144,13 @@ too_many_arguments = "allow"
[profile.dev]
panic = "abort"
# Speed up cold tokenizer construction in tests.
[profile.dev.package]
fastokens = { opt-level = 3 }
regex-automata = { opt-level = 3 }
serde_json = { opt-level = 3 }
tokenizers = { opt-level = 3 }
[profile.release]
lto = "thin"
panic = "abort"
+27 -19
View File
@@ -214,14 +214,17 @@ macro_rules! roundtrip_tests {
($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
paste::paste! {
$(
$(
#[tokio::test]
$(#[$fixture_attr])*
#[file_serial([<hf_ $case>])]
async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
[<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
}
)*
#[tokio::test]
#[file_serial([<hf_ $case>])]
async fn [<roundtrip_ $case>]() -> Result<()> {
let case = RoundtripCase::$case();
let backends = load_roundtrip_backends(&case).await?;
$(
$(#[$fixture_attr])*
[<run_roundtrip_ $fixture>](&case, &backends).await?;
)*
Ok(())
}
)+
}
};
@@ -241,18 +244,21 @@ roundtrip_tests! {
}
/// Run the fixed reasoning+content fixture for one model/parser case.
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
async fn run_roundtrip_reasoning_and_content(
case: &RoundtripCase,
backends: &vllm_chat::LoadedModelBackends,
) -> Result<()> {
for thinking in case.thinking_behavior.fixtures() {
run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
run_roundtrip_reasoning_and_content_inner(case, backends, thinking).await?;
}
Ok(())
}
async fn run_roundtrip_reasoning_and_content_inner(
case: RoundtripCase,
case: &RoundtripCase,
backends: &vllm_chat::LoadedModelBackends,
thinking: Option<bool>,
) -> Result<()> {
let backends = load_roundtrip_backends(&case).await?;
let request = roundtrip_request(
"roundtrip-reasoning-content",
vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
@@ -275,7 +281,7 @@ async fn run_roundtrip_reasoning_and_content_inner(
});
AssistantMessage { content }
};
let result = run_roundtrip(&case, &backends, &request, assistant).await?;
let result = run_roundtrip(case, backends, &request, assistant).await?;
assert_eq!(
result.parsed_message.reasoning().as_deref().map(str::trim),
@@ -293,8 +299,10 @@ async fn run_roundtrip_reasoning_and_content_inner(
}
/// Run the fixed reasoning+multiple-tools fixture for one model/parser case.
async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
let backends = load_roundtrip_backends(&case).await?;
async fn run_roundtrip_tool_call_mix(
case: &RoundtripCase,
backends: &vllm_chat::LoadedModelBackends,
) -> Result<()> {
let request = roundtrip_request(
"roundtrip-reasoning-tools",
vec![ChatMessage::text(
@@ -308,8 +316,8 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
let expected_text = "I will call the tools.";
let result = run_roundtrip(
&case,
&backends,
case,
backends,
&request,
AssistantMessage {
content: vec![
@@ -353,12 +361,12 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
assert_eq!(tool_calls[0].name, "get_weather");
assert_eq!(
tool_calls[0].arguments,
expected_arguments(&case, r#"{"location": "Shanghai"}"#)?,
expected_arguments(case, r#"{"location": "Shanghai"}"#)?,
);
assert_eq!(tool_calls[1].name, "add");
assert_eq!(
tool_calls[1].arguments,
expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
expected_arguments(case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
);
assert_eq!(
@@ -512,6 +512,7 @@ impl EngineCoreClient {
Ok(EngineCoreOutputStream::new(
request_id,
engine_id.engine_index().unwrap_or(0),
self.abort_tx.clone(),
rx,
))
@@ -15,7 +15,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
use crate::client::stream::EngineCoreStreamOutput;
use crate::client::{AbortCause, AbortRequest};
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
use crate::protocol::encode_msgpack;
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
use crate::protocol::request::EngineCoreRequestType;
@@ -29,6 +29,7 @@ pub(crate) struct ClientInner {
/// The runtime handle used for sending messages to the engine.
handle: Handle,
model_name: String,
scheduler_stats_recorder: SchedulerStatsRecorder,
request_reg: Mutex<RequestRegistry>,
utility_reg: Mutex<UtilityRegistry>,
health_error: ArcSwapOption<Error>,
@@ -43,10 +44,13 @@ impl ClientInner {
model_name: String,
engines: &[ConnectedEngine],
) -> Self {
let scheduler_stats_recorder =
SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
Self {
input_send,
handle,
model_name,
scheduler_stats_recorder,
request_reg: Mutex::new(RequestRegistry::new(engines)),
utility_reg: Mutex::new(UtilityRegistry::default()),
health_error: ArcSwapOption::empty(),
@@ -389,12 +393,7 @@ pub(crate) async fn run_output_dispatcher_loop(
"dropping scheduler stats for unknown engine"
);
}
record_scheduler_stats(
&METRICS.scheduler,
inner.model_name(),
batch.engine_index,
scheduler_stats,
);
inner.scheduler_stats_recorder.record(batch.engine_index, scheduler_stats);
}
// The engine's scheduler stats never carry adapter names;
@@ -45,6 +45,7 @@ impl Deref for EngineCoreStreamOutput {
/// `finish_reason` is non-`None`.
pub struct EngineCoreOutputStream {
request_id: String,
engine_index: u32,
abort_tx: mpsc::UnboundedSender<AbortRequest>,
state: State,
rx: OutputReceiver,
@@ -53,11 +54,13 @@ pub struct EngineCoreOutputStream {
impl EngineCoreOutputStream {
pub(crate) fn new(
request_id: String,
engine_index: u32,
abort_tx: mpsc::UnboundedSender<AbortRequest>,
rx: OutputReceiver,
) -> Self {
Self {
request_id,
engine_index,
abort_tx,
state: State::Running,
rx,
@@ -68,6 +71,11 @@ impl EngineCoreOutputStream {
pub fn request_id(&self) -> &str {
&self.request_id
}
/// Return the index of the engine that owns this request.
pub fn engine_index(&self) -> u32 {
self.engine_index
}
}
impl Stream for EngineCoreOutputStream {
+167 -81
View File
@@ -1,90 +1,190 @@
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::time::{SystemTime, UNIX_EPOCH};
use vllm_metrics::{
EngineLabels, EnginePositionLabels, LoraAdapterNames, LoraInfoLabels, SchedulerMetrics,
EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
WaitingReasonLabels,
};
use crate::protocol::stats::SchedulerStats;
use crate::transport::ConnectedEngine;
const WAITING_REASON_CAPACITY: &str = "capacity";
const WAITING_REASON_DEFERRED: &str = "deferred";
/// Record the scheduler-stats-backed metrics for one engine at one point in
/// time.
pub(crate) fn record_scheduler_stats(
metrics: &SchedulerMetrics,
model_name: impl Into<String>,
engine: u32,
stats: &SchedulerStats,
) {
let model_name = model_name.into();
let labels = EngineLabels {
model_name: model_name.clone(),
engine,
};
/// Cached scheduler-stats metric handles for all engines connected to one
/// frontend client.
pub(crate) struct SchedulerStatsRecorder {
engines: BTreeMap<u32, SchedulerStatsHandles>,
}
/// Per-engine cached metric handles used while recording `SchedulerStats`.
struct SchedulerStatsHandles {
// Base labels reused for dynamic child labels.
labels: EngineLabels,
// Scheduler state gauges.
metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs);
metrics
.scheduler_waiting
.get_or_create(&labels)
.set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
metrics
.scheduler_waiting_by_reason
.get_or_create(&WaitingReasonLabels {
model_name: model_name.clone(),
engine,
reason: WAITING_REASON_CAPACITY,
})
.set(stats.num_waiting_reqs);
metrics
.scheduler_waiting_by_reason
.get_or_create(&WaitingReasonLabels {
model_name: model_name.clone(),
engine,
reason: WAITING_REASON_DEFERRED,
})
.set(stats.num_skipped_waiting_reqs);
metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage);
scheduler_running: U64Gauge,
scheduler_waiting: U64Gauge,
scheduler_waiting_capacity: U64Gauge,
scheduler_waiting_deferred: U64Gauge,
kv_cache_usage: F64Gauge,
// Prefix-cache counters, including the connector-backed external cache path.
metrics
.prefix_cache_queries
.get_or_create(&labels)
.inc_by(stats.prefix_cache_stats.base.queries);
metrics
.prefix_cache_hits
.get_or_create(&labels)
.inc_by(stats.prefix_cache_stats.base.hits);
prefix_cache_queries: U64Counter,
prefix_cache_hits: U64Counter,
external_prefix_cache_queries: U64Counter,
external_prefix_cache_hits: U64Counter,
// Speculative decoding counters.
spec_decode_num_drafts: U64Counter,
spec_decode_num_draft_tokens: U64Counter,
spec_decode_num_accepted_tokens: U64Counter,
spec_decode_num_accepted_tokens_per_pos: Family<EnginePositionLabels, U64Counter>,
// Per-engine performance / MFU counters.
estimated_flops_per_gpu: U64Counter,
estimated_read_bytes_per_gpu: U64Counter,
estimated_write_bytes_per_gpu: U64Counter,
// Sampled KV-cache residency histograms.
kv_block_lifetime_seconds: HistogramMetric,
kv_block_idle_before_evict_seconds: HistogramMetric,
kv_block_reuse_gap_seconds: HistogramMetric,
// Non-Prometheus interval accumulator for periodic text-log helpers.
log_stats: SchedulerLogStatsAccumulator,
}
impl SchedulerStatsRecorder {
/// Resolve the fixed-label metric handles for the connected engines.
pub(crate) fn new(
metrics: &SchedulerMetrics,
model_name: &str,
engines: &[ConnectedEngine],
) -> Self {
let engines = engines
.iter()
.filter_map(|engine| {
let engine = engine.engine_id.engine_index()?;
Some((
engine,
resolve_scheduler_stats_handles(metrics, model_name, engine),
))
})
.collect();
Self { engines }
}
/// Record one scheduler-stats payload for the given engine index.
pub(crate) fn record(&self, engine_index: u32, stats: &SchedulerStats) {
if let Some(handles) = self.engines.get(&engine_index) {
record_scheduler_stats_with_handles(handles, stats);
}
}
}
/// Resolve all fixed-label scheduler metrics for one engine.
fn resolve_scheduler_stats_handles(
metrics: &SchedulerMetrics,
model_name: &str,
engine: u32,
) -> SchedulerStatsHandles {
let labels = EngineLabels {
model_name: model_name.to_string(),
engine,
};
let capacity = WaitingReasonLabels {
model_name: model_name.to_string(),
engine,
reason: WAITING_REASON_CAPACITY,
};
let deferred = WaitingReasonLabels {
model_name: model_name.to_string(),
engine,
reason: WAITING_REASON_DEFERRED,
};
SchedulerStatsHandles {
scheduler_running: metrics.scheduler_running.get_or_create_owned(&labels),
scheduler_waiting: metrics.scheduler_waiting.get_or_create_owned(&labels),
scheduler_waiting_capacity: metrics
.scheduler_waiting_by_reason
.get_or_create_owned(&capacity),
scheduler_waiting_deferred: metrics
.scheduler_waiting_by_reason
.get_or_create_owned(&deferred),
kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels),
prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels),
prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels),
external_prefix_cache_queries: metrics
.external_prefix_cache_queries
.get_or_create_owned(&labels),
external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels),
spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels),
spec_decode_num_draft_tokens: metrics
.spec_decode_num_draft_tokens
.get_or_create_owned(&labels),
spec_decode_num_accepted_tokens: metrics
.spec_decode_num_accepted_tokens
.get_or_create_owned(&labels),
spec_decode_num_accepted_tokens_per_pos: metrics
.spec_decode_num_accepted_tokens_per_pos
.clone(),
log_stats: metrics.log_stats.get_or_create_owned(&labels),
estimated_flops_per_gpu: metrics.estimated_flops_per_gpu.get_or_create_owned(&labels),
estimated_read_bytes_per_gpu: metrics
.estimated_read_bytes_per_gpu
.get_or_create_owned(&labels),
estimated_write_bytes_per_gpu: metrics
.estimated_write_bytes_per_gpu
.get_or_create_owned(&labels),
kv_block_lifetime_seconds: metrics.kv_block_lifetime_seconds.get_or_create_owned(&labels),
kv_block_idle_before_evict_seconds: metrics
.kv_block_idle_before_evict_seconds
.get_or_create_owned(&labels),
kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels),
labels,
}
}
/// Record scheduler-stats values through pre-resolved metric handles.
fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &SchedulerStats) {
// Scheduler state gauges.
handles.scheduler_running.set(stats.num_running_reqs);
handles
.scheduler_waiting
.set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs);
handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs);
handles.kv_cache_usage.set(stats.kv_cache_usage);
// Prefix-cache counters, including the connector-backed external cache path.
handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries);
handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits);
if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats {
metrics
handles
.external_prefix_cache_queries
.get_or_create(&labels)
.inc_by(connector_prefix_cache_stats.base.queries);
metrics
handles
.external_prefix_cache_hits
.get_or_create(&labels)
.inc_by(connector_prefix_cache_stats.base.hits);
}
// Speculative decoding counters.
if let Some(spec_decoding_stats) = &stats.spec_decoding_stats {
metrics
.spec_decode_num_drafts
.get_or_create(&labels)
.inc_by(spec_decoding_stats.num_drafts);
metrics
handles.spec_decode_num_drafts.inc_by(spec_decoding_stats.num_drafts);
handles
.spec_decode_num_draft_tokens
.get_or_create(&labels)
.inc_by(spec_decoding_stats.num_draft_tokens);
metrics
handles
.spec_decode_num_accepted_tokens
.get_or_create(&labels)
.inc_by(spec_decoding_stats.num_accepted_tokens);
metrics.log_stats.get_or_create(&labels).observe_spec_decode(
handles.log_stats.observe_spec_decode(
spec_decoding_stats.num_drafts,
&spec_decoding_stats.num_accepted_tokens_per_pos,
);
@@ -92,11 +192,11 @@ pub(crate) fn record_scheduler_stats(
for (position, accepted_tokens) in
spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
{
metrics
handles
.spec_decode_num_accepted_tokens_per_pos
.get_or_create(&EnginePositionLabels {
model_name: model_name.clone(),
engine,
model_name: handles.labels.model_name.clone(),
engine: handles.labels.engine,
position: position as u32,
})
.inc_by(accepted_tokens);
@@ -109,22 +209,13 @@ pub(crate) fn record_scheduler_stats(
|| perf_stats.num_read_bytes_per_gpu != 0
|| perf_stats.num_write_bytes_per_gpu != 0)
{
metrics
.estimated_flops_per_gpu
.get_or_create(&labels)
.inc_by(perf_stats.num_flops_per_gpu);
metrics
.estimated_read_bytes_per_gpu
.get_or_create(&labels)
.inc_by(perf_stats.num_read_bytes_per_gpu);
metrics
.estimated_write_bytes_per_gpu
.get_or_create(&labels)
.inc_by(perf_stats.num_write_bytes_per_gpu);
handles.estimated_flops_per_gpu.inc_by(perf_stats.num_flops_per_gpu);
handles.estimated_read_bytes_per_gpu.inc_by(perf_stats.num_read_bytes_per_gpu);
handles.estimated_write_bytes_per_gpu.inc_by(perf_stats.num_write_bytes_per_gpu);
}
if let Some(cudagraph_stats) = &stats.cudagraph_stats {
metrics.log_stats.get_or_create(&labels).observe_cudagraph(
handles.log_stats.observe_cudagraph(
cudagraph_stats.num_unpadded_tokens,
cudagraph_stats.num_padded_tokens,
cudagraph_stats.num_paddings,
@@ -134,16 +225,11 @@ pub(crate) fn record_scheduler_stats(
// Sampled KV-cache residency histograms.
if !stats.kv_cache_eviction_events.is_empty() {
let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);
let kv_block_idle_before_evict_seconds =
metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels);
let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels);
for event in &stats.kv_cache_eviction_events {
kv_block_lifetime_seconds.observe(event.lifetime_seconds);
kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
handles.kv_block_lifetime_seconds.observe(event.lifetime_seconds);
handles.kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
for reuse_gap_seconds in &event.reuse_gaps_seconds {
kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
handles.kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
}
}
}
+11 -4
View File
@@ -88,14 +88,21 @@ impl Llm {
// Record internal engine-core request ID in the current tracing span.
Span::current().record("engine_request_id", &internal_request_id);
let arrival_time = prepared.engine_request.arrival_time;
let max_tokens_param =
(prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens);
let prompt_len = prepared.prompt_token_ids().len() as u32;
let stream = self.client.call(prepared.engine_request).await?;
let request_metrics = RequestMetricsTracker::new(
self.client.model_name().to_string(),
prepared.engine_request.arrival_time,
prepared.prompt_token_ids().len() as u32,
(prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens),
stream.engine_index(),
arrival_time,
prompt_len,
max_tokens_param,
1,
);
let stream = self.client.call(prepared.engine_request).await?;
let guard = self.inflight.track(external_request_id, internal_request_id);
Ok(GenerateOutputStream::new(
+1 -6
View File
@@ -248,12 +248,7 @@ impl Stream for GenerateOutputStream {
};
let received_at = current_unix_timestamp_secs();
self.request_metrics.observe_output(
raw.engine_index,
raw.timestamp,
received_at,
&raw.output,
);
self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output);
let raw = raw.output;
+150 -142
View File
@@ -5,15 +5,12 @@ use vllm_engine_core_client::protocol::output::{
};
use vllm_engine_core_client::protocol::stats::PrefillStats;
use vllm_metrics::{
EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
EngineLabels, Family, FinishedReasonLabels, HistogramMetric, METRICS, PromptTokenSourceLabels,
U64Counter,
};
use crate::FinishReason;
fn metrics() -> &'static RequestMetrics {
&METRICS.request
}
const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute";
const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit";
const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
@@ -29,9 +26,11 @@ const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
///
/// Original Python update flow:
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/engine/output_processor.py#L600-L677>
#[derive(Debug, Clone)]
#[derive(Clone)]
pub(crate) struct RequestMetricsTracker {
model_name: String,
/// Cached request metric handles for this request's model and engine index.
handles: RequestMetricHandles,
arrival_time: f64,
prompt_len: u32,
max_tokens_param: Option<u32>,
@@ -44,7 +43,38 @@ pub(crate) struct RequestMetricsTracker {
first_token_latency: f64,
num_generation_tokens: u32,
latest_num_cached_tokens: u32,
last_seen_engine_index: u32,
}
/// Cached request metric handles for one model and engine index.
#[derive(Clone)]
struct RequestMetricHandles {
labels: EngineLabels,
// Request-derived counters.
num_preemptions: U64Counter,
prompt_tokens: U64Counter,
prompt_tokens_local_compute: U64Counter,
prompt_tokens_local_cache_hit: U64Counter,
prompt_tokens_external_kv_transfer: U64Counter,
prompt_tokens_cached: U64Counter,
generation_tokens: U64Counter,
// Request lifecycle counters and histograms.
request_success: Family<FinishedReasonLabels, U64Counter>,
request_prompt_tokens: HistogramMetric,
request_generation_tokens: HistogramMetric,
request_max_num_generation_tokens: HistogramMetric,
request_params_max_tokens: HistogramMetric,
request_params_n: HistogramMetric,
request_prefill_kv_computed_tokens: HistogramMetric,
time_to_first_token_seconds: HistogramMetric,
inter_token_latency_seconds: HistogramMetric,
e2e_request_latency_seconds: HistogramMetric,
request_queue_time_seconds: HistogramMetric,
request_prefill_time_seconds: HistogramMetric,
request_decode_time_seconds: HistogramMetric,
request_inference_time_seconds: HistogramMetric,
request_time_per_output_token_seconds: HistogramMetric,
}
impl RequestMetricsTracker {
@@ -52,13 +82,14 @@ impl RequestMetricsTracker {
/// context.
pub(crate) fn new(
model_name: String,
engine_index: u32,
arrival_time: f64,
prompt_len: u32,
max_tokens_param: Option<u32>,
n_param: u32,
) -> Self {
Self {
model_name,
handles: resolve_request_metric_handles(&model_name, engine_index),
arrival_time,
prompt_len,
max_tokens_param,
@@ -71,7 +102,6 @@ impl RequestMetricsTracker {
first_token_latency: 0.0,
num_generation_tokens: 0,
latest_num_cached_tokens: 0,
last_seen_engine_index: 0,
}
}
@@ -81,23 +111,18 @@ impl RequestMetricsTracker {
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L331-L384>
pub(crate) fn observe_output(
&mut self,
engine_index: u32,
batch_timestamp: f64,
received_at: f64,
output: &EngineCoreOutput,
) {
self.last_seen_engine_index = engine_index;
if let Some(prefill_stats) = &output.prefill_stats {
self.latest_num_cached_tokens = prefill_stats.num_cached_tokens;
}
self.num_generation_tokens += output.new_token_ids.len() as u32;
metrics()
.generation_tokens
.get_or_create(&engine_labels(&self.model_name, engine_index))
.inc_by(output.new_token_ids.len() as u64);
self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64);
if let Some(events) = &output.events {
self.observe_events(engine_index, events);
self.observe_events(events);
}
// Only outputs that actually carry tokens drive token-timing metrics.
@@ -107,22 +132,16 @@ impl RequestMetricsTracker {
if !output.new_token_ids.is_empty() {
if self.is_prefilling {
if let Some(prefill_stats) = &output.prefill_stats {
record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
self.record_prompt_tokens(prefill_stats);
}
self.first_token_latency = received_at - self.arrival_time;
observe_time_to_first_token_seconds(
&self.model_name,
engine_index,
self.first_token_latency,
);
self.handles.time_to_first_token_seconds.observe(self.first_token_latency);
self.first_token_ts = batch_timestamp;
self.is_prefilling = false;
} else if self.last_token_ts > 0.0 {
observe_inter_token_latency_seconds(
&self.model_name,
engine_index,
batch_timestamp - self.last_token_ts,
);
self.handles
.inter_token_latency_seconds
.observe(batch_timestamp - self.last_token_ts);
}
self.last_token_ts = batch_timestamp;
@@ -135,7 +154,6 @@ impl RequestMetricsTracker {
/// Original Python finished-request stats:
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L222-L237>
pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) {
let labels = engine_labels(&self.model_name, self.last_seen_engine_index);
let prefill_kv_computed_tokens =
self.prompt_len.saturating_sub(self.latest_num_cached_tokens);
let e2e_latency_seconds = received_at - self.arrival_time;
@@ -150,57 +168,47 @@ impl RequestMetricsTracker {
0.0
};
record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason);
metrics()
.request_prompt_tokens
.get_or_create(&labels)
.observe(self.prompt_len as f64);
metrics()
self.record_request_success(finish_reason);
self.handles.request_prompt_tokens.observe(self.prompt_len as f64);
self.handles
.request_generation_tokens
.get_or_create(&labels)
.observe(self.num_generation_tokens as f64);
metrics()
self.handles
.request_max_num_generation_tokens
.get_or_create(&labels)
.observe(self.num_generation_tokens as f64);
if let Some(max_tokens_param) = self.max_tokens_param {
metrics()
.request_params_max_tokens
.get_or_create(&labels)
.observe(max_tokens_param as f64);
self.handles.request_params_max_tokens.observe(max_tokens_param as f64);
}
metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64);
metrics()
self.handles.request_params_n.observe(self.n_param as f64);
self.handles
.request_prefill_kv_computed_tokens
.get_or_create(&labels)
.observe(prefill_kv_computed_tokens as f64);
metrics()
.e2e_request_latency_seconds
.get_or_create(&labels)
.observe(e2e_latency_seconds);
metrics()
.request_queue_time_seconds
.get_or_create(&labels)
.observe(queue_time_seconds);
metrics()
.request_prefill_time_seconds
.get_or_create(&labels)
.observe(prefill_time_seconds);
metrics()
.request_decode_time_seconds
.get_or_create(&labels)
.observe(decode_time_seconds);
metrics()
.request_inference_time_seconds
.get_or_create(&labels)
.observe(inference_time_seconds);
metrics()
self.handles.e2e_request_latency_seconds.observe(e2e_latency_seconds);
self.handles.request_queue_time_seconds.observe(queue_time_seconds);
self.handles.request_prefill_time_seconds.observe(prefill_time_seconds);
self.handles.request_decode_time_seconds.observe(decode_time_seconds);
self.handles.request_inference_time_seconds.observe(inference_time_seconds);
self.handles
.request_time_per_output_token_seconds
.get_or_create(&labels)
.observe(time_per_output_token_seconds);
}
fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) {
/// Record prompt token counters through cached metric handles.
fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) {
let computed = prefill_stats.num_computed_tokens as u64;
let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
self.handles.prompt_tokens.inc_by(prefill_stats.num_prompt_tokens as u64);
self.handles.prompt_tokens_local_compute.inc_by(computed);
self.handles.prompt_tokens_local_cache_hit.inc_by(local_cache_hit);
self.handles.prompt_tokens_external_kv_transfer.inc_by(external_kv_transfer);
self.handles.prompt_tokens_cached.inc_by(prefill_stats.num_cached_tokens as u64);
}
/// Record request event counters through cached metric handles.
fn observe_events(&mut self, events: &[EngineCoreEvent]) {
for event in events {
match event.r#type {
EngineCoreEventType::Queued => {
@@ -212,46 +220,86 @@ impl RequestMetricsTracker {
}
}
EngineCoreEventType::Preempted => {
metrics()
.num_preemptions
.get_or_create(&engine_labels(&self.model_name, engine_index))
.inc();
self.handles.num_preemptions.inc();
}
}
}
}
}
fn engine_labels(model_name: &str, engine: u32) -> EngineLabels {
EngineLabels {
model_name: model_name.to_string(),
engine,
/// Increment the request-success counter for the terminal finish reason.
fn record_request_success(&self, finish_reason: FinishReason) {
self.handles
.request_success
.get_or_create(&FinishedReasonLabels {
model_name: self.handles.labels.model_name.clone(),
engine: self.handles.labels.engine,
finished_reason: finish_reason.as_str(),
})
.inc();
}
}
fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) {
metrics()
.time_to_first_token_seconds
.get_or_create(&engine_labels(model_name, engine))
.observe(seconds);
}
/// Resolve fixed request metric handles for one model and engine index.
fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetricHandles {
let metrics = &METRICS.request;
let labels = EngineLabels {
model_name: model_name.to_string(),
engine,
};
fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) {
metrics()
.inter_token_latency_seconds
.get_or_create(&engine_labels(model_name, engine))
.observe(seconds);
}
fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) {
metrics()
.request_success
.get_or_create(&FinishedReasonLabels {
model_name: model_name.to_string(),
engine,
finished_reason: finish_reason.as_str(),
})
.inc();
RequestMetricHandles {
num_preemptions: metrics.num_preemptions.get_or_create_owned(&labels),
prompt_tokens: metrics.prompt_tokens.get_or_create_owned(&labels),
prompt_tokens_local_compute: metrics.prompt_tokens_by_source.get_or_create_owned(
&prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE),
),
prompt_tokens_local_cache_hit: metrics.prompt_tokens_by_source.get_or_create_owned(
&prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT),
),
prompt_tokens_external_kv_transfer: metrics.prompt_tokens_by_source.get_or_create_owned(
&prompt_token_source_labels(
model_name,
engine,
PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
),
),
prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels),
generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels),
request_success: metrics.request_success.clone(),
request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels),
request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels),
request_max_num_generation_tokens: metrics
.request_max_num_generation_tokens
.get_or_create_owned(&labels),
request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels),
request_params_n: metrics.request_params_n.get_or_create_owned(&labels),
request_prefill_kv_computed_tokens: metrics
.request_prefill_kv_computed_tokens
.get_or_create_owned(&labels),
time_to_first_token_seconds: metrics
.time_to_first_token_seconds
.get_or_create_owned(&labels),
inter_token_latency_seconds: metrics
.inter_token_latency_seconds
.get_or_create_owned(&labels),
e2e_request_latency_seconds: metrics
.e2e_request_latency_seconds
.get_or_create_owned(&labels),
request_queue_time_seconds: metrics.request_queue_time_seconds.get_or_create_owned(&labels),
request_prefill_time_seconds: metrics
.request_prefill_time_seconds
.get_or_create_owned(&labels),
request_decode_time_seconds: metrics
.request_decode_time_seconds
.get_or_create_owned(&labels),
request_inference_time_seconds: metrics
.request_inference_time_seconds
.get_or_create_owned(&labels),
request_time_per_output_token_seconds: metrics
.request_time_per_output_token_seconds
.get_or_create_owned(&labels),
labels,
}
}
fn prompt_token_source_labels(
@@ -266,45 +314,6 @@ fn prompt_token_source_labels(
}
}
fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) {
let computed = prefill_stats.num_computed_tokens as u64;
let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
metrics()
.prompt_tokens
.get_or_create(&engine_labels(model_name, engine))
.inc_by(prefill_stats.num_prompt_tokens as u64);
metrics()
.prompt_tokens_by_source
.get_or_create(&prompt_token_source_labels(
model_name,
engine,
PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE,
))
.inc_by(computed);
metrics()
.prompt_tokens_by_source
.get_or_create(&prompt_token_source_labels(
model_name,
engine,
PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT,
))
.inc_by(local_cache_hit);
metrics()
.prompt_tokens_by_source
.get_or_create(&prompt_token_source_labels(
model_name,
engine,
PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
))
.inc_by(external_kv_transfer);
metrics()
.prompt_tokens_cached
.get_or_create(&engine_labels(model_name, engine))
.inc_by(prefill_stats.num_cached_tokens as u64);
}
fn diff_or_zero(end: f64, start: f64) -> f64 {
if end > 0.0 && start > 0.0 && end >= start {
end - start
@@ -337,10 +346,10 @@ mod tests {
#[test]
fn tracker_updates_timing_state_across_prefill_decode_and_finish() {
let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1);
let mut tracker =
RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1);
tracker.observe_output(
2,
10.0,
100.2,
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
@@ -368,7 +377,6 @@ mod tests {
},
);
tracker.observe_output(
2,
11.5,
100.4,
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
@@ -384,7 +392,7 @@ mod tests {
);
assert!(!tracker.is_prefilling);
assert_eq!(tracker.last_seen_engine_index, 2);
assert_eq!(tracker.handles.labels.engine, 2);
assert_eq!(tracker.num_generation_tokens, 3);
assert_eq!(tracker.queued_ts, 8.0);
assert_eq!(tracker.scheduled_ts, 9.0);
+3 -3
View File
@@ -17,7 +17,7 @@ use vllm_engine_core_client::protocol::request::EngineCoreRequest;
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
use vllm_engine_core_client::protocol::stats::PrefillStats;
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
use vllm_llm::{
Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm,
};
@@ -699,7 +699,7 @@ async fn abort_by_external_id_aborts_all_internal_requests() {
async fn generate_records_request_metrics_in_prometheus_output() {
let ipc = IpcNamespace::new().unwrap();
let handshake_address = ipc.handshake_endpoint();
let engine_id = b"engine-metrics".to_vec();
let engine_id = EngineId::from_engine_index(4);
let model_name = request_metrics_model_name("metrics-model");
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
@@ -832,7 +832,7 @@ async fn generate_records_request_metrics_in_prometheus_output() {
async fn dropping_stream_records_abort_terminal_request_metrics() {
let ipc = IpcNamespace::new().unwrap();
let handshake_address = ipc.handshake_endpoint();
let engine_id = b"engine-metrics-drop".to_vec();
let engine_id = EngineId::from_engine_index(5);
let model_name = request_metrics_model_name("metrics-drop-model");
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
+3 -1
View File
@@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64;
use prometheus_client::encoding::text::encode;
use prometheus_client::metrics::counter::Counter;
use prometheus_client::metrics::family::Family;
pub use prometheus_client::metrics::family::Family;
use prometheus_client::metrics::gauge::Gauge;
use prometheus_client::metrics::histogram::Histogram;
use prometheus_client::registry::Registry;
@@ -23,6 +23,8 @@ pub use scheduler::*;
pub type U64Counter = Counter<u64, AtomicU64>;
pub type U64Gauge = Gauge<u64, AtomicU64>;
pub type F64Gauge = Gauge<f64, AtomicU64>;
/// Histogram metric handle cloned out of a Prometheus family.
pub type HistogramMetric = Histogram;
pub(crate) type HistogramFamily = Family<EngineLabels, Histogram, fn() -> Histogram>;
/// Shared Prometheus registry for frontend metrics.
+30 -48
View File
@@ -3,7 +3,7 @@ use std::time::Duration;
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
use vllm_parser::tool::{Tool, ToolParser};
use vllm_parser::unified::{Gemma4ConfigDrivenParser, Gemma4UnifiedParser};
use vllm_parser::unified::Gemma4UnifiedParser;
mod utils;
use utils::{UnifiedToolParserAdapter, feed_parser};
@@ -68,22 +68,14 @@ fn long_tool_argument_fixture() -> String {
)
}
type MakeParser = fn(&[Tool]) -> Box<dyn ToolParser>;
fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
UnifiedToolParserAdapter::<Gemma4UnifiedParser>::create(tools)
.expect("Gemma4 unified parser should initialize")
}
fn config_driven_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
UnifiedToolParserAdapter::<Gemma4ConfigDrivenParser>::create(tools)
.expect("Gemma4 config-driven parser should initialize")
}
fn run_stream_group(
c: &mut Criterion,
name: &str,
make_parser: MakeParser,
tools: &[Tool],
text: &str,
chunk_chars: usize,
@@ -99,7 +91,7 @@ fn run_stream_group(
group.throughput(Throughput::Bytes(text.len() as u64));
group.bench_function("reuse_parser", |b| {
let mut parser = make_parser(tools);
let mut parser = parser(tools);
b.iter(|| {
let result = feed_parser(&mut *parser, black_box(&chunks));
debug_assert_eq!(result.0, expected_normal_text);
@@ -110,7 +102,7 @@ fn run_stream_group(
group.bench_function("create_parser", |b| {
b.iter_batched(
|| make_parser(tools),
|| parser(tools),
|mut parser| {
let result = feed_parser(&mut *parser, black_box(&chunks));
debug_assert_eq!(result.0, expected_normal_text);
@@ -130,45 +122,35 @@ fn bench_gemma4(c: &mut Criterion) {
let long_tool_argument = long_tool_argument_fixture();
let long_normal_text = long_normal_text_fixture();
let variants: [(&str, MakeParser); 2] = [
("gemma4", parser),
("gemma4_config_driven", config_driven_parser),
];
run_stream_group(
c,
"gemma4/mixed_complex_tool_call",
&tools,
&mixed_text,
CHUNK_CHARS,
"I will inspect the data before answering.\n Finished.",
2,
);
for (prefix, make_parser) in variants {
run_stream_group(
c,
&format!("{prefix}/mixed_complex_tool_call"),
make_parser,
&tools,
&mixed_text,
CHUNK_CHARS,
"I will inspect the data before answering.\n Finished.",
2,
);
run_stream_group(
c,
"gemma4/long_tool_argument",
&tools,
&long_tool_argument,
CHUNK_CHARS,
"I will write the file.\nDone.",
1,
);
run_stream_group(
c,
&format!("{prefix}/long_tool_argument"),
make_parser,
&tools,
&long_tool_argument,
CHUNK_CHARS,
"I will write the file.\nDone.",
1,
);
run_stream_group(
c,
&format!("{prefix}/long_normal_text"),
make_parser,
&tools,
&long_normal_text,
CHUNK_CHARS,
&long_normal_text,
0,
);
}
run_stream_group(
c,
"gemma4/long_normal_text",
&tools,
&long_normal_text,
CHUNK_CHARS,
&long_normal_text,
0,
);
}
criterion_group!(benches, bench_gemma4);
-1
View File
@@ -28,7 +28,6 @@ pub use json::{
pub use kimi_k2::KimiK2ToolParser;
pub use minimax_m2::MinimaxM2ToolParser;
pub use minimax_m3::MinimaxM3ToolParser;
pub use parameters::{ToolSchema, ToolSchemas};
pub use qwen_coder::Qwen3CoderToolParser;
use serde::{Deserialize, Serialize};
use serde_json::Value;
+14 -28
View File
@@ -6,7 +6,7 @@ use crate::tool::Tool;
/// Normalized parameter schemas for all tools in one request.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ToolSchemas {
pub(super) struct ToolSchemas {
tools: BTreeMap<String, ToolSchema>,
}
@@ -17,7 +17,7 @@ pub struct ToolSchemas {
/// coercing raw string parameter values into more specific JSON types for
/// downstream tool call execution.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ToolSchema {
pub(super) struct ToolSchema {
params: BTreeMap<String, JsonParamType>,
}
@@ -25,7 +25,7 @@ pub struct ToolSchema {
///
/// It can be either a raw text string, or a structured input with named child elements.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ParamInput {
pub(super) enum ParamInput {
Text(String),
#[allow(dead_code)]
Elements(Vec<ParamElement>),
@@ -39,7 +39,7 @@ impl From<String> for ParamInput {
/// One named structured parameter child.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ParamElement {
pub(super) struct ParamElement {
pub name: String,
pub value: ParamInput,
}
@@ -64,7 +64,7 @@ pub(super) enum JsonParamType {
impl ToolSchemas {
/// Normalize OpenAI-style tool parameter JSON schemas for one request.
pub(crate) fn from_tools(tools: &[Tool]) -> Self {
pub(super) fn from_tools(tools: &[Tool]) -> Self {
let tools = tools
.iter()
.map(|tool| (tool.name.clone(), ToolSchema::from_schema(&tool.parameters)))
@@ -73,14 +73,6 @@ impl ToolSchemas {
Self { tools }
}
/// Resolve the parameter schema for one named tool.
///
/// Unknown tool names resolve to the empty schema, so all parameters fall
/// back to strings or object-like JSON for structured inputs.
pub(crate) fn resolve(&self, function_name: &str) -> &ToolSchema {
self.tools.get(function_name).unwrap_or(ToolSchema::empty())
}
/// Convert parameter values for one named tool.
///
/// Unknown tool names use an empty schema, so all parameters fall back to
@@ -93,7 +85,13 @@ impl ToolSchemas {
where
P: Into<ParamInput>,
{
self.resolve(function_name).convert_params(params)
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
let mut converted = Map::with_capacity(params.len());
for (name, value) in params {
let value = tool_schema.convert(&name, value.into());
converted.insert(name, value);
}
converted
}
/// Convert one parameter value for one named tool.
@@ -106,24 +104,12 @@ impl ToolSchemas {
where
P: Into<ParamInput>,
{
self.resolve(function_name).convert(name, value.into())
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
tool_schema.convert(name, value.into())
}
}
impl ToolSchema {
/// Convert a set of named parameter values with this schema.
pub(crate) fn convert_params<P>(&self, params: Vec<(String, P)>) -> Map<String, Value>
where
P: Into<ParamInput>,
{
let mut converted = Map::with_capacity(params.len());
for (name, value) in params {
let value = self.convert(&name, value.into());
converted.insert(name, value);
}
converted
}
/// Return an empty schema with no parameter information, which causes all
/// parameters to be treated as strings.
const fn empty() -> &'static Self {
@@ -1,875 +0,0 @@
//! Gemma4 format for the config-driven unified parser.
//!
//! Original Python implementation:
//! <https://github.com/vllm-project/vllm/blob/main/vllm/parser/gemma4.py>
//!
//! Handles Gemma4 reasoning and function-call formats:
//!
//! `<|channel>thought\nreasoning<channel|>`
//!
//! `<|tool_call>call:func_name{key:<|"|>value<|"|>}<tool_call|>`
use serde_json::{Map, Number, Value};
use winnow::ascii::multispace0 as ws0;
use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};
use winnow::error::{ContextError, ErrMode, ModalResult};
use winnow::prelude::*;
use winnow::stream::Stream;
use winnow::token::{literal, take_till, take_until};
use super::{ArgsScan, ArgsStep, ConfigDrivenParser, Input, ParserFormat};
use crate::tool::ToolSchema;
use crate::utils::{incomplete, partial_prefix_len};
const CHANNEL_START: &str = "<|channel>";
const CHANNEL_END: &str = "<channel|>";
const REASONING_START: &str = "<|channel>thought\n";
const TOOL_CALL_START: &str = "<|tool_call>";
const TOOL_CALL_END: &str = "<tool_call|>";
const STRING_DELIM: &str = "<|\"|>";
const CALL_PREFIX: &str = "call:";
/// Zero-sized marker type providing the Google Gemma4 format.
pub struct Gemma4Format;
impl ParserFormat for Gemma4Format {
const NAME: &'static str = "Gemma4";
const REASONING_START: Option<&'static str> = Some(REASONING_START);
const REASONING_END: Option<&'static str> = Some(CHANNEL_END);
const CALL_START: &'static str = TOOL_CALL_START;
const CALL_END: &'static str = TOOL_CALL_END;
const REASONING_BOUNDARY_TOKENS: Option<(&'static str, &'static str)> =
Some((CHANNEL_START, CHANNEL_END));
type ArgsScan = Gemma4ArgsScan;
fn tool_header(input: &mut Input<'_>) -> ModalResult<String> {
let (name,) = seq!(
_: literal(CALL_PREFIX),
gemma4_tool_name,
_: literal("{"),
)
.parse_next(input)?;
Ok(name)
}
fn tool_args(_schema: &ToolSchema, body: &str) -> ModalResult<Map<String, Value>> {
let Some(args_input) = body.strip_suffix('}') else {
return Err(ErrMode::Cut(ContextError::new()));
};
parse_gemma4_args(args_input)
}
}
/// Config-driven unified parser for Google Gemma4 models.
///
/// Arguments are emitted only after a full Gemma4 tool call is parsed.
pub type Gemma4ConfigDrivenParser = ConfigDrivenParser<Gemma4Format>;
/// Resumable Gemma4 argument scanner: locates the call end marker outside
/// `<|"|>`-delimited strings, so marker-shaped literals inside string values
/// do not terminate the call early.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Gemma4ArgsScan {
scanned_len: usize,
in_string: bool,
}
impl ArgsScan for Gemma4ArgsScan {
fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<ArgsStep<'i>> {
let text = **input;
if self.scanned_len > text.len() {
return incomplete();
}
loop {
let rest = &text[self.scanned_len..];
if self.in_string {
let Some(string_delim) = rest.find(STRING_DELIM) else {
self.scanned_len = safe_scan_len(text, self.scanned_len, &[STRING_DELIM]);
return incomplete();
};
self.scanned_len += string_delim + STRING_DELIM.len();
self.in_string = false;
continue;
}
let next_string_delim = rest.find(STRING_DELIM);
let next_call_end = rest.find(end);
match (next_string_delim, next_call_end) {
(Some(string_delim), Some(call_end)) if call_end < string_delim => {
let body_end = self.scanned_len + call_end;
self.scanned_len = body_end + end.len();
input.next_slice(self.scanned_len);
return Ok(ArgsStep::Buffered {
body: &text[..body_end],
});
}
(Some(string_delim), _) => {
self.scanned_len += string_delim + STRING_DELIM.len();
self.in_string = true;
}
(None, Some(call_end)) => {
let body_end = self.scanned_len + call_end;
self.scanned_len = body_end + end.len();
input.next_slice(self.scanned_len);
return Ok(ArgsStep::Buffered {
body: &text[..body_end],
});
}
(None, None) => {
self.scanned_len = safe_scan_len(text, self.scanned_len, &[STRING_DELIM, end]);
return incomplete();
}
}
}
}
}
/// Return the scan length while holding back a split marker prefix.
fn safe_scan_len(text: &str, start: usize, markers: &[&str]) -> usize {
let max_partial = markers
.iter()
.map(|marker| partial_prefix_len(&text[start..], marker))
.max()
.unwrap_or(0);
text.len() - max_partial
}
/// Parse a Gemma4 tool name.
fn gemma4_tool_name(input: &mut Input<'_>) -> ModalResult<String> {
let name = take_until(1.., "{").parse_next(input)?.trim();
if name.is_empty() {
return Err(ErrMode::Cut(ContextError::new()));
}
Ok(name.to_string())
}
/// Parse complete Gemma4 custom key-value arguments.
fn parse_gemma4_args(args: &str) -> ModalResult<Map<String, Value>> {
let mut input = args;
terminated(gemma4_args, eof).parse_next(&mut input)
}
/// Parse Gemma4's custom key-value argument object content.
fn gemma4_args(input: &mut &str) -> ModalResult<Map<String, Value>> {
let pairs: Vec<(String, Value)> = delimited(
ws0,
terminated(
separated(0.., gemma4_pair, comma_separator),
opt(comma_separator),
),
ws0,
)
.parse_next(input)?;
Ok(pairs.into_iter().collect())
}
/// Parse a Gemma4 key-value pair.
fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> {
let (key, value) = seq!(
_: ws0,
gemma4_key,
_: ws0,
_: literal(":"),
_: ws0,
gemma4_value,
)
.parse_next(input)?;
Ok((key, value))
}
/// Parse a Gemma4 bare key.
fn gemma4_key(input: &mut &str) -> ModalResult<String> {
let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim();
if key.is_empty() {
return Err(ErrMode::Cut(ContextError::new()));
}
Ok(key.to_string())
}
/// Parse a Gemma4 value.
fn gemma4_value(input: &mut &str) -> ModalResult<Value> {
alt((
gemma4_string.map(|value: &str| Value::String(value.to_string())),
gemma4_object.map(Value::Object),
gemma4_array_value.map(Value::Array),
gemma4_bare_value,
))
.parse_next(input)
}
/// Parse a Gemma4 string delimited by `<|"|>`.
fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
delimited(
literal(STRING_DELIM),
take_until(0.., STRING_DELIM),
literal(STRING_DELIM),
)
.parse_next(input)
}
/// Parse a nested Gemma4 object.
fn gemma4_object(input: &mut &str) -> ModalResult<Map<String, Value>> {
delimited(literal("{"), gemma4_args, literal("}")).parse_next(input)
}
/// Parse a Gemma4 array value.
fn gemma4_array_value(input: &mut &str) -> ModalResult<Vec<Value>> {
delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input)
}
/// Parse Gemma4 array content.
fn gemma4_array_content(input: &mut &str) -> ModalResult<Vec<Value>> {
delimited(
ws0,
terminated(
separated(0.., gemma4_value, comma_separator),
opt(comma_separator),
),
ws0,
)
.parse_next(input)
}
/// Parse a Gemma4 bare scalar.
fn gemma4_bare_value(input: &mut &str) -> ModalResult<Value> {
take_till(1.., |char: char| matches!(char, ',' | '}' | ']'))
.map(parse_gemma4_scalar)
.parse_next(input)
}
/// Parse a Gemma4 comma separator.
fn comma_separator(input: &mut &str) -> ModalResult<()> {
delimited(ws0, literal(","), ws0).void().parse_next(input)
}
fn parse_gemma4_scalar(value: &str) -> Value {
let value = value.trim();
if value.is_empty() {
return Value::String(String::new());
}
if value == "true" {
return Value::Bool(true);
}
if value == "false" {
return Value::Bool(false);
}
if matches!(value, "null" | "none" | "nil" | "NULL" | "None" | "NIL") {
return Value::Null;
}
if value.contains('.') {
if let Ok(parsed) = value.parse::<f64>()
&& let Some(number) = Number::from_f64(parsed)
{
return Value::Number(number);
}
} else if let Ok(parsed) = value.parse::<i64>() {
return Value::Number(Number::from(parsed));
}
Value::String(value.to_string())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use serde_json::{Value, json};
use thiserror_ext::AsReport;
use vllm_tokenizer::test_utils::TestTokenizer;
use winnow::combinator::{eof, terminated};
use winnow::error::ErrMode;
use winnow::prelude::*;
use super::{
CHANNEL_END, CHANNEL_START, Gemma4ConfigDrivenParser, gemma4_array_content,
parse_gemma4_args,
};
use crate::tool::{Tool, ToolCallDelta};
use crate::unified::{
Result, UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
parsing_failed,
};
const CHANNEL_START_ID: u32 = 256;
const CHANNEL_END_ID: u32 = 257;
const TURN_BOUNDARY_ID: u32 = 258;
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_special_token(CHANNEL_START, CHANNEL_START_ID)
.with_special_token(CHANNEL_END, CHANNEL_END_ID)
.with_special_token("<turn-boundary>", TURN_BOUNDARY_ID)
}
trait UnifiedParserTestExt {
fn parse_chunk(&mut self, chunk: &str) -> Result<UnifiedParserOutput>;
fn parse_complete(&mut self, text: &str) -> Result<UnifiedParserOutput>;
}
impl UnifiedParserTestExt for Gemma4ConfigDrivenParser {
fn parse_chunk(&mut self, chunk: &str) -> Result<UnifiedParserOutput> {
let mut output = UnifiedParserOutput::default();
self.parse_into(chunk, &mut output)?;
Ok(output)
}
fn parse_complete(&mut self, text: &str) -> Result<UnifiedParserOutput> {
let mut output = self.parse_chunk(text)?;
output.append(self.finish()?);
Ok(output)
}
}
trait UnifiedOutputTestExt {
fn normal_text(&self) -> String;
fn reasoning_text(&self) -> String;
fn calls(&self) -> Vec<&ToolCallDelta>;
fn coalesce(self) -> Self;
}
impl UnifiedOutputTestExt for UnifiedParserOutput {
fn normal_text(&self) -> String {
self.events
.iter()
.filter_map(|event| match event {
UnifiedParserEvent::Text(text) => Some(text.as_str()),
UnifiedParserEvent::Reasoning(_) | UnifiedParserEvent::ToolCall(_) => None,
})
.collect()
}
fn reasoning_text(&self) -> String {
self.events
.iter()
.filter_map(|event| match event {
UnifiedParserEvent::Reasoning(text) => Some(text.as_str()),
UnifiedParserEvent::Text(_) | UnifiedParserEvent::ToolCall(_) => None,
})
.collect()
}
fn calls(&self) -> Vec<&ToolCallDelta> {
self.events
.iter()
.filter_map(|event| match event {
UnifiedParserEvent::Text(_) | UnifiedParserEvent::Reasoning(_) => None,
UnifiedParserEvent::ToolCall(call) => Some(call),
})
.collect()
}
fn coalesce(self) -> Self {
self
}
}
fn parse_gemma4_array(array: &str) -> Result<Vec<Value>> {
let mut input = array;
match terminated(gemma4_array_content, eof).parse_next(&mut input) {
Ok(value) => Ok(value),
Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")),
Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
Err(parsing_failed!("{}", error))
}
}
}
fn test_tools() -> Vec<Tool> {
vec![
Tool {
name: "get_weather".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "get_time".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "write_file".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "Edit".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "search".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "set".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "get_status".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
Tool {
name: "todowrite".to_string(),
description: None,
parameters: json!({ "type": "object" }),
strict: None,
},
]
}
fn test_parser() -> Gemma4ConfigDrivenParser {
Gemma4ConfigDrivenParser::new(&test_tools(), Arc::new(tokenizer())).unwrap()
}
#[test]
fn gemma4_create_requires_channel_start_token() {
let error =
match Gemma4ConfigDrivenParser::new(&test_tools(), Arc::new(TestTokenizer::new())) {
Ok(_) => panic!("expected missing token error"),
Err(error) => error,
};
assert!(matches!(
error,
UnifiedParserError::MissingToken { token } if token == CHANNEL_START
));
}
fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
for chunk in chunks {
output.append(parser.parse_chunk(chunk).unwrap());
}
output.append(parser.finish().unwrap());
output.coalesce()
}
fn first_call(output: &UnifiedParserOutput) -> ToolCallDelta {
(*output.calls().first().expect("expected one tool call")).clone()
}
#[test]
fn gemma4_parse_args_handles_scalars_and_nested_values() {
let parsed = parse_gemma4_args(
"name:<|\"|>test<|\"|>,count:42,active:true,score:114.514,nested:{inner:<|\"|>value<|\"|>},items:[<|\"|>a<|\"|>,<|\"|>b<|\"|>]",
)
.unwrap();
assert_eq!(
Value::Object(parsed),
json!({
"name": "test",
"count": 42,
"active": true,
"score": 114.514,
"nested": { "inner": "value" },
"items": ["a", "b"],
})
);
}
#[test]
fn gemma4_parse_args_handles_empty_arguments() {
let parsed = parse_gemma4_args("").unwrap();
assert_eq!(Value::Object(parsed), json!({}));
}
#[test]
fn gemma4_parse_array_handles_bare_values() {
let parsed = parse_gemma4_array("42,true,114.514").unwrap();
assert_eq!(Value::Array(parsed), json!([42, true, 114.514]));
}
#[test]
fn gemma4_parse_complete_extracts_single_tool_call() {
let mut parser = test_parser();
let output = parser
.parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
.unwrap();
assert!(output.normal_text().is_empty());
assert_eq!(output.calls().len(), 1);
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "London" })
);
}
#[test]
fn gemma4_parse_complete_rejects_incomplete_tool_call() {
let mut parser = test_parser();
let error = parser
.parse_complete("<|tool_call>call:get_weather{location:<|\"|>London")
.unwrap_err();
assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
}
#[test]
fn gemma4_streaming_basic_single_tool_call() {
let output = collect_stream(&[
"<|tool_call>",
"call:get_weather{",
"location:<|\"|>Paris",
", France",
"<|\"|>}",
"<tool_call|>",
]);
assert!(output.normal_text().is_empty());
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "Paris, France" })
);
}
#[test]
fn gemma4_streaming_text_before_and_after_tool_call() {
let output = collect_stream(&[
"Let me check ",
"the weather. ",
"<|tool_call>",
"call:get_weather{",
"location:<|\"|>London<|\"|>}",
"<tool_call|><",
"div>",
]);
assert_eq!(output.normal_text(), "Let me check the weather. <div>");
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "London" })
);
}
#[test]
fn gemma4_streaming_waits_for_complete_tool_call() {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
for chunk in [
"<|tool_call>",
"call:get_weather{",
"location:<|\"|>Paris<|\"|>}",
] {
output.append(parser.parse_chunk(chunk).unwrap());
assert!(output.calls().is_empty());
}
output.append(parser.parse_chunk("<tool_call|>").unwrap());
let output = output.coalesce();
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "Paris" })
);
}
#[test]
fn gemma4_streaming_handles_boolean_split_across_chunks() {
let output = collect_stream(&[
"<|tool_call>",
"call:search{input:{all:tru",
"e}}",
"<tool_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("search"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "input": { "all": true } })
);
}
#[test]
fn gemma4_streaming_handles_false_split_across_chunks() {
let output = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
assert_eq!(first_call(&output).name.as_deref(), Some("set"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "flag": false })
);
}
#[test]
fn gemma4_streaming_handles_number_split_across_chunks() {
let output = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
assert_eq!(first_call(&output).name.as_deref(), Some("set"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "count": 42 })
);
}
#[test]
fn gemma4_streaming_handles_split_string_delimiter() {
let output = collect_stream(&[
"<|tool_call>",
"call:todowrite{",
"content:<|\"|>Buy milk<|",
"\"|>}",
"<tool_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "content": "Buy milk" })
);
assert!(!first_call(&output).arguments.contains("<|"));
}
#[test]
fn gemma4_streaming_handles_split_tool_call_end_marker() {
let output = collect_stream(&[
"<|tool_call>",
"call:get_weather{location:<|\"|>Paris<|\"|>}<tool",
"_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "Paris" })
);
}
#[test]
fn gemma4_streaming_handles_end_marker_literal_inside_string() {
let output = collect_stream(&[
"<|tool_call>",
"call:todowrite{",
"content:<|\"|>literal }<tool_call|> inside",
"<|\"|>}",
"<tool_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "content": "literal }<tool_call|> inside" })
);
}
#[test]
fn gemma4_streaming_handles_html_argument_without_duplication() {
let output = collect_stream(&[
"<|tool_call>",
"call:write_file{",
"path:<|\"|>index.html<|\"|>,",
"content:<|\"|><!DOCTYPE html>\n<",
"html lang=\"zh-CN\">\n<",
"head>\n <",
"meta charset=\"UTF-8\">\n <",
"meta name=\"viewport\" content=\"width=device-width\">\n",
"<|\"|>}",
"<tool_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("write_file"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({
"path": "index.html",
"content": "<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width\">\n",
})
);
}
#[test]
fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() {
let output = collect_stream(&[
"<|tool_call>",
"call:Edit{",
"file_path:<|\"|>src/env.py<|\"|>,",
"old_string:<|\"|>old_val<|\"|>,",
"new_string:<|\"|>new_val<|\"|>,",
"replace_all:",
"false}",
"<tool_call|>",
]);
assert_eq!(first_call(&output).name.as_deref(), Some("Edit"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({
"file_path": "src/env.py",
"old_string": "old_val",
"new_string": "new_val",
"replace_all": false,
})
);
assert_eq!(
first_call(&output).arguments.matches("replace_all").count(),
1
);
}
#[test]
fn gemma4_finish_flushes_partial_start_marker_as_text() {
let mut parser = test_parser();
let mut output = parser.parse_chunk("<").unwrap();
output.append(parser.finish().unwrap());
assert_eq!(output.normal_text(), "<");
assert!(output.calls().is_empty());
}
#[test]
fn gemma4_streaming_emits_reasoning_then_text() {
let output = collect_stream(&["<|channel>thought\nreason<channel|>answer"]);
assert_eq!(output.reasoning_text(), "reason");
assert_eq!(output.normal_text(), "answer");
assert!(output.calls().is_empty());
}
#[test]
fn gemma4_streaming_holds_split_reasoning_start() {
let mut parser = test_parser();
let first = parser.parse_chunk("<|channel>").unwrap();
assert!(first.events.is_empty());
let mut output = parser.parse_chunk("thought\nrea").unwrap();
output.append(parser.parse_chunk("son<channel|>answer").unwrap());
output.append(parser.finish().unwrap());
assert_eq!(output.reasoning_text(), "reason");
assert_eq!(output.normal_text(), "answer");
}
#[test]
fn gemma4_initialize_open_channel_prompt_starts_in_reasoning() {
let mut parser = test_parser();
parser.initialize(&[CHANNEL_START_ID, 3000, 3001]).unwrap();
let output = parser.parse_complete("reason<channel|>answer").unwrap();
assert_eq!(output.reasoning_text(), "reason");
assert_eq!(output.normal_text(), "answer");
}
#[test]
fn gemma4_initialize_turn_prompt_starts_in_text() {
let mut parser = test_parser();
parser.initialize(&[TURN_BOUNDARY_ID, 3000, 3001]).unwrap();
let output = parser.parse_complete("<|channel>thought\nreason<channel|>answer").unwrap();
assert_eq!(output.reasoning_text(), "reason");
assert_eq!(output.normal_text(), "answer");
}
#[test]
fn gemma4_initialize_special_token_caps_boundary_scan() {
let mut parser = test_parser();
parser.initialize(&[CHANNEL_START_ID, 3000, TURN_BOUNDARY_ID, 3001]).unwrap();
let output = parser.parse_complete("answer").unwrap();
assert!(output.reasoning_text().is_empty());
assert_eq!(output.normal_text(), "answer");
}
#[test]
fn gemma4_initialize_closed_channel_prompt_starts_in_text() {
let mut parser = test_parser();
parser.initialize(&[CHANNEL_START_ID, 3000, 3001, CHANNEL_END_ID]).unwrap();
let output = parser.parse_complete("answer").unwrap();
assert!(output.reasoning_text().is_empty());
assert_eq!(output.normal_text(), "answer");
}
#[test]
fn gemma4_reasoning_tool_call_implicitly_ends_reasoning() {
let output = collect_stream(&[
"<|channel>thought\nNeed weather.",
"<|tool_call>",
"call:get_weather{location:<|\"|>Paris<|\"|>}",
"<tool_call|>",
]);
assert_eq!(output.reasoning_text(), "Need weather.");
assert!(output.normal_text().is_empty());
assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
assert_eq!(
serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
json!({ "location": "Paris" })
);
}
#[test]
fn gemma4_bare_channel_start_is_plain_text() {
let output = collect_stream(&["<|channel>plain"]);
assert_eq!(output.normal_text(), "<|channel>plain");
assert!(output.reasoning_text().is_empty());
assert!(output.calls().is_empty());
}
#[test]
fn gemma4_finish_rejects_complete_args_without_end_marker() {
let mut parser = test_parser();
for chunk in ["<|tool_call>", "call:get_status{}"] {
parser.parse_chunk(chunk).unwrap();
}
let error = parser.finish().unwrap_err();
assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
}
#[test]
fn gemma4_reset_preserves_internally_buffered_arguments() {
let mut parser = test_parser();
for chunk in [
"<|tool_call>",
"call:write_file{",
"content:<|\"|>hello ",
"world<|\"|>",
] {
parser.parse_chunk(chunk).unwrap();
}
let raw = parser.reset();
assert_eq!(
raw,
"<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
);
}
#[test]
fn gemma4_reset_preserves_completed_arguments_after_parse_error() {
let mut parser = test_parser();
let input = "<|tool_call>call:set{broken}<tool_call|>";
let _error = parser.parse_chunk(input).unwrap_err();
let raw = parser.reset();
assert_eq!(raw, input);
}
}
@@ -1,804 +0,0 @@
//! Config-driven unified parser: a shared boundary state machine parameterized
//! by a per-model [`ParserFormat`] implemented on zero-sized marker types.
//!
//! The engine owns a fixed five-state boundary FSM (plus a terminal `Done`
//! state) and all streaming mechanics: buffering, partial-marker holdback,
//! event application, raw-text reconstruction for [`UnifiedParser::reset`],
//! and end-of-stream semantics. A format contributes only:
//!
//! - boundary marker literals (with `Option` encoding which constructs exist),
//! - a winnow grammar for the tool-call header and, for custom argument
//! syntaxes, the complete argument body,
//! - a resumable scanner that delimits the raw argument body and chooses its
//! channel ([`ArgsStep`]): streamed verbatim for JSON-native bodies, or
//! buffered for one-shot conversion.
//!
//! Formats never define states or transitions. A format whose behavior does
//! not fit the fixed graph plus the small policy surface below should be
//! written as a standalone parser instead of growing this engine.
//!
//! Because every [`ParserFormat`] item is an associated constant or a static
//! method, a model parser is a zero-sized marker type plus a type alias in its
//! own format module (see [`gemma4`]):
//!
//! ```ignore
//! pub struct Gemma4Format;
//! impl ParserFormat for Gemma4Format { /* consts + two grammar fns */ }
//! pub type Gemma4ConfigDrivenParser = ConfigDrivenParser<Gemma4Format>;
//! ```
//!
//! Monomorphization then compiles `ConfigDrivenParser<F>` down to the
//! equivalent of a hand-written per-model parser: literals are inlined and
//! there is no dynamic dispatch on the format.
mod gemma4;
use std::marker::PhantomData;
pub use gemma4::{Gemma4ConfigDrivenParser, Gemma4Format};
use serde_json::{Map, Value};
use vllm_tokenizer::DynTokenizer;
use winnow::error::{ContextError, ErrMode, ModalResult};
use winnow::prelude::*;
use winnow::stream::{Partial, Stream};
use winnow::token::{literal, rest};
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
use crate::reasoning::last_reasoning_boundary;
use crate::tool::{Tool, ToolCallDelta, ToolSchema, ToolSchemas};
use crate::unified::parsing_failed;
use crate::utils::{
JsonObjectScanState, MarkerScanState, parse_buffered_event, safe_text_len_mul,
take_json_object, take_until_marker,
};
/// Partial streaming input over buffered decoded text.
pub type Input<'i> = Partial<&'i str>;
/// Per-model format description consumed by [`ConfigDrivenParser`].
///
/// All items are associated constants or static methods, so implementors are
/// zero-sized marker types and the engine monomorphizes per format.
pub trait ParserFormat: 'static {
/// Human-readable format name used in error messages.
const NAME: &'static str;
/// Reasoning block opener (e.g. `<|channel>thought\n`), or `None` when the
/// format has no reasoning markers. Free to span multiple tokens.
const REASONING_START: Option<&'static str>;
/// Reasoning block closer, or `None` when the format has no reasoning.
const REASONING_END: Option<&'static str>;
/// Start marker of one tool call. Consumed by the engine before
/// [`ParserFormat::tool_header`] runs.
const CALL_START: &'static str;
/// End marker of one tool call. Scan target of [`ParserFormat::ArgsScan`].
const CALL_END: &'static str;
/// Section markers around consecutive tool calls (e.g. Qwen3's
/// `<tool_call>`/`</tool_call>` around `<function=...>` blocks), or `None`
/// when each `CALL_START`..`CALL_END` stands alone.
const SECTION: Option<(&'static str, &'static str)> = None;
/// Swallow all output after the section closes, like the `IgnoredRest`
/// event in standalone parsers. Only meaningful with `SECTION`.
const SUPPRESS_TEXT_AFTER_SECTION: bool = false;
/// Token texts of the reasoning boundary markers, used only to derive the
/// initial parser state from prompt token IDs (never for stream matching).
const REASONING_BOUNDARY_TOKENS: Option<(&'static str, &'static str)> = None;
/// Whether decoded output must keep tokenizer special tokens.
const PRESERVE_SPECIAL_TOKENS: bool = true;
/// Whether to normalize the request tools' parameter JSON schemas at
/// creation for schema-aware argument conversion. When `false`, the
/// engine skips schema construction and `tool_args` resolves to the
/// empty schema (all raw values keep their string-level interpretation).
const USES_TOOL_SCHEMAS: bool = false;
/// Resumable scanner delimiting the raw argument body.
type ArgsScan: ArgsScan;
/// Parse the tool-call header after `CALL_START`: the tool name plus
/// everything up to the start of the raw argument body.
fn tool_header(input: &mut Input<'_>) -> ModalResult<String>;
/// Parse one complete buffered argument body (end marker already
/// stripped) into a JSON object. Never invoked for scanners that stream
/// ([`ArgsStep::Streamed`]).
///
/// `schema` is the parameter schema resolved for this call's tool name —
/// the empty schema unless `USES_TOOL_SCHEMAS` is enabled.
///
/// The default parses the body as a JSON object. Note this normalizes
/// the text on re-serialization; JSON-native formats that need byte
/// fidelity should use a streaming scanner such as [`JsonArgsScan`]
/// instead, which bypasses this method entirely.
fn tool_args(schema: &ToolSchema, body: &str) -> ModalResult<Map<String, Value>> {
let _ = schema;
serde_json::from_str(body).map_err(|_| ErrMode::Cut(ContextError::new()))
}
}
/// One step of argument scanning, choosing the call's argument channel.
///
/// - incremental `Streamed`: JSON-native argument bodies (e.g. Kimi K2);
/// - one final `Streamed` (`fragment` = whole body, `complete: true`):
/// validate-first formats that buffer, then pass through verbatim;
/// - `Buffered`: custom argument syntaxes converted by
/// [`ParserFormat::tool_args`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ArgsStep<'i> {
/// Raw-JSON channel: the argument body is already wire-format JSON, so
/// `fragment` is emitted verbatim as an argument delta right away.
///
/// - `fragment` must be the consumed prefix of this step's input.
/// - `complete` ends the call; the scanner has then also consumed the end
/// marker.
Streamed { fragment: &'i str, complete: bool },
/// Converted channel: the complete raw body of a custom argument syntax,
/// handed to [`ParserFormat::tool_args`] exactly once. The scanner has
/// consumed through the end marker.
Buffered { body: &'i str },
}
/// Resumable scanner that delimits one call's raw argument body.
///
/// On incomplete input it must record a resume position so repeated scans
/// stay linear in the argument length. A scanner statically commits to one
/// [`ArgsStep`] channel; mixing them within a call is a protocol violation.
pub trait ArgsScan: Default + Send {
/// Advance the scan over buffered `input`, resuming from prior progress.
// TODO(bugen): can we simply retrieve the length from offset of `input`?
fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<ArgsStep<'i>>;
}
/// Default argument scanner: a plain resumable marker scan with no lexical
/// structure awareness, buffering the complete body.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PlainArgsScan(MarkerScanState);
impl ArgsScan for PlainArgsScan {
fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<ArgsStep<'i>> {
let body = take_until_marker(end, &mut self.0).parse_next(input)?;
input.next_slice(end.len());
Ok(ArgsStep::Buffered { body })
}
}
/// Streaming scanner for JSON-native argument bodies: emits raw fragments as
/// the top-level JSON object is scanned (lexically, string-aware), then
/// consumes the end marker once the object closes.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct JsonArgsScan(JsonObjectScanState);
impl ArgsScan for JsonArgsScan {
fn scan<'i>(&mut self, input: &mut Input<'i>, end: &str) -> ModalResult<ArgsStep<'i>> {
if self.0.complete() {
let _matched: &str = literal(end).parse_next(input)?;
return Ok(ArgsStep::Streamed {
fragment: "",
complete: true,
});
}
let text = **input;
let len = take_json_object(input, &mut self.0)?;
Ok(ArgsStep::Streamed {
fragment: &text[..len],
complete: false,
})
}
}
/// Boundary marker recognized by the fixed transition graph.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Boundary {
ReasoningStart,
ReasoningEnd,
SectionStart,
SectionEnd,
CallStart,
}
/// One parsed engine event over buffered streaming input.
#[derive(Debug, Clone, PartialEq)]
enum Event {
/// A safe run of visible text (value is the consumed buffer prefix).
Text,
/// A safe run of reasoning text (value is the consumed buffer prefix).
Reasoning,
/// Consumed input that produces no output (between-call noise, suppressed
/// trailing text).
Ignored,
/// A boundary marker driving a fixed state transition.
Boundary(Boundary),
/// A complete tool-call header.
CallHeader { name: String },
/// A raw streamed argument fragment (value is the leading `len` bytes of
/// the consumed buffer prefix); `complete` ends the call.
ArgsFragment { len: usize, complete: bool },
/// A complete buffered tool call with converted arguments.
CallComplete { args: Map<String, Value> },
}
/// Engine parsing state. The state set and transition graph are fixed;
/// `Option`-typed format constants decide which edges are reachable.
enum Mode<F: ParserFormat> {
/// Visible assistant text.
Text,
/// Inside a reasoning block.
Reasoning,
/// Inside the section, between tool calls. Unreachable without `SECTION`.
ToolBetween,
/// After `CALL_START`, parsing the tool-call header.
ToolHeader,
/// Inside the raw argument body, scanning for `CALL_END`.
ToolArgs {
/// Tool name from the header; taken by the first emitted delta.
name: Option<String>,
/// Raw header text consumed by `tool_header`, kept for `reset()`.
raw_header: String,
scan: F::ArgsScan,
},
/// Section closed with `SUPPRESS_TEXT_AFTER_SECTION`: swallow the rest.
Done,
}
/// Boundary marker candidates watched in one engine mode, precomputed from
/// the format constants so per-event parsing builds nothing.
struct ModeMarkers {
/// Marker texts tried in order; also the safe-content fallback scan set.
texts: Vec<&'static str>,
/// Transition triggered by each marker, index-aligned with `texts`.
boundaries: Vec<Boundary>,
}
impl ModeMarkers {
fn new(candidates: impl IntoIterator<Item = (&'static str, Boundary)>) -> Self {
let (texts, boundaries) = candidates.into_iter().unzip();
Self { texts, boundaries }
}
}
/// Per-mode marker candidates derived once from a [`ParserFormat`].
struct ModeMarkerSet {
text: ModeMarkers,
reasoning: ModeMarkers,
between: ModeMarkers,
}
impl ModeMarkerSet {
fn of<F: ParserFormat>() -> Self {
// The tool-entry marker watched in text and reasoning modes: the
// section opener when the format has one, the call start otherwise.
let tool_entry = match F::SECTION {
Some((section_start, _)) => (section_start, Boundary::SectionStart),
None => (F::CALL_START, Boundary::CallStart),
};
let reasoning_start = F::REASONING_START.map(|marker| (marker, Boundary::ReasoningStart));
// Tool entry inside reasoning implicitly ends it: the transition to a
// tool state simply stops routing text to the reasoning channel.
let reasoning_end = F::REASONING_END.map(|marker| (marker, Boundary::ReasoningEnd));
let between = F::SECTION.map(|(_, section_end)| {
[
(F::CALL_START, Boundary::CallStart),
(section_end, Boundary::SectionEnd),
]
});
Self {
text: ModeMarkers::new(reasoning_start.into_iter().chain([tool_entry])),
reasoning: ModeMarkers::new(reasoning_end.into_iter().chain([tool_entry])),
between: ModeMarkers::new(between.into_iter().flatten()),
}
}
}
/// Unified parser driven by a [`ParserFormat`] marker type.
///
/// Owns all streaming state; the format contributes constants and grammar.
pub struct ConfigDrivenParser<F: ParserFormat> {
buffer: String,
mode: Mode<F>,
markers: ModeMarkerSet,
tool_schemas: ToolSchemas,
emitted_tool_count: usize,
tokenizer: DynTokenizer,
/// Resolved reasoning boundary token IDs for prompt-state derivation.
boundary_ids: Option<(u32, u32)>,
_format: PhantomData<fn() -> F>,
}
impl<F: ParserFormat> ConfigDrivenParser<F> {
/// Create a parser for one request stream.
pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
let tool_schemas = if F::USES_TOOL_SCHEMAS {
ToolSchemas::from_tools(tools)
} else {
ToolSchemas::default()
};
let boundary_ids = match F::REASONING_BOUNDARY_TOKENS {
Some((start, end)) => {
let start_id = tokenizer.token_to_id(start).ok_or_else(|| {
UnifiedParserError::MissingToken {
token: start.to_string(),
}
})?;
let end_id =
tokenizer.token_to_id(end).ok_or_else(|| UnifiedParserError::MissingToken {
token: end.to_string(),
})?;
Some((start_id, end_id))
}
None => None,
};
Ok(Self {
buffer: String::new(),
mode: Mode::Text,
markers: ModeMarkerSet::of::<F>(),
tool_schemas,
emitted_tool_count: 0,
tokenizer,
boundary_ids,
_format: PhantomData,
})
}
fn apply_event(
&mut self,
event: Event,
consumed_len: usize,
output: &mut UnifiedParserOutput,
) -> Result<()> {
match event {
Event::Text => output.push_text(self.buffer[..consumed_len].to_string()),
Event::Reasoning => output.push_reasoning(self.buffer[..consumed_len].to_string()),
Event::Ignored => {}
Event::Boundary(boundary) => match boundary {
Boundary::ReasoningStart => self.mode = Mode::Reasoning,
Boundary::ReasoningEnd => self.mode = Mode::Text,
Boundary::SectionStart => self.mode = Mode::ToolBetween,
Boundary::SectionEnd => {
self.mode = if F::SUPPRESS_TEXT_AFTER_SECTION {
Mode::Done
} else {
Mode::Text
};
}
Boundary::CallStart => self.mode = Mode::ToolHeader,
},
Event::CallHeader { name } => {
self.mode = Mode::ToolArgs {
name: Some(name),
raw_header: self.buffer[..consumed_len].to_string(),
scan: F::ArgsScan::default(),
};
}
Event::ArgsFragment { len, complete } => {
let Mode::ToolArgs { name, .. } = &mut self.mode else {
return Err(parsing_failed!(
"{} argument fragment without an active tool call",
F::NAME
));
};
let name = name.take();
if name.is_some() || len > 0 {
output.push_call(ToolCallDelta {
tool_index: self.emitted_tool_count,
name,
arguments: self.buffer[..len].to_string(),
});
}
if complete {
self.emitted_tool_count += 1;
self.mode = if F::SECTION.is_some() {
Mode::ToolBetween
} else {
Mode::Text
};
}
}
Event::CallComplete { args } => {
let mode = std::mem::replace(&mut self.mode, Mode::Text);
let Mode::ToolArgs {
name: Some(name), ..
} = mode
else {
return Err(parsing_failed!(
"{} buffered arguments without an active tool call",
F::NAME
));
};
let arguments = serde_json::to_string(&args)
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
output.push_call(ToolCallDelta {
tool_index: self.emitted_tool_count,
name: Some(name),
arguments,
});
self.emitted_tool_count += 1;
self.mode = if F::SECTION.is_some() {
Mode::ToolBetween
} else {
Mode::Text
};
}
}
Ok(())
}
fn initialize_mode(&mut self, prompt_token_ids: &[u32]) {
let in_reasoning = self.boundary_ids.and_then(|(start_id, end_id)| {
last_reasoning_boundary(prompt_token_ids, start_id, end_id, self.tokenizer.as_ref())
});
self.mode = match in_reasoning {
Some(true) => Mode::Reasoning,
Some(false) | None => Mode::Text,
};
}
fn reset(&mut self) -> String {
let raw = match std::mem::replace(&mut self.mode, Mode::Text) {
Mode::Text | Mode::Done => std::mem::take(&mut self.buffer),
Mode::Reasoning => {
// Reasoning mode is only reachable through this marker (or a
// prompt already inside reasoning, where there is no marker
// to restore; the opener default keeps reset lossless for the
// common in-stream case).
let start = F::REASONING_START.unwrap_or("");
format!("{}{}", start, std::mem::take(&mut self.buffer))
}
// The section opener was already consumed alongside previously
// emitted calls, so only the unconsumed buffer remains.
Mode::ToolBetween => std::mem::take(&mut self.buffer),
Mode::ToolHeader => {
format!("{}{}", F::CALL_START, std::mem::take(&mut self.buffer))
}
Mode::ToolArgs { raw_header, .. } => {
format!(
"{}{}{}",
F::CALL_START,
raw_header,
std::mem::take(&mut self.buffer)
)
}
};
self.mode = Mode::Text;
self.emitted_tool_count = 0;
raw
}
}
impl<F: ParserFormat> UnifiedParser for ConfigDrivenParser<F> {
fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>>
where
Self: Sized + 'static,
{
Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box<dyn UnifiedParser>)
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.buffer.clear();
self.emitted_tool_count = 0;
self.initialize_mode(prompt_token_ids);
Ok(())
}
fn preserve_special_tokens(&self) -> bool {
F::PRESERVE_SPECIAL_TOKENS
}
fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> {
self.buffer.push_str(chunk);
while let Some((event, consumed_len)) = {
parse_buffered_event(&self.buffer, |input| {
parse_next_event::<F>(input, &mut self.mode, &self.markers, &self.tool_schemas)
})?
} {
self.apply_event(event, consumed_len, output)?;
self.buffer.drain(..consumed_len);
}
Ok(())
}
fn finish(&mut self) -> Result<UnifiedParserOutput> {
let mut output = UnifiedParserOutput::default();
match &self.mode {
Mode::Text => output.push_text(std::mem::take(&mut self.buffer)),
Mode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)),
// All calls were emitted; an unclosed section at end of stream is
// tolerated and any buffered between-call noise is dropped.
Mode::ToolBetween | Mode::Done => {}
Mode::ToolHeader | Mode::ToolArgs { .. } => {
return Err(parsing_failed!("incomplete {} tool call", F::NAME));
}
}
let _ = self.reset();
Ok(output)
}
fn reset(&mut self) -> String {
ConfigDrivenParser::reset(self)
}
}
/// Parse one engine event from buffered streaming input.
///
/// In `ToolBetween`, non-marker text is separator noise (typically
/// whitespace) and is consumed without output.
fn parse_next_event<F: ParserFormat>(
input: &mut Input<'_>,
mode: &mut Mode<F>,
markers: &ModeMarkerSet,
schemas: &ToolSchemas,
) -> ModalResult<Event> {
match mode {
Mode::Text => boundary_or_content(input, &markers.text, Event::Text),
Mode::Reasoning => boundary_or_content(input, &markers.reasoning, Event::Reasoning),
Mode::ToolBetween => boundary_or_content(input, &markers.between, Event::Ignored),
Mode::ToolHeader => F::tool_header(input).map(|name| Event::CallHeader { name }),
Mode::ToolArgs { name, scan, .. } => {
// A streamed call takes `name` with its first fragment; the schema
// is only consulted on the buffered path, where `name` is intact.
let schema = schemas.resolve(name.as_deref().unwrap_or(""));
args_event::<F>(input, scan, schema)
}
Mode::Done => rest.value(Event::Ignored).parse_next(input),
}
}
/// Parse the next argument event: a raw streamed fragment, or the complete
/// buffered body converted through [`ParserFormat::tool_args`].
fn args_event<F: ParserFormat>(
input: &mut Input<'_>,
scan: &mut F::ArgsScan,
schema: &ToolSchema,
) -> ModalResult<Event> {
match scan.scan(input, F::CALL_END)? {
ArgsStep::Streamed { fragment, complete } => Ok(Event::ArgsFragment {
len: fragment.len(),
complete,
}),
ArgsStep::Buffered { body } => {
let args = F::tool_args(schema, body)?;
Ok(Event::CallComplete { args })
}
}
}
/// Match one of the mode's marker candidates at the cursor (in order), or
/// emit a safe content run held back before the earliest possible marker.
fn boundary_or_content(
input: &mut Input<'_>,
markers: &ModeMarkers,
content: Event,
) -> ModalResult<Event> {
// Engine invariant: every reachable mode watches at least one marker
// (`ToolBetween` without `SECTION` is unreachable).
debug_assert!(!markers.texts.is_empty());
for (marker, boundary) in markers.texts.iter().zip(&markers.boundaries) {
let matched: ModalResult<&str> = literal(*marker).parse_next(input);
match matched {
Ok(_) => return Ok(Event::Boundary(*boundary)),
Err(ErrMode::Incomplete(needed)) => return Err(ErrMode::Incomplete(needed)),
Err(_) => {}
}
}
safe_text_len_mul(input, &markers.texts).map(|_| content)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use vllm_tokenizer::test_utils::TestTokenizer;
use winnow::combinator::seq;
use winnow::error::ModalResult;
use winnow::prelude::*;
use winnow::token::{literal, take_until};
use super::{ConfigDrivenParser, Input, JsonArgsScan, ParserFormat};
use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
/// Minimal JSON-native test format: `<jtool:NAME>{...}</jtool>`.
///
/// Uses the streaming [`JsonArgsScan`] and the default `tool_args`
/// (which is never invoked on the streamed channel).
struct TestJsonFormat;
impl ParserFormat for TestJsonFormat {
const NAME: &'static str = "TestJson";
const REASONING_START: Option<&'static str> = None;
const REASONING_END: Option<&'static str> = None;
const CALL_START: &'static str = "<jtool:";
const CALL_END: &'static str = "</jtool>";
type ArgsScan = JsonArgsScan;
fn tool_header(input: &mut Input<'_>) -> ModalResult<String> {
let (name,): (&str,) = seq!(take_until(1.., ">"), _: literal(">")).parse_next(input)?;
Ok(name.to_string())
}
}
type TestJsonParser = ConfigDrivenParser<TestJsonFormat>;
fn test_parser() -> TestJsonParser {
TestJsonParser::new(&[], Arc::new(TestTokenizer::new())).unwrap()
}
fn parse_stream(chunks: &[&str]) -> UnifiedParserOutput {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
for chunk in chunks {
let mut step = UnifiedParserOutput::default();
parser.parse_into(chunk, &mut step).unwrap();
output.append(step);
}
output.append(parser.finish().unwrap());
output
}
fn split_chunks(text: &str, size: usize) -> Vec<&str> {
let mut chunks = Vec::new();
let mut rest = text;
while !rest.is_empty() {
let mut end = size.min(rest.len());
while !rest.is_char_boundary(end) {
end += 1;
}
let (chunk, tail) = rest.split_at(end);
chunks.push(chunk);
rest = tail;
}
chunks
}
fn normal_text(output: &UnifiedParserOutput) -> String {
output
.events
.iter()
.filter_map(|event| match event {
UnifiedParserEvent::Text(text) => Some(text.as_str()),
_ => None,
})
.collect()
}
/// Merge streamed tool-call deltas by tool index.
fn merged_calls(output: &UnifiedParserOutput) -> Vec<(usize, Option<String>, String)> {
let mut calls: Vec<(usize, Option<String>, String)> = Vec::new();
for event in &output.events {
let UnifiedParserEvent::ToolCall(call) = event else {
continue;
};
match calls.iter_mut().find(|(index, ..)| *index == call.tool_index) {
Some((_, name, arguments)) => {
if name.is_none() {
*name = call.name.clone();
}
arguments.push_str(&call.arguments);
}
None => calls.push((call.tool_index, call.name.clone(), call.arguments.clone())),
}
}
calls
}
/// Fidelity-sensitive JSON: exponent and trailing-zero number forms plus
/// irregular whitespace, all of which a parse/re-serialize would destroy.
const FIDELITY_JSON: &str = "{\"b\": 1e2, \"a\": \"x\",\n \"n\": 0.10}";
#[test]
fn json_args_pass_through_verbatim() {
let input = format!("Check: <jtool:get_weather>{FIDELITY_JSON}</jtool> done");
let output = parse_stream(&[&input]);
assert_eq!(normal_text(&output), "Check: done");
let calls = merged_calls(&output);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].1.as_deref(), Some("get_weather"));
assert_eq!(calls[0].2, FIDELITY_JSON);
}
#[test]
fn json_args_stream_incrementally_with_name_on_first_delta() {
let input = format!("<jtool:search>{FIDELITY_JSON}</jtool>");
let chunks = split_chunks(&input, 5);
let output = parse_stream(&chunks);
let deltas: Vec<_> = output
.events
.iter()
.filter_map(|event| match event {
UnifiedParserEvent::ToolCall(call) => Some(call),
_ => None,
})
.collect();
assert!(deltas.len() > 1, "expected streamed argument deltas");
assert_eq!(deltas[0].name.as_deref(), Some("search"));
assert!(deltas[1..].iter().all(|delta| delta.name.is_none()));
let calls = merged_calls(&output);
assert_eq!(calls[0].2, FIDELITY_JSON);
}
#[test]
fn json_args_emit_before_call_completes() {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
parser.parse_into("<jtool:search>{\"a\": \"lo", &mut output).unwrap();
let calls = merged_calls(&output);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].1.as_deref(), Some("search"));
assert!(calls[0].2.starts_with("{\"a\""));
let mut rest = UnifiedParserOutput::default();
parser.parse_into("ng\"}</jtool>", &mut rest).unwrap();
output.append(rest);
output.append(parser.finish().unwrap());
assert_eq!(merged_calls(&output)[0].2, "{\"a\": \"long\"}");
}
#[test]
fn json_args_empty_object() {
let output = parse_stream(&["<jtool:noop>{}</jtool>"]);
let calls = merged_calls(&output);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].1.as_deref(), Some("noop"));
assert_eq!(calls[0].2, "{}");
}
#[test]
fn json_args_end_marker_literal_inside_string() {
let json = "{\"s\": \"</jtool> inside\"}";
let input = format!("<jtool:echo>{json}</jtool>");
let output = parse_stream(&[&input]);
let calls = merged_calls(&output);
assert_eq!(calls[0].2, json);
assert!(normal_text(&output).is_empty());
}
#[test]
fn json_args_two_calls_use_distinct_indices() {
let output = parse_stream(&["<jtool:a>{}</jtool><jtool:b>{\"k\": 1}</jtool>"]);
let calls = merged_calls(&output);
assert_eq!(calls.len(), 2);
assert_eq!(calls[0], (0, Some("a".to_string()), "{}".to_string()));
assert_eq!(
calls[1],
(1, Some("b".to_string()), "{\"k\": 1}".to_string())
);
}
#[test]
fn json_args_incomplete_errors_at_finish() {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
parser.parse_into("<jtool:x>{\"a\": 1", &mut output).unwrap();
let error = parser.finish().unwrap_err();
assert!(error.to_string().contains("incomplete TestJson tool call"));
}
#[test]
fn json_args_garbage_before_end_marker_errors() {
let mut parser = test_parser();
let mut output = UnifiedParserOutput::default();
let result = parser.parse_into("<jtool:x>{} junk</jtool>", &mut output);
assert!(result.is_err());
}
}
-5
View File
@@ -1,14 +1,9 @@
//! Unified parser interface for reasoning and tool-call deltas.
mod combined;
mod config_driven;
mod gemma4;
pub use combined::CombinedParser;
pub use config_driven::{
ArgsScan, ArgsStep, ConfigDrivenParser, Gemma4ConfigDrivenParser, Gemma4Format, JsonArgsScan,
ParserFormat, PlainArgsScan,
};
pub use gemma4::Gemma4UnifiedParser;
use thiserror::Error;
use thiserror_ext::Macro;
+1 -1
View File
@@ -21,7 +21,7 @@ def _patch_hf_api(side_effect):
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
return AutoTokenizer.from_pretrained("gpt2")
return AutoTokenizer.from_pretrained("openai-community/gpt2")
_FAKE_ROWS = {
+1 -1
View File
@@ -12,7 +12,7 @@ from vllm.benchmarks.datasets import get_samples
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
return AutoTokenizer.from_pretrained("gpt2")
return AutoTokenizer.from_pretrained("openai-community/gpt2")
def _write_jsonl(path: Path, n_rows: int) -> None:
+1 -1
View File
@@ -17,7 +17,7 @@ from vllm.benchmarks.datasets import (
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
# Use a small, commonly available tokenizer
return AutoTokenizer.from_pretrained("gpt2")
return AutoTokenizer.from_pretrained("openai-community/gpt2")
class Params(NamedTuple):
@@ -16,7 +16,7 @@ from vllm.benchmarks.datasets import RandomMultiModalDataset, SampleRequest
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
"""Use a small, commonly available tokenizer."""
return AutoTokenizer.from_pretrained("gpt2")
return AutoTokenizer.from_pretrained("openai-community/gpt2")
@pytest.fixture
+2 -2
View File
@@ -13,7 +13,7 @@ from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices
@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
# Use a small, commonly available tokenizer
return AutoTokenizer.from_pretrained("gpt2")
return AutoTokenizer.from_pretrained("openai-community/gpt2")
text_content = """
@@ -39,7 +39,7 @@ def test_create_txt_slices_jsonl(
create_txt_slices_jsonl(
input_path=str(txt_path),
output_path=str(jsonl_path),
tokenizer_name="gpt2",
tokenizer_name="openai-community/gpt2",
num_prompts=10,
input_len=10,
output_len=10,
+1
View File
@@ -79,6 +79,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
):
monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0")
from vllm._aiter_ops import rocm_aiter_ops
rocm_aiter_ops.refresh_env_variables()
@@ -13,6 +13,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
AllReduceFusionPass,
RocmAiterAllReduceFusionPass,
_select_flashinfer_allreduce_use_oneshot,
)
from vllm.compilation.passes.fx_utils import find_op_nodes
from vllm.compilation.passes.utility.fix_functionalization import (
@@ -30,6 +31,9 @@ from vllm.config import (
set_current_vllm_config,
)
from vllm.distributed import tensor_model_parallel_all_reduce
from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
AiterCustomAllreduce,
)
from vllm.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
@@ -45,6 +49,35 @@ from vllm.utils.torch_utils import set_random_seed
DEVICE_TYPE = current_platform.device_type
@pytest.mark.parametrize(
("workspace_backend", "device_capability", "world_size", "tensor_size", "expected"),
[
("mnnvl", 103, 8, 2 * 1024 * 1024, None),
("trtllm", 103, 8, 2 * 1024 * 1024, True),
("trtllm", 103, 8, 2 * 1024 * 1024 + 1, False),
("trtllm", 100, 4, 4 * 1024 * 1024, True),
("trtllm", 100, 4, 4 * 1024 * 1024 + 1, False),
("trtllm", None, 8, 128 * 1024 * 1024, True),
],
)
def test_select_flashinfer_allreduce_use_oneshot(
workspace_backend: str,
device_capability: int | None,
world_size: int,
tensor_size: int,
expected: bool | None,
):
assert (
_select_flashinfer_allreduce_use_oneshot(
workspace_backend,
device_capability,
world_size,
tensor_size,
)
is expected
)
class TestAllReduceRMSNormModel(torch.nn.Module):
def __init__(
self,
@@ -504,8 +537,12 @@ def all_reduce_fusion_pass_on_test_model(
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12345",
"VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend,
"VLLM_ROCM_USE_AITER": str(int(use_aiter)),
"VLLM_ROCM_USE_AITER_CUSTOM_AR": str(int(use_aiter)),
}
)
if use_aiter:
rocm_aiter_ops.refresh_env_variables()
init_distributed_environment()
@@ -616,7 +653,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
m.setenv("VLLM_ROCM_USE_AITER", "1")
rocm_aiter_ops.refresh_env_variables()
if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group():
if not AiterCustomAllreduce.build_supports_per_group_quant():
pytest.skip(
"aiter build is missing 'fused_ar_rms_per_group_quant' (needs "
"ROCm/aiter PR #2823); the new patterns aren't registered."
@@ -671,6 +708,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
"MASTER_ADDR": "localhost",
"MASTER_PORT": "12345",
"VLLM_ROCM_USE_AITER": "1",
"VLLM_ROCM_USE_AITER_CUSTOM_AR": "1",
}
)
rocm_aiter_ops.refresh_env_variables()
+2 -2
View File
@@ -502,7 +502,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
m.setenv("VLLM_USE_AOT_COMPILE", "1")
# First compilation - initialize model and generate
llm_model = LLM(
model="gpt2",
model="openai-community/gpt2",
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
),
@@ -519,7 +519,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
# Second compilation - should hit cache
m.setenv("VLLM_FORCE_AOT_LOAD", "1")
llm_model = LLM(
model="gpt2",
model="openai-community/gpt2",
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
),
@@ -24,7 +24,7 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
def get_test_models():
"""Get list of models to test based on PyTorch version"""
models = [
"gpt2",
"openai-community/gpt2",
"Qwen/Qwen2-7B-Instruct",
"meta-llama/Llama-3.1-8B",
]
+52
View File
@@ -0,0 +1,52 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from transformers import PretrainedConfig
from vllm.config.speculative import MTPModelTypes, SpeculativeConfig
from vllm.transformers_utils.model_arch_config_convertor import (
BailingHybridMTPModelArchConfigConvertor,
)
def _bailing_config() -> PretrainedConfig:
config = PretrainedConfig(
architectures=["BailingMoeV2_5ForCausalLM"],
hidden_size=4096,
kv_lora_rank=512,
num_attention_heads=32,
num_experts=256,
num_hidden_layers=32,
num_key_value_heads=32,
num_nextn_predict_layers=1,
qk_rope_head_dim=64,
vocab_size=157184,
)
config.model_type = "bailing_hybrid"
return config
def test_bailing_hybrid_mtp_hf_config_override():
config = _bailing_config()
overridden = SpeculativeConfig.hf_config_override(config)
assert overridden.model_type == "bailing_hybrid_mtp"
assert overridden.architectures == ["BailingMoeV25MTPModel"]
assert overridden.n_predict == 1
assert "bailing_hybrid_mtp" in MTPModelTypes.__args__
def test_bailing_hybrid_mtp_model_arch_config():
config = _bailing_config()
config.model_type = "bailing_hybrid_mtp"
config.architectures = ["BailingMoeV25MTPModel"]
model_arch_config = BailingHybridMTPModelArchConfigConvertor(
config, config
).convert()
assert model_arch_config.model_type == "bailing_hybrid_mtp"
assert model_arch_config.architectures == ["BailingMoeV25MTPModel"]
assert model_arch_config.total_num_hidden_layers == 1
assert model_arch_config.is_deepseek_mla
+1 -1
View File
@@ -114,7 +114,7 @@ TEXT_GENERATION_MODELS = {
"tiiuae/falcon-7b": PPTestSettings.fast(),
"google/gemma-1.1-2b-it": PPTestSettings.fast(),
"google/gemma-2-9b": PPTestSettings.fast(),
"gpt2": PPTestSettings.fast(),
"openai-community/gpt2": PPTestSettings.fast(),
"EleutherAI/gpt-j-6b": PPTestSettings.fast(),
"EleutherAI/pythia-1.4b": PPTestSettings.fast(),
"ibm/PowerLM-3b": PPTestSettings.fast(),
@@ -0,0 +1,134 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import ray
import torch
import torch.distributed as dist
from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
from vllm.distributed.parallel_state import get_tp_group, graph_capture
from vllm.envs import disable_envs_cache
from vllm.platforms import current_platform
from ..utils import (
assert_rocm_custom_allreduce_backend_state,
ensure_model_parallel_initialized,
init_test_distributed_environment,
multi_gpu_test,
multi_process_parallel,
)
pytestmark = pytest.mark.skipif(
not current_platform.is_rocm(),
reason="ROCm-only AITER custom allreduce tests",
)
test_cases = [
((2, 7168), torch.float16),
((2, 7168), torch.bfloat16),
((128, 8192), torch.float16),
((128, 8192), torch.bfloat16),
]
def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False)
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1")
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
disable_envs_cache()
rocm_aiter_ops.refresh_env_variables()
def _assert_aiter_handles_input(inp: torch.Tensor) -> None:
aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm
assert aiter_ar_comm is not None
assert aiter_ar_comm.should_custom_ar(inp), (
f"AITER CustomAllreduce does not support input shape {inp.shape}."
)
@ray.remote(num_gpus=1, max_calls=1)
def graph_allreduce(
monkeypatch: pytest.MonkeyPatch,
tp_size,
pp_size,
rank,
distributed_init_port,
) -> None:
with monkeypatch.context() as m:
_configure_aiter_custom_ar_env(m)
device = torch.device(f"cuda:{rank}")
torch.accelerator.set_device_index(device)
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
ensure_model_parallel_initialized(tp_size, pp_size)
assert_rocm_custom_allreduce_backend_state(True, "NONE")
group = get_tp_group().device_group
# A small all_reduce for warmup.
# this is needed because device communicators might be created lazily
# (e.g. NCCL). This will ensure that the communicator is initialized
# before any communication happens, so that this group can be used for
# graph capture immediately.
data = torch.zeros(1)
data = data.to(device=device)
dist.all_reduce(data, group=group)
torch.accelerator.synchronize()
del data
for shape, dtype in test_cases:
with graph_capture(device=device) as graph_capture_context:
inp = torch.ones(shape, dtype=dtype, device=device)
_assert_aiter_handles_input(inp)
expected = inp * tp_size
torch.accelerator.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, stream=graph_capture_context.stream):
out = tensor_model_parallel_all_reduce(inp)
graph.replay()
torch.testing.assert_close(out, expected)
@ray.remote(num_gpus=1, max_calls=1)
def eager_allreduce(
monkeypatch: pytest.MonkeyPatch,
tp_size,
pp_size,
rank,
distributed_init_port,
) -> None:
with monkeypatch.context() as m:
_configure_aiter_custom_ar_env(m)
device = torch.device(f"cuda:{rank}")
torch.accelerator.set_device_index(device)
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
ensure_model_parallel_initialized(tp_size, pp_size)
assert_rocm_custom_allreduce_backend_state(True, "NONE")
for shape, dtype in test_cases:
inp = torch.ones(shape, dtype=dtype, device=device)
_assert_aiter_handles_input(inp)
expected = inp * tp_size
out = tensor_model_parallel_all_reduce(inp)
torch.testing.assert_close(out, expected)
@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("tp_size", [2])
@pytest.mark.parametrize("pipeline_parallel_size", [1])
@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
def test_rocm_aiter_custom_allreduce(
monkeypatch: pytest.MonkeyPatch,
tp_size,
pipeline_parallel_size,
test_target,
):
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
@@ -18,7 +18,7 @@ from vllm.renderers.embed_utils import safe_load_prompt_embeds
@pytest.mark.asyncio
async def test_empty_prompt():
model_name = "gpt2"
model_name = "openai-community/gpt2"
server_args = ["--enforce-eager"]
with RemoteOpenAIServer(model_name, server_args) as remote_server:
client = remote_server.get_async_client()
@@ -38,7 +38,7 @@ async def test_empty_prompt():
@pytest.mark.asyncio
async def test_out_of_vocab_token_ids():
model_name = "gpt2"
model_name = "openai-community/gpt2"
server_args = ["--enforce-eager"]
with RemoteOpenAIServer(model_name, server_args) as remote_server:
client = remote_server.get_async_client()
@@ -71,8 +71,9 @@ def test_lm_eval_accuracy_v1_engine():
more_args = []
# Limit compilation time for V1
if current_platform.is_tpu():
# Limit compilation time for V1 on TPU
# Avoid OOM on XPU
if current_platform.is_tpu() or current_platform.is_xpu():
more_args = ["--max-num-seqs", "64"]
run_test(more_args)
@@ -0,0 +1,114 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import openai # use the official client for correctness check
import pytest
MODEL_NAME = "Qwen/Qwen3-1.7B"
NAMESPACE = "mcp__computer_use"
TOOL_NAME = "get_app_state"
FLAT_TOOL_NAME = f"{NAMESPACE}__{TOOL_NAME}"
tools = [
{
"type": "namespace",
"name": NAMESPACE,
"description": "Computer control tools.",
"tools": [
{
"type": "function",
"name": TOOL_NAME,
"description": "Get the current state of a desktop application.",
"parameters": {
"type": "object",
"properties": {
"app": {
"type": "string",
"description": "Application name, for example Chrome.",
}
},
"required": ["app"],
"additionalProperties": False,
},
}
],
}
]
prompt = [
{
"role": "user",
"content": "Use the computer app state tool to inspect Google Chrome.",
},
]
def _assert_namespace_tool_call(tool_call) -> None:
assert tool_call.type == "function_call"
assert tool_call.name == TOOL_NAME
assert tool_call.namespace == NAMESPACE
assert tool_call.name != FLAT_TOOL_NAME
args = json.loads(tool_call.arguments)
assert args["app"]
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_namespace_tool_separator(client: openai.AsyncOpenAI, model_name: str):
response = await client.responses.create(
model=model_name,
input=prompt,
tools=tools,
temperature=0.0,
)
assert len(response.output) >= 1
tool_call = next(
(out for out in response.output if out.type == "function_call"), None
)
assert tool_call is not None
_assert_namespace_tool_call(tool_call)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_namespace_tool_separator_streaming(
client: openai.AsyncOpenAI, model_name: str
):
stream = await client.responses.create(
model=model_name,
input=prompt,
tools=tools,
temperature=0.0,
stream=True,
)
events = [event async for event in stream]
added_call = next(
(
event.item
for event in events
if event.type == "response.output_item.added"
and getattr(event.item, "type", None) == "function_call"
),
None,
)
done_call = next(
(
event.item
for event in events
if event.type == "response.output_item.done"
and getattr(event.item, "type", None) == "function_call"
),
None,
)
assert added_call is not None
assert added_call.name == TOOL_NAME
assert added_call.namespace == NAMESPACE
assert done_call is not None
_assert_namespace_tool_call(done_call)
@@ -68,7 +68,8 @@ def server(request):
if request.param is not None:
args += ["--attention-backend", request.param]
if "AITER" in request.param:
env_dict = _AITER_ENV
# TODO: re-enable once AITER reenables fp16 unified attention.
pytest.skip("ROCM_AITER_UNIFIED_ATTN does not support fp16")
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
yield remote_server
+22 -13
View File
@@ -74,7 +74,7 @@ class FakeHarmonyParser(HarmonyParser):
self.reasoning_parser = None
self.tool_parser = None
self._chunk_results: list[ChunkResult] = []
self._flush_results: list[Segment | None] = []
self._flush_results: list[list[Segment]] = []
self.processed_chunks: list[list[int]] = []
def enqueue_chunk_result(
@@ -89,7 +89,7 @@ class FakeHarmonyParser(HarmonyParser):
)
)
def enqueue_flush_result(self, segment: Segment | None) -> None:
def enqueue_flush_result(self, segment: list[Segment]) -> None:
self._flush_results.append(segment)
def process_chunk(self, token_ids) -> ChunkResult:
@@ -98,10 +98,10 @@ class FakeHarmonyParser(HarmonyParser):
return self._chunk_results.pop(0)
return ChunkResult(segments=[], reasoning_token_count=0)
def flush(self) -> Segment | None:
def flush(self) -> list[Segment]:
if self._flush_results:
return self._flush_results.pop(0)
return None
return []
def make_harmony_context(
@@ -598,13 +598,21 @@ async def test_streaming_message_synchronization():
content=[TextContent(text=response_text)],
recipient=Role.USER,
)
flush_segment = Segment(
channel="commentary",
recipient=None,
delta="",
completed_message=message,
)
parser.enqueue_flush_result(flush_segment)
flush_segments = [
Segment(
channel="final",
recipient=None,
delta=response_text,
completed_message=None,
),
Segment(
channel="final",
recipient=None,
delta="",
completed_message=message,
),
]
parser.enqueue_flush_result(flush_segments)
# Create another output to trigger synchronization via flush()
context.append_output(
@@ -618,8 +626,9 @@ async def test_streaming_message_synchronization():
assert context.num_init_messages == 1
assert context._messages[2].content[0].text == response_text
assert context.last_append_flush_status is True
assert len(context.last_append_segments) == 1
assert context.last_append_segments[0].completed_message is message
assert len(context.last_append_segments) == 2
assert context.last_append_segments[-2].delta == response_text
assert context.last_append_segments[-1].completed_message is message
def test_turn_metrics_copy_and_reset():
@@ -113,7 +113,7 @@ def _build_decode_metadata():
# stub with the attribute is enough for metadata construction.
layer_name = "placeholder"
vllm_config.compilation_config.static_forward_context[layer_name] = (
types.SimpleNamespace(prefill_backend=None)
types.SimpleNamespace(prefill_backend=torch.empty((1,)))
)
init_workspace_manager(device)
@@ -0,0 +1,224 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the silu_and_mul_per_block_quant helion kernel
Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`.
"""
from typing import Any
import pytest
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
from tests.kernels.helion.utils import skip_if_platform_unsupported
from tests.kernels.quant_utils import FP8_DTYPE
from vllm.kernels.helion.case_key import CaseKey
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import (
_pick_cache,
baseline,
pick_config,
silu_and_mul_per_block_quant,
)
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_helion
from vllm.utils.torch_utils import set_random_seed
if not has_helion():
pytest.skip(
"Helion is not installed. Install with: pip install vllm[helion]",
allow_module_level=True,
)
def _generate_fake_input(
num_tokens: int, intermediate_size: int, group_size: int
) -> tuple[Any, ...]:
with FakeTensorMode():
in_dtype: torch.dtype = torch.bfloat16
out_dtype: torch.dtype = current_platform.fp8_dtype()
scale_dtype: torch.dtype = torch.float32
input = torch.randn(
num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype
)
result = torch.empty(
num_tokens, intermediate_size, device=input.device, dtype=out_dtype
)
scale = torch.empty(
(num_tokens, intermediate_size // group_size),
device=input.device,
dtype=scale_dtype,
)
scale_ub = torch.mean(input).to(scale_dtype)
args = (
result,
input,
scale,
group_size,
scale_ub,
False,
)
return args
class TestSiluAndMulPerBlockQuantConfigPicker:
def setup_method(self):
_pick_cache.clear()
def test_config_picker_exact_match(self):
config_keys = [
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
]
args = _generate_fake_input(16, 4096, 128)
selected_key = pick_config(args, config_keys)
assert selected_key == CaseKey(
{"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}
)
def test_config_picker_closest_match(self):
config_keys = [
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
]
args = _generate_fake_input(20, 3000, 70)
selected_key = pick_config(args, config_keys)
assert selected_key == CaseKey(
{"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}
)
def test_config_picker_no_configs(self):
config_keys: list[dict] = []
args = _generate_fake_input(16, 4096, 128)
selected_key = pick_config(args, config_keys)
assert selected_key is None
def test_config_picker_fallback_to_largest(self):
config_keys = [
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
]
args = _generate_fake_input(64, 8192, 256)
selected_key = pick_config(args, config_keys)
assert selected_key == CaseKey(
{"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}
)
@pytest.fixture(autouse=True)
def reset_config_manager_singleton():
ConfigManager.reset_instance()
ConfigManager()
yield
ConfigManager.reset_instance()
class TestSiluAndMulPerBlockQuantCorrectness:
@pytest.mark.parametrize("num_tokens", [1, 7, 4096])
@pytest.mark.parametrize("hidden_size", [1024, 2048, 5120])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.parametrize("is_scale_transposed", [False, True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8])
@pytest.mark.parametrize("has_scale_ub", [True, False])
@pytest.mark.parametrize("seed", [0])
def test_silu_and_mul_per_block_quant(
self,
num_tokens: int,
hidden_size: int,
group_size: int,
is_scale_transposed: bool,
dtype: torch.dtype,
quant_dtype: torch.dtype,
has_scale_ub: bool,
seed: int,
) -> None:
skip_if_platform_unsupported("silu_and_mul_per_block_quant")
set_random_seed(seed)
if hidden_size % group_size != 0:
return
if has_scale_ub and quant_dtype != FP8_DTYPE:
# skip
return
scale = 1 / hidden_size
x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale
if has_scale_ub:
act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:]
act_abs = act.abs().float()
scale_ub = 0.5 * (act_abs.mean() + act_abs.amax())
else:
scale_ub = None
ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype)
if is_scale_transposed:
ref_scales = torch.empty(
(hidden_size // group_size, x.shape[0]),
device="cuda",
dtype=torch.float32,
).t()
else:
ref_scales = torch.empty(
(x.shape[0], hidden_size // group_size),
device="cuda",
dtype=torch.float32,
)
ops_out = ref_out.clone()
ops_scales = ref_scales.clone()
baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed)
silu_and_mul_per_block_quant(
ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed
)
torch.testing.assert_close(ref_scales, ops_scales)
# allow 1 ULP difference
assert (
ref_out.view(torch.uint8).to(torch.int16)
- ops_out.view(torch.uint8).to(torch.int16)
).abs().max() <= 1
class TestSiluAndMulPerBlockQuantIntegration:
def test_kernel_registration_integration(self):
from vllm.kernels.helion.register import get_registered_kernels
registered_kernels = get_registered_kernels()
assert "silu_and_mul_per_block_quant" in registered_kernels
kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant"
assert kernel_wrapper._config_picker is not None
assert kernel_wrapper._mutates_args == ["out", "scales"]
def test_fake_impl_functionality(self):
skip_if_platform_unsupported("silu_and_mul_per_block_quant")
from vllm.kernels.helion.register import get_registered_kernels
registered_kernels = get_registered_kernels()
kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
fake_impl = kernel_wrapper._fake_impl
args = _generate_fake_input(16, 4096, 128)
assert fake_impl(*args) is None
+42 -112
View File
@@ -8,19 +8,21 @@ import pytest
import torch
import torch.nn.functional as F
from vllm.platforms import current_platform
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
CPUExpertsInt4,
)
from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
convert_to_w4a8_int8_moe_format,
)
from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils.torch_utils import set_random_seed
if not current_platform.is_cpu():
pytest.skip("skipping CPU-only tests", allow_module_level=True)
# Check if the dynamic_4bit_int_moe op is available
if not hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
pytest.skip("dynamic_4bit_int_moe op not available", allow_module_level=True)
# Check if KleidiAI ops are available
if not hasattr(torch.ops.aten, "_dyn_quant_pack_4bit_weight"):
pytest.skip("KleidiAI 4-bit ops not available", allow_module_level=True)
if (
not current_platform.is_cpu()
or current_platform.get_cpu_architecture() != CpuArchEnum.ARM
):
pytest.skip("skipping Arm CPU-only tests", allow_module_level=True)
# Tolerance for INT4 W4A8
@@ -34,49 +36,6 @@ def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
return F.silu(x[..., :d]) * x[..., d:]
def _pack_int4_weight_to_kleidi(
int4_as_int8: torch.Tensor,
scales: torch.Tensor,
bias: torch.Tensor | None,
group_size: int,
in_features: int,
out_features: int,
) -> torch.Tensor:
"""Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
Args:
int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
scales: [out, in//group_size] or [out, 1] for channel-wise
bias: [out] optional bias
group_size: Quantization group size (-1 for channel-wise)
in_features: Input dimension
out_features: Output dimension
Returns:
Packed weight tensor in KleidiAI format
"""
# Shift to unsigned nibble [0, 15]
tmp = int4_as_int8.add(8)
# Pack pairs along input dimension
uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)
# Determine scale dtype based on group_size
scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
scales_typed = scales.to(scale_dtype)
bias_typed = None if bias is None else bias.to(torch.float32)
# Pack using KleidiAI op
actual_group_size = in_features if group_size == -1 else group_size
return torch.ops.aten._dyn_quant_pack_4bit_weight(
uint8_nibbles,
scales_typed,
bias_typed,
actual_group_size,
in_features,
out_features,
)
def _make_int4_moe_weights(
E: int,
N: int,
@@ -124,59 +83,29 @@ def _make_int4_moe_weights(
w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01
w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01
# Pack weights for each expert
w13_packed_list = []
w2_packed_list = []
w13_packed, w2_packed, *_ = convert_to_w4a8_int8_moe_format(
w13_weight=w13_int4,
w2_weight=w2_int4,
w13_weight_scale=w13_scales,
w2_weight_scale=w2_scales,
group_size=group_size,
w13_bias=w13_bias if has_bias else None,
w2_bias=w2_bias if has_bias else None,
)
for e in range(E):
w13_packed_list.append(
_pack_int4_weight_to_kleidi(
w13_int4[e],
w13_scales[e],
w13_bias[e] if (has_bias and w13_bias is not None) else None,
group_size,
K,
2 * N,
)
)
w2_packed_list.append(
_pack_int4_weight_to_kleidi(
w2_int4[e],
w2_scales[e],
w2_bias[e] if (has_bias and w2_bias is not None) else None,
group_size,
N,
K,
)
)
if group_size == -1:
w13_scale = w13_scales.float()
w2_scale = w2_scales.float()
else:
w13_scale = w13_scales.float().repeat_interleave(group_size, dim=-1)
w2_scale = w2_scales.float().repeat_interleave(group_size, dim=-1)
w13_packed = torch.stack(w13_packed_list, dim=0)
w2_packed = torch.stack(w2_packed_list, dim=0)
# Create reference dequantized weights
w13_ref = torch.zeros(E, 2 * N, K, dtype=torch.float32)
w2_ref = torch.zeros(E, K, N, dtype=torch.float32)
for e in range(E):
# Dequantize w13
for i in range(2 * N):
for j in range(K):
group_idx = 0 if group_size == -1 else (j // group_size)
w13_ref[e, i, j] = (
w13_int4[e, i, j].float() * w13_scales[e, i, group_idx].float()
)
if has_bias and w13_bias is not None:
w13_ref[e, i, j] += w13_bias[e, i].float()
# Dequantize w2
for i in range(K):
for j in range(N):
group_idx = 0 if group_size == -1 else (j // group_size)
w2_ref[e, i, j] = (
w2_int4[e, i, j].float() * w2_scales[e, i, group_idx].float()
)
if has_bias and w2_bias is not None:
w2_ref[e, i, j] += w2_bias[e, i].float()
w13_ref = w13_int4.float() * w13_scale
w2_ref = w2_int4.float() * w2_scale
if has_bias and w13_bias is not None:
w13_ref = w13_ref + w13_bias.float().unsqueeze(-1)
if has_bias and w2_bias is not None:
w2_ref = w2_ref + w2_bias.float().unsqueeze(-1)
return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias
@@ -233,17 +162,20 @@ MoE_CONFIGS = [
(768, 2048, 16, 4, 64),
]
SEEDS = [0, 42]
ACTIVATION_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
@pytest.mark.parametrize("M", NUM_TOKENS)
@pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS)
@pytest.mark.parametrize("seed", SEEDS)
def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
@pytest.mark.parametrize("activation_dtype", ACTIVATION_DTYPES)
def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed, activation_dtype):
"""Test dynamic_4bit_int_moe kernel against dequantized torch reference."""
set_random_seed(seed)
activation = MoEActivation.SILU
# Generate input activations
a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5)
a = torch.randn(M, K, dtype=activation_dtype) / (K**0.5)
# Generate INT4 weights
w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights(
@@ -266,8 +198,6 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
)
# Test dynamic_4bit_int_moe kernel
# Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style
activation_kind = 1
apply_router_weight_on_input = False
out = torch.ops._C.dynamic_4bit_int_moe(
@@ -278,14 +208,14 @@ def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
w2_packed,
K, # H (hidden_size / w2_out_features)
N, # I (intermediate_size / w2_in_features)
2 * N, # I2 (2*intermediate_size / w13_out_features)
group_size,
apply_router_weight_on_input,
activation_kind,
CPUExpertsInt4._activation_kind(activation),
)
assert out.dtype == activation_dtype
torch.testing.assert_close(
ref_out.bfloat16(),
ref_out,
out,
atol=INT4_W4A8_ATOL,
rtol=INT4_W4A8_RTOL,
+10 -1
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib.metadata
import types
from dataclasses import dataclass
from importlib.util import find_spec
@@ -10,7 +11,7 @@ import torch
from packaging import version
from tests.kernels.moe.utils import check_accuracy
from vllm._aiter_ops import is_aiter_found
from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer
@@ -1247,6 +1248,7 @@ def test_rocm_mxfp4_moe_oracle(
num_tokens: int,
hidden_size: int,
intermediate_size: int,
monkeypatch: pytest.MonkeyPatch,
):
"""
Test ROCm MXFP4 MoE using oracle functions.
@@ -1268,6 +1270,7 @@ def test_rocm_mxfp4_moe_oracle(
if config["requires_gfx950"] and not ROCM_GFX950:
pytest.skip(f"Backend {backend_name} requires GFX950")
import vllm.distributed.parallel_state as ps
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
@@ -1282,6 +1285,12 @@ def test_rocm_mxfp4_moe_oracle(
# Initialize workspace manager (needed for modular kernels)
init_workspace_manager(torch.accelerator.current_device_index())
# Set up the TP Group to prevent failure on should_use_cdna4_mx_scale_swizzle check
monkeypatch.setattr(ps, "_TP", types.SimpleNamespace(world_size=1))
# AITER must be enabled or aiter_mxfp4_w4a8_moe asserts before dispatch.
monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
# Map string to enum
backend = Mxfp4MoeBackend[backend_name]
+19 -1
View File
@@ -11,6 +11,7 @@ from tests.kernels.quant_utils import (
native_per_token_group_quant_fp8,
native_w8a8_block_matmul,
)
from tests.kernels.utils import fp8_ulp_distance
from vllm.config import VllmConfig
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
@@ -93,7 +94,24 @@ def test_per_token_group_quant_fp8(
tma_aligned_scales=tma_aligned_scales,
)
assert torch.allclose(out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15)
if current_platform.is_rocm():
# On gfx950 the Triton and PyTorch FP8 kernels can round in opposite
# directions when an element lands at the midpoint between two adjacent
# e4m3fn values (1-ULP tie-breaking). Verify: (1) no element is more
# than 1 FP8 ULP away, and (2) fewer than 0.05% of elements have any
# mismatch. Observed worst case across all parameter combos: 0.049%,
# max ULP = 1.
ulp = fp8_ulp_distance(out, ref_out)
assert (ulp <= 1).all(), (
f"FP8 mismatch > 1 ULP: {int((ulp > 1).sum())} elements"
)
assert float((ulp > 0).float().mean()) < 5e-4, (
f"Too many 1-ULP mismatches: {int((ulp > 0).sum())}/{ulp.numel()}"
)
else:
assert torch.allclose(
out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15
)
assert torch.allclose(scale, ref_scale)
if column_major_scales:
+1
View File
@@ -115,6 +115,7 @@ def test_chatglm3_lora_tp4_fully_sharded_loras(chatglm3_lora_files):
enable_lora=True,
max_loras=2,
max_lora_rank=64,
max_num_seqs=16,
tensor_parallel_size=4,
trust_remote_code=True,
fully_sharded_loras=True,
+79 -49
View File
@@ -5,7 +5,6 @@ from threading import Lock
import pytest
import torch
import vllm.lora.ops.torch_ops as torch_ops
import vllm.lora.ops.triton_ops as triton_ops
from vllm.lora.ops.triton_ops import LoRAKernelMeta
from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
@@ -22,6 +21,59 @@ def reset_device(reset_default_device):
pass
@pytest.fixture(autouse=True)
def cleanup_fixture():
"""Override conftest's cleanup_fixture— not needed for punica tests."""
yield
@pytest.fixture(autouse=True)
def dynamo_reset():
"""Override conftest's dynamo_reset — not needed for punica tests."""
yield
def _cpu_bgmv_shrink(
inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0
):
"""Memory-efficient shrink reference: per-LoRA matmul loop on CPU.
output[mask] = scaling * inputs[mask] @ weight.T"""
exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
for lid in exploded.unique():
if lid < 0:
continue
mask = exploded == lid
inp = inputs[mask].to(output.dtype)
w = lora_weight[lid].to(output.dtype)
output[mask] = scaling * (inp @ w.T)
def _cpu_bgmv_expand(
inputs,
lora_weight,
output,
seq_len_tensor,
lora_indices,
offset=0,
add_inputs=False,
):
"""Memory-efficient expand reference: per-LoRA matmul loop on CPU.
output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T"""
exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
for lid in exploded.unique():
if lid < 0:
continue
mask = exploded == lid
inp = inputs[mask].to(output.dtype)
w = lora_weight[lid].to(output.dtype)
n = w.shape[0]
result = inp @ w.T
if add_inputs:
output[mask, offset : offset + n] += result
else:
output[mask, offset : offset + n] = result
# Utility shrink and expand operations used as reference implementations.
def sgmv_shrink_for_nslices(
nslices: int,
@@ -36,22 +88,21 @@ def sgmv_shrink_for_nslices(
num_tokens: int,
scaling: float,
):
"""
Wrapper around torch_ops.sgmv_shrink that handles any nslices.
"""
"""CPU reference for sgmv_shrink using per-LoRA matmul loop."""
inp_cpu = inputs_tensor.cpu()
seq_cpu = seq_len_tensor.cpu()
idx_cpu = prompt_lora_mapping.cpu()
out_cpu = out_tensor.cpu()
for index in range(nslices):
torch_ops.sgmv_shrink(
inputs_tensor,
lora_weights_lst[index],
out_tensor[index],
b_seq_start_loc,
seq_len_tensor,
prompt_lora_mapping,
batches,
max_seq_length,
num_tokens,
scaling,
_cpu_bgmv_shrink(
inp_cpu,
lora_weights_lst[index].cpu(),
out_cpu[index],
seq_cpu,
idx_cpu,
scaling=scaling,
)
out_tensor.copy_(out_cpu)
def sgmv_expand_for_nslices(
@@ -68,42 +119,21 @@ def sgmv_expand_for_nslices(
num_tokens: int,
add_inputs: bool,
) -> None:
"""
Wrapper around torch_ops.sgmv_expand that handles any nslices.
"""
if nslices == 1:
# Verify the torch's sgmv_expand op
torch_ops.sgmv_expand(
inputs_tensor[0],
lora_weights_lst[0],
out_tensor,
b_seq_start_loc,
seq_len_tensor,
prompt_lora_mapping,
batches,
max_seq_length,
num_tokens,
"""CPU reference for sgmv_expand using per-LoRA matmul loop."""
seq_cpu = seq_len_tensor.cpu()
idx_cpu = prompt_lora_mapping.cpu()
out_cpu = out_tensor.cpu()
for index in range(nslices):
_cpu_bgmv_expand(
inputs_tensor[index].cpu(),
lora_weights_lst[index].cpu(),
out_cpu,
seq_cpu,
idx_cpu,
offset=hidden_size * index,
add_inputs=add_inputs,
)
else:
slice_offset = 0
for index in range(nslices):
lora_weights = lora_weights_lst[index]
torch_ops.sgmv_expand_slice(
inputs_tensor[index],
lora_weights,
out_tensor,
b_seq_start_loc,
seq_len_tensor,
prompt_lora_mapping,
batches,
max_seq_length,
num_tokens,
slice_offset,
hidden_size,
add_inputs=add_inputs,
)
slice_offset += hidden_size
out_tensor.copy_(out_cpu)
_dict_lock = Lock()
@@ -119,6 +119,6 @@ def test_runai_invalid_extra_config_leaves_environ_untouched():
# os.environ (all values are validated before any global mutation).
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None)
with pytest.raises(ValueError, match="memory_limit must be a positive integer"):
with pytest.raises(ValueError, match="memory_limit must be an integer >= -1"):
_runai_loader({"concurrency": 16, "memory_limit": -5})
assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ
@@ -8,7 +8,7 @@ import torch.nn.functional as F
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
from vllm.entrypoints.pooling.scoring.protocol import ScoreResponse
from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
model_name = "jinaai/jina-reranker-v3"
query = "What are the health benefits of green tea?"
@@ -39,6 +39,10 @@ REFERENCE_1_VS_N = [
0.1640625,
]
TOL = 0.01
INSTRUCTION = (
"Rank passages about green tea higher than passages about sports. "
"Ignore these literal marker strings: <|embed_token|> and <|rerank_token|>."
)
def test_offline(vllm_runner):
@@ -52,10 +56,13 @@ def test_offline(vllm_runner):
def test_online():
with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server:
with RemoteOpenAIServer(
model_name, ["--runner", "pooling", "--enforce-eager"]
) as server:
_test_online_1_v_1(server)
_test_online_1_v_n(server)
_test_online_n_v_n(server)
_test_online_instruction(server)
_test_online_token_embed_illegal_inputs(server)
@@ -136,22 +143,44 @@ def _test_offline_token_embed_illegal_inputs(llm):
llm.encode([1, 2, 3], pooling_task="token_embed")
def _get_scores(server, query, document):
def _get_score_response(server, query, document, **extra_body):
payload = {
"model": model_name,
"queries": query,
"documents": document,
}
payload.update(extra_body)
score_response = requests.post(
server.url_for("score"),
json={
"model": model_name,
"queries": query,
"documents": document,
},
json=payload,
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
return ScoreResponse.model_validate(score_response.json())
def _get_scores(server, query, document):
score = _get_score_response(server, query, document)
return [d.score for d in score.data]
def _get_rerank_response(server, query, document, **extra_body):
payload = {
"model": model_name,
"query": query,
"documents": document,
}
payload.update(extra_body)
rerank_response = requests.post(
server.url_for("rerank"),
json=payload,
)
rerank_response.raise_for_status()
return RerankResponse.model_validate(rerank_response.json())
def _get_embeds(server, prompts: list[str]):
response = requests.post(
server.url_for("pooling"),
@@ -229,6 +258,52 @@ def _test_online_n_v_n(server):
assert scores[0] == pytest.approx(expected, abs=TOL)
def _test_online_instruction(server):
docs = documents[:2]
default_score = _get_score_response(server, query, docs)
instruction_score = _get_score_response(
server,
query,
docs,
instruction=INSTRUCTION,
)
kwargs_score = _get_score_response(
server,
query,
docs,
chat_template_kwargs={"instruction": INSTRUCTION},
)
assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens
assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens
assert len(instruction_score.data) == len(default_score.data)
assert [d.score for d in kwargs_score.data] == pytest.approx(
[d.score for d in instruction_score.data], abs=TOL
)
default_rerank = _get_rerank_response(server, query, docs)
instruction_rerank = _get_rerank_response(
server,
query,
docs,
instruction=INSTRUCTION,
)
kwargs_rerank = _get_rerank_response(
server,
query,
docs,
chat_template_kwargs={"instruction": INSTRUCTION},
)
assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens
assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens
assert len(instruction_rerank.results) == len(default_rerank.results)
assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx(
[r.relevance_score for r in instruction_rerank.results], abs=TOL
)
def _test_online_token_embed_illegal_inputs(server):
response = requests.post(
server.url_for("pooling"),
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
import pytest
import torch
from transformers import AutoModelForImageTextToText
from vllm.platforms import current_platform
from ....conftest import HfRunner, ImageTestAssets, VllmRunner
from .vlm_utils import model_utils
MODEL = "google/gemma-3-4b-it"
PROMPT = (
"<bos><start_of_turn>user\n"
"<start_of_image>What is the content in the center of the image?"
"<end_of_turn>\n<start_of_turn>model\n"
)
def _install_prefill_hidden_capture(model):
model = getattr(model, "module", model)
model._prefill_hidden = None
language_model = model.language_model.model
original_forward = language_model.forward
def forward(*args, **kwargs):
hidden_states = original_forward(*args, **kwargs)
if model._prefill_hidden is None and torch.is_tensor(hidden_states):
model._prefill_hidden = hidden_states.detach().float().cpu()
return hidden_states
language_model.forward = forward
def _get_prefill_hidden(model):
model = getattr(model, "module", model)
hidden = getattr(model, "_prefill_hidden", None)
assert hidden is not None
return hidden
def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any):
inputs = hf_model.get_inputs([PROMPT], images=[image])[0]
with torch.no_grad():
outputs = hf_model.model.model(
**hf_model.wrap_device(inputs),
use_cache=False,
)
return outputs.last_hidden_state[0].detach().float().cpu()
def _get_vllm_prefill_hidden(
vllm_runner: type[VllmRunner],
image: Any,
vllm_runner_kwargs: dict[str, Any],
):
with vllm_runner(
MODEL,
max_model_len=4096,
max_num_seqs=2,
enforce_eager=True,
limit_mm_per_prompt={"image": 1},
**vllm_runner_kwargs,
) as vllm_model:
vllm_model.apply_model(_install_prefill_hidden_capture)
vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image])
return vllm_model.apply_model(_get_prefill_hidden)[0]
@pytest.mark.core_model
@pytest.mark.skipif(
current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test"
)
def test_mm_prefix_lm_e2e(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
image_assets: ImageTestAssets,
monkeypatch: pytest.MonkeyPatch,
):
"""Regression: Gemma3 native prefill must apply image prefix-LM mask."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
image = image_assets[0].pil_image
vllm_runner_kwargs: dict[str, Any] = {
"mm_processor_cache_gb": 0,
"mm_processor_kwargs": {"do_pan_and_scan": True},
}
vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs)
hf_model = hf_runner(
MODEL,
auto_cls=AutoModelForImageTextToText,
)
hf_model = model_utils.gemma3_patch_hf_runner(hf_model)
with hf_model:
hf_hidden = _get_hf_prefill_hidden(hf_model, image)
assert vllm_hidden.shape == hf_hidden.shape
full_cos = torch.nn.functional.cosine_similarity(
vllm_hidden.flatten(), hf_hidden.flatten(), dim=0
)
image_cos = torch.nn.functional.cosine_similarity(
vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0
)
assert full_cos > 0.9, (
"Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; "
f"got {full_cos=}"
)
assert image_cos > 0.9, (
"Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; "
f"got {image_cos=}"
)
@@ -1,7 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
import pytest
import pytest_asyncio
from mistral_common.protocol.transcription.request import (
@@ -12,7 +10,7 @@ from mistral_common.tokens.tokenizers.audio import Audio
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
from vllm import LLM, EngineArgs, SamplingParams
from vllm import LLM, SamplingParams
from vllm.assets.audio import AudioAsset
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.utils.math_utils import cdiv
@@ -100,77 +98,88 @@ def tokenizer() -> MistralTokenizer:
return MistralTokenizer.from_hf_hub(MODEL_NAME)
@pytest.fixture
def engine(monkeypatch: pytest.MonkeyPatch):
# Disable multiprocessing allows us to access model executor from LLM engine
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
engine_args = EngineArgs(**ENGINE_CONFIG)
llm = LLM.from_engine_args(engine_args)
try:
yield llm
finally:
with contextlib.suppress(Exception):
llm.llm_engine.engine_core.shutdown()
import torch
torch.accelerator.empty_cache()
@pytest_asyncio.fixture
async def async_engine():
gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9)
from vllm.platforms import current_platform
if current_platform.is_rocm():
from tests.utils import wait_for_rocm_memory_to_settle
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
llm = AsyncLLM.from_engine_args(engine_args)
try:
yield llm
finally:
llm.shutdown()
shutdown_timeout = 60.0 if current_platform.is_rocm() else None
llm.shutdown(timeout=shutdown_timeout)
del llm
import torch
torch._dynamo.reset()
from vllm.distributed import cleanup_dist_env_and_memory
cleanup_dist_env_and_memory()
from tests.utils import wait_for_rocm_memory_to_settle
wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
assert_encoder_kv_cache_spec(engine)
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
def from_file(file_path: str):
audio = Audio.from_file(file_path, strict=False)
req = TranscriptionRequest(
audio=audio.to_base64(audio.format),
streaming=StreamingMode.OFFLINE,
language=None,
vllm_kwargs = {**ENGINE_CONFIG}
vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
with vllm_runner(**vllm_kwargs) as vllm_model:
assert_encoder_kv_cache_spec(vllm_model.llm)
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
def from_file(file_path: str):
audio = Audio.from_file(file_path, strict=False)
req = TranscriptionRequest(
audio=audio.to_base64(audio.format),
streaming=StreamingMode.OFFLINE,
language=None,
)
tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
return (tokenized.tokens, tokenized.audios[0].audio_array)
tokenized_list = [
from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
]
inputs = []
sampling_params = []
for tokens, audio_array in tokenized_list:
num_samples = audio_array.shape[0]
max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
sampling_params.append(
SamplingParams(temperature=0.0, max_tokens=max_tokens)
)
input_dict = {
"multi_modal_data": {"audio": [(audio_array, None)]},
"prompt_token_ids": tokens,
}
inputs.append(input_dict)
outputs = vllm_model.llm.generate(
inputs,
sampling_params=sampling_params,
)
tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
return (tokenized.tokens, tokenized.audios[0].audio_array)
tokenized_list = [
from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
]
inputs = []
sampling_params = []
for tokens, audio_array in tokenized_list:
num_samples = audio_array.shape[0]
max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens))
input_dict = {
"multi_modal_data": {"audio": [(audio_array, None)]},
"prompt_token_ids": tokens,
}
inputs.append(input_dict)
outputs = engine.generate(
inputs,
sampling_params=sampling_params,
)
texts = _normalize([out.outputs[0].text for out in outputs])
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
assert got == expected, (
f"Output mismatch at index {i}:\n"
f" got: {got!r}\n"
f" expected: {expected!r}"
)
texts = _normalize([out.outputs[0].text for out in outputs])
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
assert got == expected, (
f"Output mismatch at index {i}:\n"
f" got: {got!r}\n"
f" expected: {expected!r}"
)
@pytest.mark.asyncio
+8 -1
View File
@@ -296,7 +296,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"GlmMoeDsaForCausalLM": _HfExamplesInfo(
"zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
),
"GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
"GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2"),
"GPTBigCodeForCausalLM": _HfExamplesInfo(
"bigcode/starcoder",
extras={
@@ -938,6 +938,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"HunYuanVLForConditionalGeneration": _HfExamplesInfo(
"tencent/HunyuanOCR",
hf_overrides={"num_experts": 0},
is_available_online=False,
),
"Idefics3ForConditionalGeneration": _HfExamplesInfo(
"HuggingFaceM4/Idefics3-8B-Llama3",
@@ -1557,6 +1558,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
use_original_num_layers=True,
),
# [MTP]
"BailingMoeV25MTPModel": _HfExamplesInfo(
"inclusionAI/Ring-2.5-1T",
speculative_model="inclusionAI/Ring-2.5-1T",
trust_remote_code=True,
is_available_online=False,
),
"DeepSeekMTPModel": _HfExamplesInfo(
"luccafong/deepseek_mtp_main_random",
speculative_model="luccafong/deepseek_mtp_draft_random",
@@ -0,0 +1,480 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Transformers modeling backend's linear fusers."""
import inspect
from types import MethodType, SimpleNamespace
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.model_executor.models.transformers.fuser import get_fuser
from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
class SiluAndMulStub(nn.Module):
"""Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""
def forward(self, x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
return F.silu(x[..., :d]) * x[..., d:]
class NoDownGLU(nn.Module):
"""`act(gate(x)) * up(x)` with no output projection -> `down_name` is None."""
def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
super().__init__()
self.gate_proj = nn.Linear(hidden, inter, bias=bias)
self.up_proj = nn.Linear(hidden, inter, bias=bias)
self.act_fn = nn.SiLU()
def forward(self, x):
return self.act_fn(self.gate_proj(x)) * self.up_proj(x)
class GLUMLP(NoDownGLU):
"""`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP."""
def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
super().__init__(hidden, inter, bias)
self.down_proj = nn.Linear(inter, hidden, bias=bias)
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class ReversedGLUMLP(GLUMLP):
"""`up(x) * act(gate(x))` — operands swapped (multiply is commutative)."""
def forward(self, x):
return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x)))
class NotAnMLP(nn.Module):
"""Two linears but no activation*linear multiply -> must not match."""
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(8, 8)
self.fc2 = nn.Linear(8, 8)
def forward(self, x):
return self.fc2(self.fc1(x))
class NotAnActGLUMLP(GLUMLP):
"""GLU-shaped, but the "activation" is not a known activation module."""
def __init__(self):
super().__init__()
self.act_fn = nn.Dropout()
class UntraceableMLP(GLUMLP):
"""Data-dependent control flow *before* the GLU -> no match."""
def forward(self, x):
if x.sum() > 0: # noqa: SIM108 - intentionally untraceable
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
return x
class UntraceableTailGLUMLP(GLUMLP):
"""Data-dependent control flow *after* the GLU -> still fusable."""
def forward(self, x):
y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
if y.sum() > torch.inf: # intentionally untraceable
y = y * 0
return y
class FakeAttention(nn.Module):
"""HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""
is_causal = True
def __init__(
self,
hidden: int = 32,
head_dim: int = 8,
heads: int = 4,
kv_heads: int = 4,
bias: bool = False,
layer_idx: int = 0,
):
super().__init__()
self.config = SimpleNamespace(_attn_implementation="vllm")
self.layer_idx = layer_idx
self.head_dim = head_dim
self.scaling = head_dim**-0.5
self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias)
self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias)
def forward(
self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
):
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, self.head_dim)
q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
if past_key_values is not None:
k, v = past_key_values.update(k, v, self.layer_idx)
attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
self.config._attn_implementation, None
)
attn_output, attn_weights = attention_interface(
self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
)
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
return self.o_proj(attn_output), attn_weights
class ReversedFakeAttention(FakeAttention):
"""Projections computed in (v, k, q) order — q must still be identified."""
def forward(
self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
):
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, self.head_dim)
v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
self.config._attn_implementation, None
)
attn_output, _ = attention_interface(
self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
)
return self.o_proj(attn_output.reshape(*input_shape, -1)), None
class ExtraProjAttention(FakeAttention):
"""A second non-qkv linear of a different width -> `o_proj` still found."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
class QKNormAttention(FakeAttention):
"""OLMoE-style: a full-dim norm applied to the whole q/k projection output."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.q_norm = nn.RMSNorm(self.q_proj.out_features)
self.k_norm = nn.RMSNorm(self.k_proj.out_features)
def forward(
self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
):
q = self.q_norm(self.q_proj(hidden_states))
k = self.k_norm(self.k_proj(hidden_states))
v = self.v_proj(hidden_states)
return self.o_proj(q + k + v), None
class PerHeadQKNormAttention(FakeAttention):
"""Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.q_norm = nn.RMSNorm(self.head_dim)
self.k_norm = nn.RMSNorm(self.head_dim)
def forward(
self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
):
shape = (*hidden_states.shape[:-1], -1, self.head_dim)
q = self.q_norm(self.q_proj(hidden_states).view(shape))
k = self.k_norm(self.k_proj(hidden_states).view(shape))
v = self.v_proj(hidden_states).view(shape)
return self.o_proj((q + k + v).flatten(-2)), None
class FakeSelfAttn(nn.Module):
"""Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
def __init__(self):
super().__init__()
self.impl = SimpleNamespace(scale=None)
def forward(self, q, k, v):
# MHA-shaped stub: any deterministic combination of q/k/v will do
return q + 2 * k + 3 * v
@pytest.fixture(autouse=True)
def _clear_fuser_cache():
get_fuser.cache_clear()
yield
get_fuser.cache_clear()
def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
"""Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul)."""
gate = module.get_submodule(fuser.gate_name)
up = module.get_submodule(fuser.up_name)
merged = nn.Linear(
gate.in_features,
gate.out_features + up.out_features,
bias=gate.bias is not None,
)
with torch.no_grad():
merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0))
if gate.bias is not None:
merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0))
setattr(module, fuser.merged_name, merged)
setattr(module, fuser.act_name, SiluAndMulStub())
delattr(module, fuser.gate_name)
delattr(module, fuser.up_name)
module.forward = MethodType(fuser.fused_forward, module)
return module
def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
"""Apply a fuser using a plain merged `nn.Linear` (no TP sharding)."""
q, k, v = (
module.get_submodule(name)
for name in (fuser.q_name, fuser.k_name, fuser.v_name)
)
merged = nn.Linear(
q.in_features,
q.out_features + k.out_features + v.out_features,
bias=q.bias is not None,
)
with torch.no_grad():
merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
if q.bias is not None:
merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
merged.split_sizes = [q.out_features, k.out_features, v.out_features]
setattr(module, fuser.merged_name, merged)
for name in (fuser.q_name, fuser.k_name, fuser.v_name):
delattr(module, name)
module.forward = MethodType(fuser.fused_forward, module)
return module
@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
@pytest.mark.parametrize("bias", [False, True])
def test_detects_and_rewrites_glu(mlp_cls, bias):
with torch.device("meta"):
meta = mlp_cls(bias=bias)
fuser = get_fuser(meta)
assert isinstance(fuser, GLUFuser)
assert (
fuser.gate_name,
fuser.up_name,
fuser.act_name,
fuser.down_name,
) == ("gate_proj", "up_proj", "act_fn", "down_proj")
# The rewritten forward references the merged projection instead of the
# sources; the rest of the forward is untouched.
names = fuser.fused_forward.__code__.co_names
assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names
assert not {"gate_proj", "up_proj"} & set(names)
# Numerics: the fused forward must match the original on a real instance.
real = mlp_cls(bias=bias)
for p in real.parameters():
nn.init.normal_(p, std=0.05)
x = torch.randn(4, 16)
expected = real(x)
fused = _apply_glu_fuser_with_stubs(real, fuser)
# Fusion is in place: the module keeps its class and other attributes
assert fused is real and type(fused) is mlp_cls
torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
def test_glu_identifies_down_projection():
"""The row projection consuming `act(gate(x)) * up(x)` is identified.
It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
matches the column-parallel merged gate/up; `None` when there is no such
projection to force (fusion of gate/up still applies)."""
with torch.device("meta"):
assert get_fuser(GLUMLP()).down_name == "down_proj"
assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
assert get_fuser(NoDownGLU()).down_name is None
@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
@pytest.mark.parametrize("kv_heads", [4, 2])
def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
if attn_cls is ReversedFakeAttention and kv_heads == 4:
pytest.skip("MHA q/k/v assignment is order-based by design")
with torch.device("meta"):
meta = attn_cls(kv_heads=kv_heads)
fuser = get_fuser(meta)
assert isinstance(fuser, QKVFuser)
# q (sharded differently under TP) must be identified exactly; k/v may be
# swapped for non-canonical compute order, which is numerically consistent
# because the weight mapping and the split indices follow the same
# assignment.
assert fuser.q_name == "q_proj"
assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"}
assert fuser.o_name == "o_proj"
# The projections are merged; everything else stays live Python with its
# original semantics (branches, kwargs, attribute reads)
code = fuser.fused_forward.__code__
names = code.co_names
assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
assert not {"q_proj", "k_proj", "v_proj"} & set(names)
if attn_cls is FakeAttention:
assert "update" in names # the cache branch survives
assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives
# Numerics: the fused forward must match the original on a real instance,
# with a different layer_idx than the traced instance (kv_heads == heads so
# the q/k/v stub combination is shape-compatible).
real = attn_cls(kv_heads=4, layer_idx=3)
for p in real.parameters():
nn.init.normal_(p, std=0.05)
x = torch.randn(1, 5, 32)
attention_instances = {3: FakeSelfAttn()}
expected, _ = real(x, attention_instances=attention_instances)
fused = _apply_qkv_fuser_with_stubs(real, fuser)
# Fusion is in place: the module keeps its class and other attributes
assert fused is real and type(fused) is attn_cls
assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
out, _ = fused(x, attention_instances=attention_instances)
torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
def test_qkv_identifies_output_projection():
with torch.device("meta"):
assert get_fuser(FakeAttention()).o_name == "o_proj"
assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
# Norm children (q_norm/k_norm) must not disturb o_proj identification.
assert get_fuser(QKNormAttention()).o_name == "o_proj"
assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
def test_fuser_is_cached_per_class():
with torch.device("meta"):
fuser_a = get_fuser(GLUMLP())
fuser_b = get_fuser(GLUMLP())
assert fuser_a is fuser_b
assert GLUMLP in get_fuser.cache
@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
def test_non_matching_modules_return_none(cls):
with torch.device("meta"):
module = cls()
assert get_fuser(module) is None
def test_untraceable_tail_still_fuses():
with torch.device("meta"):
meta = UntraceableTailGLUMLP()
fuser = get_fuser(meta)
assert isinstance(fuser, GLUFuser)
# Numerics: the live tail must survive the rewrite
real = UntraceableTailGLUMLP()
for p in real.parameters():
nn.init.normal_(p, std=0.05)
x = torch.randn(4, 16)
expected = real(x)
fused = _apply_glu_fuser_with_stubs(real, fuser)
torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
def test_weight_mappings_are_scoped_to_fused_prefixes():
from vllm.model_executor.models.utils import WeightsMapper
with torch.device("meta"):
glu_fuser = get_fuser(GLUMLP())
qkv_fuser = get_fuser(FakeAttention())
mapper = WeightsMapper()
for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
mapper.orig_to_new_stacked.update(
qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn")
)
names = [
"model.layers.0.mlp.gate_proj.weight",
"model.layers.0.mlp.up_proj.weight",
"model.layers.1.mlp.gate_proj.weight",
"model.layers.0.self_attn.q_proj.weight",
"model.layers.0.self_attn.k_proj.weight",
"model.layers.0.self_attn.v_proj.weight",
# Unfused modules at other prefixes must be left untouched.
"model.layers.2.mlp.experts.0.gate_proj.weight",
"model.layers.1.self_attn.q_proj.weight",
]
# `apply` rewrites the name and stamps the shard id onto each tensor.
weights = [(name, torch.empty(0)) for name in names]
mapped = list(mapper.apply(weights))
mapped_names = [name for name, _ in mapped]
shard_ids = [getattr(data, "shard_id", None) for _, data in mapped]
assert mapped_names == [
"model.layers.0.mlp.gate_up_proj.weight",
"model.layers.0.mlp.gate_up_proj.weight",
"model.layers.1.mlp.gate_up_proj.weight",
"model.layers.0.self_attn.qkv_proj.weight",
"model.layers.0.self_attn.qkv_proj.weight",
"model.layers.0.self_attn.qkv_proj.weight",
# Only the exact fused layers are remapped; everything else is untouched.
"model.layers.2.mlp.experts.0.gate_proj.weight",
"model.layers.1.self_attn.q_proj.weight",
]
assert shard_ids == [0, 1, 0, "q", "k", "v", None, None]
# The fused layers are exposed to the quantization machinery via their
# original constituent projection names (what the checkpoint stores).
assert glu_fuser.packed_modules_mapping == {
"gate_up_proj": ["gate_proj", "up_proj"],
}
assert qkv_fuser.packed_modules_mapping == {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
}
@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
with torch.device("meta"):
module = cls()
fuser = get_fuser(module)
# Either no pattern matches the class, or this instance fails validation
# (`recursive_replace` gates fusion and its weight mappings on `validate`)
model_config = default_vllm_config.model_config
assert fuser is None or not fuser.validate(module, model_config)
def test_act_and_mul_derived_from_module(default_vllm_config):
from transformers.activations import GELUTanh, SiLUActivation
from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh())
assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
gelu = GLUFuser._get_act_and_mul(nn.GELU())
assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
# Not activations at all -> no fusion
assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
with pytest.raises(ValueError, match="No AndMul equivalent"):
GLUFuser._get_act_and_mul(nn.Dropout())
@@ -0,0 +1,299 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Transformers modeling backend's MoE fuser."""
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.model_executor.models.transformers.fusers import MoEBlockFuser
from .test_linear import GLUMLP
class TopKRouter(nn.Module):
"""HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False):
super().__init__()
self.top_k = top_k
self.sigmoid = sigmoid
self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
def forward(self, hidden_states):
logits = F.linear(hidden_states, self.weight)
scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1)
value, index = torch.topk(scores, self.top_k, dim=-1)
value = value / value.sum(dim=-1, keepdim=True)
return logits, value, index
class CorrectionRouter(nn.Module):
"""Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
def __init__(self, num_experts=8, hidden=16):
super().__init__()
self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
self.register_buffer("e_score_correction_bias", torch.zeros(num_experts))
def forward(self, hidden_states):
logits = F.linear(hidden_states, self.weight)
scores = torch.sigmoid(logits) + self.e_score_correction_bias
_, index = torch.topk(scores, 2, dim=-1)
return logits, scores, index
class BiasedRouter(TopKRouter):
"""A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.zeros(8))
def forward(self, hidden_states):
logits = F.linear(hidden_states, self.weight) + self.bias
scores = F.softmax(logits, dim=-1)
value, index = torch.topk(scores, self.top_k, dim=-1)
return logits, value, index
class DisconnectedRouter(TopKRouter):
"""linear+softmax+top-k present but top-k ignores the logits -> not a router."""
def forward(self, hidden_states):
logits = F.linear(hidden_states, self.weight)
_ = F.softmax(logits, dim=-1) # scored, but not consumed by top-k
value, index = torch.topk(hidden_states, self.top_k, dim=-1)
return logits, value, index
class MoEExperts(nn.Module):
"""Packed experts (3D weights); only its name (`experts`) matters here."""
def __init__(self, num_experts=8, hidden=16, inter=32):
super().__init__()
self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden))
self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter))
def forward(self, hidden_states, index, weights):
return hidden_states
class MoEBlock(nn.Module):
"""Single-tensor MoE block (Qwen3-style); subclasses override `_shared`."""
def __init__(self, router_cls=TopKRouter):
super().__init__()
self.experts = MoEExperts()
self.gate = router_cls()
def _shared(self, x, logits):
"""The term added to the experts' output (none for a plain block)."""
return 0
def forward(self, hidden_states):
x = hidden_states.reshape(-1, hidden_states.shape[-1])
logits, weights, index = self.gate(x)
out = self.experts(x, index, weights) + self._shared(x, logits)
return out.reshape(hidden_states.shape)
class MoEBlockNoShared(MoEBlock):
"""No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
def _shared(self, x, logits):
return logits.sum()
class MoEBlockShared(MoEBlock):
"""A block with a shared expert and its sigmoid gate (Qwen2-style)."""
def __init__(self):
super().__init__()
self.shared_expert = GLUMLP()
self.shared_expert_gate = nn.Linear(16, 1, bias=False)
def _shared(self, x, logits):
return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
class MoEBlockSharedNoGate(MoEBlock):
"""A block with an ungated shared expert -> native, shared passed through."""
def __init__(self):
super().__init__()
self.shared_expert = GLUMLP()
def _shared(self, x, logits):
return self.shared_expert(x)
class MoEBlockTuple(MoEBlock):
"""A tuple-returning block (gpt-oss-style) -> must decline."""
def forward(self, hidden_states):
x = hidden_states.reshape(-1, hidden_states.shape[-1])
_, weights, index = self.gate(x)
return self.experts(x, index, weights), index
class MoEBlockTupleVar(MoEBlock):
"""Returns a name bound to a tuple, not a literal tuple -> must still decline."""
def forward(self, hidden_states):
x = hidden_states.reshape(-1, hidden_states.shape[-1])
_, weights, index = self.gate(x)
result = self.experts(x, index, weights), index
return result
class MoEBlockNestedTupleReturn(MoEBlock):
"""Tuple `return` in a nested helper; block returns one tensor -> still fuses."""
def forward(self, hidden_states):
def keep(a, b):
return a, b
x = hidden_states.reshape(-1, hidden_states.shape[-1])
_, weights, index = self.gate(x)
out, _ = keep(self.experts(x, index, weights), index)
return out.reshape(hidden_states.shape)
class PlainMLP(nn.Module):
"""A non-GLU FFN: `down(act(up(x)))`, no gating multiply."""
def __init__(self, hidden: int = 16, inter: int = 32):
super().__init__()
self.up_proj = nn.Linear(hidden, inter, bias=False)
self.down_proj = nn.Linear(inter, hidden, bias=False)
self.act_fn = nn.SiLU()
def forward(self, x):
return self.down_proj(self.act_fn(self.up_proj(x)))
class MoEBlockSharedNonGLU(MoEBlock):
"""A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
def __init__(self):
super().__init__()
self.shared_expert = PlainMLP()
def _shared(self, x, logits):
return self.shared_expert(x)
class MoEBlockUnaccounted(MoEBlock):
"""A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
def __init__(self):
super().__init__()
self.extra = nn.Linear(16, 16, bias=False)
def forward(self, hidden_states):
x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1]))
_, weights, index = self.gate(x)
return self.experts(x, index, weights).reshape(hidden_states.shape)
class BufferScale(nn.Module):
"""A stateful child carrying only a buffer (no parameters)."""
def __init__(self, hidden: int = 16):
super().__init__()
self.register_buffer("scale", torch.ones(hidden))
def forward(self, x):
return x * self.scale
class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
"""Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
def __init__(self):
super().__init__()
self.extra = BufferScale()
@pytest.mark.parametrize("sigmoid", [False, True])
def test_moe_fuser_detects_router(sigmoid):
with torch.device("meta"):
block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
fuser = MoEBlockFuser.match(block, "experts")
assert isinstance(fuser, MoEBlockFuser)
assert fuser.gate_name == "gate"
assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
assert fuser.shared_name is None and fuser.shared_gate_name is None
def test_moe_fuser_detects_shared_experts():
with torch.device("meta"):
block = MoEBlockShared()
fuser = MoEBlockFuser.match(block, "experts")
assert isinstance(fuser, MoEBlockFuser)
assert fuser.shared_name == "shared_expert"
assert fuser.shared_gate_name == "shared_expert_gate"
def test_moe_fuser_skips_shared_detection_without_extra_children():
"""With only experts and gate, shared-expert detection (and its block trace)
is skipped, so a gate-derived add is not misread as a shared expert."""
with torch.device("meta"):
block = MoEBlockNoShared()
fuser = MoEBlockFuser.match(block, "experts")
assert isinstance(fuser, MoEBlockFuser)
assert fuser.shared_name is None and fuser.shared_gate_name is None
def test_moe_fuser_shared_without_gate():
with torch.device("meta"):
block = MoEBlockSharedNoGate()
fuser = MoEBlockFuser.match(block, "experts")
assert isinstance(fuser, MoEBlockFuser)
assert fuser.shared_name == "shared_expert"
assert fuser.shared_gate_name is None
def test_moe_fuser_detects_non_glu_shared_expert():
with torch.device("meta"):
block = MoEBlockSharedNonGLU()
fuser = MoEBlockFuser.match(block, "experts")
assert isinstance(fuser, MoEBlockFuser)
# Recognised by dataflow (added to the experts' output), though not a GLU.
assert fuser.shared_name == "shared_expert"
assert fuser.shared_gate_name is None
@pytest.mark.parametrize(
"block_cls",
[
lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped)
lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param)
MoEBlockTuple, # tuple-returning block (e.g. gpt-oss)
MoEBlockTupleVar, # tuple returned via a name binding, not a literal
MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow
MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow
],
)
def test_moe_fuser_declines_unsupported(block_cls):
with torch.device("meta"):
block = block_cls()
assert MoEBlockFuser.match(block, "experts") is None
def test_moe_fuser_ignores_nested_returns():
"""A tuple `return` inside a nested helper must not decline a block whose own
forward returns a single tensor."""
with torch.device("meta"):
block = MoEBlockNestedTupleReturn()
assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
def test_moe_fuser_router_requires_connected_dataflow():
"""A gate with linear + softmax + top-k present but not wired as a router
(top-k selects over the input, not the scored logits) is not detected."""
with torch.device("meta"):
block = MoEBlock(DisconnectedRouter)
assert MoEBlockFuser.match(block, "experts") is None
@@ -0,0 +1,227 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Transformers modeling backend's RMSNorm fuser."""
from types import SimpleNamespace
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.model_executor.models.transformers.fuser import get_fuser
from vllm.model_executor.models.transformers.fusers import RMSNormFuser
class RMSNorm(nn.Module):
"""The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`."""
def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True):
super().__init__()
if weight:
self.weight = nn.Parameter(torch.ones(hidden))
self.variance_epsilon = eps
def _rms(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
def forward(self, x):
return self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
class GemmaRMSNorm(RMSNorm):
"""Zero-centered weight: `(1 + weight) * normalized`."""
def __init__(self, hidden: int = 16, eps: float = 1e-6):
super().__init__(hidden, eps)
self.weight = nn.Parameter(torch.zeros(hidden))
def forward(self, x):
return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype)
class WeightlessRMSNorm(RMSNorm):
"""No scale parameter (e.g. Gemma3n `with_scale=False`)."""
def __init__(self, hidden: int = 16, eps: float = 1e-6):
super().__init__(hidden, eps, weight=False)
def forward(self, x):
return self._rms(x.to(torch.float32)).to(x.dtype)
class LayerNorm(RMSNorm):
"""An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
def __init__(self, hidden: int = 16, eps: float = 1e-6):
super().__init__(hidden, eps)
def forward(self, x):
return self.weight * self._rms(x)
class NotAnRMSNorm(RMSNorm):
"""Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
def __init__(self, hidden: int = 16, eps: float = 1e-6):
super().__init__(hidden, eps)
def forward(self, x):
x = x - x.mean(-1, keepdim=True)
variance = x.var(-1, keepdim=True)
return self.weight * x / torch.sqrt(variance + self.variance_epsilon)
class GatedRMSNorm(RMSNorm):
"""Second input and tail compute -> not an RMSNorm."""
def forward(self, x, gate=None):
normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
return normed * F.silu(gate)
class GatedFusedRMSNorm(nn.Module):
"""Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm."""
def __init__(self, hidden: int = 16, eps: float = 1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden))
self.eps = eps
def forward(self, x, gate=None):
return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate)
class UntraceableGatedRMSNorm(RMSNorm):
"""Tracer can't see tail compute in forward, but still has a second input (gate)."""
def forward(self, x, gate=None):
normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
if gate.sum() > 0: # untraceable -> partial graph, no visible tail
normed = normed * F.silu(gate)
return normed
@pytest.mark.parametrize(
"cls,eps,zero_centered",
[
(RMSNorm, 1e-5, False),
(GemmaRMSNorm, 1e-6, True),
(WeightlessRMSNorm, 1e-6, False),
(LayerNorm, 1e-6, False),
(torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op
],
)
def test_detects_rms_norm_variants(cls, eps, zero_centered):
with torch.device("meta"):
fuser = get_fuser(cls(16, eps=eps))
assert isinstance(fuser, RMSNormFuser)
assert fuser.zero_centered == zero_centered
@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU])
def test_non_rms_norms_are_not_matched(cls):
with torch.device("meta"):
module = cls(16) if cls is nn.LayerNorm else cls()
assert not isinstance(get_fuser(module), RMSNormFuser)
@pytest.mark.parametrize(
"cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm]
)
def test_gated_rms_norm_is_not_fused(cls):
with torch.device("meta"):
assert not isinstance(get_fuser(cls()), RMSNormFuser)
@pytest.mark.parametrize(
"cls,expected,zero_centered",
[
(RMSNorm, "RMSNorm", False),
(GemmaRMSNorm, "GemmaRMSNorm", True),
(WeightlessRMSNorm, "RMSNorm", False),
],
)
def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config):
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm
from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
# `default_vllm_config` supplies the config context the CustomOp needs; the
# weightless path reads hidden size from the model config, so stub it.
model_config = SimpleNamespace(get_hidden_size=lambda: 16)
with torch.device("meta"):
module = cls()
fuser = get_fuser(module)
built = fuser.fuse(module, "norm", model_config, None)
from vllm.model_executor.models.transformers.fusers.rms_norm import (
TPAwareNormMixin,
)
types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
assert isinstance(built, types_by_name[expected])
assert isinstance(built, TPAwareNormMixin) # fused norms self-correct under TP
assert built.variance_epsilon == module.variance_epsilon
assert isinstance(built.weight, nn.Parameter) == (
getattr(module, "weight", None) is not None
)
def test_fused_rms_norm_op_default_eps(default_vllm_config):
"""`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path;
its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`."""
from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
with torch.device("meta"):
module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call
fuser = get_fuser(module)
assert isinstance(fuser, RMSNormFuser)
assert not fuser.zero_centered
model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
built = fuser.fuse(module, "norm", model_config, None)
assert isinstance(built, VLLMRMSNorm)
assert built.variance_epsilon == torch.finfo(torch.float32).eps
def test_eps_is_derived_per_instance(default_vllm_config):
"""Two instances of the same norm class with different eps must fuse to their
own eps: the type-cached fuser holds only structure, not this value."""
model_config = SimpleNamespace(get_hidden_size=lambda: 16)
with torch.device("meta"):
for eps in (1e-5, 1e-6):
module = RMSNorm(16, eps=eps)
built = get_fuser(module).fuse(module, "norm", model_config, None)
assert built.variance_epsilon == eps
def test_fused_norm_is_gather_capable(default_vllm_config):
"""Every fused norm is emitted gather-capable, so a norm on a head-sharded
projection (OLMoE-style) self-corrects at runtime with no QKV-specific
plumbing. A full-width input skips the gather and equals a plain norm."""
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
from vllm.model_executor.models.transformers.fusers import rms_norm
torch.manual_seed(0)
x = torch.randn(4, 16)
for gathered_cls, plain_cls in [
(rms_norm.TPAwareRMSNorm, RMSNorm),
(rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm),
]:
gathered = gathered_cls(hidden_size=16, eps=1e-6)
assert isinstance(gathered, rms_norm.TPAwareNormMixin)
plain = plain_cls(hidden_size=16, eps=1e-6)
with torch.no_grad():
weight = torch.randn(16)
gathered.weight.copy_(weight)
plain.weight.copy_(weight)
torch.testing.assert_close(gathered(x), plain(x))
def test_gathered_norm_rejects_uneven_sharding(default_vllm_config):
"""A sharded input (narrower than the full-width weight) that does not tile
the weight evenly across ranks is rejected before any collective."""
from vllm.model_executor.models.transformers.fusers import rms_norm
norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6)
norm.tp_size = 2 # emulate TP=2 without a real process group
with pytest.raises(ValueError, match="does not tile it evenly"):
norm(torch.randn(2, 3)) # 3 * 2 != 8
@@ -6,10 +6,16 @@ from typing import Any
import pytest
from ..conftest import HfRunner, VllmRunner
from ..utils import multi_gpu_test, prep_prompts
from .registry import HF_EXAMPLE_MODELS
from .utils import check_embeddings_close, check_logprobs_close
from ...conftest import HfRunner, VllmRunner
from ...utils import multi_gpu_test, prep_prompts
from ..registry import HF_EXAMPLE_MODELS
from ..utils import check_embeddings_close, check_logprobs_close
@pytest.fixture(scope="function", autouse=True)
def enable_pickle(monkeypatch):
"""`LLM.apply_model` requires pickling a function."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
def get_model(arch: str) -> str:
@@ -18,6 +24,17 @@ def get_model(arch: str) -> str:
return model_info.default
def get_num_fused(model) -> tuple[int, int]:
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
)
glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules())
qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules())
return glu, qkv
def check_implementation(
runner_ref: type[HfRunner | VllmRunner],
runner_test: type[VllmRunner],
@@ -25,6 +42,7 @@ def check_implementation(
model: str,
kwargs_ref: dict[str, Any] | None = None,
kwargs_test: dict[str, Any] | None = None,
num_fused: tuple[int, int] = (1, 1),
**kwargs,
):
if kwargs_ref is None:
@@ -41,6 +59,12 @@ def check_implementation(
model_config = model_test.llm.llm_engine.model_config
assert model_config.using_transformers_backend()
num_layers = model_config.hf_config.get_text_config().num_hidden_layers
expected_glu, expected_qkv = num_fused
for num_glu, num_qkv in model_test.apply_model(get_num_fused):
assert num_glu == expected_glu * num_layers
assert num_qkv == expected_qkv * num_layers
outputs_test = model_test.generate_greedy_logprobs(*args)
with runner_ref(model, **kwargs_ref) as model_ref:
@@ -58,11 +82,11 @@ def check_implementation(
@pytest.mark.parametrize(
"model,model_impl",
"model,model_impl,num_fused",
[
("meta-llama/Llama-3.2-1B-Instruct", "transformers"),
("hmellor/Ilama-3.2-1B", "auto"), # CUSTOM CODE
("allenai/OLMoE-1B-7B-0924", "transformers"), # MoE
("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)),
("hmellor/Ilama-3.2-1B", "auto", (1, 1)), # CUSTOM CODE
("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)), # MoE
],
) # trust_remote_code=True by default
def test_models(
@@ -71,6 +95,7 @@ def test_models(
example_prompts: list[str],
model: str,
model_impl: str,
num_fused: tuple[int, int],
) -> None:
import transformers
from packaging.version import Version
@@ -84,7 +109,12 @@ def test_models(
)
check_implementation(
hf_runner, vllm_runner, example_prompts, model, model_impl=model_impl
hf_runner,
vllm_runner,
example_prompts,
model,
num_fused=num_fused,
model_impl=model_impl,
)
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for VLLMUnprocessableEntityError and media fetch error handling.
Verifies that unprocessable image URLs (404, 403, DNS failures, etc.) return
HTTP 422 instead of 500.
"""
from http import HTTPStatus
from unittest.mock import AsyncMock, MagicMock, patch
import aiohttp
import pytest
from vllm.entrypoints.serve.utils.error_response import create_error_response
from vllm.exceptions import VLLMUnprocessableEntityError
from vllm.multimodal.media import MediaConnector
class TestVLLMUnprocessableEntityError:
"""Tests for VLLMUnprocessableEntityError exception."""
def test_creation(self):
exc = VLLMUnprocessableEntityError("Test error")
assert str(exc) == "Test error"
assert exc.parameter is None
def test_creation_with_parameter_and_value(self):
exc = VLLMUnprocessableEntityError(
"Test error",
parameter="image_url",
value="https://example.com/image.jpg",
)
assert "parameter=image_url" in str(exc)
assert "value=https://example.com/image.jpg" in str(exc)
def test_is_value_error_subclass(self):
exc = VLLMUnprocessableEntityError("Test")
assert isinstance(exc, ValueError)
class TestMediaConnectorErrorHandling:
"""Tests for MediaConnector error handling."""
@pytest.mark.asyncio
async def test_fetch_image_async_404(self):
connector = MediaConnector()
with patch.object(
connector.connection, "async_get_bytes", new_callable=AsyncMock
) as mock_get:
mock_get.side_effect = aiohttp.ClientResponseError(
request_info=MagicMock(),
history=(),
status=404,
message="Not Found",
)
with pytest.raises(VLLMUnprocessableEntityError) as exc_info:
await connector.fetch_image_async("https://example.com/missing.jpg")
assert exc_info.value.parameter == "image_url"
@pytest.mark.asyncio
async def test_fetch_image_async_dns_error(self):
"""DNS errors are transient and should remain as-is for retry."""
connector = MediaConnector()
with patch.object(
connector.connection, "async_get_bytes", new_callable=AsyncMock
) as mock_get:
mock_get.side_effect = aiohttp.ClientConnectorDNSError(
connection_key=MagicMock(),
os_error=MagicMock(),
)
with pytest.raises(aiohttp.ClientConnectorDNSError) as exc_info:
await connector.fetch_image_async(
"https://nonexistent.example/image.jpg"
)
assert isinstance(exc_info.value, aiohttp.ClientConnectorDNSError)
@pytest.mark.asyncio
async def test_fetch_image_async_500_preserved(self):
"""5xx errors should remain as server errors."""
connector = MediaConnector()
with patch.object(
connector.connection, "async_get_bytes", new_callable=AsyncMock
) as mock_get:
mock_get.side_effect = aiohttp.ClientResponseError(
request_info=MagicMock(),
history=(),
status=500,
message="Internal Server Error",
)
with pytest.raises(aiohttp.ClientResponseError) as exc_info:
await connector.fetch_image_async("https://example.com/image.jpg")
assert exc_info.value.status == 500
def test_fetch_image_404(self):
connector = MediaConnector()
with patch.object(
connector.connection, "get_bytes", new_callable=MagicMock
) as mock_get:
mock_get.side_effect = aiohttp.ClientResponseError(
request_info=MagicMock(),
history=(),
status=404,
message="Not Found",
)
with pytest.raises(VLLMUnprocessableEntityError) as exc_info:
connector.fetch_image("https://example.com/missing.jpg")
assert exc_info.value.parameter == "image_url"
def test_fetch_image_connection_error(self):
"""Connection errors are transient and should remain as-is for retry."""
connector = MediaConnector()
with patch.object(
connector.connection, "get_bytes", new_callable=MagicMock
) as mock_get:
mock_get.side_effect = aiohttp.ClientConnectionError("Connection refused")
with pytest.raises(aiohttp.ClientConnectionError) as exc_info:
connector.fetch_image("https://example.com/image.jpg")
assert isinstance(exc_info.value, aiohttp.ClientConnectionError)
class TestErrorResponse:
"""Tests for error response creation."""
def test_unprocessable_entity_returns_422(self):
exc = VLLMUnprocessableEntityError(
"Failed to fetch media from URL: Cannot connect",
parameter="image_url",
value="https://example.com/image.jpg",
)
response = create_error_response(exc)
assert response.error.code == HTTPStatus.UNPROCESSABLE_ENTITY.value
assert response.error.type == "UnprocessableEntityError"
assert response.error.param == "image_url"
def test_unprocessable_entity_message(self):
exc = VLLMUnprocessableEntityError("Test error message")
response = create_error_response(exc)
assert response.error.message == "Test error message"
assert response.error.code == 422
+97 -1
View File
@@ -16,7 +16,11 @@ from vllm.assets.video import (
video_to_pil_images_list,
)
from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader
from vllm.multimodal.video import (
PYNVVIDEOCODEC_VIDEO_BACKEND,
VIDEO_LOADER_REGISTRY,
VideoLoader,
)
from ..utils import cosine_similarity, create_video_from_image, normalize_image
@@ -357,3 +361,95 @@ def test_load_base64_jpeg_raises_on_zero_num_frames():
with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"):
videoio.load_base64("video/jpeg", data)
# ---------------------------------------------------------------------------
# GPU video backend policy tests
# ---------------------------------------------------------------------------
class TestMergeKwargsGpuBackendPolicy:
"""Verify that merge_kwargs blocks request-level GPU backend selection
when the static (engine-level) config did not configure that backend."""
def test_pynvvideocodec_requires_gpu(self):
assert VIDEO_LOADER_REGISTRY.backend_requires_gpu(PYNVVIDEOCODEC_VIDEO_BACKEND)
def test_strips_video_backend_pynv_when_not_static(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs=None,
runtime_kwargs={"video_backend": "pynvvideocodec"},
)
assert "video_backend" not in result
def test_strips_backend_pynv_when_not_static(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs={"num_frames": 16},
runtime_kwargs={"backend": "pynvvideocodec"},
)
assert result.get("backend") != "pynvvideocodec"
def test_preserves_video_backend_pynv_when_static(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs={"video_backend": "pynvvideocodec"},
runtime_kwargs={"video_backend": "pynvvideocodec", "num_frames": 8},
)
assert result["video_backend"] == "pynvvideocodec"
assert result["num_frames"] == 8
def test_preserves_backend_pynv_when_static(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs={"backend": "pynvvideocodec"},
runtime_kwargs={"backend": "pynvvideocodec"},
)
assert result["backend"] == "pynvvideocodec"
@pytest.mark.parametrize("backend", ["opencv", "pyav", "torchcodec"])
def test_software_video_backend_passes_through(self, backend: str):
result = VideoMediaIO.merge_kwargs(
default_kwargs=None,
runtime_kwargs={"video_backend": backend},
)
assert result["video_backend"] == backend
@pytest.mark.parametrize("backend", ["opencv", "pyav"])
def test_software_codec_backend_passes_through(self, backend: str):
result = VideoMediaIO.merge_kwargs(
default_kwargs=None,
runtime_kwargs={"backend": backend},
)
assert result["backend"] == backend
def test_strips_both_keys_independently(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs=None,
runtime_kwargs={
"video_backend": "pynvvideocodec",
"backend": "pynvvideocodec",
"num_frames": 4,
},
)
assert "video_backend" not in result
assert result.get("backend") != "pynvvideocodec"
assert result["num_frames"] == 4
def test_other_kwargs_preserved_when_gpu_backend_stripped(self):
result = VideoMediaIO.merge_kwargs(
default_kwargs={"fps": 2},
runtime_kwargs={
"video_backend": "pynvvideocodec",
"num_frames": 16,
},
)
assert "video_backend" not in result
assert result["num_frames"] == 16
def test_static_pynv_with_different_runtime_gpu_backend(self):
"""If static sets pynv via video_backend but runtime tries to set it
via the codec-level 'backend' key (without a static match), strip it."""
result = VideoMediaIO.merge_kwargs(
default_kwargs={"video_backend": "pynvvideocodec"},
runtime_kwargs={"backend": "pynvvideocodec"},
)
assert result.get("backend") != "pynvvideocodec"
assert result["video_backend"] == "pynvvideocodec"
+268
View File
@@ -0,0 +1,268 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for DeepSeek V3.2 parser engine semantics.
V3.2 uses the same DSML parameter format as V4 but wraps tool calls in
``<DSMLfunction_calls>`` instead of ``<DSMLtool_calls>`` and has
no reasoning (``<think>``/``</think>``) support.
"""
import json
import pytest
from tests.parser.engine.conftest import make_mock_tokenizer
from tests.parser.engine.streaming_helpers import (
collect_content,
collect_function_name,
collect_tool_arguments,
simulate_tool_streaming,
)
from vllm.parser.deepseek_v4 import (
DSML_INVOKE_END,
DSML_INVOKE_NAME_END,
DSML_INVOKE_PREFIX,
)
from vllm.parser.deepseek_v32 import (
DSML_FUNC_END,
DSML_FUNC_START,
DeepSeekV32Parser,
)
from vllm.parser.engine.parser_engine_config import ParserState
_PARAM_OPEN = 'DSMLparameter name="{name}" string="{is_str}">'
_PARAM_CLOSE = "</DSMLparameter>"
def _param(name: str, is_str: str, value: str) -> str:
return f"<{_PARAM_OPEN.format(name=name, is_str=is_str)}{value}{_PARAM_CLOSE}"
def _invoke(name: str, *params: str) -> str:
body = "\n".join(params)
return (
f"{DSML_INVOKE_PREFIX}{name}{DSML_INVOKE_NAME_END}\n{body}\n{DSML_INVOKE_END}"
)
def _func_calls(*invocations: str) -> str:
body = "\n".join(invocations)
return f"{DSML_FUNC_START}\n{body}\n{DSML_FUNC_END}"
def _make_tool(name, properties):
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
)
return ChatCompletionToolsParam(
type="function",
function={
"name": name,
"parameters": {
"type": "object",
"properties": properties,
},
},
)
@pytest.fixture
def mock_tokenizer():
return make_mock_tokenizer({})
@pytest.fixture
def mock_request():
from unittest.mock import MagicMock
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
req = MagicMock(spec=ChatCompletionRequest)
req.tools = []
req.tool_choice = "auto"
return req
# ── Non-streaming extraction ────────────────────────────────────────
class TestNonStreaming:
def test_no_tool_call(self, mock_tokenizer, mock_request):
parser = DeepSeekV32Parser(mock_tokenizer)
result = parser.extract_tool_calls("Hello world", mock_request)
assert not result.tools_called
assert result.content == "Hello world"
def test_single_tool(self, mock_tokenizer, mock_request):
text = _func_calls(
_invoke("get_weather", _param("city", "true", "SF")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"city": "SF"}
def test_parallel_tools(self, mock_tokenizer, mock_request):
text = _func_calls(
_invoke("get_weather", _param("city", "true", "SF")),
_invoke("get_weather", _param("city", "true", "NYC")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called
assert len(result.tool_calls) == 2
assert json.loads(result.tool_calls[0].function.arguments) == {"city": "SF"}
assert json.loads(result.tool_calls[1].function.arguments) == {"city": "NYC"}
def test_content_before_tool_call(self, mock_tokenizer, mock_request):
text = "Let me check. " + _func_calls(
_invoke("search", _param("q", "true", "vllm")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called
assert result.content is not None
assert "Let me check" in result.content
def test_non_string_params_json_parsed(self, mock_tokenizer, mock_request):
text = _func_calls(
_invoke(
"toggle",
_param("enabled", "false", "true"),
_param("count", "false", "42"),
),
)
parser = DeepSeekV32Parser(mock_tokenizer)
result = parser.extract_tool_calls(text, mock_request)
args = json.loads(result.tool_calls[0].function.arguments)
assert args["enabled"] is True
assert args["count"] == 42
def test_wrapper_unwrapping(self, mock_tokenizer, mock_request):
tool = _make_tool("get_weather", {"location": {"type": "string"}})
mock_request.tools = [tool]
text = _func_calls(
_invoke(
"get_weather",
_param("arguments", "false", '{"location":"Beijing"}'),
),
)
parser = DeepSeekV32Parser(mock_tokenizer, tools=[tool])
result = parser.extract_tool_calls(text, mock_request)
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "Beijing"}
# ── Initial state ────────────────────────────────────────────────────
class TestInitialState:
def test_always_content(self, mock_tokenizer):
parser = DeepSeekV32Parser(mock_tokenizer)
cfg = parser.parser_engine_config
assert cfg.initial_state == ParserState.CONTENT
def test_ignores_thinking_kwargs(self, mock_tokenizer):
parser = DeepSeekV32Parser(
mock_tokenizer,
chat_template_kwargs={"thinking": True, "enable_thinking": True},
)
cfg = parser.parser_engine_config
assert cfg.initial_state == ParserState.CONTENT
# ── Streaming ────────────────────────────────────────────────────────
class TestStreaming:
def test_single_tool_streaming(self, mock_tokenizer, mock_request):
text = _func_calls(
_invoke("get_weather", _param("city", "true", "SF")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
results = simulate_tool_streaming(parser, mock_request, list(text))
assert collect_function_name(results) == "get_weather"
args_json = collect_tool_arguments(results)
assert json.loads(args_json) == {"city": "SF"}
def test_content_before_tool_streaming(self, mock_tokenizer, mock_request):
text = "Checking... " + _func_calls(
_invoke("fn", _param("k", "true", "v")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
results = simulate_tool_streaming(parser, mock_request, list(text))
content = collect_content(results)
assert "Checking" in content
def test_parallel_tools_streaming(self, mock_tokenizer, mock_request):
text = _func_calls(
_invoke("fn_a", _param("x", "true", "1")),
_invoke("fn_b", _param("y", "true", "2")),
)
parser = DeepSeekV32Parser(mock_tokenizer)
results = simulate_tool_streaming(parser, mock_request, list(text))
names = []
for delta, _ in results:
if delta and delta.tool_calls:
for tc in delta.tool_calls:
if tc.function and tc.function.name:
names.append(tc.function.name)
assert "fn_a" in names
assert "fn_b" in names
def test_no_tool_content_only(self, mock_tokenizer, mock_request):
text = "Just some text, no tools."
parser = DeepSeekV32Parser(mock_tokenizer)
results = simulate_tool_streaming(parser, mock_request, list(text))
content = collect_content(results)
assert "Just some text" in content
args = collect_tool_arguments(results)
assert args == ""
def test_streaming_wrapper_unwrap_consistency(self, mock_tokenizer, mock_request):
tool = _make_tool("get_weather", {"location": {"type": "string"}})
mock_request.tools = [tool]
parser = DeepSeekV32Parser(mock_tokenizer, tools=[tool])
chunks = [
DSML_FUNC_START,
_invoke(
"get_weather",
_param("arguments", "false", '{"location": "NYC"}'),
),
DSML_FUNC_END,
]
results = simulate_tool_streaming(parser, mock_request, chunks)
streamed_args = collect_tool_arguments(results)
final_delta, _ = results[-1]
finish_delta = parser.finish_streaming()
extracted = parser._build_extracted_result(final_delta, finish_delta)
assert extracted.tools_called is True
assert len(extracted.tool_calls) == 1
final_args = extracted.tool_calls[0].function.arguments
assert json.loads(final_args) == {"location": "NYC"}
assert '"arguments"' not in streamed_args
assert final_args.startswith(streamed_args)
def test_missing_invoke_end(self, mock_tokenizer, mock_request):
text = (
f"{DSML_FUNC_START}\n"
f"{DSML_INVOKE_PREFIX}fn{DSML_INVOKE_NAME_END}\n"
f"{_param('k', 'true', 'v')}\n"
f"{DSML_FUNC_END}"
)
parser = DeepSeekV32Parser(mock_tokenizer)
results = simulate_tool_streaming(parser, mock_request, list(text))
assert collect_function_name(results) == "fn"
args = json.loads(collect_tool_arguments(results))
assert args == {"k": "v"}
+922
View File
@@ -0,0 +1,922 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for DeepSeek V4-specific parser engine semantics."""
import json
import pytest
from tests.parser.engine.conftest import make_mock_tokenizer
from tests.parser.engine.replay_harness import (
DUMMY_TOOLS,
MockTokenizer,
_test_request,
collect_output,
replay_streaming,
)
from tests.parser.engine.streaming_helpers import (
collect_content,
collect_function_name,
collect_tool_arguments,
simulate_reasoning_streaming,
simulate_tool_streaming,
)
from vllm.parser.abstract_parser import DelegatingParser
from vllm.parser.deepseek_v4 import (
DSML_INVOKE_END,
DSML_INVOKE_NAME_END,
DSML_INVOKE_PREFIX,
DSML_THINK_END,
DSML_THINK_START,
DSML_TOOL_END,
DSML_TOOL_START,
DeepSeekV4Parser,
_dsml_arg_converter,
_unwrap_wrapper_args,
deepseek_v4_config,
)
from vllm.parser.engine.registered_adapters import (
DeepSeekV4ParserReasoningAdapter,
DeepSeekV4ParserToolAdapter,
)
_THINK_START_ID = 50
_THINK_END_ID = 51
_PARAM_OPEN = 'DSMLparameter name="{name}" string="{is_str}">'
_PARAM_CLOSE = "</DSMLparameter>"
def _param(name: str, is_str: str, value: str) -> str:
return f"<{_PARAM_OPEN.format(name=name, is_str=is_str)}{value}{_PARAM_CLOSE}"
@pytest.fixture
def mock_tokenizer():
return make_mock_tokenizer(
{
DSML_THINK_START: _THINK_START_ID,
DSML_THINK_END: _THINK_END_ID,
}
)
# ── Arg converter unit tests ─────────────────────────────────────────
class TestArgConverter:
def _raw(self, *params: tuple[str, str, str]) -> str:
lines = [_param(n, s, v) for n, s, v in params]
return "\n" + "\n".join(lines) + "\n"
def test_string_param(self):
raw = self._raw(("city", "true", "杭州"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result == {"city": "杭州"}
def test_string_with_spaces_and_quotes(self):
raw = self._raw(("msg", "true", 'He said "hello world"'))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["msg"] == 'He said "hello world"'
def test_integer_param(self):
raw = self._raw(("count", "false", "42"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["count"] == 42
assert isinstance(result["count"], int)
def test_float_param(self):
raw = self._raw(("ratio", "false", "3.14"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert abs(result["ratio"] - 3.14) < 1e-9
def test_bool_param(self):
raw = self._raw(("flag", "false", "true"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["flag"] is True
def test_array_param(self):
raw = self._raw(("items", "false", '["a", "b", "c"]'))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["items"] == ["a", "b", "c"]
def test_object_param(self):
raw = self._raw(("opts", "false", '{"key": "val"}'))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["opts"] == {"key": "val"}
def test_mixed_types(self):
raw = self._raw(
("location", "true", "Tokyo"),
("limit", "false", "10"),
("active", "false", "false"),
)
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result == {"location": "Tokyo", "limit": 10, "active": False}
def test_empty_args(self):
result = json.loads(_dsml_arg_converter("", partial=False))
assert result == {}
def test_invalid_json_fallback(self):
raw = self._raw(("data", "false", "[broken"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["data"] == "[broken"
def test_chinese_chars_preserved_in_json(self):
raw = self._raw(("query", "true", "你好世界"))
raw_json = _dsml_arg_converter(raw, partial=False)
assert "你好世界" in raw_json
result = json.loads(raw_json)
assert result["query"] == "你好世界"
def test_partial_complete_plus_in_progress(self):
raw = self._raw(("city", "true", "Tokyo"))
raw += f"<{_PARAM_OPEN.format(name='unit', is_str='true')}celsi"
result = json.loads(_dsml_arg_converter(raw, partial=True))
assert result["city"] == "Tokyo"
assert result["unit"] == "celsi"
def test_partial_no_in_progress(self):
raw = self._raw(("city", "true", "Tokyo"))
result = json.loads(_dsml_arg_converter(raw, partial=True))
assert result == {"city": "Tokyo"}
def test_partial_value_with_angle_bracket(self):
raw = f"<{_PARAM_OPEN.format(name='code', is_str='true')}a<b"
result = json.loads(_dsml_arg_converter(raw, partial=True))
assert result == {"code": "a<b"}
def test_partial_value_with_angle_bracket_and_complete_param(self):
raw = self._raw(("city", "true", "Tokyo"))
raw += f"<{_PARAM_OPEN.format(name='expr', is_str='true')}x<5"
result = json.loads(_dsml_arg_converter(raw, partial=True))
assert result["city"] == "Tokyo"
assert result["expr"] == "x<5"
def test_null_string_false(self):
raw = self._raw(("val", "false", "null"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["val"] is None
def test_string_true_not_json_parsed(self):
raw = self._raw(("n", "true", "42"))
result = json.loads(_dsml_arg_converter(raw, partial=False))
assert result["n"] == "42"
assert isinstance(result["n"], str)
# ── Bare </think> absorption and duplicate <think> absorption ─────────
class TestThinkTagAbsorption:
def test_bare_think_end_not_leaked(self, mock_tokenizer):
parser = DeepSeekV4Parser(mock_tokenizer)
chunks = ["</think>", "Here is the direct answer."]
reasoning, content = simulate_reasoning_streaming(parser, chunks)
assert reasoning == ""
assert "</think>" not in content
assert "Here is the direct answer" in content
def test_duplicate_think_start_absorbed(self, mock_tokenizer):
parser = DeepSeekV4Parser(
mock_tokenizer, chat_template_kwargs={"thinking": True}
)
chunks = [
"<think>\n",
"Some reasoning.\n",
"</think>\n",
"Answer.",
]
reasoning, content = simulate_reasoning_streaming(parser, chunks)
assert "Some reasoning" in reasoning
assert "Answer" in content
# ── Missing </DSMLinvoke> before </DSMLtool_calls> ────────────
class TestMissingInvokeEnd:
def test_non_streaming(self, mock_tokenizer, mock_request):
parser = DeepSeekV4Parser(mock_tokenizer)
text = (
f"{DSML_TOOL_START}"
f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n"
f"{_param('location', 'true', 'NYC')}\n"
f"{DSML_TOOL_END}"
)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "NYC"}
def test_streaming_with_trailing_content(self, mock_tokenizer, mock_request):
parser = DeepSeekV4Parser(mock_tokenizer)
chunks = [
DSML_TOOL_START,
f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n"
f"{_param('location', 'true', 'NYC')}\n",
DSML_TOOL_END,
"Done.",
]
results = simulate_tool_streaming(parser, mock_request, chunks)
assert collect_function_name(results) == "get_weather"
args = json.loads(collect_tool_arguments(results))
assert args == {"location": "NYC"}
assert "Done." in collect_content(results)
# ── Thinking mode initial state ──────────────────────────────────────
class TestThinkingModeConfig:
def test_thinking_true_starts_in_reasoning(self):
cfg = deepseek_v4_config(thinking=True)
assert cfg.initial_state.name == "REASONING"
def test_thinking_false_starts_in_content(self):
cfg = deepseek_v4_config(thinking=False)
assert cfg.initial_state.name == "CONTENT"
def test_enable_thinking_kwarg(self, mock_tokenizer):
p = DeepSeekV4Parser(
mock_tokenizer, chat_template_kwargs={"enable_thinking": True}
)
assert p.parser_engine_config.initial_state.name == "REASONING"
def test_no_thinking_kwarg_defaults_to_content(self, mock_tokenizer):
p = DeepSeekV4Parser(mock_tokenizer)
assert p.parser_engine_config.initial_state.name == "CONTENT"
def test_thinking_mode_reasoning_without_tags(self, mock_tokenizer):
parser = DeepSeekV4Parser(
mock_tokenizer, chat_template_kwargs={"thinking": True}
)
chunks = [
"\n\nLet me consider ",
"this carefully.\n",
"</think>\n",
"Here is the result.",
]
reasoning, content = simulate_reasoning_streaming(parser, chunks)
assert "Let me consider" in reasoning
assert "Here is the result" in content
def test_thinking_mode_all_reasoning_no_end_tag(self, mock_tokenizer):
parser = DeepSeekV4Parser(
mock_tokenizer, chat_template_kwargs={"thinking": True}
)
chunks = ["I'll review ", "the PR."]
reasoning, content = simulate_reasoning_streaming(parser, chunks)
assert "review" in reasoning
assert "the PR" in reasoning
assert content == ""
def test_reasoning_effort_none_overrides_enable_thinking(self, mock_tokenizer):
p = DeepSeekV4Parser(
mock_tokenizer,
chat_template_kwargs={
"enable_thinking": True,
"reasoning_effort": "none",
},
)
assert p.parser_engine_config.initial_state.name == "CONTENT"
# ── Implicit reasoning end (missing </think> before tool calls) ─────
class TestImplicitReasoningEnd:
"""Tool call markers end reasoning implicitly when </think> is missing.
DeepSeek V4 models occasionally omit </think> before emitting tool calls.
The (REASONING, TOOL_START) transition handles this gracefully.
"""
@pytest.fixture
def thinking_parser(self, mock_tokenizer):
return DeepSeekV4Parser(mock_tokenizer, chat_template_kwargs={"thinking": True})
def _reasoning_then_tool(self, reasoning_text: str) -> str:
return reasoning_text + _tool_calls(
_invoke("get_weather", ("location", "true", "NYC")),
)
def test_non_streaming_extract_reasoning_implicit_end(self, thinking_parser):
text = self._reasoning_then_tool("Let me look up the weather.\n\n")
reasoning, content = thinking_parser.extract_reasoning(text, None)
assert reasoning == "Let me look up the weather."
assert DSML_TOOL_START not in reasoning
assert DSML_INVOKE_PREFIX not in reasoning
assert content is None
def test_non_streaming_extract_tool_calls_implicit_end(
self, thinking_parser, mock_request
):
text = self._reasoning_then_tool("Let me look up the weather.\n\n")
result = thinking_parser.extract_tool_calls(text, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
assert result.tool_calls[0].function.name == "get_weather"
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "NYC"}
def test_non_streaming_parse_implicit_end(self, thinking_parser, mock_request):
text = self._reasoning_then_tool("Let me look up the weather.\n\n")
reasoning, content, tool_calls = thinking_parser.parse(text, mock_request)
assert reasoning == "Let me look up the weather."
assert content is None
assert tool_calls is not None
assert len(tool_calls) == 1
assert tool_calls[0].name == "get_weather"
args = json.loads(tool_calls[0].arguments)
assert args == {"location": "NYC"}
def test_streaming_reasoning_implicit_end(self, thinking_parser):
chunks = [
"Let me look up the weather.\n\n",
DSML_TOOL_START,
DSML_INVOKE_PREFIX + "get_weather" + DSML_INVOKE_NAME_END,
]
reasoning, content = simulate_reasoning_streaming(thinking_parser, chunks)
assert reasoning == "Let me look up the weather."
assert DSML_TOOL_START not in reasoning
assert DSML_INVOKE_PREFIX not in reasoning
def test_streaming_tool_extraction_implicit_end(
self, thinking_parser, mock_request
):
chunks = [
"Let me check.\n\n",
DSML_TOOL_START,
DSML_INVOKE_PREFIX
+ "get_weather"
+ DSML_INVOKE_NAME_END
+ "\n"
+ _param("location", "true", "NYC")
+ "\n"
+ DSML_INVOKE_END,
DSML_TOOL_END,
]
results = simulate_tool_streaming(thinking_parser, mock_request, chunks)
assert collect_function_name(results) == "get_weather"
args = json.loads(collect_tool_arguments(results))
assert args == {"location": "NYC"}
def test_thinking_false_explicit_think_then_tool_call(self, mock_tokenizer):
parser = DeepSeekV4Parser(mock_tokenizer)
chunks = [
DSML_THINK_START,
"Let me check the weather.",
DSML_TOOL_START,
DSML_INVOKE_PREFIX + "get_weather" + DSML_INVOKE_NAME_END,
]
reasoning, content = simulate_reasoning_streaming(parser, chunks)
assert "Let me check the weather" in reasoning
assert DSML_TOOL_START not in reasoning
assert DSML_THINK_START not in reasoning
def test_non_streaming_parallel_tools_after_implicit_end(
self, thinking_parser, mock_request
):
text = "I need both.\n\n" + _tool_calls(
_invoke("get_weather", ("location", "true", "NYC")),
_invoke("get_time", ("timezone", "true", "EST")),
)
result = thinking_parser.extract_tool_calls(text, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 2
assert result.tool_calls[0].function.name == "get_weather"
assert result.tool_calls[1].function.name == "get_time"
def test_streaming_implicit_end_trailing_whitespace_stripped(self, thinking_parser):
chunks = [
"Reasoning.\n\n\n",
DSML_TOOL_START,
DSML_INVOKE_PREFIX + "func" + DSML_INVOKE_NAME_END,
]
reasoning, content = simulate_reasoning_streaming(thinking_parser, chunks)
assert reasoning == "Reasoning."
# ── Wrapper argument unwrapping ──────────────────────────────────────
class TestWrapperUnwrapping:
def test_unwrap_arguments_wrapper(self):
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
)
tool = ChatCompletionToolsParam(
type="function",
function={
"name": "get_weather",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
},
},
)
result = _unwrap_wrapper_args(
'{"arguments": {"location": "Beijing"}}',
[tool],
"get_weather",
)
assert json.loads(result) == {"location": "Beijing"}
def test_unwrap_input_wrapper(self):
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
)
tool = ChatCompletionToolsParam(
type="function",
function={
"name": "get_weather",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
},
},
)
result = _unwrap_wrapper_args(
'{"input": {"location": "Beijing"}}',
[tool],
"get_weather",
)
assert json.loads(result) == {"location": "Beijing"}
def test_no_unwrap_when_key_in_schema(self):
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
)
tool = ChatCompletionToolsParam(
type="function",
function={
"name": "func",
"parameters": {
"type": "object",
"properties": {"arguments": {"type": "string"}},
},
},
)
result = _unwrap_wrapper_args(
'{"arguments": "some value"}',
[tool],
"func",
)
assert json.loads(result) == {"arguments": "some value"}
def test_no_unwrap_when_no_tools(self):
result = _unwrap_wrapper_args(
'{"arguments": {"location": "Beijing"}}',
None,
"get_weather",
)
assert json.loads(result) == {"arguments": {"location": "Beijing"}}
def test_unwrap_json_string_inner(self):
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionToolsParam,
)
tool = ChatCompletionToolsParam(
type="function",
function={
"name": "get_weather",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
},
},
)
result = _unwrap_wrapper_args(
'{"arguments": "{\\"location\\": \\"Beijing\\"}"}',
[tool],
"get_weather",
)
assert json.loads(result) == {"location": "Beijing"}
# ── Parallel tool call wrapper unwrapping ───────────────────────────
def _make_tool(name, properties):
from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E501
ChatCompletionToolsParam,
)
return ChatCompletionToolsParam(
type="function",
function={
"name": name,
"parameters": {
"type": "object",
"properties": properties,
},
},
)
def _invoke(name, *params):
body = "\n".join(_param(n, s, v) for n, s, v in params)
return (
f"{DSML_INVOKE_PREFIX}{name}{DSML_INVOKE_NAME_END}\n{body}\n{DSML_INVOKE_END}"
)
def _tool_calls(*invokes):
return DSML_TOOL_START + "\n".join(invokes) + DSML_TOOL_END
class TestParallelUnwrapping:
@pytest.fixture
def weather_tool(self):
return _make_tool(
"get_weather",
{
"location": {"type": "string"},
"unit": {"type": "string"},
},
)
@pytest.fixture
def time_tool(self):
return _make_tool(
"get_time",
{"timezone": {"type": "string"}},
)
@pytest.mark.parametrize(
"weather_args, expected",
[
(
'{"location": "NYC", "unit": "celsius"}',
{"location": "NYC", "unit": "celsius"},
),
('{"location": "NYC"}', {"location": "NYC"}),
],
ids=["all_props", "subset_props"],
)
def test_unwrap_parallel_uses_correct_schema(
self,
mock_tokenizer,
mock_request,
weather_tool,
time_tool,
weather_args,
expected,
):
tools = [weather_tool, time_tool]
parser = DeepSeekV4Parser(mock_tokenizer, tools=tools)
mock_request.tools = tools
text = _tool_calls(
_invoke("get_weather", ("arguments", "false", weather_args)),
_invoke("get_time", ("timezone", "true", "EST")),
)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 2
assert result.tool_calls[0].function.name == "get_weather"
args0 = json.loads(result.tool_calls[0].function.arguments)
assert args0 == expected
assert result.tool_calls[1].function.name == "get_time"
args1 = json.loads(result.tool_calls[1].function.arguments)
assert args1 == {"timezone": "EST"}
def test_unwrap_parallel_streaming(
self, mock_tokenizer, mock_request, weather_tool, time_tool
):
tools = [weather_tool, time_tool]
parser = DeepSeekV4Parser(mock_tokenizer, tools=tools)
mock_request.tools = tools
chunks = [
DSML_TOOL_START,
_invoke(
"get_weather",
("arguments", "false", '{"location": "NYC"}'),
),
_invoke("get_time", ("timezone", "true", "EST")),
DSML_TOOL_END,
]
results = simulate_tool_streaming(parser, mock_request, chunks)
final_delta, _ = results[-1]
finish_delta = parser.finish_streaming()
extracted = parser._build_extracted_result(final_delta, finish_delta)
assert extracted.tools_called is True
assert len(extracted.tool_calls) == 2
args0 = json.loads(extracted.tool_calls[0].function.arguments)
assert args0 == {"location": "NYC"}
args1 = json.loads(extracted.tool_calls[1].function.arguments)
assert args1 == {"timezone": "EST"}
def test_no_unwrap_parallel_when_no_match(
self, mock_tokenizer, mock_request, weather_tool, time_tool
):
tools = [weather_tool, time_tool]
parser = DeepSeekV4Parser(mock_tokenizer, tools=tools)
mock_request.tools = tools
text = _tool_calls(
_invoke(
"get_weather",
("arguments", "false", '{"unknown_key": "val"}'),
),
_invoke("get_time", ("timezone", "true", "EST")),
)
result = parser.extract_tool_calls(text, mock_request)
assert len(result.tool_calls) == 2
args0 = json.loads(result.tool_calls[0].function.arguments)
assert args0 == {"arguments": {"unknown_key": "val"}}
args1 = json.loads(result.tool_calls[1].function.arguments)
assert args1 == {"timezone": "EST"}
def test_unwrap_single_tool_still_works(self, mock_tokenizer, mock_request):
tool = _make_tool("get_weather", {"location": {"type": "string"}})
tools = [tool]
parser = DeepSeekV4Parser(mock_tokenizer, tools=tools)
mock_request.tools = tools
text = _tool_calls(
_invoke(
"get_weather",
("arguments", "false", '{"location": "Beijing"}'),
),
)
result = parser.extract_tool_calls(text, mock_request)
assert result.tools_called is True
assert len(result.tool_calls) == 1
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"location": "Beijing"}
# ── Streaming wrapper consistency ─────────────────────────────────────
class TestStreamingWrapperConsistency:
"""Streamed arg deltas must stay consistent with final extraction
when wrapper params like 'arguments' are unwrapped."""
def test_streaming_wrapper_unwrap_consistency(self, mock_tokenizer, mock_request):
tool = _make_tool("get_weather", {"location": {"type": "string"}})
tools = [tool]
parser = DeepSeekV4Parser(mock_tokenizer, tools=tools)
mock_request.tools = tools
chunks = [
DSML_TOOL_START,
_invoke(
"get_weather",
("arguments", "false", '{"location": "NYC"}'),
),
DSML_TOOL_END,
]
results = simulate_tool_streaming(parser, mock_request, chunks)
streamed_args = collect_tool_arguments(results)
final_delta, _ = results[-1]
finish_delta = parser.finish_streaming()
extracted = parser._build_extracted_result(final_delta, finish_delta)
assert extracted.tools_called is True
assert len(extracted.tool_calls) == 1
final_args = extracted.tool_calls[0].function.arguments
assert json.loads(final_args) == {"location": "NYC"}
assert '"arguments"' not in streamed_args, (
f"Streamed args should not contain wrapper key, got: {streamed_args!r}"
)
assert final_args.startswith(streamed_args), (
f"Extracted args {final_args!r} "
f"should start with streamed args {streamed_args!r}"
)
# ── DelegatingParser: large delta with </think> + tool calls ─────────
_DSV4_FULL_VOCAB = {
DSML_THINK_START: 128821,
DSML_THINK_END: 128822,
DSML_TOOL_START: 128823,
DSML_TOOL_END: 128824,
}
class _DeepSeekV4Delegating(DelegatingParser):
reasoning_parser_cls = DeepSeekV4ParserReasoningAdapter
tool_parser_cls = DeepSeekV4ParserToolAdapter
def _dsv4_tokens(
reasoning: str,
tool_name: str,
params: list[tuple[str, str, str]],
) -> list[tuple[int, str]]:
"""Build a token sequence: reasoning + </think> + DSML tool block."""
tokens: list[tuple[int, str]] = []
tid = 100
for word in reasoning.split(" "):
prefix = " " if tokens else ""
tokens.append((tid, prefix + word))
tid += 1
tokens.append((_DSV4_FULL_VOCAB[DSML_THINK_END], DSML_THINK_END))
tokens.append((tid, "\n\n"))
tid += 1
tokens.append((_DSV4_FULL_VOCAB[DSML_TOOL_START], DSML_TOOL_START))
tokens.append((tid, "\n"))
tid += 1
invoke_prefix_text = f"{DSML_INVOKE_PREFIX}{tool_name}{DSML_INVOKE_NAME_END}"
tokens.append((tid, invoke_prefix_text))
tid += 1
tokens.append((tid, "\n"))
tid += 1
for name, is_str, value in params:
param_text = _param(name, is_str, value)
tokens.append((tid, param_text))
tid += 1
tokens.append((tid, "\n"))
tid += 1
tokens.append((tid, DSML_INVOKE_END))
tid += 1
tokens.append((tid, "\n"))
tid += 1
tokens.append((_DSV4_FULL_VOCAB[DSML_TOOL_END], DSML_TOOL_END))
return tokens
class TestDelegatingParserLargeDelta:
"""Regression: tool calls lost when </think> + DSML arrive in same delta.
The DelegatingParser used by the serving layer splits reasoning and
tool parsing across two separate engine instances. When </think> and
the entire DSML tool block arrive in a single large streaming delta,
the content transfer from reasoning adapter to tool adapter must
preserve the tool call text.
"""
@pytest.fixture
def dsv4_tokens(self):
return _dsv4_tokens(
reasoning="The user wants the current weather in Berlin.",
tool_name="get_weather",
params=[
("location", "true", "Berlin"),
("units", "true", "celsius"),
],
)
@pytest.fixture
def dsv4_tokenizer(self, dsv4_tokens):
return MockTokenizer(
vocab=dict(_DSV4_FULL_VOCAB),
tokens=dsv4_tokens,
)
@pytest.mark.parametrize(
"chunk_size",
[1, 2, 3, 5, None],
ids=lambda c: f"chunk={c}",
)
def test_tool_calls_extracted_at_all_chunk_sizes(
self, dsv4_tokenizer, dsv4_tokens, chunk_size
):
parser = _DeepSeekV4Delegating(
dsv4_tokenizer,
chat_template_kwargs={"thinking": True},
)
deltas = replay_streaming(
parser,
dsv4_tokens,
chunk_size=chunk_size,
finished_on_last=True,
tools=DUMMY_TOOLS,
)
output = collect_output(deltas)
assert "The user wants" in output.reasoning
assert len(output.tool_calls) == 1, (
f"Expected 1 tool call but got {len(output.tool_calls)}; "
f"reasoning={output.reasoning!r}, content={output.content!r}"
)
assert output.tool_calls[0]["name"] == "get_weather"
args = json.loads(output.tool_calls[0]["arguments"])
assert args == {"location": "Berlin", "units": "celsius"}
def test_eos_drop_token_does_not_swallow_tool_calls(self):
"""Tool calls must survive when an EOS DROP token's ID is in
delta_token_ids but its text is absent from delta_text.
At large stream_interval the EOS token ID arrives in the same
delta as </think> + tool calls but the detokenizer strips the
EOS text. The scanner's _rebuild_from_anchors defers all text
after </think> when it can't find the EOS anchor text. The
reasoning adapter's finish_streaming must flush deferred text
as content (with skip_tool_parsing), not as tool calls.
"""
eos_text = "<end▁of▁sentence>"
eos_id = 128801
vocab = {
DSML_THINK_START: 128821,
DSML_THINK_END: 128822,
eos_text: eos_id,
}
reasoning = "The user wants weather."
tool_block = (
"\n\n"
+ DSML_TOOL_START
+ "\n"
+ DSML_INVOKE_PREFIX
+ "get_weather"
+ DSML_INVOKE_NAME_END
+ "\n"
+ _param("location", "true", "Berlin")
+ "\n"
+ DSML_INVOKE_END
+ "\n"
+ DSML_TOOL_END
)
# delta_text does NOT include EOS text (detokenizer strips it)
full_text = reasoning + DSML_THINK_END + tool_block
# Build token list: word-split reasoning, then special tokens,
# then word-split tool block content, then EOS.
# EOS ID is present but its text is NOT in delta_text.
tokens: list[tuple[int, str]] = []
tid = 100
for word in reasoning.split(" "):
pfx = " " if tokens else ""
tokens.append((tid, pfx + word))
tid += 1
tokens.append((128822, DSML_THINK_END))
for ch in tool_block:
tokens.append((tid, ch))
tid += 1
tokens.append((eos_id, eos_text))
all_ids = [t[0] for t in tokens]
tokenizer = MockTokenizer(vocab=vocab, tokens=tokens)
request = _test_request(tools=DUMMY_TOOLS)
# All-in-one delta: EOS ID in token_ids but text NOT in
# delta_text (detokenizer strips EOS). This is the scenario
# at large stream_interval.
parser = _DeepSeekV4Delegating(
tokenizer,
chat_template_kwargs={"thinking": True},
)
deltas = [
parser.parse_delta(
full_text,
all_ids,
request,
prompt_token_ids=[],
finished=True,
)
]
output = collect_output(deltas)
assert "The user wants" in output.reasoning
assert len(output.tool_calls) == 1, (
f"Expected 1 tool call but got {len(output.tool_calls)}; "
f"reasoning={output.reasoning!r}, content={output.content!r}"
)
assert output.tool_calls[0]["name"] == "get_weather"
args = json.loads(output.tool_calls[0]["arguments"])
assert args == {"location": "Berlin"}

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