Compare commits

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

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

Signed-off-by: dependabot[bot] <support@github.com>
2026-06-02 18:31:03 +00:00
dependabot[bot]GitHubdependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
b8b49e2395 Bump actions/github-script from 8.0.0 to 9.0.0 (#39667)
Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-06-02 11:26:57 -07:00
da107a59e5 [MRV2] Also enable MRV2 for Llama and Mistral dense models (#43458)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: yewentao256 <zhyanwentao@126.com>
2026-06-02 11:18:46 -07:00
ChaunceyGitHubAleksandar YanakievAng Kah Min, Kelvin <syraxius@hotmail.com>
ed9a7526b6 [Anthropic] Support system role messages inside messages array (#44283)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Co-authored-by: Aleksandar Yanakiev <alexander.yanakiev@discretestack.com>
Co-authored-by: Ang Kah Min, Kelvin <syraxius@hotmail.com>
2026-06-02 18:13:54 +00:00
2427094152 [Feature] Support EPLB for DeepSeek v4 Mega Moe (#43339)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: Wei Zhao (Engrg-Hardware 1) <weizha@login-lyris01.lyris.clusters.nvidia.com>
2026-06-02 10:56:44 -07:00
Kartavya sonarandGitHub fe32e7830b [Bugfix] flashinfer: fail fast when --kv-cache-dtype nvfp4 used on unsupported arch (#43669)
Signed-off-by: Kartavya Sonar <sonarkartavya@gmail.com>
2026-06-02 10:50:00 -07:00
afcb580715 [BugFix] Fix Humming MoE deploy error (#43100)
Signed-off-by: Alireza Dadgarnia <dadgarnia@Alirezas-MacBook-Pro-2.local>
Signed-off-by: Alireza Dadgarnia <49554709+adotdad@users.noreply.github.com>
Co-authored-by: Alireza Dadgarnia <dadgarnia@Alirezas-MacBook-Pro-2.local>
Co-authored-by: Jinzhen Lin <linjinzhen@hotmail.com>
2026-06-02 09:32:50 -07:00
3f3e2702c2 [XPU] Enable rms_norm/act quant fusions (#43963)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-02 16:14:41 +00:00
Flora FengandGitHub 478b49ddec [Refactor] Remove dead code from parser infrastructure (#44279)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-02 12:08:27 -04:00
Nick HillandGitHub cab5c9a2a9 [Core] Move max_concurrent_batches to VllmConfig (#44274)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-06-02 08:57:25 -07:00
Brian DellabettaandGitHub 774e552397 [compressed-tensors] Asymmetric support for MoE WNA16 marlin (#44025)
Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>
2026-06-02 08:51:45 -07:00
XiaoZandGitHub 53fa09d085 [Misc] Support local image encoding in benchmarks (#43843)
Signed-off-by: xiaoz <Sukra1@outlook.com>
2026-06-02 15:15:06 +00:00
Chris LeonardandGitHub 4d93bc35c9 Migrate header files to torch stable abi (#44013) 2026-06-02 08:09:52 -07:00
Bugen ZhaoandGitHub 586201ebdc [Rust Frontend] Cover different thinking modes in roundtrip tests (#44320)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-02 07:51:25 -07:00
pschlan-amdandGitHub 88f172188b [ROCm] Fix AITER RMSNormQuantFusion for Kimi-Linear (#44308)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
2026-06-02 14:50:21 +00:00
Bugen ZhaoandGitHub 880fc032f4 [Rust Frontend] Support recursive tool parameter conversion (#44299)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-02 07:45:35 -07:00
6314de8bad [XPU] [Bug] remove xpuw4a16 output size check (#44168)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-02 22:26:20 +08:00
IdoAtadTDandGitHub c91a87f01a [BugFix] [GDN] Read linear_key_head_dim from hf_text_config for multimodal models (#43978)
Signed-off-by: IdoAtadTD <ido.atad@twodelta.com>
2026-06-02 17:17:55 +03:00
Matthew BonanniandGitHub ea0d045a05 [FlashAttention] Sync FA with upstream (#44065)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-06-02 07:15:37 -07:00
0bdfd5eb84 [Bugfix] Vendor MiniCPMV/MiniCPMO processors to unblock Transformers v5 (#44282)
Signed-off-by: guanwei-wu <b08901019@ntu.edu.tw>
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: guanwei-wu <b08901019@ntu.edu.tw>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-02 07:14:38 -07:00
0cbc48c4f9 Support ModelOpt MXFP8 non-gated MoE (#42958)
Signed-off-by: tbarnatan <tbarnatan@nvidia.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-06-02 13:56:03 +00:00
2fd0e52252 [Bugfix] Fix Gemma4 startup crash with recent transformers multimodal processor (#44232)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
2026-06-02 13:42:40 +00:00
654bd2bca4 [Bugfix] Sync block_size from EngineCore to frontend for hybrid Mamba… (#42967)
Signed-off-by: Amit Gruner <agruner@crusoe.ai>
Co-authored-by: Amit Gruner <agruner@crusoe.ai>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
2026-06-02 13:41:00 +00:00
wang.yuqiandGitHub b623f7ea95 [Frontend] Consolidate dev entrypoints. (#44170)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-02 06:30:21 -07:00
Shreyas KulkarniandGitHub 0eeba5eec1 Fix DFlash prefix cache corruption due to missing lookahead block (#42971)
Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>
2026-06-02 12:06:33 +00:00
f69ede495b [XPU][Mamba] Triton-based selective scan forward op for XPU (#43421)
Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-02 03:50:26 -07:00
Ronen SchafferandGitHub 2a2b5ca791 [KV Offload] Add on_schedule_end() hook to separate step lifecycle from event draining (#44206)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
2026-06-02 13:42:52 +03:00
Rukhaiya2004GitHubRukhaiyaAkash kaothalkarLi, Jiang <jiang1.li@intel.com>
689b0eeb9e [HARDWARE][POWER] Enable SHM communicator support for PowerPC (#43754)
Signed-off-by: Rukhaiya <rukhaiya@c643n08aix1-lp1.pok.stglabs.ibm.com>
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com>
Co-authored-by: Rukhaiya <rukhaiya@c643n08aix1-lp1.pok.stglabs.ibm.com>
Co-authored-by: Akash kaothalkar <61960177+Akashcodes732@users.noreply.github.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-06-02 18:06:32 +08:00
Isotr0pyandGitHub f8e9c56d15 [Multimodal] Automatically select registered video loader for VLM (#44126)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-02 09:09:47 +00:00
albertoandGitHub e30313220c [Parser] Migrate ResponsesParser to unified Parser interface (#42977)
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>
2026-06-02 08:50:05 +00:00
omerpaz95GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d247a9dc13 [EC Connector] Non blocking EC Connector lookup (#41627)
Signed-off-by: omerpaz95 <omerpaz95@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-02 08:48:25 +00:00
Yifan QiaoandGitHub 7c37096620 [Core][Refactor]: thread scheduler_block_size into KVCacheManager and KVCacheCoordinator (#44165)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-06-02 01:14:44 -07:00
Maria GuevaraandGitHub b817b23f7b [Rust Frontend] add --enable-request-id-headers flag support. (#43883)
Signed-off-by: Maria Guevara <kawaiiplush14@gmail.com>
2026-06-02 16:08:37 +08:00
Ronen SchafferandGitHub 93da882e73 [kv_offload] Add @override decorators to subclass method implementations (#44177)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
2026-06-02 08:07:47 +00:00
Fadi ArafehGitHubLi, Jiang <jiang1.li@intel.com>
0b25cf4419 [CPU][Perf] Enable fused kernels for GDN's gated delta rules (#43534)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-06-02 08:00:48 +00:00
Jiangyun ZhuandGitHub dcdfe66bfa [Perf] use triton moe backend on hopper by default (#44220)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-06-02 15:52:30 +08:00
Flora FengandGitHub 68dafcca75 [Refactor] Unify reasoning + tool-call parsing behind Parser.parse() (#44267)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-02 15:11:42 +08:00
zhrrrandGitHub 1edfd09ffd [Model Runner V2] Use actual batch max_seq_len for attn metadata (#43991)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
2026-06-02 06:07:56 +00:00
zhrrrandGitHub 8a9eb40808 [Model Runner V2] Support zeroing freshly allocated KV blocks for hybrid + fp8 KVCache (#43990)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
2026-06-02 05:56:53 +00:00
f91fb2fcf3 [Bugfix] Convert Gemma4-MM ViT linear layers to vllm native impl (#43798)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Co-authored-by: ZiTian Zhao <zitian.zhao@tencentmusic.com>
Co-authored-by: B-201 <Joy25810@foxmail.com>
2026-06-01 21:41:16 -07:00
JooHo LeeandGitHub a045c7425f [MM][CG] Profile encoder CUDA graph pool memory (#41714)
Signed-off-by: JooHo Lee <jooho414@gmail.com>
2026-06-02 12:27:34 +08:00
a3a5a5ece5 [XPU][Bugfix] Fix per_token_group_fp8_quant missing dummy args on XPU (#43930)
Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-02 03:09:21 +00:00
Or OzeriandGitHub 480fadab1b [BugFix][kv_offload]: Prevent offloading stale sliding window blocks (#42959)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
2026-06-02 05:59:48 +03:00
279d25f5cb [BugFix] Fix TypeError in MiniCPM-O audio feature unpadding (#38053)
Signed-off-by: Krishna Chaitanya Balusu <krishnabkc15@gmail.com>
Signed-off-by: wjinxu <1299461899@qq.com>
Signed-off-by: Kc Balusu <kcbalusu@users.noreply.github.com>
Co-authored-by: wjinxu <1299461899@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Co-authored-by: Kc Balusu <kcbalusu@users.noreply.github.com>
2026-06-01 19:57:28 -07:00
Andreas KaratzasandGitHub 54d0c36fff [CI] Stabilize OpenAI schema fuzzing for malformed structural tags (#44131)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-06-01 19:56:15 -07:00
Flora FengandGitHub 9affc17a05 [Refactor] Move unstreamed tool-arg flush from serving layer to parser (#44017)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-02 10:37:43 +08:00
AlecandGitHub 816cc73a9b [Bugfix][CI] Normalize NIXL connector CUDA wheel installs (#44266)
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
2026-06-01 19:34:05 -07:00
Micah WilliamsonandGitHub 2588ec4f0a [ROCm] Upgrade AITER to v0.1.13.post1 (#44265)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-02 01:48:59 +00:00
d68f0b220e [Bugfix][Mooncake] Release GPU pin on failed store in MooncakeStoreConnector (#43742)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-06-01 18:29:18 -07:00
Woosuk KwonandGitHub 517e74a964 [DSV4] Refactor RoPE initialization (#44262)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-02 01:26:58 +00:00
JartXandGitHub 48c0d13e65 [ROCm][CI] Skip unbacked dynamic shapes tests on PyTorch < 2.11 (#44256)
Signed-off-by: JartX <sagformas@epdcenter.es>
2026-06-01 19:09:01 -05:00
Woosuk KwonandGitHub 8c3cc98cff [DSV4] Remove unncessary classes & functions (#44246)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-01 14:43:00 -07:00
Nick HillandGitHub e4cbc4385d [Test][BugFix] Fix double-BOS in PD+specdec acceptance test (#44234)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-06-01 14:31:12 -07:00
Nick HillandGitHub 6f8b40a23f [BugFix][CI] Fix added _has_module tests (#44248)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-06-01 14:23:12 -07:00
266b9d9c64 [Frontend][Core] Add sparse NCCL weight transfer support for in-place updates (#40096)
Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-06-01 15:37:30 -04:00
182c67daf1 [Rust Frontend] Support streaming generate endpoint (#43779)
Signed-off-by: xunzhuo <xunzhuo@vllm-semantic-router.ai>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
2026-06-01 19:30:55 +00:00
fd9e91d7e4 [ROCm][CI] Fix and stabilize EAGLE3 acceptance tests (#41294)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Co-authored-by: Micah Williamson <micah.williamson@amd.com>
2026-06-01 12:40:01 -05:00
Yongye ZhuandGitHub 035733515f [Kernel][DSv4] Optimize sparse FP8 compressor kernels (#44161)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-02 00:18:32 +08:00
023808c23d [Feature] Add support for JetBrains' Mellum v2 code generation model (#43992)
Signed-off-by: Madeesh Kannan <madeeswaran.kannan@jetbrains.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-06-01 10:11:35 -04:00
Wentao YeGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
985c97a6a8 [Perf] Optimize cutlass fp8 scaled mm bypassing padding, 20% kernel performance improvement (#43706)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-01 09:05:21 -04:00
Chaojun ZhangandGitHub bd0aecdc08 [XPU][CI] Fix test_audio_in_video flake by using module-scoped server fixture (#44146)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
2026-06-01 11:21:36 +00:00
8796838910 [Bugfix] fix wrong partial_rotary_factor calculation for bailing_moe model. (#43770)
Signed-off-by: zzt <zengzetang.zzt@antgroup.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
2026-06-01 02:42:49 -07:00
de21863419 [Rust Frontend] Add InternLM2 tool parser (#43481)
Signed-off-by: Will.hou <1205157517@qq.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
2026-06-01 08:58:46 +00:00
wang.yuqiandGitHub 0910f7e0e1 [Frontend] Resettle generative scoring entrypoint. (#44153)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-06-01 07:54:59 +00:00
UranusandGitHub 1f6048abe5 fix: glm5.1 pp model loading (#42944)
Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
2026-06-01 15:14:47 +08:00
wcyGitHubLi, Jiang <jiang1.li@intel.com>
98f1279815 [CPU][RISC-V] Add missing RVV cpu_types helpers for WNA16 (#42730)
Signed-off-by: wcy <233313160abc@gmail.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
2026-06-01 14:56:41 +08:00
Isotr0pyandGitHub 1fd8bd02a4 [Docs] Replace broken video url in examples (#44159)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-01 06:01:10 +00:00
29d69332aa [BugFix] Fix _has_module to verify native deps via trial import (#44035)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: esmeetu <jasonailu87@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-05-31 22:06:33 -07:00
Lucas WilkinsonandGitHub 4721bb3aa4 [MRV2] Remove Eagle's dedicated CUDA graph pool (#44078)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
2026-05-31 22:00:33 -07:00
Umut PolatandGitHub f46e6be169 [Misc] Use VLLMValidationError consistently in chat completion and completion protocol validators (#36254)
Signed-off-by: umut-polat <52835619+umut-polat@users.noreply.github.com>
2026-06-01 04:04:11 +00:00
8b8546da1c docs: fix MLA attention docstring examples (#44118)
Co-authored-by: nightcityblade <nightcityblade@gmail.com>
2026-05-31 12:28:38 -07:00
257 changed files with 10369 additions and 2849 deletions
+3 -1
View File
@@ -16,6 +16,7 @@ steps:
- tests/kernels/test_onednn.py
- tests/kernels/test_awq_int4_to_int8.py
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -24,7 +25,8 @@ steps:
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
- label: CPU-Compatibility Tests
depends_on: []
@@ -37,7 +37,8 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/test_onednn.py
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_moe.py -k test_cpu_fused_moe_basic"
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
+39
View File
@@ -0,0 +1,39 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
set -euo pipefail
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
NIXL_METADATA=$(python3 - <<'PY'
import importlib.metadata as metadata
import torch
cuda_version = torch.version.cuda
if cuda_version is None:
raise SystemExit("torch.version.cuda is not set")
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
PY
)
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
# matching this CI image so nixl_ep_cpp links against the available libcudart.
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
python3 - <<'PY'
import importlib.metadata as metadata
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
try:
version = metadata.version(package_name)
except metadata.PackageNotFoundError:
version = "not installed"
print(f"{package_name}: {version}")
PY
+15 -13
View File
@@ -1238,14 +1238,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -1275,10 +1272,14 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
@@ -1368,7 +1369,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- label: OpenAI API correctness # TBD
timeout_in_minutes: 180
@@ -2745,14 +2746,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -2782,10 +2780,14 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
+9 -9
View File
@@ -11,7 +11,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+8 -8
View File
@@ -11,7 +11,7 @@ steps:
- tests/entrypoints/
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
@@ -60,9 +60,13 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
mirror:
amd:
device: mi325_1
@@ -98,14 +102,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
- tests/entrypoints/serve
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
mirror:
amd:
device: mi325_1
@@ -153,6 +154,5 @@ steps:
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
commands: # LMEval
- pytest -s entrypoints/openai/correctness/
+1 -1
View File
@@ -86,7 +86,7 @@ steps:
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# split the test to avoid interference
- pytest -v -s -m 'not cpu_test' v1/core
+4 -4
View File
@@ -45,19 +45,19 @@ steps:
- vllm/entrypoints/serve/
- vllm/v1/engine/
- tests/utils.py
# - tests/entrypoints/rpc/test_collective_rpc.py
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
- tests/entrypoints/serve/disagg/test_serving_tokens.py
- tests/entrypoints/serve/instrumentator/test_basic.py
- tests/entrypoints/serve/instrumentator/test_metrics.py
# - tests/entrypoints/serve/instrumentator/test_sleep.py
# - tests/entrypoints/serve/dev/test_sleep.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s entrypoints/rpc/test_collective_rpc.py
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
# - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
- label: Rust Frontend Core Correctness
timeout_in_minutes: 30
+1 -1
View File
@@ -10,7 +10,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Add label
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
github.rest.issues.addLabels({
+3 -3
View File
@@ -14,7 +14,7 @@ jobs:
steps:
- name: Label issues based on keywords
id: label-step
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
// Configuration: Add new labels and keywords here
@@ -315,7 +315,7 @@ jobs:
- name: CC users for labeled issues
if: steps.label-step.outputs.labels_added != '[]'
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
// Configuration: Map labels to GitHub users to CC
@@ -392,7 +392,7 @@ jobs:
- name: Request missing ROCm info from issue author
if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const body = (context.payload.issue.body || '').toLowerCase();
+2 -2
View File
@@ -12,7 +12,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Update PR description
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { owner, repo } = context.repo;
@@ -55,7 +55,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Post welcome comment for first-time contributors
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { owner, repo } = context.repo;
+2 -2
View File
@@ -20,7 +20,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
with:
script: |
const { data: pr } = await github.rest.pulls.get({
@@ -49,7 +49,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: "3.12"
- run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
+1 -1
View File
@@ -21,7 +21,7 @@ repos:
rev: v21.1.2
hooks:
- id: clang-format
exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
types_or: [c++, cuda]
args: [--style=file, --verbose]
- repo: https://github.com/DavidAnson/markdownlint-cli2
+31 -1
View File
@@ -369,6 +369,18 @@ else()
add_compile_definitions(-DVLLM_NUMA_DISABLED)
endif()
# check if the pytorch wheel ships libopenblas.so.
set(VLLM_OPENBLAS_LIB "")
if (NOT ENABLE_X86_ISA)
file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
"${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
# Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
if (_VLLM_TORCH_OPENBLAS_LIBS)
list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
endif()
endif()
#
# Generate CPU attention dispatch header
#
@@ -387,6 +399,7 @@ endif()
#
set(VLLM_EXT_SRC
"csrc/cpu/activation.cpp"
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/layernorm.cpp"
@@ -396,6 +409,13 @@ set(VLLM_EXT_SRC
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
set(VLLM_EXT_SRC
"csrc/cpu/cpu_wna16.cpp"
${VLLM_EXT_SRC})
endif()
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
@@ -403,6 +423,12 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
${VLLM_EXT_SRC})
endif()
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
${VLLM_EXT_SRC})
endif()
if(USE_ONEDNN)
set(VLLM_EXT_SRC
"csrc/cpu/dnnl_kernels.cpp"
@@ -411,7 +437,6 @@ endif()
if (ENABLE_X86_ISA)
set(VLLM_EXT_SRC_SGL
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/sgl-kernels/conv.cpp"
"csrc/cpu/sgl-kernels/gemm.cpp"
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
@@ -423,6 +448,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
set(VLLM_EXT_SRC_AVX512
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/shm.cpp"
"csrc/cpu/cpu_wna16.cpp"
"csrc/cpu/cpu_fused_moe.cpp"
@@ -439,6 +465,7 @@ if (ENABLE_X86_ISA)
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
set(VLLM_EXT_SRC_AVX2
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
@@ -512,6 +539,9 @@ else()
USE_SABI 3
WITH_SOABI
)
if (VLLM_OPENBLAS_LIB)
target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
endif()
endif()
message(STATUS "Enabling C extension.")
@@ -31,7 +31,7 @@ endif()
if(VLLM_FLASH_ATTN_SRC_DIR)
FetchContent_Declare(
vllm-flash-attn SOURCE_DIR
vllm-flash-attn SOURCE_DIR
${VLLM_FLASH_ATTN_SRC_DIR}
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
)
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+51
View File
@@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
explicit FP16Vec16(const c10::Half v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
@@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
explicit BF16Vec16(const c10::BFloat16 v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
explicit BF16Vec16(const c10::BFloat16 v)
: reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
VEC_ELEM_NUM)) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
@@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
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);
}
explicit FP32Vec16(const FP16Vec16& v);
#ifdef __riscv_zvfbfmin
@@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
// FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
// load_b_pair_vec template to compile on all platforms.
explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
template <typename VecT>
static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
vec0.save(values0);
vec1.save(values1);
auto* packed = reinterpret_cast<uint32_t*>(ptr);
for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
packed[i] = static_cast<uint32_t>(values0[i]) |
(static_cast<uint32_t>(values1[i]) << 16);
}
}
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
#ifdef __riscv_zvfbfmin
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+106 -1
View File
@@ -89,6 +89,35 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
ss16x8x2_t reg;
explicit FP16Vec16(const void* ptr) {
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
explicit FP16Vec16(const FP32Vec16&);
void save(void* ptr) const {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst(reg.val[1], 16, (signed short*)ptr);
}
void save(void* ptr, int elem_num) const {
int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
if (num <= 8) {
vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
} else {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
@@ -100,6 +129,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -379,6 +410,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = vec_xl(48, ptr);
}
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
explicit FP32Vec16(const FP32Vec16& data) {
@@ -402,6 +435,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = data.reg.val[1];
}
explicit FP32Vec16(const FP16Vec16& v);
explicit FP32Vec16(const BF16Vec16& v) {
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -735,6 +769,40 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
#endif
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
alignas(16) float temp_fp32[16];
alignas(16) c10::Half temp_fp16[16];
vec_xst(v.reg.val[0], 0, temp_fp32);
vec_xst(v.reg.val[1], 16, temp_fp32);
vec_xst(v.reg.val[2], 32, temp_fp32);
vec_xst(v.reg.val[3], 48, temp_fp32);
for (int i = 0; i < 16; i++) {
temp_fp16[i] = c10::Half(temp_fp32[i]);
}
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
alignas(16) c10::Half temp_fp16[16];
alignas(16) float temp_fp32[16];
vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
for (int i = 0; i < 16; i++) {
temp_fp32[i] = float(temp_fp16[i]);
}
reg.val[0] = vec_xl(0, temp_fp32);
reg.val[1] = vec_xl(16, temp_fp32);
reg.val[2] = vec_xl(32, temp_fp32);
reg.val[3] = vec_xl(48, temp_fp32);
}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
#ifdef _ARCH_PWR10
__vector signed short ret[4];
@@ -794,6 +862,43 @@ inline void prefetch(const void* addr) {
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
}
}; // namespace vec_op
struct INT8Vec64 {
__vector signed char data[4];
INT8Vec64() = default;
explicit INT8Vec64(const int8_t* ptr) {
data[0] = vec_xl(0, ptr);
data[1] = vec_xl(16, ptr);
data[2] = vec_xl(32, ptr);
data[3] = vec_xl(48, ptr);
}
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
void save(int8_t* ptr) const {
vec_xst(data[0], 0, ptr);
vec_xst(data[1], 16, ptr);
vec_xst(data[2], 32, ptr);
vec_xst(data[3], 48, ptr);
}
void save(int8_t* ptr, int elem_num) const {
if (elem_num <= 0) return;
int full_vecs = elem_num / 16;
for (int i = 0; i < full_vecs && i < 4; i++) {
vec_xst(data[i], i * 16, ptr);
}
int remaining = elem_num % 16;
if (remaining > 0 && full_vecs < 4) {
vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
}
}
void nt_save(int8_t* ptr) const { save(ptr); }
};
} // namespace vec_op
#endif
+82
View File
@@ -0,0 +1,82 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include <ATen/native/CPUBlas.h>
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
#if defined(VLLM_HAS_OPENBLAS)
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const at::BFloat16* a, int* lda,
const at::BFloat16* b, int* ldb, float* beta, float* c,
int* ldc);
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const float* a, int* lda, const float* b,
int* ldb, float* beta, float* c, int* ldc);
inline char blas_transpose(at::native::TransposeType trans) {
switch (trans) {
case at::native::TransposeType::NoTranspose:
return 'n';
case at::native::TransposeType::Transpose:
return 't';
case at::native::TransposeType::ConjTranspose:
return 'c';
}
return 'n';
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const at::BFloat16* a,
int64_t lda, const at::BFloat16* b, int64_t ldb,
float beta, float* c, int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const float* a, int64_t lda,
const float* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
int64_t, int64_t, int64_t, float, const at::Half*,
int64_t, const at::Half*, int64_t, float, float*,
int64_t) {
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
}
#else
template <typename scalar_t>
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const scalar_t* a, int64_t lda,
const scalar_t* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
}
#endif
+278 -141
View File
@@ -301,25 +301,42 @@ void chunk_gated_delta_rule_kernel_impl(
// attn = k_beta @ key.transpose(-1, -2)
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
// transpose and pack for key
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
curr_k_pad,
qk_head_size,
curr_k_beta,
qk_head_size,
0.0f,
curr_attn,
chunk_size);
}
// attn = attn * decay_mask
for (int64_t m = 0; m < chunk_size; m++) {
at::vec::map2<float>(
@@ -413,25 +430,42 @@ void chunk_gated_delta_rule_kernel_impl(
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
// pack for value
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
curr_v_beta,
v_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
curr_value,
v_head_size);
}
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
for (int64_t j = 0; j < chunk_size; j++) {
int64_t i = 0;
@@ -445,25 +479,42 @@ void chunk_gated_delta_rule_kernel_impl(
}
}
// pack for k_beta_g
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
qk_head_size,
chunk_size,
chunk_size,
1.0f,
k_beta_g,
qk_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
k_cumdecay,
qk_head_size);
}
for (int i = 0; i < chunk_size; i++) {
at::vec::map<scalar_t>(
[](fVec x) { return x; },
@@ -551,25 +602,42 @@ void chunk_gated_delta_rule_kernel_impl(
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
// k_transpose_i = k_i.transpose(-1, -2)
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
k_i,
qk_head_size,
q_i,
qk_head_size,
0.0f,
attn_i,
chunk_size);
}
// attn_i = attn_i * decay_mask_i
for (int64_t m = 0; m < chunk_size; m++) {
auto attn_i_m = attn_i + m * chunk_size;
@@ -609,28 +677,45 @@ void chunk_gated_delta_rule_kernel_impl(
}
// pack for curr_last_recurrent_state
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
k_cumdecay_i_reduced,
qk_head_size,
0.0f,
v_prime,
v_head_size);
}
// v_new = v_prime = v_i - v_prime
// v_i: [chunk_size, EV]
@@ -663,41 +748,75 @@ void chunk_gated_delta_rule_kernel_impl(
}
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
qg,
qk_head_size,
0.0f,
attn_inter,
v_head_size);
}
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
// pack for v_prime
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
attn_i_reduced,
chunk_size,
1.0f,
attn_inter,
v_head_size);
}
// core_attn_out[:, :, i] = attn_inter
for (int64_t m = 0; m < chunk_size; m++) {
@@ -762,17 +881,34 @@ void chunk_gated_delta_rule_kernel_impl(
/* ld_dst */ chunk_size);
// kgv = kg.transpose(-1, -2) @ v_new
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
qk_head_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
kg_transpose,
chunk_size,
0.0f,
kgv,
v_head_size);
}
// last_recurrent_state = 1) + 2)
for (int64_t m = 0; m < qk_head_size; m++) {
at::vec::map2<float>(
@@ -921,7 +1057,8 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
// See: https://github.com/sgl-project/sglang/pull/26634
float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
fVec beta_vec = fVec(beta_val);
int64_t dvi = 0;
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
+18 -7
View File
@@ -4,9 +4,12 @@
// clang-format off
#pragma once
#include <ATen/native/CPUBlas.h>
#include "common.h"
#include "blas_gemm.h"
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
#define CPU_CAPABILITY_AVX512
#endif
// amx-bf16
#define TILE_M 16
@@ -21,31 +24,39 @@ constexpr int block_size_n() {
return 2 * TILE_N;
}
constexpr bool brgemm_supported() {
#if defined(CPU_CAPABILITY_AVX512)
return true;
#else
return false;
#endif
}
// define threshold using brgemm (intel AMX)
template <typename T>
inline bool can_use_brgemm(int M);
template <>
inline bool can_use_brgemm<at::BFloat16>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<at::Half>(int M) {
return true;
return brgemm_supported();
}
// this requires PyTorch 2.7 or above
template <>
inline bool can_use_brgemm<int8_t>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<uint8_t>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
template <>
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
return M > 4;
return brgemm_supported() && M > 4;
}
// work around compiler internal error
+2
View File
@@ -11,7 +11,9 @@
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
#if defined(CPU_CAPABILITY_AVX512)
#include <immintrin.h>
#endif
namespace {
using namespace at::vec;
+10 -8
View File
@@ -5,7 +5,7 @@
#include <sys/stat.h>
#include <unistd.h>
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
#include <atomic>
#endif
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
};
struct ThreadSHMContext {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
// memory model is weaker on AArch64, so we use atomic variables for
// consumer (load-acquire) and producer (store-release) to make sure
// that a stamp cannot be ready before the corresponding data is ready.
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
TORCH_CHECK((size_t)this % 64 == 0);
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
}
char get_curr_stamp(int idx) const {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _curr_thread_stamp[idx];
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
}
char get_ready_stamp(int idx) const {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _ready_thread_stamp[idx];
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
}
void next_stamp() {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
1, std::memory_order_release);
#else
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
}
void commit_ready_stamp() {
#ifdef __aarch64__
#if defined(__aarch64__) || defined(__powerpc64__)
_ready_thread_stamp[local_stamp_buffer_idx].store(
_curr_thread_stamp[local_stamp_buffer_idx].load(
std::memory_order_relaxed),
@@ -186,8 +186,10 @@ struct ThreadSHMContext {
break;
}
++_spinning_count;
#ifdef __aarch64__
#if defined(__aarch64__)
__asm__ __volatile__("yield");
#elif defined(__powerpc64__)
__asm__ __volatile__("or 1,1,1");
#else
_mm_pause();
#endif // __aarch64__
+22 -21
View File
@@ -378,7 +378,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
// SHM CCL
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
defined(__powerpc64__)
ops.def(
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
"int",
@@ -447,6 +448,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"bool is_vnni) -> Tensor");
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// Adapted from sglang: GDN kernels
ops.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -470,25 +490,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"-> (Tensor, Tensor)");
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// CPU attention kernels
ops.def(
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
@@ -518,7 +519,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
// WNA16
#if defined(__AVX512F__)
#if defined(__AVX512F__) || defined(__riscv_v)
ops.def(
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
@@ -2,7 +2,7 @@
#include <torch/csrc/stable/tensor.h>
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
#include "broadcast_load_epilogue_c2x.hpp"
/*
This file defines custom epilogues for fusing channel scales, token scales,
@@ -20,7 +20,7 @@
#include "torch_utils.h"
#include "../async_util.cuh"
#include "async_util.cuh"
#include "../cuda_compat.h"
#include "../type_convert.cuh"
#include "dispatch_utils.h"
@@ -25,7 +25,7 @@
#include <cuda_fp8.h>
#include "cuda_utils.h"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
@@ -34,7 +34,7 @@
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
namespace vllm {
@@ -26,7 +26,7 @@
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
namespace vllm {
@@ -26,7 +26,7 @@
#include "../../cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "launch_bounds_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
@@ -7,14 +7,11 @@
#include <torch/csrc/stable/ops.h>
// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
// moved to libtorch_stable) to avoid unnecessary reformatting that would break
// git rename detection and pollute blame history.
#include "../../../quantization/gguf/ggml-common.h"
#include "../../../quantization/gguf/vecdotq.cuh"
#include "../../../quantization/gguf/dequantize.cuh"
#include "../../../quantization/gguf/mmvq.cuh"
#include "../../../quantization/gguf/mmq.cuh"
#include "ggml-common.h"
#include "vecdotq.cuh"
#include "dequantize.cuh"
#include "mmvq.cuh"
#include "mmq.cuh"
#include "moe.cuh"
#include "moe_vec.cuh"
+1 -1
View File
@@ -7,7 +7,7 @@
#include "torch_utils.h"
#ifndef USE_ROCM
#include "../persistent_topk.cuh"
#include "persistent_topk.cuh"
#endif
namespace {
+1 -1
View File
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.13"
ARG AITER_BRANCH="v0.1.13.post1"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
+8 -2
View File
@@ -246,6 +246,12 @@ Every image listed in "image_files" is added to the request in the listed order
The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
By default, image references are sent to the serving endpoint as provided, with local image paths converted to `file://` URLs.
If the benchmark client should load local and HTTP(S) images before sending requests, pass `--custom-ensure-client-side-data` to encode them as base64 data URLs on the client side.
Existing `data:image/...` URLs are already self-contained and are kept unchanged.
```bash
# need a model with vision capability here
vllm serve Qwen/Qwen2-VL-7B-Instruct
@@ -253,13 +259,13 @@ vllm serve Qwen/Qwen2-VL-7B-Instruct
```bash
# run benchmarking script
vllm bench serve--save-result --save-detailed \
vllm bench serve --save-result --save-detailed \
--backend openai-chat \
--model Qwen/Qwen2-VL-7B-Instruct \
--endpoint /v1/chat/completions \
--dataset-name custom_image \
--dataset-path <path-to-your-image-data-jsonl> \
--allowed-local-media-path /path/to/image/folder
--custom-ensure-client-side-data
```
Note that we need to use the `openai-chat` backend and `/v1/chat/completions` endpoint for multimodal inputs.
+1 -1
View File
@@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows:
base_url=openai_api_base,
)
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
## Use video url in the payload
chat_completion_from_url = client.chat.completions.create(
+1
View File
@@ -437,6 +437,7 @@ th {
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
+32 -2
View File
@@ -100,14 +100,44 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
- `/version` - Version information
- `/load` - Server load metrics
## Sleep Mode APIs
## Server in development mode
When using the flag VLLM_SERVER_DEV_MODE=1, you enable development endpoints.
**SECURITY WARNING: These endpoints should NOT be used in production!**
### Cache Management APIs
- `/reset_prefix_cache` - Reset prefix cache (can disrupt service)
- `/reset_mm_cache` - Reset multimodal cache (can disrupt service)
- `/reset_encoder_cache` - Reset encoder cache (can disrupt service)
### Weight Transfer APIs (RL Training)
For further details on Weight Transfer, please refer to [this page](../../training/weight_transfer/README.md).
- `/pause` - Pause generation (causes denial of service)
- `/resume` - Resume generation
- `/is_paused` - Check if generation is paused
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
- `/update_weights` - Update model weights (can alter model behavior)
- `/get_world_size` - Get distributed world size
### Collective RPC
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
### Server info
- `/server_info` - Get detailed server configuration
### Sleep Mode APIs
For further details on sleep mode, please refer to [this page](../../features/sleep_mode.md).
- `/sleep` - Put engine to sleep (causes denial of service)
- `/wake_up` - Wake engine from sleep
- `/is_sleeping` - Check if engine is sleeping
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
## Chat Template
+18 -3
View File
@@ -84,7 +84,10 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
## Receiving Weights (Inference Side)
The inference side triggers weight reception using the four-phase protocol`init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
The inference side triggers weight reception using the four-phase protocol:
`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
`finish_weight_update`. The init phase is shown [above](#initialization). The
remaining three steps are:
```python
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
@@ -108,12 +111,24 @@ llm.update_weights(
llm.finish_weight_update()
```
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
must match the order in which the trainer iterates over its parameters.
`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
`start_weight_update` must be called before `update_weights`, and
`finish_weight_update` must be called after all weight chunks have been
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
processing is applied (`True` for checkpoint-format weights, `False` for
pre-processed kernel-format weights).
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
`update_info`, but they should be wrapped in
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
directly to runtime/kernel-format parameters. The current sparse MVP requires
`TP=1` and `PP=1`.
## Examples
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
@@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None:
# Video input inference
def run_video(model: str, max_completion_tokens: int) -> None:
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
video_base64 = encode_base64_content_from_url(video_url)
## Use video url in the payload
+526
View File
@@ -0,0 +1,526 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
This example mirrors the validation story used for the sparse NCCL MVP:
both the dense update path and the sparse patch path start from the same real
checkpoint and apply the same deterministic trainer-side patch. The script then
checks that greedy 1-token outputs match between the dense and sparse vLLM
engines after the update.
The example performs the following steps:
* Load a training model on one GPU via a Ray actor.
* Launch a vLLM engine with the same real model on a second GPU.
* Verify trainer vs vLLM baseline agreement before any update.
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
outputs.
* Reset the trainer back to the baseline checkpoint.
* Apply the same deterministic patch again.
* Run a sparse NCCL update into another fresh vLLM engine and collect
post-update outputs.
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
outputs, estimated payload sizes, and trainer-side send times.
Current sparse weight transfer MVP limitations:
* ``TP=1`` and ``PP=1`` only
* sparse updates use runtime/kernel-format parameter names
* sparse updates are not composable with checkpoint-format or packed updates
This example assumes a single-node cluster with two GPUs.
"""
import hashlib
import os
import time
from collections.abc import Sequence
import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from transformers import AutoModelForCausalLM, AutoTokenizer
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.base import SparseWeightPatch
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
from vllm.utils.network_utils import get_ip, get_open_port
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
MAX_PATCH_ROWS = 32
PROMPTS = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
class MyLLM(LLM):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, *args, **kwargs):
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
super().__init__(*args, **kwargs)
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that owns the trainer-side model and deterministic patch state."""
def __init__(self, model_name: str):
self.model_name = model_name
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = None
self.patched_param = None
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
self.model_update_group = None
self.master_address = get_ip()
self.port = get_open_port()
self.reset_model()
def reset_model(self) -> None:
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.bfloat16,
).to("cuda:0")
self.model.eval()
try:
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
except AttributeError as exc:
raise RuntimeError(
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
) from exc
self.pending_sparse_patches = None
def create_rendezvous(self) -> tuple[str, int]:
self.port = get_open_port()
return self.master_address, self.port
def init_weight_transfer_group(self, world_size: int) -> None:
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
dict(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
)
)
def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
names = []
dtype_names = []
shapes = []
payload_bytes = 0
for name, param in self.model.named_parameters():
names.append(name)
dtype_names.append(str(param.dtype).split(".")[-1])
shapes.append(list(param.shape))
payload_bytes += param.numel() * param.element_size()
return (
dict(
names=names,
dtype_names=dtype_names,
shapes=shapes,
packed=packed,
),
payload_bytes,
)
@torch.inference_mode()
def generate(
self,
prompts: Sequence[str],
max_new_tokens: int = 1,
) -> list[dict[str, object]]:
generations = []
for prompt in prompts:
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
output = self.model.generate(
**model_inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
pad_token_id=self.tokenizer.pad_token_id,
)
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
generations.append(
{
"token_ids": new_token_ids,
"text": self.tokenizer.decode(
new_token_ids,
skip_special_tokens=False,
),
}
)
return generations
def prepare_sparse_patch(
self,
prompts: Sequence[str],
max_patch_rows: int = MAX_PATCH_ROWS,
) -> tuple[dict[str, object], list[int], str, int]:
selected_token_ids: list[int] = []
special_ids = set(self.tokenizer.all_special_ids)
for prompt in prompts:
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
for token_id in token_ids:
if token_id in special_ids or token_id in selected_token_ids:
continue
selected_token_ids.append(token_id)
if len(selected_token_ids) == max_patch_rows:
break
if len(selected_token_ids) == max_patch_rows:
break
if not selected_token_ids:
raise ValueError("Could not derive any non-special token IDs to patch")
vocab_size = self.patched_param.shape[0]
next_token_id = selected_token_ids[-1]
while len(selected_token_ids) < max_patch_rows:
next_token_id = (next_token_id + 1) % vocab_size
if next_token_id in special_ids or next_token_id in selected_token_ids:
continue
selected_token_ids.append(next_token_id)
row_ids = torch.tensor(
selected_token_ids,
device=self.patched_param.device,
dtype=torch.long,
)
hidden_size = self.patched_param.shape[1]
column_offsets = torch.arange(
hidden_size,
device=self.patched_param.device,
dtype=torch.long,
)
with torch.no_grad():
# Rotate the selected embedding rows instead of zeroing them so the
# patch remains deterministic while avoiding a degenerate collapse
# to the same special token after the update.
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
self.patched_param[row_ids] = replacement_rows
flat_indices = (
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
)
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
self.pending_sparse_patches = [
SparseWeightPatch(
name=PATCHED_PARAM_NAME,
indices=flat_indices.to(torch.int32),
values=flat_values,
)
]
patch_digest = hashlib.sha256(
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
+ self.pending_sparse_patches[0]
.values.detach()
.float()
.cpu()
.numpy()
.tobytes()
).hexdigest()
sparse_payload_bytes = (
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
+ flat_values.numel() * flat_values.element_size()
)
update_info = dict(
names=[PATCHED_PARAM_NAME],
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
shapes=[list(self.patched_param.shape)],
num_updates_list=[flat_indices.numel()],
update_kind="sparse_flat",
)
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
def broadcast_weights(self, packed: bool = False) -> float:
if self.model_update_group is None:
raise RuntimeError("Weight transfer group is not initialized")
trainer_args = NCCLTrainerSendWeightsArgs(
group=self.model_update_group,
packed=packed,
)
start = time.perf_counter()
NCCLWeightTransferEngine.trainer_send_weights(
iterator=self.model.named_parameters(),
trainer_args=trainer_args,
)
torch.accelerator.synchronize()
return (time.perf_counter() - start) * 1000.0
def broadcast_pending_sparse_patch(self) -> float:
if self.model_update_group is None:
raise RuntimeError("Weight transfer group is not initialized")
if self.pending_sparse_patches is None:
raise RuntimeError("Sparse patch has not been prepared")
start = time.perf_counter()
NCCLWeightTransferEngine.trainer_send_sparse_weights(
iter(self.pending_sparse_patches),
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
)
torch.accelerator.synchronize()
self.pending_sparse_patches = None
return (time.perf_counter() - start) * 1000.0
def launch_llm(
scheduling_inference: PlacementGroupSchedulingStrategy,
):
return ray.remote(
num_cpus=0,
num_gpus=0,
scheduling_strategy=scheduling_inference,
)(MyLLM).remote(
model=MODEL_NAME,
enforce_eager=True,
tensor_parallel_size=1,
distributed_executor_backend="ray",
gpu_memory_utilization=0.7,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
generations = []
for output in outputs:
generations.append(
{
"token_ids": output.outputs[0].token_ids,
"text": output.outputs[0].text,
}
)
return generations
def token_sequences_match(
left: Sequence[dict[str, object]],
right: Sequence[dict[str, object]],
) -> bool:
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
print(f"\n{label}")
print("-" * 50)
for prompt, generation in zip(prompts, generations):
print(f"Prompt: {prompt!r}")
print(f"Token IDs: {generation['token_ids']}")
print(f"Text: {generation['text']!r}")
print("-" * 50)
def run_dense_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
try:
dense_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
world_size = ray.get(llm.get_world_size.remote()) + 1
inference_init = llm.init_weight_transfer_engine.remote(
dict(
init_info=dict(
master_address=master_address,
master_port=master_port,
rank_offset=1,
world_size=world_size,
)
)
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
dense_update_info, dense_payload_bytes = ray.get(
train_model.get_dense_update_info.remote()
)
_, selected_token_ids, patch_digest, _ = ray.get(
train_model.prepare_sparse_patch.remote(PROMPTS)
)
inference_update = llm.update_weights.remote(
dict(update_info=dense_update_info)
)
dense_send_ms, _ = ray.get(
[
train_model.broadcast_weights.remote(packed=False),
inference_update,
]
)
ray.get(llm.finish_weight_update.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
dense_after = collect_vllm_generations(llm)
return {
"dense_before": dense_before,
"dense_after": dense_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"dense_payload_bytes": dense_payload_bytes,
"dense_send_ms": dense_send_ms,
}
finally:
ray.kill(llm)
def run_sparse_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
try:
sparse_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
world_size = ray.get(llm.get_world_size.remote()) + 1
inference_init = llm.init_weight_transfer_engine.remote(
dict(
init_info=dict(
master_address=master_address,
master_port=master_port,
rank_offset=1,
world_size=world_size,
)
)
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
)
inference_update = llm.update_weights.remote(
dict(update_info=sparse_update_info)
)
sparse_send_ms, _ = ray.get(
[
train_model.broadcast_pending_sparse_patch.remote(),
inference_update,
]
)
ray.get(llm.finish_weight_update.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
sparse_after = collect_vllm_generations(llm)
return {
"sparse_before": sparse_before,
"sparse_after": sparse_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"sparse_payload_bytes": sparse_payload_bytes,
"sparse_send_ms": sparse_send_ms,
}
finally:
ray.kill(llm)
ray.init()
try:
train_model = TrainModel.remote(MODEL_NAME)
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
placement_group=pg_inference,
placement_group_capture_child_tasks=True,
placement_group_bundle_index=0,
)
dense_results = run_dense_phase(train_model, scheduling_inference)
sparse_results = run_sparse_phase(train_model, scheduling_inference)
baseline_equal = token_sequences_match(
dense_results["dense_before"],
sparse_results["sparse_before"],
)
patch_selection_equal = (
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
)
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
after_equal = token_sequences_match(
dense_results["dense_after"],
sparse_results["sparse_after"],
)
any_output_changed = any(
before["token_ids"] != after["token_ids"]
for before, after in zip(
dense_results["dense_before"],
dense_results["dense_after"],
)
)
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
print_generations(
"Dense baseline outputs",
PROMPTS,
dense_results["dense_before"],
)
print_generations(
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
)
print_generations(
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
)
print_generations(
"Sparse outputs after update",
PROMPTS,
sparse_results["sparse_after"],
)
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
print(f"patch_selection_equal = {patch_selection_equal}")
print(f"dense_patch_digest = {dense_results['patch_digest']}")
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
print(f"patch_digest_equal = {patch_digest_equal}")
print(f"baseline_equal = {baseline_equal}")
print(f"after_equal = {after_equal}")
print(f"any_output_changed = {any_output_changed}")
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
if not baseline_equal:
raise RuntimeError(
"Dense and sparse phases did not start from the same baseline"
)
if not patch_selection_equal:
raise RuntimeError("Dense and sparse phases used different sparse patches")
if not patch_digest_equal:
raise RuntimeError("Dense and sparse phases produced different patch values")
if not after_equal:
raise RuntimeError("Dense and sparse updates produced different outputs")
if not any_output_changed:
raise RuntimeError("Patch did not change the observed outputs")
finally:
ray.shutdown()
+1 -1
View File
@@ -233,7 +233,7 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
}
#[test]
+13 -3
View File
@@ -5,9 +5,9 @@ use std::sync::LazyLock;
pub use vllm_tool_parser::{
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
ToolParserOutput,
Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
ToolParserError, ToolParserOutput,
};
use crate::parser::ParserFactory;
@@ -24,6 +24,9 @@ pub mod names {
pub const GEMMA4: &str = "gemma4";
pub const HERMES: &str = "hermes";
pub const HY_V3: &str = "hy_v3";
// Matches the Python CLI name `--tool-call-parser internlm`, which Python
// also routes to `Internlm2ToolParser` despite the version-agnostic name.
pub const INTERNLM: &str = "internlm";
pub const KIMI_K2: &str = "kimi_k2";
pub const LLAMA3_JSON: &str = "llama3_json";
pub const LLAMA4_JSON: &str = "llama4_json";
@@ -62,6 +65,7 @@ impl ToolParserFactory {
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
.register_parser::<HermesToolParser>(names::HERMES)
.register_parser::<HyV3ToolParser>(names::HY_V3)
.register_parser::<Internlm2ToolParser>(names::INTERNLM)
.register_parser::<KimiK2ToolParser>(names::KIMI_K2)
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
@@ -80,6 +84,12 @@ impl ToolParserFactory {
.register_pattern("hermes", names::HERMES)
.register_pattern("hy3", names::HY_V3)
.register_pattern("hy_v3", names::HY_V3)
// Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
// and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
// routes to Llama), `internlm3-*` (also Llama-architecture per
// vllm/model_executor/models/registry.py:146), or `Intern-S1` /
// `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
.register_pattern("internlm2", names::INTERNLM)
.register_pattern("llama-4", names::LLAMA4_JSON)
.register_pattern("llama-3.2", names::LLAMA3_JSON)
.register_pattern("llama-3.1", names::LLAMA3_JSON)
+29
View File
@@ -161,4 +161,33 @@ fn factory_new_resolves_default_patterns() {
factory.resolve_name_for_model("org/mm-m2-base"),
Some(names::MINIMAX_M2)
);
// InternLM2 positive: both dashed and underscored versioned names route.
assert_eq!(
factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
Some(names::INTERNLM)
);
assert_eq!(
factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
Some(names::INTERNLM)
);
// Negative: other internlm-org models do NOT route to the InternLM2 parser,
// since they use unrelated prompt formats.
// - InternLM v1 (`internlm-chat-7b`) routes to Llama
// - InternLM3 (`internlm3-8b-instruct`) routes to Llama
// - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
assert_eq!(
factory.resolve_name_for_model("internlm/internlm-chat-7b"),
None
);
assert_eq!(
factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
None
);
assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
assert_eq!(
factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
None
);
}
+75 -21
View File
@@ -20,6 +20,7 @@ use vllm_chat::{
use vllm_text::{DecodedTextEvent, Finished, Prompt};
/// One model/parser configuration used to run the fixed roundtrip fixtures.
#[derive(Clone)]
struct RoundtripCase {
/// Hugging Face model id resolved through the production backend loader.
model_id: &'static str,
@@ -31,11 +32,45 @@ struct RoundtripCase {
tool_call_parser: ParserSelection,
/// Reasoning parser selection used by the output processor.
reasoning_parser: ParserSelection,
/// How this model's chat template handles thinking mode.
thinking_behavior: ThinkingBehavior,
/// JSON formatting expected after this model's template has materialized
/// tool-call arguments.
json_fmt: JsonFmt,
}
#[derive(Clone, Copy)]
enum ThinkingBehavior {
/// The chat template accepts explicit thinking on/off kwargs, and uses
/// `default` when the request does not specify either kwarg.
Toggleable { default: bool },
/// The chat template always behaves as `value` for this fixture.
Always { value: bool },
}
impl ThinkingBehavior {
fn default(self) -> bool {
match self {
Self::Toggleable { default } => default,
Self::Always { value } => value,
}
}
fn fixtures(self) -> Vec<Option<bool>> {
match self {
Self::Toggleable { .. } => vec![
Some(true), // explicitly enable thinking
Some(false), // explicitly disable thinking
None, // use default template behavior
],
Self::Always { value } => vec![
Some(value), // explicitly request the supported thinking behavior
None, // use default template behavior
],
}
}
}
impl RoundtripCase {
/// Qwen3 XML tool-call format with `qwen3` reasoning tags.
fn qwen3() -> Self {
@@ -44,6 +79,7 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: spaced_json_fmt(),
}
}
@@ -55,6 +91,7 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: compact_json_fmt(),
}
}
@@ -66,6 +103,7 @@ impl RoundtripCase {
assistant_stop_suffix: "[e~[\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
@@ -77,6 +115,7 @@ impl RoundtripCase {
assistant_stop_suffix: "<end▁of▁sentence>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: false },
json_fmt: compact_json_fmt(),
}
}
@@ -88,6 +127,7 @@ impl RoundtripCase {
assistant_stop_suffix: "",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: compact_json_fmt(),
}
}
@@ -100,6 +140,7 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: spaced_json_fmt(),
}
}
@@ -135,35 +176,44 @@ roundtrip_tests! {
/// Run the fixed reasoning+content fixture for one model/parser case.
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
for thinking in case.thinking_behavior.fixtures() {
run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
}
Ok(())
}
async fn run_roundtrip_reasoning_and_content_inner(
case: RoundtripCase,
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?")],
Vec::new(),
thinking,
);
let expected_reasoning = "Need compute 2 + 2 directly.";
let expected_text = "The answer is 4.";
let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default());
let result = run_roundtrip(
&case,
&backends,
&request,
AssistantMessage {
content: vec![
AssistantContentBlock::Reasoning {
text: expected_reasoning.to_string(),
},
AssistantContentBlock::Text {
text: expected_text.to_string(),
},
],
},
)
.await?;
let assistant = {
let mut content = Vec::new();
if effective_thinking {
content.push(AssistantContentBlock::Reasoning {
text: expected_reasoning.to_string(),
});
}
content.push(AssistantContentBlock::Text {
text: expected_text.to_string(),
});
AssistantMessage { content }
};
let result = run_roundtrip(&case, &backends, &request, assistant).await?;
assert_eq!(
result.parsed_message.reasoning().as_deref().map(str::trim),
Some(expected_reasoning)
effective_thinking.then_some(expected_reasoning)
);
assert_eq!(result.parsed_message.text().trim(), expected_text);
assert_eq!(result.parsed_message.tool_calls().count(), 0);
@@ -186,6 +236,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
"Check Shanghai weather and add 1.00 plus 2.",
)],
test_tools(),
Some(true), // always enable thinking in this fixture
);
let expected_reasoning = "Need call the weather and add tools.";
let expected_text = "I will call the tools.";
@@ -487,6 +538,7 @@ fn roundtrip_request(
request_id: impl Into<String>,
messages: Vec<ChatMessage>,
tools: Vec<ChatTool>,
thinking: Option<bool>,
) -> ChatRequest {
let mut request = ChatRequest {
request_id: request_id.into(),
@@ -500,10 +552,12 @@ fn roundtrip_request(
..ChatRequest::for_test()
};
// Enable thinking for some models so that rendering and parsing the reasoning block is
// exercised in the roundtrip.
for key in ["thinking", "enable_thinking"] {
request.chat_options.template_kwargs.insert(key.to_string(), true.into());
// Explicitly enable or disable thinking so that rendering and parsing the reasoning block is
// exercised or skipped in the roundtrip. If unspecified, use the default template behavior.
if let Some(thinking) = thinking {
for key in ["thinking", "enable_thinking"] {
request.chat_options.template_kwargs.insert(key.to_string(), thinking.into());
}
}
request
+11
View File
@@ -165,6 +165,15 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub enable_log_requests: bool,
/// If specified, API server will add X-Request-Id header to responses.
#[arg(
long,
default_missing_value = "true",
num_args = 0..=1
)]
#[serde(default)]
pub enable_request_id_headers: bool,
/// Disable periodic logging of engine statistics (throughput, queue depth,
/// cache usage).
#[arg(long)]
@@ -238,6 +247,7 @@ impl SharedRuntimeArgs {
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
@@ -278,6 +288,7 @@ impl SharedRuntimeArgs {
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
+46
View File
@@ -43,6 +43,7 @@ fn serve_args_forward_python_flags_with_separator() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -116,6 +117,46 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() {
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
}
#[test]
fn serve_passes_enable_request_id_headers_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--enable-request-id-headers",
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert!(config.enable_request_id_headers);
}
#[test]
fn frontend_args_json_passes_enable_request_id_headers_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"frontend",
"--listen-fd",
"3",
"--input-address",
"ipc:///tmp/input.sock",
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#,
])
.unwrap();
let Command::Frontend(args) = cli.command else {
panic!("expected frontend args");
};
let config = args.into_config();
assert!(config.enable_request_id_headers);
}
#[test]
fn serve_args_reject_unknown_renderer_value() {
let error = Cli::try_parse_from([
@@ -218,6 +259,7 @@ fn frontend_args_accept_json() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -616,6 +658,7 @@ fn serve_args_accept_handshake_aliases() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -733,6 +776,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -795,6 +839,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -872,6 +917,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
-9
View File
@@ -620,15 +620,6 @@ pub struct ServerUnsupportedArgs {
#[arg(long)]
pub middleware: Option<Unsupported>,
/// If specified, API server will add X-Request-Id header to responses.
#[arg(
long,
visible_alias = "no-enable-request-id-headers",
default_missing_value = "true",
num_args = 0..=1
)]
pub enable_request_id_headers: Option<Unsupported>,
/// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
#[arg(
long,
@@ -68,6 +68,7 @@ async fn main() -> Result<()> {
default_chat_template_kwargs: None,
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: Duration::ZERO,
+2
View File
@@ -61,6 +61,8 @@ pub struct Config {
pub chat_template_content_format: ChatTemplateContentFormatOption,
/// Log a summary line for each completed request.
pub enable_log_requests: bool,
/// When `true`, set `X-Request-Id` on every HTTP response.
pub enable_request_id_headers: bool,
/// When `true`, suppress periodic stats logging (throughput, queue depth,
/// cache usage).
pub disable_log_stats: bool,
+3 -1
View File
@@ -85,7 +85,9 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
};
Ok(Arc::new(
AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
AppState::new(served_model_names, chat)
.with_log_requests(config.enable_log_requests)
.with_request_id_headers(config.enable_request_id_headers),
))
}
+2
View File
@@ -1,5 +1,7 @@
mod load;
mod metrics;
mod request_id;
pub use load::track_server_load;
pub use metrics::track_http_metrics;
pub use request_id::set_request_id_header;
@@ -0,0 +1,24 @@
use axum::extract::Request;
use axum::http::HeaderValue;
use axum::http::header::HeaderName;
use axum::middleware::Next;
use axum::response::Response;
use uuid::Uuid;
const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
/// Echo the request's `X-Request-Id` on the response, or generate a fresh
/// `uuid4` hex if the request did not provide one.
///
/// Original Python:
/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`.
pub async fn set_request_id_header(req: Request, next: Next) -> Response {
let incoming = req.headers().get(&X_REQUEST_ID).cloned();
let mut response = next.run(req).await;
let value = incoming.unwrap_or_else(|| {
HeaderValue::from_str(&Uuid::new_v4().simple().to_string())
.expect("uuid hex is valid header value")
});
response.headers_mut().insert(X_REQUEST_ID, value);
response
}
+9 -2
View File
@@ -56,11 +56,18 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
.route("/is_sleeping", get(sleep::is_sleeping))
}
router
let enable_request_id_headers = state.enable_request_id_headers;
let mut router = router
.with_state(state.clone())
.layer(from_fn_with_state(state, middleware::track_server_load))
.layer(from_fn(middleware::track_http_metrics))
.layer(TraceLayer::new_for_http())
.layer(TraceLayer::new_for_http());
if enable_request_id_headers {
router = router.layer(from_fn(middleware::set_request_id_header));
}
router
}
#[cfg(test)]
@@ -3,23 +3,33 @@ mod types;
mod validate;
use std::collections::HashMap;
use std::convert::Infallible;
use std::result::Result;
use std::sync::Arc;
use asynk_strim_attr::{TryYielder, try_stream};
use axum::Json;
use axum::extract::State;
use axum::http::HeaderMap;
use axum::response::sse::{Event, Sse};
use axum::response::{IntoResponse, Response};
use futures::{Stream, StreamExt as _, pin_mut};
use thiserror_ext::AsReport as _;
use tracing::info;
use tracing::{error, info, trace};
use tracing_futures::Instrument as _;
use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
use vllm_llm::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
};
use self::convert::prepare_generate_request;
use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
use crate::error::{ApiError, server_error};
use self::types::{
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
GenerateResponseStreamChoice, GenerateStreamResponse,
};
use crate::error::{ApiError, bail_server_error, server_error};
use crate::routes::openai::utils::logprobs::clamp_logprob;
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
use crate::routes::openai::utils::validated_json::ValidatedJson;
use crate::state::AppState;
use crate::utils::resolve_request_context;
@@ -46,6 +56,7 @@ pub async fn generate(
let log_request = state.enable_log_requests;
let include_logprobs = prepared.include_logprobs;
let include_prompt_logprobs = prepared.include_prompt_logprobs;
let stream = prepared.stream;
let raw_stream = match state
.chat
@@ -64,6 +75,20 @@ pub async fn generate(
}
};
if stream {
let chunk_stream = generate_chunk_stream(
raw_stream,
prepared.request_id,
log_request,
prepared.include_usage,
prepared.include_continuous_usage,
include_logprobs,
);
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
return Sse::new(sse_stream).into_response();
}
let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
Ok(collected) => collected,
Err(error) => {
@@ -98,6 +123,102 @@ pub async fn generate(
Json(response).into_response()
}
#[try_stream]
async fn generate_chunk_stream(
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
request_id: String,
log_request: bool,
include_usage: bool,
include_continuous_usage: bool,
include_logprobs: bool,
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
let mut prompt_tokens: Option<u32> = None;
let mut output_tokens = 0_u32;
while let Some(next) = stream.next().await {
match next {
Ok(output) => {
if prompt_tokens.is_none() {
prompt_tokens =
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
}
let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
let token_ids = output.token_ids;
output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
let finish_reason = output.finish_reason;
if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
bail_server_error!("Internal server error");
}
if let Some(finish_reason) = finish_reason.as_ref()
&& log_request
{
info!(
stream = true,
prompt_tokens = usage_prompt_tokens,
output_tokens,
finish_reason = finish_reason.as_str(),
"generate finished"
);
}
if token_ids.is_empty() && finish_reason.is_none() {
continue;
}
let logprobs = if include_logprobs && !token_ids.is_empty() {
let logprobs = output.logprobs.as_ref().ok_or_else(|| {
server_error!(
"raw generate stream requested logprobs but generation returned none"
)
})?;
Some(raw_logprobs_to_openai_chat(logprobs)?)
} else {
None
};
y.yield_ok(GenerateStreamResponse {
request_id: request_id.clone(),
choices: vec![GenerateResponseStreamChoice {
index: 0,
logprobs,
finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
token_ids,
}],
usage: include_continuous_usage
.then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
})
.await;
}
Err(error) => {
error!(
error = %error.as_report(),
"raw generate stream failed"
);
bail_server_error!("{}", error.to_report_string());
}
}
}
if include_usage {
y.yield_ok(GenerateStreamResponse {
request_id,
choices: Vec::new(),
usage: Some(Usage::from_counts(
prompt_tokens.unwrap_or_default(),
output_tokens,
)),
})
.await;
}
Ok(())
}
fn collect_generate(
collected: CollectedGenerateOutput,
request_id: String,
@@ -213,3 +334,94 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
fn format_token_id(token_id: u32) -> String {
format!("token_id:{token_id}")
}
/// Convert one raw-generate chunk stream into SSE events.
#[try_stream]
async fn generate_sse_stream(
stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
mut y: TryYielder<Event, Infallible>,
) -> Result<(), Infallible> {
pin_mut!(stream);
while let Some(next) = stream.next().await {
match next {
Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
Err(error) => {
y.yield_ok(to_error_sse_event(&error)).await;
break;
}
}
}
y.yield_ok(done_sse_event()).await;
Ok(())
}
fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
trace!(payload, "generate emitting chunk");
Event::default().data(payload)
}
fn to_error_sse_event(error: &ApiError) -> Event {
let payload = serde_json::to_string(&error.to_error_response())
.expect("ErrorResponse must serialize to JSON");
trace!(payload, "generate emitting error");
Event::default().data(payload)
}
fn done_sse_event() -> Event {
trace!("generate emitting done");
Event::default().data("[DONE]")
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use futures::{TryStreamExt as _, stream};
use vllm_llm::GeneratePromptInfo;
use super::*;
#[tokio::test]
async fn generate_chunk_stream_captures_late_prompt_info() {
let stream = stream::iter(vec![
Ok(GenerateOutput {
request_id: String::new(),
prompt_info: None,
token_ids: Vec::new(),
logprobs: None,
finish_reason: None,
kv_transfer_params: None,
}),
Ok(GenerateOutput {
request_id: String::new(),
prompt_info: Some(GeneratePromptInfo {
prompt_token_ids: Arc::from([11_u32, 22_u32]),
prompt_logprobs: None,
}),
token_ids: vec![33],
logprobs: None,
finish_reason: Some(FinishReason::stop_eos()),
kv_transfer_params: None,
}),
]);
let chunks: Vec<_> =
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
.try_collect()
.await
.expect("collect chunks");
assert_eq!(chunks.len(), 2);
assert_eq!(
chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
2
);
assert_eq!(
chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
2
);
}
}
@@ -10,6 +10,9 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
pub struct PreparedRequest {
pub request_id: String,
pub text_request: TextRequest,
pub stream: bool,
pub include_usage: bool,
pub include_continuous_usage: bool,
pub include_logprobs: bool,
pub include_prompt_logprobs: bool,
}
@@ -23,6 +26,18 @@ pub fn prepare_generate_request(
) -> Result<PreparedRequest, ApiError> {
validate::validate_request_compat(&request, served_model_names)?;
let stream = request.stream;
let include_usage = request
.stream_options
.as_ref()
.and_then(|options| options.include_usage)
.unwrap_or(false);
let include_continuous_usage = include_usage
&& request
.stream_options
.as_ref()
.and_then(|options| options.continuous_usage_stats)
.unwrap_or(false);
let include_logprobs = request.sampling_params.logprobs.is_some();
let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
let mut sampling_params = request.sampling_params;
@@ -47,6 +62,9 @@ pub fn prepare_generate_request(
Ok(PreparedRequest {
request_id: ctx.request_id,
text_request,
stream,
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
})
@@ -109,4 +127,28 @@ mod tests {
Some(json!({"connector": "x"}))
);
}
#[test]
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
let request: GenerateRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"continuous_usage_stats": true
},
"sampling_params": {}
}))
.expect("parse request");
let prepared = prepare_generate_request(
request,
&["Qwen/Qwen1.5-0.5B-Chat".to_string()],
ResolvedRequestContext::default(),
)
.expect("prepare");
assert!(!prepared.include_usage);
assert!(!prepared.include_continuous_usage);
}
}
@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
use validator::Validate;
use vllm_text::SamplingParams;
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
/// vLLM-compatible request type for the token-in/token-out generate API.
#[serde_with::skip_serializing_none]
@@ -17,6 +17,7 @@ pub struct GenerateRequest {
pub sampling_params: SamplingParams,
#[serde(default)]
pub stream: bool,
pub stream_options: Option<StreamOptions>,
pub cache_salt: Option<String>,
#[serde(default)]
pub priority: i32,
@@ -37,6 +38,25 @@ pub(super) struct GenerateResponseChoice {
pub token_ids: Vec<u32>,
}
/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub(super) struct GenerateResponseStreamChoice {
pub index: u32,
pub logprobs: Option<ChatLogProbs>,
pub finish_reason: Option<String>,
pub token_ids: Vec<u32>,
}
/// Mirrors the Python vLLM `GenerateStreamResponse` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub(super) struct GenerateStreamResponse {
pub request_id: String,
pub choices: Vec<GenerateResponseStreamChoice>,
pub usage: Option<Usage>,
}
/// Mirrors the Python vLLM `GenerateResponse` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
@@ -13,8 +13,11 @@ pub(super) fn validate_request_compat(
return Err(ApiError::model_not_found(model.clone()));
}
if request.stream {
bail_invalid_request!(param = "stream", "stream=true is not supported.");
if request.stream_options.is_some() && !request.stream {
bail_invalid_request!(
param = "stream_options",
"stream_options are only supported when stream=true."
);
}
if request.token_ids.is_empty() {
@@ -65,11 +68,24 @@ mod tests {
}
#[test]
fn validate_request_compat_rejects_streaming() {
fn validate_request_compat_accepts_streaming() {
let request = GenerateRequest {
stream: true,
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
}
#[test]
fn validate_request_compat_rejects_stream_options_without_streaming() {
let request: GenerateRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"token_ids": [11, 22],
"stream": false,
"stream_options": {"include_usage": true},
"sampling_params": {}
}))
.expect("parse request");
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
}
+327 -6
View File
@@ -746,6 +746,20 @@ async fn test_app() -> axum::Router {
)))
}
async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
b"engine-openai-request-id",
default_stream_output_specs(),
Arc::new(FakeChatBackend::new()),
)
.await;
let app = build_router(Arc::new(
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
.with_request_id_headers(true),
));
(app, engine_task)
}
async fn test_health_app_with_engine_script<F>(
script: F,
) -> (axum::Router, Arc<AppState>, MockEngineTask)
@@ -947,6 +961,18 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
(status, body)
}
async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response {
let mut builder = Request::builder().method("GET").uri("/health");
if let Some(request_id) = request_id {
builder = builder.header("X-Request-Id", request_id);
}
app.clone()
.call(builder.body(Body::empty()).expect("build request"))
.await
.expect("call app")
}
fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
rendered.lines().find_map(|line| {
let rest = line.strip_prefix(metric)?;
@@ -994,6 +1020,43 @@ async fn list_models_returns_configured_model() {
assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_is_absent_by_default() {
let app = test_app().await;
let response = health_response(&app, None).await;
assert_eq!(response.status(), StatusCode::OK);
assert!(!response.headers().contains_key("x-request-id"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_generates_uuid_hex_when_enabled() {
let (app, _engine_task) = test_app_with_request_id_headers().await;
let response = health_response(&app, None).await;
assert_eq!(response.status(), StatusCode::OK);
let request_id = response
.headers()
.get("x-request-id")
.expect("x-request-id header")
.to_str()
.expect("header is ascii");
assert_eq!(request_id.len(), 32);
assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit()));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_echoes_incoming_header_when_enabled() {
let (app, _engine_task) = test_app_with_request_id_headers().await;
let response = health_response(&app, Some("req-123")).await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn version_returns_engine_vllm_version() {
@@ -2425,8 +2488,72 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn raw_generate_rejects_streaming() {
let mut app = test_app().await;
async fn stream_raw_generate_returns_sse_chunks_and_usage() {
let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint();
let engine_id = b"engine-raw-generate-stream".to_vec();
let engine_task = MockEngineTask::new(spawn_mock_engine_task(
handshake_address.clone(),
engine_id.clone(),
|dealer, push| {
boxed_test_future(async move {
let add = recv_engine_message(dealer).await;
let request: EngineCoreRequest =
rmp_serde::from_slice(&add[1]).expect("decode request");
assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
outputs: vec![
request_output_with_logprobs(
&request.request_id,
vec![33],
None,
None,
Some(sample_logprobs_for_token(33, 34)),
None,
),
request_output_with_logprobs(
&request.request_id,
vec![44],
Some(EngineCoreFinishReason::Stop),
None,
Some(sample_logprobs_for_token(44, 45)),
None,
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
)
.await;
})
},
));
let client = EngineCoreClient::connect(
EngineCoreClientConfig::new_single(handshake_address)
.with_model_name("test-model")
.with_local_input_output_addresses(
Some(ipc.input_endpoint()),
Some(ipc.output_endpoint()),
),
)
.await
.expect("connect client");
let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
let mut app = build_router(Arc::new(AppState::new(
vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
chat,
)));
let response = app
.call(
@@ -2437,9 +2564,17 @@ async fn raw_generate_rejects_streaming() {
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream",
"token_ids": [11, 22],
"stream": true,
"sampling_params": {}
"stream_options": {
"include_usage": true,
"continuous_usage_stats": true
},
"sampling_params": {
"max_tokens": 2,
"logprobs": 1
}
})
.to_string(),
))
@@ -2448,10 +2583,196 @@ async fn raw_generate_rejects_streaming() {
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get("content-type").and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
assert_eq!(json["error"]["param"], "stream");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 4, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["request_id"], "raw-stream");
assert_eq!(first["choices"][0]["index"], 0);
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert_eq!(
first["choices"][0]["logprobs"]["content"][0]["token"],
"token_id:33"
);
assert_eq!(first["usage"]["prompt_tokens"], 2);
assert_eq!(first["usage"]["completion_tokens"], 1);
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert_eq!(second["usage"]["completion_tokens"], 2);
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
assert_eq!(usage["choices"], json!([]));
assert_eq!(usage["usage"]["prompt_tokens"], 2);
assert_eq!(usage["usage"]["completion_tokens"], 2);
assert_eq!(usage["usage"]["total_tokens"], 4);
assert_eq!(payloads[3], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
(vec![33], None),
(vec![44], Some(EngineCoreFinishReason::Stop)),
])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-final-usage",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"include_usage": true
},
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 4, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert!(first.get("usage").is_none());
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert!(second.get("usage").is_none());
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
assert_eq!(usage["choices"], json!([]));
assert_eq!(usage["usage"]["prompt_tokens"], 2);
assert_eq!(usage["usage"]["completion_tokens"], 2);
assert_eq!(usage["usage"]["total_tokens"], 4);
assert_eq!(payloads[3], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_emits_empty_finish_chunk() {
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
(vec![33], None),
(vec![], Some(EngineCoreFinishReason::Stop)),
])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-empty-finish",
"token_ids": [11, 22],
"stream": true,
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 3, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert!(first["choices"][0].get("finish_reason").is_none());
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert_eq!(payloads[2], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_error_finish_returns_sse_error() {
let (mut app, engine_task) =
test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-error",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"include_usage": true
},
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
assert!(text.contains("\"type\":\"server_error\""), "{text}");
assert!(text.contains("Internal server error"), "{text}");
assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
assert!(!text.contains("\"usage\":"), "{text}");
assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+9
View File
@@ -17,6 +17,8 @@ pub struct AppState {
pub chat: ChatLlm,
/// Whether to log a summary line for each completed request.
pub enable_log_requests: bool,
/// Whether to set X-Request-Id on every HTTP response.
pub enable_request_id_headers: bool,
/// Number of in-flight inference requests currently owned by this frontend.
server_load: AtomicU64,
}
@@ -39,6 +41,7 @@ impl AppState {
served_model_names,
chat,
enable_log_requests: false,
enable_request_id_headers: false,
server_load: AtomicU64::new(0),
}
}
@@ -49,6 +52,12 @@ impl AppState {
self
}
/// Enable X-Request-Id response headers.
pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
self.enable_request_id_headers = enabled;
self
}
/// The primary model name echoed back in API responses (the first served
/// name).
pub fn primary_model_name(&self) -> &str {
@@ -104,7 +104,7 @@ impl DeepSeekDsmlToolParser {
self.tool_parameters.convert_param_with_schema(
&name,
&param.name,
&param.value,
param.value,
)
};
arguments.insert(param.name, value);
+1 -1
View File
@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: None,
name_key: "name",
arguments_key: "arguments",
arguments_key: &["arguments"],
};
/// Tool parser for Hermes XML-wrapped JSON tool calls.
+351
View File
@@ -0,0 +1,351 @@
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
use crate::{Result, Tool, ToolParser, ToolParserOutput};
const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
parser_name: "InternLM2",
start_marker: "<|action_start|><|plugin|>",
end_marker: "<|action_end|>",
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: None,
name_key: "name",
// The Python parser's `get_arguments()` accepts either `parameters` or
// `arguments` and prefers `parameters` when both are present. This Rust
// parser uses first-encountered semantics because the header parser only
// permits one args key per tool-call object; if a future model emits
// both keys in the same object, the Rust port will accept the first one
// and reject the trailing one as a syntax error rather than silently
// shadowing it.
arguments_key: &["parameters", "arguments"],
};
/// Tool parser for InternLM2 special-token wrapped JSON tool calls.
///
/// Example tool call content:
///
/// ```text
/// <|action_start|><|plugin|>{"name": "get_weather", "parameters": {"location":"Tokyo"}}<|action_end|>
/// ```
///
/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
/// argument deltas without schema conversion or JSON normalization.
///
/// # Divergences from the Python reference
///
/// This Rust port intentionally diverges from
/// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways:
///
/// - **Parallel tool calls are supported.** Python silently drops every
/// `<|action_start|>` block after the first (`current_tool_id > 0` returns
/// an empty delta); this parser emits every well-formed block with
/// incrementing `tool_index`. Models that legitimately emit multiple action
/// blocks therefore produce more tool calls under Rust than under Python.
/// - **End-marker bytes inside JSON string values are preserved.** Python
/// does `action.split("<|action_end|>")[0]` which truncates regardless of
/// JSON context; this parser scans matched braces and quotes so a literal
/// `<|action_end|>` inside an arguments string is forwarded intact.
/// - **Only whitespace is allowed before the `{`.** Python's non-streaming
/// `action[action.find("{"):]` drops any bytes before the first `{`, but
/// its streaming path has no equivalent and the model format always emits
/// `<|plugin|>{...`; this parser allows only whitespace there, matching the
/// other JSON parsers in this crate.
/// - **Truncated tool calls error rather than silently dropping.** Python's
/// streaming wrapper swallows mid-stream errors with `except Exception:
/// return None` (logging a traceback) while its non-streaming path raises
/// `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool
/// call` error from `finish()`, matching the other JSON parsers and Python's
/// non-streaming behavior.
///
/// # Known unaddressed divergences (TODO)
///
/// The following Python behaviors are NOT yet matched. They are deferred to
/// follow-up work because they require non-local changes to the shared
/// `JsonToolCallParser` core that would affect Hermes / Llama / Mistral /
/// Qwen as well. If a real-world InternLM2 deployment hits one of these,
/// prioritize the corresponding fix.
///
/// - **Arguments value type.** The shared core requires the arguments value
/// to be a JSON object (`take_json_object` rejects anything not starting
/// with `{`). Python's `json.dumps(action_dict.get("parameters", ...))`
/// accepts `null`, arrays, strings, and numbers and round-trips them
/// verbatim. Models that legitimately emit `"parameters":null` will hard-
/// fail under Rust.
/// - **Unknown arguments key.** Python falls back to `{}` via
/// `action_dict.get("parameters", action_dict.get("arguments", {}))` when
/// neither key is present; the Rust header parser raises
/// `parsing failed: invalid InternLM2` for any unrecognized key. A model
/// that emits a typo (e.g. `"params"`) breaks the whole response.
/// - **Field order independence.** The header parser requires the JSON keys
/// to appear in the order `name` then arguments key. Python's
/// `json.loads` + `dict.get` is order-independent, so a model emitting
/// `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
pub struct Internlm2ToolParser {
inner: JsonToolCallParser,
}
impl Internlm2ToolParser {
/// Create an InternLM2 tool parser.
fn new(_tools: &[Tool]) -> Self {
Self {
inner: JsonToolCallParser::new(INTERNLM2_CONFIG),
}
}
}
impl ToolParser for Internlm2ToolParser {
/// Create a boxed InternLM2 tool parser.
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tools)))
}
/// Preserve special-token markers while decoding, since
/// `<|action_start|>`, `<|plugin|>`, and `<|action_end|>` are tokenizer
/// special tokens in InternLM2 models.
fn preserve_special_tokens(&self) -> bool {
true
}
/// Feed one decoded text chunk through the InternLM2 parser.
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
self.inner.parse_into(chunk, output)
}
/// Flush any buffered partial state at end of stream.
fn finish(&mut self) -> Result<ToolParserOutput> {
self.inner.finish()
}
/// Clear parser state and return currently uncommitted buffered text.
fn reset(&mut self) -> String {
self.inner.reset()
}
}
#[cfg(test)]
mod tests {
use expect_test::expect;
use thiserror_ext::AsReport;
use super::Internlm2ToolParser;
use crate::test_utils::{collect_stream, split_by_chars, test_tools};
use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
const ACTION_START: &str = "<|action_start|><|plugin|>";
const ACTION_END: &str = "<|action_end|>";
fn build_tool_call(function_name: &str, args_key: &str, arguments: &str) -> String {
format!(
r#"{ACTION_START}{{"name":"{function_name}","{args_key}":{arguments}}}{ACTION_END}"#
)
}
#[test]
fn internlm2_parse_complete_without_tool_call_keeps_text() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let result = parser.parse_complete("Hello, world!").unwrap();
assert_eq!(result.normal_text, "Hello, world!");
assert!(result.calls.is_empty());
}
#[test]
fn internlm2_parse_complete_extracts_parameters_key() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let arguments = r#"{"location":"Tokyo","days":"3"}"#;
let result = parser
.parse_complete(&format!(
"Let me check.\n{}",
build_tool_call("get_weather", "parameters", arguments)
))
.unwrap();
assert_eq!(result.normal_text, "Let me check.\n");
assert_eq!(result.calls.len(), 1);
assert_eq!(result.calls[0].tool_index, 0);
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
assert_eq!(result.calls[0].arguments, arguments);
}
#[test]
fn internlm2_parse_complete_extracts_arguments_key_fallback() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let arguments = r#"{"location":"Tokyo"}"#;
let result = parser
.parse_complete(&build_tool_call("get_weather", "arguments", arguments))
.unwrap();
assert_eq!(result.calls.len(), 1);
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
assert_eq!(result.calls[0].arguments, arguments);
}
#[test]
fn internlm2_accepts_whitespace_after_plugin_marker() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let result = parser
.parse_complete(&format!(
r#"{ACTION_START}
{{"name":"get_weather","parameters":{{}}}}{ACTION_END}"#
))
.unwrap();
assert_eq!(result.calls.len(), 1);
assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
}
#[test]
fn internlm2_does_not_validate_or_normalize_arguments() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let arguments = r#"{"location":"Tokyo",}"#;
let result = parser
.parse_complete(&build_tool_call("get_weather", "parameters", arguments))
.unwrap();
assert_eq!(result.calls[0].arguments, arguments);
}
#[test]
fn internlm2_streaming_emits_argument_deltas() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let chunks = [
"preface <|action",
"_start|><|plugin|>",
r#"{"name":"get_weather","parameters":"#,
r#"{"location":"#,
r#""Beijing""#,
r#"}"#,
r#"}<|action_end|> suffix"#,
];
let mut result = ToolParserOutput::default();
let mut observed_arguments = Vec::new();
for chunk in chunks {
let next = parser.parse_chunk(chunk).unwrap();
observed_arguments.extend(
next.calls
.iter()
.filter(|call| call.name.is_none())
.map(|call| call.arguments.clone()),
);
result.append(next);
}
result.append(parser.finish().unwrap());
assert_eq!(
observed_arguments,
[r#"{"location":"#, r#""Beijing""#, r#"}"#]
);
assert_eq!(result.normal_text, "preface suffix");
assert_eq!(
result.coalesce_calls().calls[0].arguments,
r#"{"location":"Beijing"}"#
);
}
#[test]
fn internlm2_streaming_handles_split_markers() {
let input = format!(
"hello {}",
build_tool_call("get_weather", "parameters", r#"{"location":"Tokyo"}"#)
);
let chunks = split_by_chars(&input, 5);
let mut parser = Internlm2ToolParser::new(&test_tools());
let result = collect_stream(&mut parser, &chunks);
assert_eq!(result.normal_text, "hello ");
assert_eq!(result.calls.len(), 1);
assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
}
#[test]
fn internlm2_streaming_extracts_multiple_blocks() {
let input = format!(
"{}{}",
build_tool_call("get_weather", "parameters", r#"{"location":"Shanghai"}"#),
build_tool_call("add", "arguments", r#"{"x":1,"y":2}"#),
);
let chunks = split_by_chars(&input, 7);
let mut parser = Internlm2ToolParser::new(&test_tools());
let result = collect_stream(&mut parser, &chunks);
expect![[r#"
ToolParserOutput {
normal_text: "",
calls: [
ToolCallDelta {
tool_index: 0,
name: Some(
"get_weather",
),
arguments: "{\"location\":\"Shanghai\"}",
},
ToolCallDelta {
tool_index: 1,
name: Some(
"add",
),
arguments: "{\"x\":1,\"y\":2}",
},
],
}
"#]]
.assert_debug_eq(&result);
}
#[test]
fn internlm2_keeps_end_marker_literal_inside_json_string() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let arguments = format!(r#"{{"text":"literal {ACTION_END} inside"}}"#);
let input = build_tool_call("echo", "parameters", &arguments);
let result = parser.parse_complete(&input).unwrap();
assert_eq!(result.calls.len(), 1);
assert_eq!(result.calls[0].arguments, arguments);
}
#[test]
fn internlm2_finish_errors_on_truncated_tool_call() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let pre_finish = parser
.parse_chunk(&format!(
r#"{ACTION_START}{{"name":"get_weather","parameters":{{"location""#
))
.unwrap();
let error = parser.finish().unwrap_err();
assert_eq!(
pre_finish.calls[0].name.as_deref(),
Some("get_weather"),
"name delta is still emitted from parse_chunk() before truncation",
);
assert!(
error.to_report_string().contains("incomplete InternLM2 tool call"),
"finish() reports the truncated tool call as incomplete: {}",
error.to_report_string(),
);
}
#[test]
fn internlm2_unknown_arguments_key_fails() {
let mut parser = Internlm2ToolParser::new(&test_tools());
let input = build_tool_call("get_weather", "params", r#"{"location":"Tokyo"}"#);
let error = parser.parse_chunk(&input).unwrap_err();
expect![[r#"
tool parser parsing failed: invalid InternLM2
expected `parameters`, `arguments`"#]]
.assert_eq(&error.to_report_string());
}
#[test]
fn internlm2_preserve_special_tokens_is_true() {
let parser = Internlm2ToolParser::new(&test_tools());
assert!(parser.preserve_special_tokens());
}
}
+1 -1
View File
@@ -203,7 +203,7 @@ fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<Ll
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: Some(";"),
name_key: "name",
arguments_key: "parameters",
arguments_key: &["parameters"],
};
match tool_call_header_event(input, CONFIG)? {
+1 -1
View File
@@ -8,7 +8,7 @@ const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: Some(","),
name_key: "name",
arguments_key: "arguments",
arguments_key: &["arguments"],
};
/// Tool parser for Mistral JSON-array tool calls.
+43 -5
View File
@@ -1,20 +1,22 @@
//! Shared parser core for JSON tool calls wrapped by text markers.
pub use hermes::HermesToolParser;
pub use internlm2::Internlm2ToolParser;
pub use llama::Llama3JsonToolParser;
pub use mistral::MistralToolParser;
pub use qwen::Qwen3XmlToolParser;
mod hermes;
mod internlm2;
mod llama;
mod mistral;
mod qwen;
use winnow::ascii::multispace0 as ws0;
use winnow::combinator::{alt, seq};
use winnow::error::{ModalResult, StrContext, StrContextValue};
use winnow::error::{AddContext, ModalResult, StrContext, StrContextValue};
use winnow::prelude::*;
use winnow::stream::Partial;
use winnow::stream::{Partial, Stream};
use winnow::token::literal;
use super::utils::{
@@ -32,7 +34,10 @@ struct JsonToolCallConfig {
marker_whitespace: JsonToolCallWhitespace,
delimiter: Option<&'static str>,
name_key: &'static str,
arguments_key: &'static str,
/// Candidate JSON keys naming the arguments payload, tried in order.
/// Most parsers use a single key like `["arguments"]`, but some accept
/// multiple (e.g. InternLM2 accepts `parameters` or `arguments`).
arguments_key: &'static [&'static str],
}
#[derive(Debug, Clone, Copy)]
@@ -224,7 +229,7 @@ fn tool_call_header_event(
_: ws0,
_: literal(","),
_: ws0,
_: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key),
_: |input: &mut JsonToolInput<'_>| json_arguments_key(input, config.arguments_key),
_: ws0,
_: literal(":"),
_: ws0,
@@ -246,6 +251,39 @@ fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()>
.parse_next(input)
}
/// Parse a JSON object key accepting any of `candidates`.
///
/// The full quoted key is consumed and compared against the candidate list,
/// so this works correctly under partial input regardless of key lengths.
///
/// On mismatch, each candidate is attached as its own `Expected` context so the
/// error enumerates every valid key ("expected `a`, expected `b`"). Because
/// `StrContextValue::StringLiteral` carries a single `&'static str`, the
/// contexts are added in a loop over `candidates` rather than through chained
/// `.context(...)` calls, which keeps the diagnostics complete for any number
/// of candidates.
fn json_arguments_key(
input: &mut JsonToolInput<'_>,
candidates: &'static [&'static str],
) -> ModalResult<()> {
let start = input.checkpoint();
json_str
.verify(|key: &String| candidates.contains(&key.as_str()))
.void()
.parse_next(input)
.map_err(|err| {
err.map(|context_error| {
candidates.iter().fold(context_error, |context_error, candidate| {
context_error.add_context(
&*input,
&start,
StrContext::Expected(StrContextValue::StringLiteral(candidate)),
)
})
})
})
}
/// Parse one event inside a marker-wrapped JSON tool-call arguments payload.
fn parse_arguments_event(
input: &mut JsonToolInput<'_>,
@@ -341,7 +379,7 @@ mod tests {
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: Some("<"),
name_key: "function",
arguments_key: "parameters",
arguments_key: &["parameters"],
};
fn build_tool_call(function_name: &str, arguments: &str) -> String {
+1 -1
View File
@@ -8,7 +8,7 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
marker_whitespace: JsonToolCallWhitespace::Exact("\n"),
delimiter: None,
name_key: "name",
arguments_key: "arguments",
arguments_key: &["arguments"],
};
/// Tool parser for Qwen XML-wrapped JSON tool calls.
+4 -1
View File
@@ -24,7 +24,10 @@ pub use error::{Result, ToolParserError};
pub use gemma4::Gemma4ToolParser;
pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
pub use hy_v3::HyV3ToolParser;
pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
pub use json::{
HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser,
Qwen3XmlToolParser,
};
pub use kimi_k2::KimiK2ToolParser;
pub use minimax_m2::MinimaxM2ToolParser;
pub use qwen_coder::Qwen3CoderToolParser;
+440 -84
View File
@@ -1,6 +1,6 @@
use std::collections::BTreeMap;
use serde_json::{Number, Value};
use serde_json::{Map, Number, Value};
use crate::Tool;
@@ -21,6 +21,29 @@ pub(super) struct ToolSchema {
params: BTreeMap<String, JsonParamType>,
}
/// Parameter input for schema-aware conversion.
///
/// It can be either a raw text string, or a structured input with named child elements.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum ParamInput {
Text(String),
#[allow(dead_code)]
Elements(Vec<ParamElement>),
}
impl From<String> for ParamInput {
fn from(value: String) -> Self {
Self::Text(value)
}
}
/// One named structured parameter child.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct ParamElement {
pub name: String,
pub value: ParamInput,
}
/// Normalized JSON parameter type used for raw string coercion.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum JsonParamType {
@@ -28,8 +51,13 @@ pub(super) enum JsonParamType {
Integer,
Number,
Boolean,
Object,
Array,
Object {
properties: BTreeMap<String, JsonParamType>,
additional_properties: Option<Box<JsonParamType>>,
},
Array {
items: Option<Box<JsonParamType>>,
},
Null,
OneOf(Vec<JsonParamType>),
}
@@ -45,33 +73,39 @@ impl ToolSchemas {
Self { tools }
}
/// Convert raw string parameter values for one named tool.
/// Convert parameter values for one named tool.
///
/// Unknown tool names use an empty schema, so all parameters fall back to
/// strings.
pub(super) fn convert_params_with_schema(
/// strings or object-like JSON for structured inputs.
pub(super) fn convert_params_with_schema<P>(
&self,
function_name: &str,
params: Vec<(String, String)>,
) -> serde_json::Map<String, Value> {
params: Vec<(String, P)>,
) -> Map<String, Value>
where
P: Into<ParamInput>,
{
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
let mut converted = serde_json::Map::with_capacity(params.len());
let mut converted = Map::with_capacity(params.len());
for (name, value) in params {
let value = tool_schema.convert(&name, &value);
let value = tool_schema.convert(&name, value.into());
converted.insert(name, value);
}
converted
}
/// Convert one raw string parameter value for one named tool.
pub(super) fn convert_param_with_schema(
/// Convert one parameter value for one named tool.
pub(super) fn convert_param_with_schema<P>(
&self,
function_name: &str,
name: &str,
value: &str,
) -> Value {
value: P,
) -> Value
where
P: Into<ParamInput>,
{
let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
tool_schema.convert(name, value)
tool_schema.convert(name, value.into())
}
}
@@ -101,21 +135,13 @@ impl ToolSchema {
Self { params }
}
/// Convert one raw parameter value using its normalized schema type.
/// Convert one parameter value using its normalized schema type.
///
/// If the parameter name is unknown, or we don't have a schema for it, or
/// the value fails to convert, this falls back to returning the raw
/// string as a JSON string value.
fn convert(&self, name: &str, value: &str) -> Value {
if value.eq_ignore_ascii_case("null") {
return Value::Null;
}
let Some(param_type) = self.params.get(name) else {
return Value::String(value.to_string());
};
convert_value(param_type, value).unwrap_or_else(|| Value::String(value.to_string()))
/// string as a JSON string value, or object-like JSON for structured input.
fn convert(&self, name: &str, input: ParamInput) -> Value {
convert_with_optional_schema(self.params.get(name), &input)
}
}
@@ -125,7 +151,7 @@ impl JsonParamType {
let schema = schema.as_object()?;
if let Some(type_value) = schema.get("type") {
return Self::from_type_value(type_value);
return Self::from_type_value(type_value, schema);
}
if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) {
@@ -134,32 +160,34 @@ impl JsonParamType {
.map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::<Vec<_>>())
.filter(|types| !types.is_empty())
.map(Self::one_of)
.unwrap_or(Self::Object);
.unwrap_or_else(|| Self::object_from_schema(Some(schema)));
return Some(param_type);
}
// Typically, these types are already handled by checking the "type" field, but
// we can also infer them from their characteristic fields if "type" is missing.
if schema.contains_key("enum") {
return Some(Self::String);
}
if schema.contains_key("items") {
return Some(Self::Array);
return Some(Self::array_from_schema(Some(schema)));
}
if schema.contains_key("properties") {
return Some(Self::Object);
if schema.contains_key("properties") || schema.contains_key("additionalProperties") {
return Some(Self::object_from_schema(Some(schema)));
}
None
}
/// Normalize a JSON schema `type` value.
fn from_type_value(type_value: &Value) -> Option<Self> {
fn from_type_value(type_value: &Value, schema: &Map<String, Value>) -> Option<Self> {
match type_value {
Value::String(kind) => Self::from_type_name(kind),
Value::String(kind) => Self::from_type_name(kind, Some(schema)),
Value::Array(kinds) => {
let types = kinds
.iter()
.filter_map(Value::as_str)
.filter_map(Self::from_type_name)
.filter_map(|kind| Self::from_type_name(kind, Some(schema)))
.collect::<Vec<_>>();
if types.is_empty() {
None
@@ -172,15 +200,15 @@ impl JsonParamType {
}
/// Normalize one JSON schema type name.
fn from_type_name(kind: &str) -> Option<Self> {
fn from_type_name(kind: &str, schema: Option<&Map<String, Value>>) -> Option<Self> {
let kind = kind.trim().to_ascii_lowercase();
match kind.as_str() {
"string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String),
"integer" | "int" => Some(Self::Integer),
"number" | "float" | "double" => Some(Self::Number),
"boolean" | "bool" | "binary" => Some(Self::Boolean),
"object" | "dict" | "map" => Some(Self::Object),
"array" | "arr" | "list" | "sequence" => Some(Self::Array),
"object" | "dict" | "map" => Some(Self::object_from_schema(schema)),
"array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)),
"null" => Some(Self::Null),
_ if kind.starts_with("int")
|| kind.starts_with("uint")
@@ -191,12 +219,52 @@ impl JsonParamType {
Some(Self::Integer)
}
_ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number),
_ if kind.starts_with("dict") => Some(Self::Object),
_ if kind.starts_with("list") => Some(Self::Array),
_ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)),
_ if kind.starts_with("list") => Some(Self::array_from_schema(schema)),
_ => None,
}
}
/// Normalize object schema fields.
fn object_from_schema(schema: Option<&Map<String, Value>>) -> Self {
let properties = schema
.and_then(|schema| schema.get("properties"))
.and_then(Value::as_object)
.map(|properties| {
properties
.iter()
.filter_map(|(name, schema)| {
Self::from_schema(schema).map(|param_type| (name.clone(), param_type))
})
.collect()
})
.unwrap_or_default();
let additional_properties =
schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| {
if schema.is_object() {
Self::from_schema(schema).map(Box::new)
} else {
None
}
});
Self::Object {
properties,
additional_properties,
}
}
/// Normalize array schema fields.
fn array_from_schema(schema: Option<&Map<String, Value>>) -> Self {
let items = schema
.and_then(|schema| schema.get("items"))
.and_then(Self::from_schema)
.map(Box::new);
Self::Array { items }
}
/// Collapse a candidate type list into one normalized type.
fn one_of(mut types: Vec<Self>) -> Self {
if types.len() == 1 {
@@ -207,23 +275,126 @@ impl JsonParamType {
}
}
/// Convert one raw string value to a normalized JSON type.
fn convert_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
match param_type {
JsonParamType::String => Some(Value::String(value.to_string())),
JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
JsonParamType::Number => convert_number(value),
JsonParamType::Boolean => convert_boolean(value),
JsonParamType::Object | JsonParamType::Array => serde_json::from_str(value).ok(),
JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
JsonParamType::OneOf(types) => {
types.iter().find_map(|param_type| convert_value(param_type, value))
/// Convert one parameter input to a normalized JSON value.
fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value {
// For literal `null`, always convert to JSON null value.
if let ParamInput::Text(value) = input
&& value.eq_ignore_ascii_case("null")
{
return Value::Null;
}
// If we have a schema, try to convert the value using it.
if let Some(param_type) = param_type
&& let Some(value) = try_convert_value(param_type, input)
{
return value;
}
// We don't have a schema, or conversion failed, use fallback logic.
match input {
ParamInput::Text(value) => Value::String(value.clone()),
ParamInput::Elements(elements) => {
// Convert structured input to object without a schema.
Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None))
}
}
}
/// Convert one parameter input to a normalized JSON type.
fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option<Value> {
match input {
ParamInput::Text(value) => try_convert_text_value(param_type, value),
ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements),
}
}
/// Convert one raw string value to a normalized JSON type.
fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
match param_type {
JsonParamType::String => Some(Value::String(value.to_string())),
JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
JsonParamType::Number => try_convert_number(value),
JsonParamType::Boolean => try_convert_boolean(value),
JsonParamType::Object { .. } if value.is_empty() => Some(Value::Object(Map::new())),
JsonParamType::Array { .. } if value.is_empty() => Some(Value::Array(Vec::new())),
JsonParamType::Object { .. } | JsonParamType::Array { .. } => {
// For composite types with string input, simply interpret the string as JSON.
serde_json::from_str(value).ok()
}
JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
JsonParamType::OneOf(types) => {
types.iter().find_map(|param_type| try_convert_text_value(param_type, value))
}
}
}
/// Convert one structured parameter input to a normalized JSON type.
fn try_convert_elements_value(
param_type: &JsonParamType,
elements: &[ParamElement],
) -> Option<Value> {
match param_type {
JsonParamType::Object {
properties,
additional_properties,
} => Some(Value::Object(convert_elements_to_object(
elements,
properties,
additional_properties.as_deref(),
))),
JsonParamType::Array { items } => Some(Value::Array(
// Collect all child elements into an array, regardless of their names.
elements
.iter()
.map(|element| convert_with_optional_schema(items.as_deref(), &element.value))
.collect(),
)),
JsonParamType::OneOf(types) => types
.iter()
.find_map(|param_type| try_convert_elements_value(param_type, elements)),
// Primitive types can't be converted from structured input.
JsonParamType::String
| JsonParamType::Integer
| JsonParamType::Number
| JsonParamType::Boolean
| JsonParamType::Null => None,
}
}
/// Convert structured elements to an object, using field schemas when present.
fn convert_elements_to_object(
elements: &[ParamElement],
properties: &BTreeMap<String, JsonParamType>,
additional_properties: Option<&JsonParamType>,
) -> Map<String, Value> {
let mut object = Map::with_capacity(elements.len());
for element in elements {
let param_type = properties.get(&element.name).or(additional_properties);
let value = convert_with_optional_schema(param_type, &element.value);
insert_object_value(&mut object, element.name.clone(), value);
}
object
}
/// Insert an object field while preserving duplicate keys as arrays.
fn insert_object_value(object: &mut Map<String, Value>, key: String, value: Value) {
if let Some(existing) = object.get_mut(&key) {
match existing {
// Collect values under the same key into an array.
Value::Array(values) => values.push(value),
existing => {
let first = std::mem::replace(existing, Value::Null);
*existing = Value::Array(vec![first, value]);
}
}
} else {
object.insert(key, value);
}
}
/// Convert one raw string value to a JSON number.
fn convert_number(value: &str) -> Option<Value> {
fn try_convert_number(value: &str) -> Option<Value> {
serde_json::from_str::<Number>(value)
.or_else(|_| value.parse::<i64>().map(Number::from))
.or_else(|_| value.parse::<f64>().ok().and_then(Number::from_f64).ok_or(()))
@@ -232,7 +403,7 @@ fn convert_number(value: &str) -> Option<Value> {
}
/// Convert one raw string value to a boolean.
fn convert_boolean(value: &str) -> Option<Value> {
fn try_convert_boolean(value: &str) -> Option<Value> {
match value.trim().to_ascii_lowercase().as_str() {
"true" | "1" => Some(Value::Bool(true)),
"false" | "0" => Some(Value::Bool(false)),
@@ -242,9 +413,9 @@ fn convert_boolean(value: &str) -> Option<Value> {
#[cfg(test)]
mod tests {
use serde_json::json;
use serde_json::{Value, json};
use super::{ToolSchema, ToolSchemas};
use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas};
use crate::Tool;
fn test_tool(name: &str, parameters: serde_json::Value) -> Tool {
@@ -260,8 +431,8 @@ mod tests {
fn invalid_schema_converts_everything_as_string() {
let params = ToolSchema::from_schema(&json!({ "type": "object" }));
assert_eq!(params.convert("count", "42"), json!("42"));
assert_eq!(params.convert("count", "null"), json!(null));
assert_eq!(params.convert("count", text("42")), json!("42"));
assert_eq!(params.convert("count", text("null")), json!(null));
}
#[test]
@@ -275,9 +446,9 @@ mod tests {
}
}));
assert_eq!(params.convert("unknown_schema", "42"), json!("42"));
assert_eq!(params.convert("unknown_type", "42"), json!("42"));
assert_eq!(params.convert("known", "42"), json!(42));
assert_eq!(params.convert("unknown_schema", text("42")), json!("42"));
assert_eq!(params.convert("unknown_type", text("42")), json!("42"));
assert_eq!(params.convert("known", text("42")), json!(42));
}
#[test]
@@ -298,16 +469,25 @@ mod tests {
}
}));
assert_eq!(params.convert("text", "42"), json!("42"));
assert_eq!(params.convert("count", "42"), json!(42));
assert_eq!(params.convert("size", "5.0"), json!(5.0));
assert_eq!(params.convert("ratio", "2.5"), json!(2.5));
assert_eq!(params.convert("enabled", "1"), json!(true));
assert_eq!(params.convert("payload", r#"{"k":1}"#), json!({ "k": 1 }));
assert_eq!(params.convert("mapping", r#"{"k":1}"#), json!({ "k": 1 }));
assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
assert_eq!(params.convert("names", r#"["a","b"]"#), json!(["a", "b"]));
assert_eq!(params.convert("nothing", "null"), json!(null));
assert_eq!(params.convert("text", text("42")), json!("42"));
assert_eq!(params.convert("count", text("42")), json!(42));
assert_eq!(params.convert("size", text("5.0")), json!(5.0));
assert_eq!(params.convert("ratio", text("2.5")), json!(2.5));
assert_eq!(params.convert("enabled", text("1")), json!(true));
assert_eq!(
params.convert("payload", text(r#"{"k":1}"#)),
json!({ "k": 1 })
);
assert_eq!(
params.convert("mapping", text(r#"{"k":1}"#)),
json!({ "k": 1 })
);
assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
assert_eq!(
params.convert("names", text(r#"["a","b"]"#)),
json!(["a", "b"])
);
assert_eq!(params.convert("nothing", text("null")), json!(null));
}
#[test]
@@ -333,7 +513,22 @@ mod tests {
}
fn converted_number_text(params: &ToolSchema, value: &str) -> String {
serde_json::to_string(&params.convert("value", value)).unwrap()
serde_json::to_string(&params.convert("value", text(value))).unwrap()
}
fn text(value: &str) -> ParamInput {
ParamInput::Text(value.to_string())
}
fn elem(name: &str, value: ParamInput) -> ParamElement {
ParamElement {
name: name.to_string(),
value,
}
}
fn elements(elements: Vec<ParamElement>) -> ParamInput {
ParamInput::Elements(elements)
}
#[test]
@@ -350,12 +545,12 @@ mod tests {
}
}));
assert_eq!(params.convert("s", "x"), json!("x"));
assert_eq!(params.convert("i", "7"), json!(7));
assert_eq!(params.convert("n", "7.5"), json!(7.5));
assert_eq!(params.convert("b", "true"), json!(true));
assert_eq!(params.convert("a", "[1]"), json!([1]));
assert_eq!(params.convert("o", r#"{"x":1}"#), json!({ "x": 1 }));
assert_eq!(params.convert("s", text("x")), json!("x"));
assert_eq!(params.convert("i", text("7")), json!(7));
assert_eq!(params.convert("n", text("7.5")), json!(7.5));
assert_eq!(params.convert("b", text("true")), json!(true));
assert_eq!(params.convert("a", text("[1]")), json!([1]));
assert_eq!(params.convert("o", text(r#"{"x":1}"#)), json!({ "x": 1 }));
}
#[test]
@@ -373,8 +568,8 @@ mod tests {
}
}));
assert_eq!(integer_first.convert("value", "42"), json!(42));
assert_eq!(string_first.convert("value", "42"), json!("42"));
assert_eq!(integer_first.convert("value", text("42")), json!(42));
assert_eq!(string_first.convert("value", text("42")), json!("42"));
}
#[test]
@@ -396,9 +591,9 @@ mod tests {
}
}));
assert_eq!(params.convert("choice", "42"), json!(42));
assert_eq!(params.convert("choice", text("42")), json!(42));
assert_eq!(
params.convert("fallback_object", r#"{"x":1}"#),
params.convert("fallback_object", text(r#"{"x":1}"#)),
json!({ "x": 1 })
);
}
@@ -414,9 +609,12 @@ mod tests {
}
}));
assert_eq!(params.convert("choice", "a"), json!("a"));
assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
assert_eq!(params.convert("payload", r#"{"x":1}"#), json!({ "x": 1 }));
assert_eq!(params.convert("choice", text("a")), json!("a"));
assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
assert_eq!(
params.convert("payload", text(r#"{"x":1}"#)),
json!({ "x": 1 })
);
}
#[test]
@@ -518,4 +716,162 @@ mod tests {
assert_eq!(converted.get("topn"), Some(&json!("5")));
assert_eq!(converted.get("nullish"), Some(&json!(null)));
}
#[test]
fn converts_structured_inputs_with_recursive_schema() {
let schemas = ToolSchemas::from_tools(&[test_tool(
"create_order",
json!({
"type": "object",
"properties": {
"user_id": { "type": "integer" },
"urgent": { "type": "boolean" },
"note": { "type": "string" },
"nil": { "type": "string" },
"shipping": {
"type": "object",
"properties": {
"city": { "type": "string" },
"zip": { "type": "integer" }
}
},
"items": {
"type": "array",
"items": {
"type": "object",
"properties": {
"sku": { "type": "string" },
"qty": { "type": "integer" }
}
}
},
"metadata": {
"type": "object",
"additionalProperties": { "type": "integer" }
},
"duplicate_demo": {
"type": "object",
"properties": {
"tag": { "type": "string" }
}
},
"schema_mismatch_array": {
"type": "array",
"items": { "type": "integer" }
},
"closed_object": {
"type": "object",
"additionalProperties": false
},
"open_object": {
"type": "object",
"additionalProperties": true
},
"payload_text": { "type": "object" },
"items_text": { "type": "array" }
}
}),
)]);
let converted = schemas.convert_params_with_schema(
"create_order",
vec![
("user_id".to_string(), text("42")),
("urgent".to_string(), text("true")),
("note".to_string(), text("Please leave at front desk.")),
("nil".to_string(), text("NULL")),
(
"shipping".to_string(),
elements(vec![
elem("city", text("Singapore")),
elem("zip", text("018956")),
]),
),
(
"items".to_string(),
elements(vec![
elem(
"item1",
elements(vec![elem("sku", text("book-001")), elem("qty", text("2"))]),
),
elem(
"item2",
elements(vec![elem("sku", text("pen-007")), elem("qty", text("5"))]),
),
]),
),
(
"metadata".to_string(),
elements(vec![elem("score", text("42")), elem("rank", text("7"))]),
),
(
"duplicate_demo".to_string(),
elements(vec![elem("tag", text("a")), elem("tag", text("b"))]),
),
(
"closed_object".to_string(),
elements(vec![elem("unknown", text("x"))]),
),
(
"open_object".to_string(),
elements(vec![elem("unknown", text("y"))]),
),
("payload_text".to_string(), text(r#"{"x":1}"#)),
("items_text".to_string(), text("[1,2]")),
(
"unknown_struct".to_string(),
elements(vec![
elem("a", text("1")),
elem("a", text("2")),
elem("nil", text("null")),
]),
),
],
);
assert_eq!(
Value::Object(converted),
json!({
"user_id": 42,
"urgent": true,
"note": "Please leave at front desk.",
"nil": null,
"shipping": {
"city": "Singapore",
"zip": 18956
},
"items": [
{
"sku": "book-001",
"qty": 2
},
{
"sku": "pen-007",
"qty": 5
}
],
"metadata": {
"score": 42,
"rank": 7
},
"duplicate_demo": {
"tag": ["a", "b"]
},
"closed_object": {
"unknown": "x"
},
"open_object": {
"unknown": "y"
},
"payload_text": {
"x": 1
},
"items_text": [1, 2],
"unknown_struct": {
"a": ["1", "2"],
"nil": null
}
})
);
}
}
@@ -2,11 +2,15 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from argparse import Namespace
from io import BytesIO
from pathlib import Path
from typing import Any
import pybase64 as base64
import pytest
from PIL import Image
import vllm.benchmarks.datasets.datasets as datasets_module
from vllm.benchmarks.datasets import CustomImageDataset, get_samples
from vllm.benchmarks.lib.endpoint_request_func import (
RequestFuncInput,
@@ -33,6 +37,22 @@ def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
f.write(json.dumps(row) + "\n")
def _write_png(path: Path, color: tuple[int, int, int] = (255, 0, 0)) -> None:
Image.new("RGB", (1, 1), color=color).save(path)
def _decode_data_url(data_url: str) -> tuple[str, bytes]:
prefix, image_base64 = data_url.split(",", 1)
return prefix, base64.b64decode(image_base64)
def _assert_png_data_url(data_url: str) -> None:
prefix, image_bytes = _decode_data_url(data_url)
assert prefix == "data:image/png;base64"
with Image.open(BytesIO(image_bytes)) as image:
image.verify()
def _args_for_custom_image(dataset_path: Path) -> Namespace:
return Namespace(
dataset_name="custom_image",
@@ -42,6 +62,7 @@ def _args_for_custom_image(dataset_path: Path) -> Namespace:
num_prompts=2,
custom_output_len=32,
enable_multimodal_chat=False,
custom_ensure_client_side_data=False,
request_id_prefix="req-",
no_oversample=False,
)
@@ -230,6 +251,125 @@ def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat(
assert _get_chat_messages(request_input) == sample.prompt
@pytest.mark.benchmark
def test_custom_image_dataset_encodes_image_media_when_requested(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
image_a = tmp_path / "chart_a.png"
image_b = tmp_path / "chart b.png"
_write_png(image_a, color=(255, 0, 0))
_write_png(image_b, color=(0, 255, 0))
data_url = "data:image/png;base64,Zm9v"
remote_url = "https://example.com/chart.png"
original_fetch_image = datasets_module.fetch_image
def fake_fetch_image(image_url: str) -> Image.Image:
if image_url == remote_url:
return Image.new("RGB", (1, 1), color=(0, 0, 255))
return original_fetch_image(image_url)
monkeypatch.setattr(datasets_module, "fetch_image", fake_fetch_image)
jsonl = tmp_path / "images.jsonl"
_write_jsonl(
jsonl,
[
{
"prompt": "Compare the charts.",
"image_files": [
str(image_a),
image_b.as_uri(),
remote_url,
data_url,
],
}
],
)
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
samples = dataset.sample(
tokenizer=_Tokenizer(),
num_requests=1,
output_len=32,
ensure_client_side_data=True,
)
assert len(samples) == 1
assert isinstance(samples[0].multi_modal_data, list)
image_urls = [part["image_url"]["url"] for part in samples[0].multi_modal_data]
_assert_png_data_url(image_urls[0])
_assert_png_data_url(image_urls[1])
_assert_png_data_url(image_urls[2])
assert image_urls[3] == data_url
@pytest.mark.benchmark
def test_custom_image_dataset_encodes_interleaved_image_media(
tmp_path: Path,
) -> None:
image_a = tmp_path / "chart_a.png"
image_b = tmp_path / "chart_b.png"
_write_png(image_a, color=(255, 0, 0))
_write_png(image_b, color=(0, 255, 0))
jsonl = tmp_path / "images.jsonl"
_write_jsonl(
jsonl,
[
{
"content": [
{"type": "text", "text": "Compare "},
{"type": "image", "image": str(image_a)},
{
"type": "image_url",
"image_url": {
"url": image_b.as_uri(),
"detail": "low",
},
},
],
}
],
)
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
samples = dataset.sample(
tokenizer=_Tokenizer(),
num_requests=1,
output_len=32,
ensure_client_side_data=True,
)
sample = samples[0]
assert isinstance(sample.prompt, list)
_assert_png_data_url(sample.prompt[1]["image_url"]["url"])
_assert_png_data_url(sample.prompt[2]["image_url"]["url"])
assert sample.prompt[2]["image_url"]["detail"] == "low"
@pytest.mark.benchmark
def test_custom_image_dataset_rejects_invalid_image_media(
tmp_path: Path,
) -> None:
invalid_image = tmp_path / "not_an_image.png"
invalid_image.write_text("not an image")
jsonl = tmp_path / "images.jsonl"
_write_jsonl(
jsonl,
[{"prompt": "Describe the image.", "image_files": [str(invalid_image)]}],
)
dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
with pytest.raises(ValueError, match="Invalid image URL"):
dataset.sample(
tokenizer=_Tokenizer(),
num_requests=1,
output_len=32,
ensure_client_side_data=True,
)
@pytest.mark.benchmark
def test_custom_image_dataset_rejects_invalid_content_part(
tmp_path: Path,
@@ -55,6 +55,15 @@ def test_dynamic_shapes_compilation(
evaluate_guards,
):
"""Test that all dynamic shapes types compile successfully"""
if shapes_type == DynamicShapesType.UNBACKED and not is_torch_equal_or_newer(
"2.11.0"
):
# NOTE[ROCm]: shape_id (used by Qwen2/Llama to relate input dims) only
# landed in torch 2.11, but the ROCm CI still runs torch 2.10.x. On
# older torch there's no way to express it, so unbacked shapes go
# data-dependent and compilation blows up -- nothing to test.
pytest.skip("unbacked dynamic shapes with shape_id require torch>=2.11")
if evaluate_guards and shapes_type == DynamicShapesType.UNBACKED:
pytest.skip("unbacked dynamic shapes do not add guards")
+1 -1
View File
@@ -284,7 +284,7 @@ def test_multiproc_executor_pipeline_parallel():
assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)"
# Verify max_concurrent_batches for pipeline parallel
assert executor.max_concurrent_batches == 2, (
assert vllm_config.max_concurrent_batches == 2, (
"Max concurrent batches should equal PP size"
)
+1 -1
View File
@@ -83,7 +83,7 @@ def assert_executor(executor, tp_size, pp_size):
assert executor._get_output_rank() == expected_output_rank
if pp_size > 1:
assert executor.max_concurrent_batches == pp_size
assert executor.vllm_config.max_concurrent_batches == pp_size
executor.check_health()
assert not executor.is_failed
+330 -3
View File
@@ -18,6 +18,7 @@ from torch.multiprocessing.reductions import reduce_tensor
from vllm.config.parallel import ParallelConfig
from vllm.config.weight_transfer import WeightTransferConfig
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
from vllm.distributed.weight_transfer.base import SparseWeightPatch
from vllm.distributed.weight_transfer.ipc_engine import (
IPCWeightTransferEngine,
IPCWeightTransferInitInfo,
@@ -89,6 +90,67 @@ class TestNCCLWeightTransferUpdateInfoValidation:
)
assert len(info.names) == 0
def test_valid_sparse_update_info(self):
"""Test creating valid sparse NCCL update info."""
info = NCCLWeightTransferUpdateInfo(
names=["layer.weight", "layer.bias"],
dtype_names=["float32", "bfloat16"],
shapes=[[10, 10], [10]],
num_updates_list=[4, 2],
update_kind="sparse_flat",
)
assert info.update_kind == "sparse_flat"
assert info.num_updates_list == [4, 2]
def test_sparse_update_requires_num_updates_list(self):
with pytest.raises(ValueError, match="`num_updates_list` is required"):
NCCLWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
update_kind="sparse_flat",
)
def test_sparse_update_rejects_empty_num_updates_list(self):
with pytest.raises(ValueError, match="cannot be empty"):
NCCLWeightTransferUpdateInfo(
names=[],
dtype_names=[],
shapes=[],
num_updates_list=[],
update_kind="sparse_flat",
)
def test_sparse_update_rejects_packed(self):
with pytest.raises(ValueError, match="cannot be combined with `packed=True`"):
NCCLWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
num_updates_list=[3],
update_kind="sparse_flat",
packed=True,
)
def test_sparse_update_rejects_mismatched_num_updates(self):
with pytest.raises(ValueError, match="`num_updates_list`"):
NCCLWeightTransferUpdateInfo(
names=["layer.weight", "layer.bias"],
dtype_names=["float32", "float32"],
shapes=[[10, 10], [10]],
num_updates_list=[3],
update_kind="sparse_flat",
)
def test_dense_update_rejects_sparse_metadata(self):
with pytest.raises(ValueError, match="Sparse metadata"):
NCCLWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
num_updates_list=[3],
)
# --- Unit Tests: Engine Parsing ---
@@ -222,6 +284,27 @@ def test_nccl_receive_weights_without_init_raises():
engine.receive_weights(update_info, lambda x: None)
def test_nccl_receive_sparse_weights_without_init_raises():
"""Test that sparse receive raises if init_transfer_engine wasn't called."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
config = WeightTransferConfig(backend="nccl")
parallel_config = create_mock_parallel_config()
engine = NCCLWeightTransferEngine(config, parallel_config)
update_info = NCCLWeightTransferUpdateInfo(
names=["w"],
dtype_names=["float32"],
shapes=[[10]],
num_updates_list=[2],
update_kind="sparse_flat",
)
with pytest.raises(RuntimeError, match="not initialized"):
engine.receive_sparse_weights(update_info, lambda x: None)
# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
@@ -379,6 +462,136 @@ def test_nccl_weight_transfer_between_processes():
)
@ray.remote(num_gpus=1)
def trainer_broadcast_sparse_tensor(
master_address: str,
master_port: int,
world_size: int,
) -> bool:
"""Trainer task that broadcasts sparse patches via NCCL."""
import torch
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
from vllm.distributed.weight_transfer.base import SparseWeightPatch
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
pg = StatelessProcessGroup.create(
host=master_address,
port=master_port,
rank=0,
world_size=world_size,
)
comm = PyNcclCommunicator(pg, device=0)
patch = SparseWeightPatch(
name="test.weight",
indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"),
values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"),
)
NCCLWeightTransferEngine.trainer_send_sparse_weights(
iter([patch]),
NCCLTrainerSendWeightsArgs(group=comm),
)
torch.accelerator.synchronize()
return True
@ray.remote(num_gpus=1)
def inference_receive_sparse_tensor(
master_address: str,
master_port: int,
world_size: int,
) -> dict:
"""Inference task that receives sparse patches via NCCLWeightTransferEngine."""
from unittest.mock import MagicMock
import torch
from vllm.config.parallel import ParallelConfig
from vllm.config.weight_transfer import WeightTransferConfig
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLWeightTransferEngine,
NCCLWeightTransferInitInfo,
NCCLWeightTransferUpdateInfo,
)
config = WeightTransferConfig(backend="nccl")
parallel_config = MagicMock(spec=ParallelConfig)
parallel_config.rank = 0
parallel_config.world_size = 1
parallel_config.data_parallel_rank = 0
parallel_config.data_parallel_index = 0
engine = NCCLWeightTransferEngine(config, parallel_config)
engine.init_transfer_engine(
NCCLWeightTransferInitInfo(
master_address=master_address,
master_port=master_port,
rank_offset=1,
world_size=world_size,
)
)
target = torch.zeros(30, dtype=torch.float32, device="cuda")
def apply_sparse_patches(patches: list[SparseWeightPatch]):
for patch in patches:
target.index_copy_(0, patch.indices.to(torch.long), patch.values)
update_info = NCCLWeightTransferUpdateInfo(
names=["test.weight"],
dtype_names=["float32"],
shapes=[[30]],
num_updates_list=[3],
update_kind="sparse_flat",
)
engine.receive_sparse_weights(update_info, apply_sparse_patches)
torch.accelerator.synchronize()
expected = torch.zeros(30, dtype=torch.float32, device="cuda")
expected[[1, 7, 25]] = torch.tensor(
[10.0, 20.0, 30.0], dtype=torch.float32, device="cuda"
)
success = torch.equal(target, expected)
engine.shutdown()
return {
"success": success,
"selected_values": target[[1, 7, 25]].cpu().tolist(),
}
@pytest.mark.skipif(
torch.accelerator.device_count() < 2,
reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
)
def test_nccl_sparse_weight_transfer_between_processes():
"""Test NCCL sparse weight transfer from trainer to inference process."""
ray.init(ignore_reinit_error=True)
master_address = "127.0.0.1"
master_port = get_open_port()
world_size = 2
inference_future = inference_receive_sparse_tensor.remote(
master_address, master_port, world_size
)
trainer_future = trainer_broadcast_sparse_tensor.remote(
master_address, master_port, world_size
)
trainer_result, result = ray.get([trainer_future, inference_future])
assert trainer_result, "Trainer should complete successfully"
assert result["success"], (
"Sparse weight transfer failed. "
f"Received selected values: {result['selected_values']}"
)
# --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
@@ -461,9 +674,101 @@ class TestIPCWeightTransferUpdateInfoValidation:
ipc_handles=ipc_handles,
)
def test_missing_ipc_handles_raises(self):
"""Test that omitting ipc_handles raises TypeError."""
with pytest.raises(TypeError):
def test_sparse_update_kind_rejected(self):
"""Test that IPC backend rejects sparse update metadata."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
dummy_tensor = torch.ones(10, 10, device="cuda:0")
ipc_handle = reduce_tensor(dummy_tensor)
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
ipc_handles = [{gpu_uuid: ipc_handle}]
with pytest.raises(NotImplementedError, match="dense updates"):
IPCWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
num_updates_list=[1],
ipc_handles=ipc_handles,
update_kind="sparse_flat",
)
def test_sparse_methods_not_supported(self):
"""Test that IPC engine inherits sparse rejection from the base class."""
config = WeightTransferConfig(backend="ipc")
parallel_config = create_mock_parallel_config()
engine = IPCWeightTransferEngine(
config, parallel_config, MagicMock(spec=torch.nn.Module)
)
with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
engine.receive_sparse_weights(MagicMock(), lambda _: None)
with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
engine.trainer_send_sparse_weights(
iter([]),
{"mode": "http", "url": "http://localhost:8000"},
)
def test_valid_update_info_from_pickled(self, monkeypatch):
"""Test creating IPCWeightTransferUpdateInfo from pickled handles."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
dummy_tensor = torch.ones(10, 10, device="cuda:0")
ipc_handle = reduce_tensor(dummy_tensor)
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
ipc_handles = [{gpu_uuid: ipc_handle}]
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
info = IPCWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
ipc_handles_pickled=pickled,
)
assert info.ipc_handles == ipc_handles
assert info.ipc_handles_pickled is None
def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
"""Test that pickled handles are rejected unless env flag is enabled."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
IPCWeightTransferUpdateInfo(
names=[],
dtype_names=[],
shapes=[],
ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
)
def test_both_handles_and_pickled_raises(self):
"""Test that providing both ipc_handles and ipc_handles_pickled raises."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
dummy_tensor = torch.ones(10, 10, device="cuda:0")
ipc_handle = reduce_tensor(dummy_tensor)
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
ipc_handles = [{gpu_uuid: ipc_handle}]
pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
with pytest.raises(ValueError, match="Cannot specify both"):
IPCWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
shapes=[[10, 10]],
ipc_handles=ipc_handles,
ipc_handles_pickled=pickled,
)
def test_neither_handles_nor_pickled_raises(self):
"""Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
with pytest.raises(ValueError, match="must be provided"):
IPCWeightTransferUpdateInfo(
names=["layer.weight"],
dtype_names=["float32"],
@@ -558,6 +863,28 @@ class TestIPCEngineParsing:
assert gpu_uuid in update_info.ipc_handles[0]
assert gpu_uuid in update_info.ipc_handles[1]
def test_parse_update_info_ignores_none_pickled_handles(self):
"""Test Ray/asdict payloads with a null pickled field use ipc_handles."""
config = WeightTransferConfig(backend="ipc")
parallel_config = create_mock_parallel_config()
engine = IPCWeightTransferEngine(
config, parallel_config, MagicMock(spec=torch.nn.Module)
)
ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
update_info = engine.parse_update_info(
{
"names": ["w1"],
"dtype_names": ["float32"],
"shapes": [[1]],
"ipc_handles": ipc_handles,
"ipc_handles_pickled": None,
}
)
assert isinstance(update_info, IPCWeightTransferUpdateInfo)
assert update_info.ipc_handles == ipc_handles
def test_parse_update_info_both_handles_and_pickled_raises(self):
"""Test that providing both ipc_handles and ipc_handles_pickled raises."""
if torch.accelerator.device_count() < 1:
@@ -635,3 +635,143 @@ class TestThinkingBlockConversion:
# Redacted thinking is ignored, normal thinking still becomes reasoning.
assert asst.get("reasoning") == "Thinking..."
assert asst.get("content") == "Hi!"
class TestInlineSystemMessageInMessagesArray:
"""Verify that ``role: system`` messages embedded inside the ``messages``
array are accepted and merged with the top-level ``system`` prompt.
This handles clients that place system messages inside the messages array
instead of the Anthropic-standard top-level ``system`` field.
"""
def test_inline_system_merged_with_top_level_system(self):
"""Full integration: inline system + top-level system + user message."""
request = _make_request(
[
{
"role": "user",
"content": [
{
"type": "text",
"text": "<system-reminder>\n.....\n</system-reminder>\n\n",
},
{
"type": "text",
"text": "help?",
"cache_control": {"type": "ephemeral"},
},
],
},
{
"role": "system",
"content": ".....",
},
],
system=[
{
"type": "text",
"text": "x-anthropic-billing-header: "
"cc_version=2.1.160.bca; cc_entrypoint=cli; cch=d1d48;",
},
{
"type": "text",
"text": "You are Claude Code, Anthropic's official CLI for Claude.",
"cache_control": {"type": "ephemeral"},
},
{
"type": "text",
"text": "....",
"cache_control": {"type": "ephemeral"},
},
],
tools=[],
)
result = _convert(request)
# First message should be the merged system prompt.
assert result.messages[0]["role"] == "system"
# Billing header stripped, inline system appended.
assert (
result.messages[0]["content"]
== "You are Claude Code, Anthropic's official CLI for Claude."
"...."
"....."
)
# Second message should be the user message, content preserved.
assert result.messages[1]["role"] == "user"
user_content = result.messages[1]["content"]
assert len(user_content) == 2
assert user_content[0] == {
"type": "text",
"text": "<system-reminder>\n.....\n</system-reminder>\n\n",
}
assert user_content[1] == {
"type": "text",
"text": "help?",
}
def test_inline_system_string_only(self):
"""Only an inline system string, no top-level system."""
request = _make_request(
[
{"role": "user", "content": "Hello"},
{"role": "system", "content": "Be concise."},
]
)
result = _convert(request)
assert result.messages[0]["role"] == "system"
assert result.messages[0]["content"] == "Be concise."
assert result.messages[1]["role"] == "user"
def test_inline_system_list_content(self):
"""Inline system with list content blocks."""
request = _make_request(
[
{"role": "user", "content": "Hi"},
{
"role": "system",
"content": [
{"type": "text", "text": "Part one. "},
{"type": "text", "text": "Part two."},
],
},
]
)
result = _convert(request)
assert result.messages[0]["role"] == "system"
assert result.messages[0]["content"] == "Part one. Part two."
def test_multiple_inline_system_messages(self):
"""Multiple inline system messages should all be merged."""
request = _make_request(
[
{"role": "system", "content": "First system."},
{"role": "user", "content": "Hello"},
{"role": "system", "content": "Second system."},
]
)
result = _convert(request)
assert result.messages[0]["role"] == "system"
assert result.messages[0]["content"] == "First system.Second system."
assert result.messages[1]["role"] == "user"
def test_inline_system_with_top_level_string(self):
"""Top-level system is a string, inline system is also present."""
request = _make_request(
[
{"role": "user", "content": "Hello"},
{"role": "system", "content": "Inline hint."},
],
system="Top-level prompt.",
)
result = _convert(request)
assert result.messages[0]["role"] == "system"
assert result.messages[0]["content"] == "Top-level prompt.Inline hint."
assert result.messages[1]["role"] == "user"
@@ -18,13 +18,13 @@ from unittest.mock import MagicMock
import pytest
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.generative_scoring.serving import (
from vllm.entrypoints.generate.generative_scoring.serving import (
GenerativeScoringItemResult,
GenerativeScoringRequest,
GenerativeScoringResponse,
OpenAIServingGenerativeScoring,
ServingGenerativeScoring,
)
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.models.protocol import BaseModelPath
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.logprobs import Logprob
@@ -86,13 +86,13 @@ def _create_mock_engine():
return mock_engine
def _create_serving(mock_engine) -> OpenAIServingGenerativeScoring:
"""Create an OpenAIServingGenerativeScoring instance with mocks."""
def _create_serving(mock_engine) -> ServingGenerativeScoring:
"""Create an ServingGenerativeScoring instance with mocks."""
models = OpenAIServingModels(
engine_client=mock_engine,
base_model_paths=BASE_MODEL_PATHS,
)
return OpenAIServingGenerativeScoring(mock_engine, models, request_logger=None)
return ServingGenerativeScoring(mock_engine, models, request_logger=None)
def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput:
@@ -186,7 +186,7 @@ class TestProbabilityComputation:
self, label_logprobs, apply_softmax, should_sum_to_one
):
"""Test probability computation for softmax and true probability modes."""
serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
probs = serving._compute_probabilities(
label_logprobs, apply_softmax=apply_softmax
)
@@ -211,7 +211,7 @@ class TestProbabilityComputation:
def test_score_formula(self):
"""Test the score formula: P(token[0]) / (P(token[0]) + P(token[1]))."""
serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
# With logprobs -0.5 and -2.0, softmax gives higher prob to first token
logprobs = {9454: -0.5, 2753: -2.0}
@@ -8,7 +8,7 @@ Tests verify the full HTTP request/response flow using RemoteOpenAIServer.
import pytest
import requests
from ....utils import RemoteOpenAIServer
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen3-0.6B"
@@ -14,8 +14,12 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
@pytest.fixture
@pytest.fixture(scope="module")
def server():
# Use module scope so the server is started once and shared across all
# tests in this file. Starting a new vLLM server per test on XPU can
# cause the second server startup to hang silently and exceed the
# wait-for-server timeout, resulting in RuntimeError.
args = [
"--max-model-len",
"16384",
@@ -6,9 +6,13 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic import ValidationError
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionRequest,
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine.protocol import GenerationError
from vllm.entrypoints.openai.models.protocol import BaseModelPath
@@ -444,3 +448,45 @@ def test_json_schema_response_format_missing_schema():
messages=[{"role": "user", "content": "hello"}],
response_format={"type": "json_schema"},
)
@pytest.mark.parametrize("format_value", [None, {}])
def test_structural_tag_response_format_invalid(format_value):
"""Malformed structural tags should be rejected during request validation."""
with pytest.raises(
ValidationError,
match="Invalid response_format structural_tag",
):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
response_format={"type": "structural_tag", "format": format_value},
)
@pytest.mark.parametrize("format_value", [None, {}])
def test_batch_structural_tag_response_format_invalid(format_value):
"""Batch chat should reject malformed structural tags at request parsing."""
with pytest.raises(
ValidationError,
match="Invalid response_format structural_tag",
):
BatchChatCompletionRequest(
model=MODEL_NAME,
messages=[[{"role": "user", "content": "hello"}]],
response_format={"type": "structural_tag", "format": format_value},
)
@pytest.mark.parametrize("structural_tag", ["not json", ""])
def test_structured_outputs_structural_tag_invalid(structural_tag):
"""Malformed direct structured_outputs structural tags should be rejected."""
with pytest.raises(
ValidationError,
match="Invalid structured_outputs structural_tag",
):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
structured_outputs={"structural_tag": structural_tag},
)
@@ -1935,8 +1935,10 @@ async def test_streaming_n_gt1_independent_tool_parsers():
finished=True,
)
# Collect tool-call deltas per choice from the SSE stream.
# Collect tool-call deltas and finish_reasons per choice from the SSE
# stream.
tc_deltas_by_choice: dict[int, list[dict]] = {i: [] for i in range(num_choices)}
finish_reasons_by_choice: dict[int, list[str]] = {i: [] for i in range(num_choices)}
async for chunk_str in serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -1959,6 +1961,8 @@ async def test_streaming_n_gt1_independent_tool_parsers():
if delta.get("tool_calls"):
for tc in delta["tool_calls"]:
tc_deltas_by_choice[idx].append(tc)
if choice.get("finish_reason") is not None:
finish_reasons_by_choice[idx].append(choice["finish_reason"])
# Both choices must independently produce the correct tool call.
for choice_idx in range(num_choices):
@@ -1984,141 +1988,11 @@ async def test_streaming_n_gt1_independent_tool_parsers():
f"Choice {choice_idx}: expected {{'city': 'Tokyo'}}, got {parsed_args}"
)
class TestCreateRemainingArgsDelta:
"""Tests for _create_remaining_args_delta helper function.
This helper is used when streaming tool calls to preserve id/type/name
fields in the finish chunk, which would otherwise be lost.
"""
def test_preserves_id_type_name(self):
"""Test that id, type, and name are preserved from original delta."""
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
reasons = finish_reasons_by_choice[choice_idx]
assert len(reasons) == 1, (
f"Choice {choice_idx}: expected exactly 1 finish_reason, got {reasons}"
)
original_delta = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=0,
id="call_abc123",
type="function",
function=DeltaFunctionCall(
name="get_weather",
arguments='{"location": "Paris"}',
),
)
]
assert reasons[0] == "tool_calls", (
f"Choice {choice_idx}: expected finish_reason='tool_calls', "
f"got '{reasons[0]}'"
)
result = OpenAIServingChat._create_remaining_args_delta(
original_delta, '", "unit": "celsius"}', 0
)
assert len(result.tool_calls) == 1
tc = result.tool_calls[0]
assert tc.index == 0
assert tc.id == "call_abc123"
assert tc.type == "function"
assert tc.function.name == "get_weather"
assert tc.function.arguments == '", "unit": "celsius"}'
def test_matches_by_index(self):
"""Test that the correct tool call is matched by index."""
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
)
original_delta = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=0,
id="call_first",
type="function",
function=DeltaFunctionCall(name="func_a", arguments="{}"),
),
DeltaToolCall(
index=1,
id="call_second",
type="function",
function=DeltaFunctionCall(name="func_b", arguments="{}"),
),
]
)
result = OpenAIServingChat._create_remaining_args_delta(
original_delta, '{"extra": true}', 1
)
assert len(result.tool_calls) == 1
tc = result.tool_calls[0]
assert tc.index == 1
assert tc.id == "call_second"
assert tc.function.name == "func_b"
def test_no_matching_tool_call(self):
"""Test graceful handling when no matching tool call is found."""
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
)
original_delta = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=0,
id="call_zero",
type="function",
function=DeltaFunctionCall(name="func", arguments="{}"),
)
]
)
result = OpenAIServingChat._create_remaining_args_delta(
original_delta, '{"arg": 1}', 5
)
assert len(result.tool_calls) == 1
tc = result.tool_calls[0]
assert tc.index == 5
assert tc.id is None
assert tc.type is None
assert tc.function.name is None
assert tc.function.arguments == '{"arg": 1}'
def test_function_is_none(self):
"""Test handling when original tool call has no function."""
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall
original_delta = DeltaMessage(
tool_calls=[
DeltaToolCall(
index=0,
id="call_nofunc",
type="function",
function=None,
)
]
)
result = OpenAIServingChat._create_remaining_args_delta(
original_delta, '{"data": "value"}', 0
)
assert len(result.tool_calls) == 1
tc = result.tool_calls[0]
assert tc.index == 0
assert tc.id == "call_nofunc"
assert tc.type == "function"
assert tc.function.name is None
assert tc.function.arguments == '{"data": "value"}'
@@ -6,6 +6,7 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from pydantic import ValidationError
from vllm.config.multimodal import MultiModalConfig
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
@@ -302,6 +303,36 @@ def test_json_schema_response_format_missing_schema():
)
@pytest.mark.parametrize("format_value", [None, {}])
def test_structural_tag_response_format_invalid(format_value):
"""Malformed structural tags should be rejected during request validation."""
with pytest.raises(
ValidationError,
match="Invalid response_format structural_tag",
):
CompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
max_tokens=10,
response_format={"type": "structural_tag", "format": format_value},
)
@pytest.mark.parametrize("structural_tag", ["not json", ""])
def test_structured_outputs_structural_tag_invalid(structural_tag):
"""Malformed direct structured_outputs structural tags should be rejected."""
with pytest.raises(
ValidationError,
match="Invalid structured_outputs structural_tag",
):
CompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
max_tokens=10,
structured_outputs={"structural_tag": structural_tag},
)
def test_negative_prompt_token_ids_nested():
"""Negative token IDs in prompt (nested list) should raise validation error."""
with pytest.raises(Exception, match="greater than or equal to 0"):
@@ -0,0 +1,382 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ResponsesParser with the unified Parser interface.
These tests verify that ResponsesParser correctly delegates to the unified
Parser (via extract_response_outputs) instead of calling separate
ReasoningParser / ToolParser instances directly.
"""
from collections.abc import Sequence
from unittest.mock import MagicMock
import pytest
from vllm.entrypoints.openai.engine.protocol import (
DeltaMessage,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.parser.responses_parser import (
ResponsesParser,
get_responses_parser_for_simple_context,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.outputs import CompletionOutput
from vllm.parser.abstract_parser import DelegatingParser
pytestmark = pytest.mark.skip_global_cleanup
# ---------------------------------------------------------------------------
# Test parser stubs
# ---------------------------------------------------------------------------
class _NoOpParser(DelegatingParser):
"""Parser that extracts no reasoning and no tool calls."""
def is_reasoning_end(self, input_ids: list[int]) -> bool:
return False
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
return input_ids
def extract_reasoning(self, model_output, request):
return None, model_output
def extract_reasoning_streaming(self, *args, **kwargs):
return None
def extract_tool_calls(self, model_output, request):
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
def extract_tool_calls_streaming(self, *args, **kwargs):
return None
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
return None
class _ReasoningOnlyParser(DelegatingParser):
"""Parser that extracts reasoning but no tool calls."""
def is_reasoning_end(self, input_ids: list[int]) -> bool:
return False
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
return input_ids
def extract_reasoning(self, model_output, request):
if "<think>" in model_output and "</think>" in model_output:
start = model_output.index("<think>") + len("<think>")
end = model_output.index("</think>")
reasoning = model_output[start:end]
content = model_output[end + len("</think>") :]
return reasoning, content.strip() or None
return None, model_output
def extract_reasoning_streaming(self, *args, **kwargs):
return None
def extract_tool_calls(self, model_output, request):
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
def extract_tool_calls_streaming(self, *args, **kwargs):
return None
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
return None
class _StubToolParser:
"""Minimal tool parser stub that always returns a hardcoded tool call."""
supports_required_and_named = False
def __init__(self, tokenizer=None, tools=None):
pass
def extract_tool_calls(self, model_output, request):
return ExtractedToolCallInformation(
tools_called=True,
tool_calls=[
ToolCall(
id="call_123",
type="function",
function=FunctionCall(
name="get_weather",
arguments='{"location": "Paris"}',
),
)
],
content=None,
)
def extract_tool_calls_streaming(self, *args, **kwargs):
return None
def adjust_request(self, request):
return request
class _ToolCallingParser(DelegatingParser):
"""Parser that extracts a hardcoded tool call from any input."""
def __init__(self, tokenizer, *args, **kwargs):
super().__init__(tokenizer)
self._tool_parser = _StubToolParser()
def is_reasoning_end(self, input_ids: list[int]) -> bool:
return False
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
return input_ids
def extract_reasoning(self, model_output, request):
return None, model_output
def extract_reasoning_streaming(self, *args, **kwargs):
return None
def extract_tool_calls_streaming(self, *args, **kwargs):
return None
def parse_delta(self, *args, **kwargs) -> DeltaMessage | None:
return None
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_request(**overrides) -> ResponsesRequest:
defaults = {"model": "test-model", "input": "test"}
defaults.update(overrides)
return ResponsesRequest.model_validate(defaults)
def _make_output(
text: str = "Hello, world!",
token_ids: Sequence[int] = (1, 2, 3),
finish_reason: str = "stop",
) -> CompletionOutput:
return CompletionOutput(
index=0,
text=text,
token_ids=list(token_ids),
cumulative_logprob=None,
logprobs=None,
finish_reason=finish_reason,
)
def _make_parser(parser_cls, **overrides):
defaults = dict(
tokenizer=MagicMock(),
parser_cls=parser_cls,
response_messages=[],
request=_make_request(),
chat_template=None,
chat_template_content_format="auto",
)
defaults.update(overrides)
return ResponsesParser(**defaults)
# ---------------------------------------------------------------------------
# Tests: basic text passthrough
# ---------------------------------------------------------------------------
def test_process_text_with_parser():
"""Parser with no reasoning/tools returns a single message item."""
parser = _make_parser(_NoOpParser)
parser.process(_make_output(text="Hello!"))
assert len(parser.response_messages) == 1
msg = parser.response_messages[0]
assert msg.type == "message"
assert msg.content[0].text == "Hello!"
def test_process_text_without_parser():
"""parser_cls=None falls back to plain text wrapping."""
parser = _make_parser(None)
parser.process(_make_output(text="Hello!"))
assert len(parser.response_messages) == 1
msg = parser.response_messages[0]
assert msg.type == "message"
assert msg.content[0].text == "Hello!"
# ---------------------------------------------------------------------------
# Tests: empty / whitespace output
# ---------------------------------------------------------------------------
def test_process_empty_text_without_parser():
"""Empty text with no parser produces no output items."""
parser = _make_parser(None)
parser.process(_make_output(text=""))
assert len(parser.response_messages) == 0
def test_process_empty_text_with_parser():
"""Empty text with parser produces no output items."""
parser = _make_parser(_NoOpParser)
parser.process(_make_output(text=""))
assert len(parser.response_messages) == 0
# ---------------------------------------------------------------------------
# Tests: reasoning extraction
# ---------------------------------------------------------------------------
def test_process_extracts_reasoning():
"""Parser that finds reasoning produces both reasoning and message items."""
parser = _make_parser(_ReasoningOnlyParser)
parser.process(_make_output(text="<think>Let me check</think>The answer is 42"))
types = [m.type for m in parser.response_messages]
assert "reasoning" in types
assert "message" in types
reasoning_item = next(m for m in parser.response_messages if m.type == "reasoning")
assert reasoning_item.content[0].text == "Let me check"
message_item = next(m for m in parser.response_messages if m.type == "message")
assert message_item.content[0].text == "The answer is 42"
def test_process_reasoning_only_no_content():
"""When reasoning consumes all text, only a reasoning item is produced."""
parser = _make_parser(_ReasoningOnlyParser)
parser.process(_make_output(text="<think>Just thinking</think>"))
types = [m.type for m in parser.response_messages]
assert "reasoning" in types
assert "message" not in types
# ---------------------------------------------------------------------------
# Tests: tool call extraction
# ---------------------------------------------------------------------------
def test_process_extracts_tool_calls():
"""Parser that finds tool calls produces function_call items."""
request = _make_request(
tool_choice="auto",
tools=[
{
"type": "function",
"name": "get_weather",
"parameters": {"type": "object", "properties": {}},
}
],
)
parser = _make_parser(_ToolCallingParser, request=request, enable_auto_tools=True)
parser.process(_make_output(text="calling tool"))
types = [m.type for m in parser.response_messages]
assert "function_call" in types
tool_item = next(m for m in parser.response_messages if m.type == "function_call")
assert tool_item.name == "get_weather"
assert tool_item.arguments == '{"location": "Paris"}'
assert tool_item.status == "completed"
# ---------------------------------------------------------------------------
# Tests: finish_reason tracking
# ---------------------------------------------------------------------------
def test_finish_reason_tracked():
"""finish_reason from CompletionOutput is stored on the parser."""
parser = _make_parser(_NoOpParser)
assert parser.finish_reason is None
parser.process(_make_output(finish_reason="stop"))
assert parser.finish_reason == "stop"
parser.process(_make_output(finish_reason="length"))
assert parser.finish_reason == "length"
# ---------------------------------------------------------------------------
# Tests: multi-turn accumulation
# ---------------------------------------------------------------------------
def test_multi_turn_accumulation():
"""Multiple process() calls accumulate response_messages."""
parser = _make_parser(_NoOpParser)
parser.process(_make_output(text="First turn"))
parser.process(_make_output(text="Second turn"))
assert len(parser.response_messages) == 2
texts = [m.content[0].text for m in parser.response_messages]
assert texts == ["First turn", "Second turn"]
def test_num_init_messages_offset():
"""Initial messages are preserved and offset works correctly."""
init_messages = [MagicMock(type="message")]
parser = _make_parser(_NoOpParser, response_messages=init_messages)
assert parser.num_init_messages == 1
parser.process(_make_output(text="New output"))
assert len(parser.response_messages) == 2
items = parser.make_response_output_items_from_parsable_context()
assert len(items) == 1
assert items[0].type == "message"
# ---------------------------------------------------------------------------
# Tests: factory function
# ---------------------------------------------------------------------------
def test_factory_function_creates_parser():
"""get_responses_parser_for_simple_context returns a working parser."""
rp = get_responses_parser_for_simple_context(
tokenizer=MagicMock(),
parser_cls=_NoOpParser,
response_messages=[],
request=_make_request(),
chat_template=None,
chat_template_content_format="auto",
)
assert isinstance(rp, ResponsesParser)
rp.process(_make_output(text="Works!"))
assert len(rp.response_messages) == 1
def test_factory_function_none_parser():
"""Factory function works with parser_cls=None."""
rp = get_responses_parser_for_simple_context(
tokenizer=MagicMock(),
parser_cls=None,
response_messages=[],
request=_make_request(),
chat_template=None,
chat_template_content_format="auto",
)
assert isinstance(rp, ResponsesParser)
assert rp.parser_instance is None
@@ -4,7 +4,6 @@
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.engine.serving import OpenAIServing
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.parser.abstract_parser import DelegatingParser
@@ -32,11 +31,8 @@ class _DummyDelegatingParser(DelegatingParser):
):
return None
def extract_tool_calls(self, model_output: str, request):
return None
def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_content():
def test_chat_completion_named_tool_choice_with_none_content():
request = ChatCompletionRequest.model_validate(
{
"model": "test-model",
@@ -53,17 +49,15 @@ def test_parse_tool_calls_from_content_allows_named_tool_choice_with_none_conten
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
}
)
parser = _DummyDelegatingParser(tokenizer=None)
tool_calls, content = OpenAIServing._parse_tool_calls_from_content(
request=request,
tokenizer=None,
enable_auto_tools=True,
tool_parser_cls=None,
tool_calls, content = parser._extract_tool_calls(
content=None,
request=request,
enable_auto_tools=True,
)
assert content is None
assert tool_calls is not None
assert tool_calls == []
@@ -37,7 +37,7 @@ def server():
"--max-num-seqs",
"128",
"--worker-extension-cls",
"tests.entrypoints.rpc.test_collective_rpc.TestWorkerExtension",
"tests.entrypoints.serve.dev.rpc.test_collective_rpc.TestWorkerExtension",
]
with RemoteOpenAIServer(
MODEL_NAME,
@@ -48,6 +48,7 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
names: list[str] | None = None
dtype_names: list[str] | None = None
shapes: list[list[int]] | None = None
num_updates_list: list[int] | None = None
class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
@@ -87,6 +88,15 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
# (In real implementation, this would receive and load actual weights)
load_weights([])
def receive_sparse_weights(
self,
update_info: MockUpdateInfo,
apply_patches: Callable[[list], None],
) -> None:
MockWeightTransferEngine.receive_weights_called = True
MockWeightTransferEngine.last_update_info = update_info
apply_patches([])
def shutdown(self) -> None:
MockWeightTransferEngine.shutdown_called = True
@@ -198,8 +208,6 @@ def test_update_weights_calls_engine():
llm.init_weight_transfer_engine(
WeightTransferInitRequest(init_info={"test_param": "init"})
)
# Start weight update (required before update_weights)
llm.start_weight_update(is_checkpoint_format=True)
# Call update_weights
@@ -232,14 +240,67 @@ def test_update_weights_calls_engine():
assert dtypes == test_dtypes
assert shapes == test_shapes
# Finish weight update
llm.finish_weight_update()
@create_new_process_for_each_test()
def test_update_weights_passes_sparse_metadata():
"""Test sparse update metadata is forwarded unchanged to the engine."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
with patch(
"vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
mock_create_engine,
):
llm = LLM(
model=MODEL_NAME,
enforce_eager=True,
load_format="dummy",
tensor_parallel_size=1,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
llm.init_weight_transfer_engine(
WeightTransferInitRequest(init_info={"test_param": "init"})
)
llm.start_weight_update(is_checkpoint_format=False)
llm.update_weights(
WeightTransferUpdateRequest(
update_info={
"names": ["layer.weight"],
"dtype_names": ["bfloat16"],
"shapes": [[100]],
"num_updates_list": [3],
"update_kind": "sparse_flat",
}
)
)
def check_sparse_update_called(self):
engine = self.weight_transfer_engine
if not engine.receive_weights_called:
return None
info = engine.last_update_info
return (
info.update_kind,
info.num_updates_list,
)
results = llm.collective_rpc(check_sparse_update_called)
for result in results:
assert result == ("sparse_flat", [3])
llm.finish_weight_update()
@create_new_process_for_each_test()
def test_full_weight_transfer_flow():
"""Test the complete weight transfer flow:
init -> start -> update -> finish."""
"""Test the complete weight transfer flow: init -> start -> update -> finish."""
if torch.accelerator.device_count() < 1:
pytest.skip("Need at least 1 GPU for this test")
+314
View File
@@ -0,0 +1,314 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools
import pytest
import torch
import torch.nn.functional as F
import vllm._custom_ops as ops
from vllm.platforms import 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)
set_random_seed(12345)
NUM_HEADS = [
(2, 4),
(4, 4),
]
HEAD_DIMS = [
(32, 32),
(64, 32),
]
CHUNK_SIZE = 64
PREFILL_SEQ_LENS = [
[1],
[1, 2, 3],
[CHUNK_SIZE - 1],
[CHUNK_SIZE],
[CHUNK_SIZE + 1],
[CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE + 1],
[2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
[4 * CHUNK_SIZE + 17],
]
DECODE_BATCH_SIZES = [1, 3, 5]
@functools.lru_cache(maxsize=128, typed=False)
def tensor_cache(
elem_num: int,
dtype: torch.dtype,
) -> torch.Tensor:
tensor = torch.rand(elem_num, dtype=dtype)
return tensor
def ref_l2norm(
x: torch.Tensor,
dim: int = -1,
eps: float = 1e-5,
) -> torch.Tensor:
inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
return x * inv_norm
def ref_gdn_gating(
A_log: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
dt_bias: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
softplus_x = F.softplus(a.float() + dt_bias.float(), beta=1.0, threshold=20.0)
g = -torch.exp(A_log.float()) * softplus_x
beta = torch.sigmoid(b.float()).to(dtype=b.dtype)
return g, beta
def ref_gated_delta_rule(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
initial_state: torch.Tensor,
cu_seqlens: torch.Tensor,
use_qk_l2norm_in_kernel: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
out = torch.empty_like(value)
final_state = torch.empty_like(initial_state)
for seq_idx in range(cu_seqlens.numel() - 1):
begin = int(cu_seqlens[seq_idx].item())
end = int(cu_seqlens[seq_idx + 1].item())
q_seq = query[:, begin:end]
k_seq = key[:, begin:end]
v_seq = value[:, begin:end]
g_seq = g[begin:end].unsqueeze(0)
beta_seq = beta[begin:end].unsqueeze(0)
initial_dtype = q_seq.dtype
if use_qk_l2norm_in_kernel:
q_seq = ref_l2norm(q_seq, dim=-1)
k_seq = ref_l2norm(k_seq, dim=-1)
if q_seq.shape[2] != v_seq.shape[2]:
repeat_factor = v_seq.shape[2] // q_seq.shape[2]
q_seq = q_seq.repeat_interleave(repeat_factor, dim=2)
k_seq = k_seq.repeat_interleave(repeat_factor, dim=2)
q_seq, k_seq, v_seq, beta_seq, g_seq = [
x.transpose(1, 2).contiguous().to(torch.float32)
for x in (q_seq, k_seq, v_seq, beta_seq, g_seq)
]
batch_size, num_heads, seq_len, head_dim = q_seq.shape
v_head_dim = v_seq.shape[-1]
q_seq = q_seq * (1 / (head_dim**0.5))
out_seq = torch.empty(
batch_size,
num_heads,
seq_len,
v_head_dim,
dtype=v_seq.dtype,
)
state = initial_state[seq_idx : seq_idx + 1].to(v_seq)
for token_idx in range(seq_len):
q_t = q_seq[:, :, token_idx]
k_t = k_seq[:, :, token_idx]
v_t = v_seq[:, :, token_idx]
g_t = g_seq[:, :, token_idx].exp().unsqueeze(-1).unsqueeze(-1)
beta_t = beta_seq[:, :, token_idx].unsqueeze(-1)
state = state * g_t
kv_mem = (state * k_t.unsqueeze(-2)).sum(dim=-1)
delta = (v_t - kv_mem) * beta_t
state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
final_state[seq_idx] = state.squeeze(0)
return out, final_state
def gdn_inputs(
num_tokens: int,
num_heads: tuple[int, int],
head_dims: tuple[int, int],
) -> tuple[torch.Tensor, ...]:
num_qk_heads, num_v_heads = num_heads
head_dim, v_head_dim = head_dims
q_shape = (1, num_tokens, num_qk_heads, head_dim)
q_numel = num_tokens * num_qk_heads * head_dim
q = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
k = tensor_cache(q_numel, torch.bfloat16).view(q_shape)
v_shape = (1, num_tokens, num_v_heads, v_head_dim)
v = tensor_cache(num_tokens * num_v_heads * v_head_dim, torch.bfloat16).view(
v_shape
)
gate_shape = (num_tokens, num_v_heads)
gate_numel = num_tokens * num_v_heads
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
A_log = tensor_cache(num_v_heads, torch.float32)
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
return q, k, v, a, b, A_log, dt_bias
@pytest.mark.parametrize("num_tokens", [1, 9])
@pytest.mark.parametrize("num_v_heads", [4, 8])
@torch.inference_mode()
def test_fused_gdn_gating_cpu(
num_tokens: int,
num_v_heads: int,
) -> None:
gate_shape = (num_tokens, num_v_heads)
gate_numel = num_tokens * num_v_heads
a = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
b = tensor_cache(gate_numel, torch.bfloat16).view(gate_shape)
A_log = tensor_cache(num_v_heads, torch.float32)
dt_bias = tensor_cache(num_v_heads, torch.bfloat16)
g_ref, beta_ref = ref_gdn_gating(A_log, a, b, dt_bias)
g, beta = ops.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
torch.testing.assert_close(g, g_ref.unsqueeze(0), atol=1e-4, rtol=1e-4)
torch.testing.assert_close(
beta.float(), beta_ref.unsqueeze(0).float(), atol=5e-3, rtol=5e-3
)
# decode path
@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
@torch.inference_mode()
def test_fused_sigmoid_gating_delta_rule_update_cpu(
batch_size: int,
num_heads: tuple[int, int],
head_dims: tuple[int, int],
) -> None:
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
num_tokens=batch_size,
num_heads=num_heads,
head_dims=head_dims,
)
_, num_v_heads = num_heads
head_dim, v_head_dim = head_dims
state_indices = torch.arange(batch_size, dtype=torch.int32)
cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
state = tensor_cache(
batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
).view(state_shape)
state_ref = state[state_indices].transpose(-1, -2).contiguous()
out_ref, final_state_ref = ref_gated_delta_rule(
query=q,
key=k,
value=v,
a=a,
b=b,
A_log=A_log,
dt_bias=dt_bias,
initial_state=state_ref,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=True,
)
out_ref = out_ref.transpose(0, 1).contiguous()
state_out = state.clone()
out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=state_out,
initial_state_indices=state_indices,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=True,
)
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(
state_out[state_indices].transpose(-1, -2),
final_state_ref,
atol=1e-2,
rtol=1e-2,
)
# prefill path
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
@torch.inference_mode()
def test_chunk_gated_delta_rule_cpu(
seq_lens: list[int],
num_heads: tuple[int, int],
head_dims: tuple[int, int],
) -> None:
total_tokens = sum(seq_lens)
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
num_tokens=total_tokens,
num_heads=num_heads,
head_dims=head_dims,
)
_, num_v_heads = num_heads
head_dim, v_head_dim = head_dims
cu_seqlens = torch.tensor(
[0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
)
initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
initial_state = tensor_cache(
len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
).view(initial_state_shape)
initial_state_ref = initial_state.transpose(-1, -2).contiguous()
out_ref, final_state_ref = ref_gated_delta_rule(
query=q,
key=k,
value=v,
a=a,
b=b,
A_log=A_log,
dt_bias=dt_bias,
initial_state=initial_state_ref,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=True,
)
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
out, final_state = ops.chunk_gated_delta_rule_cpu(
query=q,
key=k,
value=v,
g=g.unsqueeze(0),
beta=beta.unsqueeze(0),
initial_state=initial_state,
output_final_state=True,
cu_seqlens=cu_seqlens,
head_first=False,
use_qk_l2norm_in_kernel=True,
)
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(
final_state.transpose(-1, -2),
final_state_ref,
atol=1e-2,
rtol=1e-2,
)
-11
View File
@@ -1,10 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from importlib.metadata import version
import pytest
from packaging.version import Version
import vllm
from vllm.assets.image import ImageAsset
@@ -13,14 +10,6 @@ from vllm.platforms import current_platform
from ..utils import multi_gpu_test
pytestmark = pytest.mark.skipif(
Version("5.0") <= Version(version("transformers")),
reason=(
"MiniCPMV custom processor uses tokenizer.im_start_id which is not "
"available on TokenizersBackend in transformers v5.0+"
),
)
MODEL_PATH = "openbmb/MiniCPM-Llama3-V-2_5"
PROMPT_TEMPLATE = (

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