Compare commits

...
74 Commits
Author SHA1 Message Date
Andreas KaratzasandGitHub 54a62a79f7 [ROCm] Fix AttributeError for torch.compiler.skip_all_guards_unsafe on older PyTorch (#37219)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-17 11:34:49 +08:00
Flora FengandGitHub 384dc7f77b [Refactor] Relocate completion and chat completion tests (#37125)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-03-17 11:31:23 +08:00
Flora FengandGitHub f04d5226f8 [CI] Fix flaky tool_use chat completion tests with deterministic seed (#37027)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-03-17 03:24:34 +00:00
Kyuyeun KimandGitHub 0a0a1a198b Add ability to replace oot ops when using lora (#37181)
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
2026-03-16 18:04:15 -07:00
6c1cfbad32 Support non-contiguous KV cache in TRTLLM fp8 dequant kernel (#36867)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Signed-off-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Co-authored-by: Pavani Majety <pavanimajety@gmail.com>
2026-03-16 17:48:42 -07:00
Harry HuangandGitHub 45f526d652 [BugFix] Correct max memory usage for multiple KV-cache groups (#36030)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
2026-03-17 00:38:52 +00:00
Julien DenizeandGitHub 5db91f0aaf Fix some Mistral parser issues (#37209)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
2026-03-17 00:08:56 +00:00
Walter Beller-MoralesandGitHub 061980c36a [Feature][Frontend] add support for Cohere Embed v2 API (#37074)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
2026-03-16 19:55:53 -04:00
Ben BrowningandGitHub 7a49742b88 [CI/Build] Add common tool call parser test suite (#27599)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-03-16 19:46:20 -04:00
Terry GaoandGitHub 3e6a1e1686 [Custom Ops] Add functional + out variant for scaled_fp4_quant (#34389)
Signed-off-by: tianrengao <terrygao87@gmail.com>
2026-03-16 18:51:46 -04:00
Julien DenizeandGitHub 7961486a9b Fix EagleMistralLarge3Model initialization (#37232)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
2026-03-16 15:41:00 -07:00
Andreas KaratzasandGitHub 4f9b14c21c [CI] Stabilize multinode DP internal LB completion tests (#36356)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-16 15:40:23 -07:00
Yuchen FamaandGitHub 31a458c091 [Doc] Clarify schema enforcement behavior for tool_choice modes (#37064)
Signed-off-by: yfama <yuchengu@gmail.com>
2026-03-16 22:27:42 +00:00
Wei ZhaoandGitHub a3a51d20e7 [Benchmark] Improvements to attention benchmark script (#37115)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
2026-03-16 22:22:40 +00:00
EdalatiAliandGitHub e5b807607c [Quant][Feature] Support online MXFP8 quantization for MoE and dense models (#35448)
Signed-off-by: EdalatiAli <aliedalati@cohere.com>
2026-03-16 18:07:39 -04:00
fd4d96302a Fix eplb nvfp4 experts hook (#37217)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
Signed-off-by: Elvir Crncevic <elvir@anthropic.com>
Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-16 22:03:54 +00:00
Krish GuptaandGitHub c0f011918d [Bugfix] opcheck false mutation error in rms_norm_per_block_quant (#36688) (#36779)
Signed-off-by: Krish Gupta <krishom70@gmail.com>
2026-03-16 21:11:33 +00:00
Zhengxu ChenandGitHub e6ae4b1be1 [compile] Enable mega aot artifact for torch 2.12+. (#37198)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
2026-03-16 21:05:51 +00:00
zhanqiuhuandGitHub 2dccb38f73 [Bugfix][MultiConnector] Fix MultiConnector for SupportsHMA sub-connectors (#36549) 2026-03-16 20:51:04 +00:00
Kunshang JiandGitHub d157216093 [BUGFIX][Mamba] Use uint64 for address in KVBlockZeroer (#37197)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-03-16 21:39:56 +01:00
Matthew BonanniandGitHub 93f3c8e531 [Misc] Add float16 to CacheDType (#37199)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-03-16 13:24:48 -07:00
rasmithandGitHub 2cc26c3a99 [CI][BugFix][MORI][AMD] Add transfer_id to kv transfer params for test (#37213)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
2026-03-16 13:22:57 -07:00
Flora FengGitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
dfa8852db2 [Refactor] Consolidate GPT-OSS reasoning parser tests (#36915)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Signed-off-by: Flora Feng <4florafeng@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-03-16 15:53:07 -04:00
Lucas KabelaandGitHub 714c6e0eab [torch.compile][BE] Modify cudagraph callable to check for is_forward_context_set (#36288)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
2026-03-16 19:42:34 +00:00
SageandGitHub 0fefd00e6c [Bugfix] Fix render server crash for quantized models on CPU-only hosts (#37215)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
2026-03-16 18:59:01 +00:00
Nicolò LucchesiandGitHub f5c081d432 [PD][Nixl] Add support for hybrid SSM-FA models (#36687) 2026-03-16 19:58:06 +01:00
Matthew BonanniandGitHub c88ea8338b [MTP][Sparse MLA] Take advantage of native MTP support in indexer when possible (#36982)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
2026-03-16 13:51:21 -04:00
Max de BayserandGitHub 9f9ecff4cd Add simple granite4 tool parser (#36827)
Signed-off-by: Max de Bayser <maxdebayser@gmail.com>
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
2026-03-16 10:49:09 -07:00
ca1954d58c [Bugfix] Disable cross-layer KV cache for MLA attention backends (#37090)
Signed-off-by: haosdent <haosdent@gmail.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
2026-03-16 19:03:10 +02:00
Raushan TurganbayandGitHub 55e6d3d5c0 [Bugfix] Make siglip/clip compatible with transformers v5 (#37200)
Signed-off-by: raushan <raushan@huggingface.co>
2026-03-16 16:48:18 +00:00
ChaunceyandGitHub 6682c231fa [Bugfix] Add error handling for FINISHED_ERROR in OpenAIServing (#37148)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-03-16 16:27:47 +00:00
5ae685c1c8 [Bugfix] Relax TRTLLM KV cache contiguity assertion for cross-layer layout (#34158)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
2026-03-16 11:20:51 -04:00
Wentao YeandGitHub ce8cf9161d [Compile] Fix compile warning st256_cs in cuda_vec_utils.cuh (#36693)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-03-16 11:12:15 -04:00
xjxandGitHub 18be11fd59 [BUGFIX]fix CUDA OOM ERROR : invalid argument at cumem_allocator.cpp:119 (#35594)
Signed-off-by: xjx <493337577@qq.com>
2026-03-16 15:10:42 +00:00
8d8855fdae [Bugfix] Add safety check and fallback for null scaling factor (#36106)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 14:27:29 +00:00
Wentao YeandGitHub e855d380fa [Compile] Fix compile warning in moe_permute (#36529)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-03-16 10:16:14 -04:00
0e5a9382af [Bugfix] accept redacted thinking blocks in Anthropic messages (#36992)
Signed-off-by: Benjamin Bartels <benjaminba@tiglab-ubuntu.ilab.local>
Signed-off-by: bbartels <benjamin@bartels.dev>
Co-authored-by: Benjamin Bartels <benjaminba@tiglab-ubuntu.ilab.local>
2026-03-16 22:01:57 +08:00
Fynn Schmitt-UlmsandGitHub 04bf5a35fa [Spec Decode] Update extract_hidden_states to use deferred kv_connector clear (#37013) 2026-03-16 14:53:45 +01:00
Tianyu GuoandGitHub 43a73f853b Remove unused EVS functions in qwen3_vl.py (#37183)
Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
2026-03-16 13:09:09 +00:00
Julien DenizeandGitHub ffbc2e5bdb Patch Mistral config (#37104)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
2026-03-16 12:22:18 +00:00
f9e6db3034 [Models][Qwen3 ViT] Keep max_seqlen on CPU to prevent D2H sync (#37139)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-16 12:11:59 +00:00
d61d2b08e9 [Build] Fix API rate limit exceeded when using VLLM_USE_PRECOMPILED=1 (#36229)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 12:09:27 +00:00
Artem PerevedentsevandGitHub f5e59ee7a6 [Performance] Add prefetch for checkpoints to OS page cache (#36012)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
2026-03-16 11:32:02 +00:00
Harry MellorandGitHub 9b005edc48 [Docs] Make the link to hardware plugins clearer (#37174)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 04:12:58 -07:00
bf9a185395 GLM4 tool parser: fix streaming mode (#35208)
Signed-off-by: Robin Nabel <opensource@nabel.co>
Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
2026-03-16 18:48:52 +08:00
Harry MellorandGitHub ad041c79db Fix text only inputs for MRoPE models with the Transformers modelling backend (#37055)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 10:31:16 +00:00
Kunshang JiandGitHub 747b068136 [Hardware] Replace memory related torch.cuda APIs (#37031)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
2026-03-16 10:24:48 +00:00
Harry MellorandGitHub 122f75d939 Fix pipeline parallel with multimodal models with the Transformers modelling backend (#37057)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 10:20:37 +00:00
d8f8a7aad2 [Misc] Sync pre-commit to 4.5.1 in workflows and docs (#36675)
Signed-off-by: SoluMilken <ypiheyn.imm02g@g2.nctu.edu.tw>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-03-16 10:03:21 +00:00
Roy WangandGitHub 0115e957d4 [Frontend][Misc] Remove unused log in /is_sleeping (#37093)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-16 17:46:28 +08:00
haosdentandGitHub 116ed130f4 [Bugfix] Fix GDN attention crash with mixed decode/spec-decode batches (#34871)
Signed-off-by: haosdent <haosdent@gmail.com>
2026-03-16 10:30:23 +01:00
Vadim GimpelsonandGitHub 8374387bd8 [FlashInfer] Revert block_size 16 + head_size 256 workaround on Blackwell (#36987)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-03-16 09:04:29 +00:00
Isotr0pyandGitHub 912fbe9555 [Bugfix] Fix Qwen2.5-Omni/Qwen3-Omni use_audio_in_video with multi-video inputs (#37147)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-16 08:56:06 +00:00
Laith SakkaandGitHub 52131f88d9 use skip_all_guards_unsafe to drop global_state and torch_function_mode_stack guards instead of previous hacks (#36204)
Signed-off-by: Laith Sakka <lsakka@meta.com>
2026-03-16 08:52:31 +00:00
Roy WangandGitHub 821eb80c0d [Performance][Model Loader] Skip non-local expert weights during EP model loading (#37136)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-16 01:33:36 -07:00
Andreas KaratzasandGitHub a2956a0f8e [ROCm][CI] Retrying in case of batch variance effects and reducing flakiness (#36442)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-16 16:08:51 +08:00
Andreas KaratzasandGitHub 911355e216 [ROCm] Fix KV copy methods and auto-select attention backend for ROCm (#36845)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-16 16:07:27 +08:00
ChaunceyandGitHub 8d3f8f485e [Bugfix] fix Qwen3.5 tool calling bug (#36774)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
2026-03-16 15:38:42 +08:00
Woosuk KwonandGitHub 96efb91480 [Model Runner V2] Fix processed logits in sample() (#37144)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-03-16 00:35:49 -07:00
2754231ba3 [Kernel] Add FlashInfer MoE A2A Kernel (#36022)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Signed-off-by: Leo Tian <lctian@nvidia.com>
Co-authored-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Co-authored-by: root <root@lyris0267.lyris.clusters.nvidia.com>
2026-03-15 23:45:32 -07:00
bigshanedoggandGitHub 2390d44209 [Model] Add HyperCLOVAX-SEED-Think-14B language model support (#37107)
Signed-off-by: bigshanedogg <bigshane319@gmail.com>
2026-03-16 06:40:05 +00:00
Li, JiangandGitHub 7362b4450a [Bugfix] Avoid LD_PRELOAD check on MacOS (#37145)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
2026-03-15 23:31:44 -07:00
Andreas KaratzasandGitHub 57a314d155 [CI][Bugfix] Fix 500 errors from priority overflow and TemplateError subclasses in schema fuzz tests (#37127)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-16 05:27:21 +00:00
Andreas KaratzasandGitHub d4c57863f7 [ROCm][CI] Fix engine teardown and text normalization to stabilize voxtral test (#37138)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
2026-03-16 04:49:31 +00:00
Wang, YitingandGitHub 68e1b711f1 [XPU] Add deepseek_scaling_rope fused kernel (#36612)
Signed-off-by: yitingw1 <yiting.wang@intel.com>
2026-03-16 12:35:08 +08:00
rasmithandGitHub 0024f39a32 [ROCm][P/D][MORI][BugFix] Add transfer_id for moriio_connector so moriio_connector to restore P/D functionality (#34907)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
2026-03-16 10:36:51 +08:00
Andrew XiaandGitHub e9163b536e [responsesAPI][ez] add a unit test for SimpleContext logprobs (#37126)
Signed-off-by: Andrew Xia <axia@meta.com>
2026-03-15 17:12:26 -07:00
7acaea634c In-Tree AMD Zen CPU Backend via zentorch [1/N] (#35970)
Signed-off-by: Lalithnarayan C <Lalithnarayan.C@amd.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Chinmay-Kulkarni-AMD <Chinmay.Kulkarni@amd.com>
Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 23:35:35 +00:00
697e4ff352 [GDN] add a config for gdn kernel selection (#36647)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-03-16 00:40:17 +08:00
HariandGitHub a3e2e250f0 [Feature] Add Azure Blob Storage support for RunAI Model Streamer (#34614)
Signed-off-by: hasethuraman <hsethuraman@microsoft.com>
2026-03-15 19:38:21 +08:00
Isotr0pyandGitHub 143e4dccdf [Misc] Add online audio_in_video test (#36775)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-15 00:14:11 -07:00
Isotr0pyandGitHub 6590a3ecda [Frontend] Remove torchcodec from audio dependency (#37061)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-03-15 05:15:59 +00:00
Russell BryantandGitHub b3debb7e77 [Build] Upgrade xgrammar to get a security fix (#36168)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
2026-03-15 03:13:48 +00:00
Nick HillandGitHub 458c1a4b2d [Frontend] Reduce chat template warmup logging levels (#37062)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-03-14 13:48:59 -07:00
287 changed files with 10365 additions and 1998 deletions
@@ -333,15 +333,15 @@ apply_rocm_test_overrides() {
# --- Entrypoint ignores ---
if [[ $cmds == *" entrypoints/openai "* ]]; then
cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
--ignore=entrypoints/openai/test_audio.py \
--ignore=entrypoints/openai/test_shutdown.py \
--ignore=entrypoints/openai/chat_completion/test_audio.py \
--ignore=entrypoints/openai/completion/test_shutdown.py \
--ignore=entrypoints/openai/test_completion.py \
--ignore=entrypoints/openai/test_models.py \
--ignore=entrypoints/openai/test_lora_adapters.py \
--ignore=entrypoints/openai/test_return_tokens_as_ids.py \
--ignore=entrypoints/openai/test_root_path.py \
--ignore=entrypoints/openai/chat_completion/test_root_path.py \
--ignore=entrypoints/openai/test_tokenization.py \
--ignore=entrypoints/openai/test_prompt_validation.py "}
--ignore=entrypoints/openai/completion/test_prompt_validation.py "}
fi
if [[ $cmds == *" entrypoints/llm "* ]]; then
+12 -12
View File
@@ -162,7 +162,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration Test (API Server 2)
@@ -674,12 +674,12 @@ steps:
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: Benchmarks # 11min
timeout_in_minutes: 20
@@ -1143,7 +1143,7 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/test_oot_registration.py
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
- pytest -v -s models/test_oot_registration.py
- pytest -v -s plugins/lora_resolvers
@@ -1502,7 +1502,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration Test (API Server 2)
@@ -2133,12 +2133,12 @@ steps:
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: Benchmarks # 11min
timeout_in_minutes: 20
@@ -2735,7 +2735,7 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
- pytest -v -s models/test_oot_registration.py # it needs a clean process
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
@@ -3257,7 +3257,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/test_chat_utils.py
- label: Entrypoints Integration Test (API Server 2)
@@ -3872,12 +3872,12 @@ steps:
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
- label: Benchmarks # 11min
timeout_in_minutes: 20
@@ -4508,7 +4508,7 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
- pytest -v -s models/test_oot_registration.py # it needs a clean process
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
+1 -1
View File
@@ -34,7 +34,7 @@ steps:
- tests/entrypoints/test_chat_utils
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
- pytest -v -s entrypoints/test_chat_utils.py
mirror:
amd:
+2 -2
View File
@@ -9,9 +9,9 @@ steps:
- vllm/config/model.py
- vllm/model_executor
- tests/model_executor
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
commands:
- apt-get update && apt-get install -y curl libsodium23
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s model_executor
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+1 -1
View File
@@ -36,6 +36,6 @@ steps:
- pytest -v -s plugins_tests/test_scheduler_plugins.py
- pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s distributed/test_distributed_oot.py
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
- pytest -v -s models/test_oot_registration.py # it needs a clean process
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
+2 -2
View File
@@ -27,7 +27,7 @@ pull_request_rules:
Hi @{{author}}, the pre-commit checks have failed. Please run:
```bash
uv pip install pre-commit
uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files
```
@@ -381,7 +381,7 @@ pull_request_rules:
- or:
- files~=^vllm/model_executor/model_loader/tensorizer.py
- files~=^vllm/model_executor/model_loader/tensorizer_loader.py
- files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
- files~=^tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
- files~=^tests/model_executor/model_loader/tensorizer_loader/
actions:
assign:
+54 -16
View File
@@ -47,6 +47,8 @@ from common import (
is_mla_backend,
)
from vllm.v1.worker.workspace import init_workspace_manager
def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
"""Run standard attention benchmark (Flash/Triton/FlashInfer)."""
@@ -462,7 +464,7 @@ def main():
parser.add_argument(
"--batch-specs",
nargs="+",
default=["q2k", "8q1s1k"],
default=None,
help="Batch specifications using extended grammar",
)
@@ -478,6 +480,21 @@ def main():
parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
parser.add_argument(
"--kv-cache-dtype",
default="auto",
choices=["auto", "fp8"],
help="KV cache dtype: auto or fp8",
)
parser.add_argument(
"--cuda-graphs",
action=argparse.BooleanOptionalAction,
default=True,
help=(
"Launch kernels with CUDA graphs to eliminate CPU overhead"
"in measurements (default: True)"
),
)
# Parameter sweep (use YAML config for advanced sweeps)
parser.add_argument(
@@ -536,21 +553,24 @@ def main():
# Batch specs and sizes
# Support both explicit batch_specs and generated batch_spec_ranges
if "batch_spec_ranges" in yaml_config:
# Generate batch specs from ranges
generated_specs = generate_batch_specs_from_ranges(
yaml_config["batch_spec_ranges"]
)
# Combine with any explicit batch_specs
if "batch_specs" in yaml_config:
args.batch_specs = yaml_config["batch_specs"] + generated_specs
else:
args.batch_specs = generated_specs
console.print(
f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
)
elif "batch_specs" in yaml_config:
args.batch_specs = yaml_config["batch_specs"]
# CLI --batch-specs takes precedence over YAML when provided.
cli_batch_specs_provided = args.batch_specs is not None
if not cli_batch_specs_provided:
if "batch_spec_ranges" in yaml_config:
# Generate batch specs from ranges
generated_specs = generate_batch_specs_from_ranges(
yaml_config["batch_spec_ranges"]
)
# Combine with any explicit batch_specs
if "batch_specs" in yaml_config:
args.batch_specs = yaml_config["batch_specs"] + generated_specs
else:
args.batch_specs = generated_specs
console.print(
f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
)
elif "batch_specs" in yaml_config:
args.batch_specs = yaml_config["batch_specs"]
if "batch_sizes" in yaml_config:
args.batch_sizes = yaml_config["batch_sizes"]
@@ -575,6 +595,10 @@ def main():
args.warmup_iters = yaml_config["warmup_iters"]
if "profile_memory" in yaml_config:
args.profile_memory = yaml_config["profile_memory"]
if "kv_cache_dtype" in yaml_config:
args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
if "cuda_graphs" in yaml_config:
args.cuda_graphs = yaml_config["cuda_graphs"]
# Parameter sweep configuration
if "parameter_sweep" in yaml_config:
@@ -629,12 +653,18 @@ def main():
# Determine backends
backends = args.backends or ([args.backend] if args.backend else ["flash"])
prefill_backends = getattr(args, "prefill_backends", None)
if not args.batch_specs:
args.batch_specs = ["q2k", "8q1s1k"]
console.print(f"Backends: {', '.join(backends)}")
if prefill_backends:
console.print(f"Prefill backends: {', '.join(prefill_backends)}")
console.print(f"Batch specs: {', '.join(args.batch_specs)}")
console.print(f"KV cache dtype: {args.kv_cache_dtype}")
console.print(f"CUDA graphs: {args.cuda_graphs}")
console.print()
init_workspace_manager(args.device)
# Run benchmarks
all_results = []
@@ -687,6 +717,8 @@ def main():
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
use_cuda_graphs=args.cuda_graphs,
)
# Add decode pipeline config
@@ -839,6 +871,8 @@ def main():
"repeats": args.repeats,
"warmup_iters": args.warmup_iters,
"profile_memory": args.profile_memory,
"kv_cache_dtype": args.kv_cache_dtype,
"use_cuda_graphs": args.cuda_graphs,
}
all_results = run_model_parameter_sweep(
backends,
@@ -861,6 +895,8 @@ def main():
"repeats": args.repeats,
"warmup_iters": args.warmup_iters,
"profile_memory": args.profile_memory,
"kv_cache_dtype": args.kv_cache_dtype,
"use_cuda_graphs": args.cuda_graphs,
}
all_results = run_parameter_sweep(
backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -891,6 +927,8 @@ def main():
repeats=args.repeats,
warmup_iters=args.warmup_iters,
profile_memory=args.profile_memory,
kv_cache_dtype=args.kv_cache_dtype,
use_cuda_graphs=args.cuda_graphs,
)
result = run_benchmark(config)
@@ -213,6 +213,9 @@ class BenchmarkConfig:
profile_memory: bool = False
use_cuda_graphs: bool = False
# "auto" or "fp8"
kv_cache_dtype: str = "auto"
# MLA-specific
prefill_backend: str | None = None
kv_lora_rank: int | None = None
@@ -369,6 +372,7 @@ class ResultsFormatter:
"backend",
"batch_spec",
"num_layers",
"kv_cache_dtype",
"mean_time",
"std_time",
"throughput",
@@ -382,6 +386,7 @@ class ResultsFormatter:
"backend": r.config.backend,
"batch_spec": r.config.batch_spec,
"num_layers": r.config.num_layers,
"kv_cache_dtype": r.config.kv_cache_dtype,
"mean_time": r.mean_time,
"std_time": r.std_time,
"throughput": r.throughput_tokens_per_sec or 0,
@@ -30,9 +30,9 @@ batch_specs:
- "2q16k_32q1s4k" # 2 very large prefill + 32 decode
# Context extension + decode
- "2q1kkv2k_16q1s1k" # 2 extend + 16 decode
- "4q2kkv4k_32q1s2k" # 4 extend + 32 decode
- "2q1kkv8k_32q1s2k" # 2 large extend + 32 decode
- "2q1ks2k_16q1s1k" # 2 extend + 16 decode
- "4q2ks4k_32q1s2k" # 4 extend + 32 decode
- "2q1ks8k_32q1s2k" # 2 large extend + 32 decode
# Explicitly chunked prefill
- "q8k" # 8k prefill with chunking hint
@@ -0,0 +1,58 @@
# MLA decode-only benchmark configuration
model:
name: "deepseek-v3"
num_layers: 60
num_q_heads: 128 # Base value, can be swept for TP simulation
num_kv_heads: 1 # MLA uses single latent KV
head_dim: 576
kv_lora_rank: 512
qk_nope_head_dim: 128
qk_rope_head_dim: 64
v_head_dim: 128
block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128
# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
model_parameter_sweep:
param_name: "num_q_heads"
values: [128, 64, 32, 16]
label_format: "{backend}_{value}h"
batch_specs:
# Small batches, varying sequence lengths
- "16q1s512" # 16 requests, 512 KV cache
- "16q1s1k" # 16 requests, 1k KV cache
- "16q1s2k" # 16 requests, 2k KV cache
- "16q1s4k" # 16 requests, 4k KV cache
# Medium batches
- "32q1s1k" # 32 requests, 1k KV cache
- "32q1s2k" # 32 requests, 2k KV cache
- "32q1s4k" # 32 requests, 4k KV cache
- "32q1s8k" # 32 requests, 8k KV cache
# Large batches
- "64q1s1k" # 64 requests, 1k KV cache
- "64q1s2k" # 64 requests, 2k KV cache
- "64q1s4k" # 64 requests, 4k KV cache
- "64q1s8k" # 64 requests, 8k KV cache
# Very large batches
- "128q1s1k" # 128 requests, 1k KV cache
- "128q1s2k" # 128 requests, 2k KV cache
- "128q1s4k" # 128 requests, 4k KV cache
- "128q1s8k" # 128 requests, 8k KV cache
# Long context
- "32q1s16k" # 32 requests, 16k KV cache
- "32q1s32k" # 32 requests, 32k KV cache
backends:
- FLASHMLA_SPARSE
- FLASHINFER_MLA_SPARSE
device: "cuda:0"
repeats: 100
warmup_iters: 10
profile_memory: true
+132 -33
View File
@@ -60,9 +60,11 @@ def create_minimal_vllm_config(
model_name: str = "deepseek-v3",
block_size: int = 128,
max_num_seqs: int = 256,
max_num_batched_tokens: int = 8192,
mla_dims: dict | None = None,
index_topk: int | None = None,
prefill_backend: str | None = None,
kv_cache_dtype: str = "auto",
) -> VllmConfig:
"""
Create minimal VllmConfig for MLA benchmarks.
@@ -149,13 +151,13 @@ def create_minimal_vllm_config(
cache_config = CacheConfig(
block_size=block_size,
gpu_memory_utilization=0.9,
cache_dtype="auto",
cache_dtype=kv_cache_dtype,
enable_prefix_caching=False,
)
scheduler_config = SchedulerConfig(
max_num_seqs=max_num_seqs,
max_num_batched_tokens=8192,
max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
max_model_len=32768,
is_encoder_decoder=False,
enable_chunked_prefill=True,
@@ -535,6 +537,7 @@ def _create_backend_impl(
device: torch.device,
max_num_tokens: int = 8192,
index_topk: int | None = None,
kv_cache_dtype: str = "auto",
):
"""
Create backend implementation instance.
@@ -583,7 +586,7 @@ def _create_backend_impl(
"num_kv_heads": mla_dims["num_kv_heads"],
"alibi_slopes": None,
"sliding_window": None,
"kv_cache_dtype": "auto",
"kv_cache_dtype": kv_cache_dtype,
"logits_soft_cap": None,
"attn_type": "decoder",
"kv_sharing_target_layer_name": None,
@@ -701,6 +704,7 @@ def _run_single_benchmark(
mla_dims: dict,
device: torch.device,
indexer=None,
kv_cache_dtype: str | None = None,
) -> BenchmarkResult:
"""
Run a single benchmark iteration.
@@ -734,49 +738,124 @@ def _run_single_benchmark(
)
# Create KV cache
kv_cache = torch.zeros(
num_blocks,
block_size,
mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
device=device,
dtype=torch.bfloat16,
)
if kv_cache_dtype is None:
kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
if kv_cache_dtype == "fp8_ds_mla":
# FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
# Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
# + 2*rope_dim bf16 bytes
# = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
kv_cache = torch.zeros(
num_blocks,
block_size,
656,
device=device,
dtype=torch.uint8,
)
elif kv_cache_dtype == "fp8":
from vllm.platforms import current_platform
# Create input tensors for both decode and prefill modes
decode_inputs, prefill_inputs = _create_input_tensors(
total_q,
mla_dims,
backend_cfg["query_format"],
device,
torch.bfloat16,
)
kv_cache = torch.zeros(
num_blocks,
block_size,
head_size,
device=device,
dtype=torch.uint8,
).view(current_platform.fp8_dtype())
else:
kv_cache = torch.zeros(
num_blocks,
block_size,
head_size,
device=device,
dtype=torch.bfloat16,
)
# Fill indexer with random indices for sparse backends
is_sparse = backend_cfg.get("is_sparse", False)
if is_sparse and indexer is not None:
indexer.fill_random_indices(total_q, max_kv_len)
# Determine which forward method to use based on metadata
if metadata.decode is not None:
forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
elif metadata.prefill is not None:
forward_fn = lambda: impl.forward_mha(
prefill_inputs["q"],
prefill_inputs["k_c_normed"],
prefill_inputs["k_pe"],
kv_cache,
metadata,
prefill_inputs["k_scale"],
prefill_inputs["output"],
)
else:
# Determine which forward methods to use based on metadata.
# Sparse MLA backends always use forward_mqa
has_decode = is_sparse or getattr(metadata, "decode", None) is not None
has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
if not has_decode and not has_prefill:
raise RuntimeError("Metadata has neither decode nor prefill metadata")
num_decode = (
metadata.num_decode_tokens
if (has_decode and has_prefill)
else total_q
if has_decode
else 0
)
num_prefill = total_q - num_decode
# Some backends requires fp8 queries when using fp8 KV cache.
is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
quantize_query = is_fp8_kvcache and getattr(
impl, "supports_quant_query_input", False
)
# quantize_query forces concat format
query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
# Create decode query tensors
if has_decode:
decode_inputs, _ = _create_input_tensors(
num_decode, mla_dims, query_fmt, device, torch.bfloat16
)
# Cast decode query to fp8 if the backend supports it
if quantize_query:
from vllm.platforms import current_platform
if isinstance(decode_inputs, tuple):
decode_inputs = torch.cat(list(decode_inputs), dim=-1)
decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
# Create prefill input tensors
if has_prefill:
_, prefill_inputs = _create_input_tensors(
num_prefill, mla_dims, query_fmt, device, torch.bfloat16
)
# Build forward function
def forward_fn():
results = []
if has_decode:
results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
if has_prefill:
results.append(
impl.forward_mha(
prefill_inputs["q"],
prefill_inputs["k_c_normed"],
prefill_inputs["k_pe"],
kv_cache,
metadata,
prefill_inputs["k_scale"],
prefill_inputs["output"],
)
)
return results[0] if len(results) == 1 else tuple(results)
# Warmup
for _ in range(config.warmup_iters):
forward_fn()
torch.accelerator.synchronize()
# Optionally capture a CUDA graph after warmup.
# Graph replay eliminates CPU launch overhead so timings reflect pure
# kernel time.
if config.use_cuda_graphs:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
forward_fn()
benchmark_fn = graph.replay
else:
benchmark_fn = forward_fn
# Benchmark
times = []
for _ in range(config.repeats):
@@ -785,7 +864,7 @@ def _run_single_benchmark(
start.record()
for _ in range(config.num_layers):
forward_fn()
benchmark_fn()
end.record()
torch.accelerator.synchronize()
@@ -852,13 +931,30 @@ def _run_mla_benchmark_batched(
# Determine if this is a sparse backend
is_sparse = backend_cfg.get("is_sparse", False)
# Extract kv_cache_dtype from the first config
kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
# FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
# Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
kv_cache_dtype = "fp8_ds_mla"
# Compute max total_q across all configs so the metadata builder buffer
# and scheduler config are large enough for all batch specs.
max_total_q = max(
sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
for cfg, *_ in configs_with_params
)
# Create and set vLLM config for MLA (reused across all benchmarks)
vllm_config = create_minimal_vllm_config(
model_name="deepseek-v3", # Used only for model path
block_size=block_size,
max_num_batched_tokens=max_total_q,
mla_dims=mla_dims, # Use custom dims from config or default
index_topk=index_topk if is_sparse else None,
prefill_backend=prefill_backend,
kv_cache_dtype=kv_cache_dtype,
)
results = []
@@ -883,7 +979,9 @@ def _run_mla_benchmark_batched(
mla_dims,
vllm_config,
device,
max_num_tokens=max_total_q,
index_topk=index_topk if is_sparse else None,
kv_cache_dtype=kv_cache_dtype,
)
# Verify the actual prefill backend matches what was requested
@@ -942,6 +1040,7 @@ def _run_mla_benchmark_batched(
mla_dims,
device,
indexer=indexer,
kv_cache_dtype=kv_cache_dtype,
)
results.append(result)
+61 -18
View File
@@ -140,7 +140,7 @@ def _create_vllm_config(
cache_config = CacheConfig(
block_size=config.block_size,
cache_dtype="auto",
cache_dtype=config.kv_cache_dtype,
)
cache_config.num_gpu_blocks = max_num_blocks
cache_config.num_cpu_blocks = 0
@@ -215,7 +215,7 @@ def _create_backend_impl(
num_kv_heads=config.num_kv_heads,
alibi_slopes=None,
sliding_window=None,
kv_cache_dtype="auto",
kv_cache_dtype=config.kv_cache_dtype,
)
kv_cache_spec = FullAttentionSpec(
@@ -288,12 +288,22 @@ def _create_input_tensors(
total_q: int,
device: torch.device,
dtype: torch.dtype,
quantize_query: bool = False,
) -> tuple:
"""Create Q, K, V input tensors for all layers."""
"""Create Q, K, V input tensors for all layers.
When quantize_query is True, queries are cast to fp8 to match backends
that require query/key/value dtype consistency.
"""
q_dtype = dtype
if quantize_query:
from vllm.platforms import current_platform
q_dtype = current_platform.fp8_dtype()
q_list = [
torch.randn(
total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
)
).to(q_dtype)
for _ in range(config.num_layers)
]
k_list = [
@@ -344,10 +354,17 @@ def _create_kv_cache(
# Compute inverse permutation to get back to logical view
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
# Use fp8 dtype for cache when requested.
cache_dtype = dtype
if config.kv_cache_dtype == "fp8":
from vllm.platforms import current_platform
cache_dtype = current_platform.fp8_dtype()
cache_list = []
for _ in range(config.num_layers):
# Allocate in physical layout order (contiguous in memory)
cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
# Permute to logical view
cache = cache.permute(*inv_order)
cache_list.append(cache)
@@ -392,6 +409,37 @@ def _run_single_benchmark(
)
torch.accelerator.synchronize()
# Optionally capture a CUDA graph after warmup.
# Graph replay eliminates CPU launch overhead so timings reflect pure
# kernel time.
if config.use_cuda_graphs:
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
for i in range(config.num_layers):
impl.forward(
layer,
q_list[i],
k_list[i],
v_list[i],
cache_list[i],
attn_metadata,
output=out,
)
benchmark_fn = graph.replay
else:
def benchmark_fn():
for i in range(config.num_layers):
impl.forward(
layer,
q_list[i],
k_list[i],
v_list[i],
cache_list[i],
attn_metadata,
output=out,
)
# Benchmark
times = []
for _ in range(config.repeats):
@@ -399,16 +447,7 @@ def _run_single_benchmark(
end = torch.cuda.Event(enable_timing=True)
start.record()
for i in range(config.num_layers):
impl.forward(
layer,
q_list[i],
k_list[i],
v_list[i],
cache_list[i],
attn_metadata,
output=out,
)
benchmark_fn()
end.record()
torch.accelerator.synchronize()
@@ -418,8 +457,8 @@ def _run_single_benchmark(
mem_stats = {}
if config.profile_memory:
mem_stats = {
"allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
"reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
"allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
"reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
}
return times, mem_stats
@@ -502,8 +541,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
common_attn_metadata=common_metadata,
)
# Only quantize queries when the impl supports it
quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
impl, "supports_quant_query_input", False
)
q_list, k_list, v_list = _create_input_tensors(
config, total_q, device, dtype
config, total_q, device, dtype, quantize_query=quantize_query
)
cache_list = _create_kv_cache(
+5 -2
View File
@@ -95,13 +95,16 @@ def create_logits(
def measure_memory() -> tuple[int, int]:
"""Return (allocated, reserved) memory in bytes."""
torch.accelerator.synchronize()
return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
return (
torch.accelerator.memory_allocated(),
torch.accelerator.max_memory_allocated(),
)
def reset_memory_stats():
"""Reset peak memory statistics."""
reset_buffer_cache()
torch.cuda.reset_peak_memory_stats()
torch.accelerator.reset_peak_memory_stats()
torch.accelerator.empty_cache()
gc.collect()
+1
View File
@@ -196,6 +196,7 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
return val;
#else
assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
return u32x8_t{};
#endif
}
+8 -6
View File
@@ -109,16 +109,18 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
#ifndef USE_ROCM
int flag = 0;
CUDA_CHECK(cuDeviceGetAttribute(
CUresult rdma_result = cuDeviceGetAttribute(
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
device));
if (flag) { // support GPUDirect RDMA if possible
device);
if (rdma_result == CUDA_SUCCESS &&
flag) { // support GPUDirect RDMA if possible
prop.allocFlags.gpuDirectRDMACapable = 1;
}
int fab_flag = 0;
CUDA_CHECK(cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
if (fab_flag) { // support fabric handle if possible
CUresult fab_result = cuDeviceGetAttribute(
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
if (fab_result == CUDA_SUCCESS &&
fab_flag) { // support fabric handle if possible
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
}
#endif
+3 -4
View File
@@ -73,10 +73,9 @@ void moe_permute(
MOE_DISPATCH(input.scalar_type(), [&] {
expandInputRowsKernelLauncher<scalar_t>(
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
n_hidden, topk, n_local_expert, stream);
get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
});
}
@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
template <typename T>
void expandInputRowsKernelLauncher(
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
T const* unpermuted_input, T* permuted_output,
int const* expanded_dest_row_to_expanded_source_row,
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -2,7 +2,7 @@
template <typename T, bool CHECK_SKIPPED>
__global__ void expandInputRowsKernel(
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
T const* unpermuted_input, T* permuted_output,
int const* expanded_dest_row_to_expanded_source_row,
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -16,7 +16,6 @@ __global__ void expandInputRowsKernel(
int64_t expanded_dest_row = blockIdx.x;
int64_t const expanded_source_row =
expanded_dest_row_to_expanded_source_row[expanded_dest_row];
int expert_id = sorted_experts[expanded_dest_row];
if (threadIdx.x == 0) {
assert(expanded_dest_row <= INT32_MAX);
@@ -54,7 +53,7 @@ __global__ void expandInputRowsKernel(
template <typename T>
void expandInputRowsKernelLauncher(
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
T const* unpermuted_input, T* permuted_output,
int const* expanded_dest_row_to_expanded_source_row,
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -70,12 +69,12 @@ void expandInputRowsKernelLauncher(
bool is_check_skip = num_valid_tokens_ptr != nullptr;
auto func = func_map[is_check_skip];
func<<<blocks, threads, 0, stream>>>(
unpermuted_input, permuted_output, sorted_experts,
expanded_dest_row_to_expanded_source_row,
expanded_source_row_to_expanded_dest_row, permuted_idx,
expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
num_local_experts);
func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
expanded_dest_row_to_expanded_source_row,
expanded_source_row_to_expanded_dest_row,
permuted_idx, expert_first_token_offset,
num_rows, num_valid_tokens_ptr, cols, k,
num_local_experts);
}
template <class T, class U>
+8 -4
View File
@@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
torch::Tensor& output_scale,
torch::Tensor const& input_scale,
bool is_sf_swizzled_layout);
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
torch::Tensor const& input, torch::Tensor const& input_scale,
bool is_sf_swizzled_layout);
void scaled_fp4_quant_out(torch::Tensor const& input,
torch::Tensor const& input_scale,
bool is_sf_swizzled_layout, torch::Tensor& output,
torch::Tensor& output_scale);
void scaled_fp4_experts_quant(
torch::Tensor& output, torch::Tensor& output_scale,
+34 -3
View File
@@ -16,6 +16,8 @@
#include <torch/all.h>
#include "nvfp4_utils.cuh"
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
@@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
torch::Tensor const& output_scale_offset_by_experts);
#endif
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
torch::Tensor& output_sf, torch::Tensor const& input_sf,
bool is_sf_swizzled_layout) {
void scaled_fp4_quant_out(torch::Tensor const& input,
torch::Tensor const& input_sf,
bool is_sf_swizzled_layout, torch::Tensor& output,
torch::Tensor& output_sf) {
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
@@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
}
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
torch::Tensor const& input, torch::Tensor const& input_sf,
bool is_sf_swizzled_layout) {
int64_t n = input.size(-1);
int64_t m = input.numel() / n;
auto device = input.device();
// Two fp4 values packed into a uint8
auto output = torch::empty(
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
torch::Tensor output_sf;
if (is_sf_swizzled_layout) {
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
output_sf = torch::empty(
{sf_m, sf_n},
torch::TensorOptions().device(device).dtype(torch::kInt32));
} else {
output_sf = torch::empty(
{m, n / CVT_FP4_SF_VEC_SIZE},
torch::TensorOptions().device(device).dtype(torch::kUInt8));
}
scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
output_sf);
return {output, output_sf};
}
void scaled_fp4_experts_quant(
torch::Tensor& output, torch::Tensor& output_scale,
torch::Tensor const& input, torch::Tensor const& input_global_scale,
+13
View File
@@ -18,6 +18,7 @@
#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <utility>
#include "../../cuda_vec_utils.cuh"
@@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) {
return round_up(m, ROW_TILE);
}
// Compute the shape of the swizzled SF output tensor.
// Returns (rounded_m, rounded_n / 4) where:
// rounded_m = round_up(m, 128)
// rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
int64_t n) {
int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
return {rounded_m, rounded_n / 4};
}
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
uint32_t val;
@@ -286,6 +286,15 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
"Outer scale stride must be 1 when scales are not transposed");
}
int64_t hidden_size = input.size(-1);
TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
"hidden_size must be a positive multiple of group_size");
int64_t num_tokens = input.numel() / hidden_size;
int64_t num_groups = hidden_size / group_size;
TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
"scales buffer too small: need ", num_tokens * num_groups,
" elements, got ", scales.numel());
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
var_epsilon, scale_ub, residual,
is_scale_transposed);
+1 -1
View File
@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
// The range of logits within the row.
int rowStart = 0;
int seq_len = seqLens[rowIdx / next_n];
int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
// Local pointers to this block
if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
+15 -4
View File
@@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Compute NVFP4 block quantized tensor.
ops.def(
"scaled_fp4_quant(Tensor! output, Tensor input,"
" Tensor! output_scale, Tensor input_scale, bool "
"is_sf_swizzled_layout) -> ()");
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
"scaled_fp4_quant(Tensor input,"
" Tensor input_scale, bool "
"is_sf_swizzled_layout) -> (Tensor, Tensor)");
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
// Out variant
// TODO: Add {at::Tag::out_variant} tag and update all call sites
// to use the functional variant once vLLM upgrades PyTorch.
// See pytorch/pytorch#176117.
ops.def(
"scaled_fp4_quant.out(Tensor input,"
" Tensor input_scale, bool "
"is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
"-> ()");
ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
// Compute NVFP4 experts quantization.
ops.def(
+2 -2
View File
@@ -620,7 +620,7 @@ RUN set -eux; \
ARG BITSANDBYTES_VERSION_X86=0.46.1
ARG BITSANDBYTES_VERSION_ARM64=0.42.0
ARG TIMM_VERSION=">=1.0.17"
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
RUN --mount=type=cache,target=/root/.cache/uv \
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
@@ -628,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
fi; \
uv pip install --system accelerate hf_transfer modelscope \
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
# ============================================================
# VLLM INSTALLATION (depends on build stage)
+17
View File
@@ -9,6 +9,7 @@
#
# Build targets:
# vllm-openai (default): used for serving deployment
# vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
# vllm-test: used for CI tests
# vllm-dev: used for development
#
@@ -222,3 +223,19 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
ENTRYPOINT ["vllm", "serve"]
######################### ZEN CPU PYPI IMAGE #########################
FROM vllm-openai AS vllm-openai-zen
ARG TARGETARCH
RUN if [ "$TARGETARCH" != "amd64" ]; then \
echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
exit 1; \
fi
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install "vllm[zen]"
ENTRYPOINT ["vllm", "serve"]
+1 -1
View File
@@ -83,7 +83,7 @@
"default": ">=1.0.17"
},
"RUNAI_MODEL_STREAMER_VERSION": {
"default": ">=0.15.3"
"default": ">=0.15.7"
}
}
}
+1 -1
View File
@@ -75,7 +75,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment
vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
```bash
uv pip install pre-commit
uv pip install pre-commit>=4.5.1
pre-commit install
```
+19 -19
View File
@@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
>
@@ -204,14 +204,14 @@ configuration.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
+2 -1
View File
@@ -35,7 +35,8 @@ th {
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
-3
View File
@@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the
with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
components (see Compile Range Integration).
3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
for torch.compile integration.
### CompilationConfig
With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend
+25 -1
View File
@@ -107,6 +107,27 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
!!! note
When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option.
## Constrained Decoding Behavior
Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode:
| `tool_choice` value | Schema-constrained decoding | Behavior |
| --- | --- | --- |
| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. |
| `"none"` | N/A | No tool calls are produced. |
When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`.
### Strict Mode (`strict` parameter)
The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode.
vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic.
Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313).
## Automatic Function Calling
To enable this feature, you should set the following flags:
@@ -124,6 +145,9 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!
!!! note
With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details.
### Hermes Models (`hermes`)
All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported.
@@ -219,7 +243,7 @@ Supported models:
* `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models
Recommended flags: `--tool-call-parser hermes`
Recommended flags: `--tool-call-parser granite4`
* `ibm-granite/granite-3.0-8b-instruct`
+3 -1
View File
@@ -16,4 +16,6 @@ vLLM supports the following hardware platforms:
vLLM supports third-party hardware plugins that live **outside** the main `vllm` repository. These follow the [Hardware-Pluggable RFC](../../design/plugin_system.md).
A list of all supported hardware can be found on the [vllm.ai website](https://vllm.ai/#compatibility). If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
A list of all supported hardware can be found on the vLLM website, see [Universal Compatibility - Hardware](https://vllm.ai/#compatibility).
If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
@@ -31,6 +31,16 @@ vllm serve gs://core-llm/Llama-3-8b \
--load-format runai_streamer
```
To run model from Azure Blob Storage run:
```bash
AZURE_STORAGE_ACCOUNT_NAME=<account> \
vllm serve az://<container>/<model-path> \
--load-format runai_streamer
```
Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.
To run model from a S3 compatible object store run:
```bash
+1
View File
@@ -418,6 +418,7 @@ th {
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
+2 -1
View File
@@ -21,7 +21,8 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
| `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
| `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes |
| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
| `naive` | Testing/debugging | Simple broadcast-based implementation | Debugging, not recommended for production |
## Single Node Deployment
+134
View File
@@ -72,6 +72,9 @@ In addition, we have the following custom APIs:
- Only applicable to [classification models](../models/pooling_models.md).
- [Score API](#score-api) (`/score`)
- Applicable to [embedding models and cross-encoder models](../models/pooling_models.md).
- [Cohere Embed API](#cohere-embed-api) (`/v2/embed`)
- Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
- Works with any [embedding model](../models/pooling_models.md), including multimodal models.
- [Re-rank API](#re-rank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
- Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
- Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
@@ -429,6 +432,137 @@ these extra parameters are supported instead:
--8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
```
### Cohere Embed API
Our API is also compatible with [Cohere's Embed v2 API](https://docs.cohere.com/reference/embed) which adds support for some modern embedding feature such as truncation, output dimensions, embedding types, and input types. This endpoint works with any embedding model (including multimodal models).
#### Cohere Embed API request parameters
| Parameter | Type | Required | Description |
| --------- | ---- | -------- | ----------- |
| `model` | string | Yes | Model name |
| `input_type` | string | No | Prompt prefix key (model-dependent, see below) |
| `texts` | list[string] | No | Text inputs (use one of `texts`, `images`, or `inputs`) |
| `images` | list[string] | No | Base64 data URI images |
| `inputs` | list[object] | No | Mixed text and image content objects |
| `embedding_types` | list[string] | No | Output types (default: `["float"]`) |
| `output_dimension` | int | No | Truncate embeddings to this dimension (Matryoshka) |
| `truncate` | string | No | `END`, `START`, or `NONE` (default: `END`) |
#### Text embedding
```bash
curl -X POST "http://localhost:8000/v2/embed" \
-H "Content-Type: application/json" \
-d '{
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
"input_type": "query",
"texts": ["Hello world", "How are you?"],
"embedding_types": ["float"]
}'
```
??? console "Response"
```json
{
"id": "embd-...",
"embeddings": {
"float": [
[0.012, -0.034, ...],
[0.056, 0.078, ...]
]
},
"texts": ["Hello world", "How are you?"],
"meta": {
"api_version": {"version": "2"},
"billed_units": {"input_tokens": 12}
}
}
```
#### Mixed text and image inputs
For multimodal models, you can embed images by passing base64 data URIs. The `inputs` field accepts a list of objects with mixed text and image content:
```bash
curl -X POST "http://localhost:8000/v2/embed" \
-H "Content-Type: application/json" \
-d '{
"model": "google/siglip-so400m-patch14-384",
"inputs": [
{
"content": [
{"type": "text", "text": "A photo of a cat"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR..."}}
]
}
],
"embedding_types": ["float"]
}'
```
#### Embedding types
The `embedding_types` parameter controls the output format. Multiple types can be requested in a single call:
| Type | Description |
| ---- | ----------- |
| `float` | Raw float32 embeddings (default) |
| `binary` | Bit-packed signed binary |
| `ubinary` | Bit-packed unsigned binary |
| `base64` | Little-endian float32 encoded as base64 |
```bash
curl -X POST "http://localhost:8000/v2/embed" \
-H "Content-Type: application/json" \
-d '{
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
"input_type": "query",
"texts": ["What is machine learning?"],
"embedding_types": ["float", "binary"]
}'
```
??? console "Response"
```json
{
"id": "embd-...",
"embeddings": {
"float": [[0.012, -0.034, ...]],
"binary": [[42, -117, ...]]
},
"texts": ["What is machine learning?"],
"meta": {
"api_version": {"version": "2"},
"billed_units": {"input_tokens": 8}
}
}
```
#### Truncation
The `truncate` parameter controls how inputs exceeding the model's maximum sequence length are handled:
| Value | Behavior |
| ----- | --------- |
| `END` (default) | Keep the first tokens, drop the end |
| `START` | Keep the last tokens, drop the beginning |
| `NONE` | Return an error if the input is too long |
#### Input type and prompt prefixes
The `input_type` field selects a prompt prefix to prepend to each text input. The available values
depend on the model:
- **Models with `task_instructions` in `config.json`**: The keys from the `task_instructions` dict are
the valid `input_type` values and the corresponding value is prepended to each text.
- **Models with `config_sentence_transformers.json` prompts**: The keys from the `prompts` dict are
the valid `input_type` values. For example, `Snowflake/snowflake-arctic-embed-xs` defines `"query"`,
so setting `input_type: "query"` prepends `"Represent this sentence for searching relevant passages: "`.
- **Other models**: `input_type` is not accepted and will raise a validation error if passed.
### Transcriptions API
Our Transcriptions API is compatible with [OpenAI's Transcriptions API](https://platform.openai.com/docs/api-reference/audio/createTranscription);
@@ -14,6 +14,10 @@ import regex as re
import zmq
from quart import Quart, make_response, request
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
MoRIIOConstants,
)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
prefill_instances: list[dict] = []
@@ -213,6 +217,8 @@ async def handle_request():
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
req_data_to_prefill = copy.deepcopy(req_data)
req_data_to_prefill["kv_transfer_params"] = {}
req_data["kv_transfer_params"] = {}
@@ -222,6 +228,7 @@ async def handle_request():
req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
decode_instance_endpoint["tp_size"]
)
req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
send_prefill_task = asyncio.create_task(
send_request_to_prefill(
@@ -267,6 +274,7 @@ async def handle_request():
if selected_prefill_dp_rank is not None:
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
req_data["kv_transfer_params"]["transfer_id"] = transfer_id
decode_request_task = asyncio.create_task(
start_decode_request(
+1 -1
View File
@@ -24,7 +24,7 @@ outlines_core == 0.2.11
# required for outlines backend disk cache
diskcache == 5.6.3
lark == 1.2.2
xgrammar == 0.1.29; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
typing_extensions >= 4.10
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
partial-json-parser # used for parsing partial JSON outputs
+1 -1
View File
@@ -1,2 +1,2 @@
# formatting
pre-commit==4.0.1
pre-commit>=4.5.1
+1 -1
View File
@@ -42,7 +42,7 @@ tritonclient>=2.51.0
numba == 0.61.2 # Required for N-gram speculative decoding
numpy
runai-model-streamer[s3,gcs]==0.15.3
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.2.2
instanttensor>=0.1.5
pydantic>=2.12 # 2.11 leads to error on python 3.13
+3 -1
View File
@@ -45,10 +45,12 @@ pystemmer==3.0.0
# via mteb
# Multi-modal processing
av==16.1.0
# required for audio_in_video tests
blobfile==3.0.0
# Multi-Modal Models Test
decord==0.6.0
# video processing, required by entrypoints/openai/test_video.py
# video processing, required by entrypoints/openai/chat_completion/test_video.py
rapidfuzz==3.12.1
# OpenAI compatibility and testing
+1 -1
View File
@@ -15,7 +15,7 @@ tensorizer==2.10.1
packaging>=24.2
setuptools>=77.0.3,<80.0.0
setuptools-scm>=8
runai-model-streamer[s3,gcs]==0.15.3
runai-model-streamer[s3,gcs,azure]==0.15.7
conch-triton-kernels==1.2.1
timm>=1.0.17
# amd-quark: required for Quark quantization on ROCm
+2 -1
View File
@@ -10,6 +10,7 @@ pytest-cov
# testing utils
albumentations # required for Nemotron Parse in test_common.py
av # required for audio_in_video tests
backoff # required for phi4mm test
blobfile # required for kimi-vl test
einops # required for MPT, qwen-vl
@@ -55,7 +56,7 @@ grpcio-reflection==1.78.0
arctic-inference == 0.1.1 # Required for suffix decoding test
numba == 0.61.2 # Required for N-gram speculative decoding
numpy
runai-model-streamer[s3,gcs]==0.15.3
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
instanttensor>=0.1.5
pydantic>=2.12 # 2.11 leads to error on python 3.13
+40 -5
View File
@@ -62,6 +62,16 @@ attrs==24.2.0
# referencing
audioread==3.0.1
# via librosa
av==16.1.0
# via -r requirements/test.in
azure-core==1.38.2
# via
# azure-identity
# azure-storage-blob
azure-identity==1.25.2
# via runai-model-streamer-azure
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test.in
@@ -101,8 +111,10 @@ certifi==2024.8.30
# rasterio
# requests
# sentry-sdk
cffi==1.17.1
# via soundfile
cffi==2.0.0
# via
# cryptography
# soundfile
chardet==5.2.0
# via mbstrdecoder
charset-normalizer==3.4.0
@@ -146,6 +158,12 @@ coverage==7.10.6
# via pytest-cov
cramjam==2.9.0
# via fastparquet
cryptography==46.0.5
# via
# azure-identity
# azure-storage-blob
# msal
# pyjwt
cuda-bindings==12.9.4
# via torch
cuda-pathfinder==1.3.3
@@ -377,6 +395,8 @@ iniconfig==2.0.0
# via pytest
instanttensor==0.1.5
# via -r requirements/test.in
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==5.13.2
@@ -490,6 +510,12 @@ more-itertools==10.5.0
# via lm-eval
mpmath==1.3.0
# via sympy
msal==1.34.0
# via
# azure-identity
# msal-extensions
msal-extensions==1.3.1
# via azure-identity
msgpack==1.1.0
# via
# librosa
@@ -826,6 +852,8 @@ pydantic-extra-types==2.10.5
# via mistral-common
pygments==2.18.0
# via rich
pyjwt==2.11.0
# via msal
pyogrio==0.11.0
# via geopandas
pyparsing==3.2.0
@@ -943,6 +971,7 @@ regex==2024.9.11
# transformers
requests==2.32.3
# via
# azure-core
# buildkite-test-collector
# datasets
# diffusers
@@ -955,6 +984,7 @@ requests==2.32.3
# lightly
# lm-eval
# mistral-common
# msal
# mteb
# pooch
# ray
@@ -991,11 +1021,13 @@ rsa==4.9.1
# via google-auth
rtree==1.4.0
# via torchgeo
runai-model-streamer==0.15.3
runai-model-streamer==0.15.7
# via -r requirements/test.in
runai-model-streamer-gcs==0.15.3
runai-model-streamer-azure==0.15.7
# via runai-model-streamer
runai-model-streamer-s3==0.15.3
runai-model-streamer-gcs==0.15.7
# via runai-model-streamer
runai-model-streamer-s3==0.15.7
# via runai-model-streamer
s3transfer==0.10.3
# via boto3
@@ -1264,6 +1296,9 @@ typing-extensions==4.15.0
# aiosignal
# albumentations
# alembic
# azure-core
# azure-identity
# azure-storage-blob
# chz
# fastapi
# grpcio
+14 -8
View File
@@ -657,13 +657,18 @@ class precompiled_wheel_utils:
def get_base_commit_in_main_branch() -> str:
try:
# Get the latest commit hash of the upstream main branch.
resp_json = subprocess.check_output(
[
"curl",
"-s",
"https://api.github.com/repos/vllm-project/vllm/commits/main",
curl_cmd = [
"curl",
"-s",
"https://api.github.com/repos/vllm-project/vllm/commits/main",
]
github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN"))
if github_token:
curl_cmd += [
"-H",
f"Authorization: token {github_token}",
]
).decode("utf-8")
resp_json = subprocess.check_output(curl_cmd).decode("utf-8")
upstream_main_commit = json.loads(resp_json)["sha"]
print(f"Upstream main branch latest commit: {upstream_main_commit}")
@@ -966,18 +971,19 @@ setup(
ext_modules=ext_modules,
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": ["zentorch"],
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
"instanttensor": ["instanttensor >= 0.1.5"],
"runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
"audio": [
"librosa",
"scipy",
"soundfile",
"mistral_common[audio]",
"av",
"torchcodec",
], # Required for audio processing
"video": [], # Kept for backwards compatibility
"flashinfer": [], # Kept for backwards compatibility
@@ -179,7 +179,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
def ops_in_model_before(self):
return [
torch.ops.vllm.all_reduce.default,
torch.ops._C.scaled_fp4_quant.default,
torch.ops._C.scaled_fp4_quant.out,
]
+3 -1
View File
@@ -1,7 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from ..entrypoints.openai.test_oot_registration import run_and_test_dummy_opt_api_server
from tests.entrypoints.openai.chat_completion.test_oot_registration import (
run_and_test_dummy_opt_api_server,
)
def test_distributed_oot(dummy_opt_path: str):
@@ -4,6 +4,9 @@
Tests the image source handling and tool_result content parsing in
AnthropicServingMessages._convert_anthropic_to_openai_request().
Also covers extended-thinking edge cases such as ``redacted_thinking``
blocks echoed back by Anthropic clients.
"""
from vllm.entrypoints.anthropic.protocol import (
@@ -373,3 +376,262 @@ class TestAttributionHeaderStripping:
result = _convert(request)
system_msg = result.messages[0]
assert system_msg["content"] == "You are a helpful assistant."
# ======================================================================
# Thinking block conversion (Anthropic → OpenAI)
# ======================================================================
class TestThinkingBlockConversion:
"""Verify that thinking blocks in assistant messages are correctly
moved to the ``reasoning`` field and stripped from ``content`` during
the Anthropic→OpenAI conversion.
This is the Anthropic-endpoint path: the client echoes back the full
assistant message (including thinking blocks emitted by vllm) in
subsequent requests.
"""
def test_thinking_plus_text_in_assistant_message(self):
"""thinking + text → reasoning field + plain-string content."""
request = _make_request(
[
{"role": "user", "content": "Write me some code."},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "I should write a simple example.",
"signature": "sig_abc123",
},
{"type": "text", "text": "Sure! Here is the code."},
],
},
{"role": "user", "content": "Can you fix the bug?"},
]
)
result = _convert(request)
# Find the assistant message in the converted output.
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
# Thinking content must be in reasoning, NOT in content.
assert asst.get("reasoning") == "I should write a simple example."
assert asst.get("content") == "Sure! Here is the code."
def test_thinking_only_in_assistant_message(self):
"""Assistant message with only a thinking block (no visible text).
This can happen when the model emits reasoning but no final answer
yet (e.g. a mid-turn reasoning step). Content should be None.
"""
request = _make_request(
[
{"role": "user", "content": "Hello"},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "Just thinking...",
"signature": "sig_xyz",
}
],
},
{"role": "user", "content": "Go on."},
]
)
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
assert asst.get("reasoning") == "Just thinking..."
# No visible text → content should be absent or None.
assert asst.get("content") is None
def test_thinking_plus_tool_use_in_assistant_message(self):
"""thinking + tool_use: reasoning field set, tool_calls populated."""
request = _make_request(
[
{"role": "user", "content": "What is 2+2?"},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "I need to call the calculator.",
"signature": "sig_tool",
},
{
"type": "tool_use",
"id": "call_001",
"name": "calculator",
"input": {"expression": "2+2"},
},
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "call_001",
"content": "4",
}
],
},
]
)
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
assert asst.get("reasoning") == "I need to call the calculator."
tool_calls = list(asst.get("tool_calls", []))
assert len(tool_calls) == 1
assert tool_calls[0]["function"]["name"] == "calculator"
# No text content alongside reasoning + tool_use.
assert asst.get("content") is None
def test_multiple_thinking_blocks_concatenated(self):
"""Multiple thinking blocks should be joined in order."""
request = _make_request(
[
{"role": "user", "content": "Think hard."},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "First thought. ",
"signature": "s1",
},
{
"type": "thinking",
"thinking": "Second thought.",
"signature": "s2",
},
{"type": "text", "text": "Done."},
],
},
]
)
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
assert asst.get("reasoning") == "First thought. Second thought."
assert asst.get("content") == "Done."
def test_no_thinking_blocks_unchanged(self):
"""Messages without thinking blocks must not be modified."""
request = _make_request(
[
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello!"},
]
)
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
assert asst.get("content") == "Hello!"
assert "reasoning" not in asst
def test_multi_turn_with_thinking_blocks(self):
"""Full multi-turn conversation: previous assistant messages that
include thinking blocks must all be converted without a 400 error.
This is the primary regression scenario from the bug report:
upgrading vllm from v0.15.1 → v0.17.0 introduced thinking-block
support in responses, but echoing those responses back in subsequent
requests caused a Pydantic validation failure.
"""
request = _make_request(
[
{"role": "user", "content": "Turn 1 question"},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "Reasoning for turn 1.",
"signature": "s_t1",
},
{"type": "text", "text": "Answer for turn 1."},
],
},
{"role": "user", "content": "Turn 2 question"},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "Reasoning for turn 2.",
"signature": "s_t2",
},
{"type": "text", "text": "Answer for turn 2."},
],
},
{"role": "user", "content": "Turn 3 question"},
]
)
# Must not raise a ValidationError / 400.
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 2
assert asst_msgs[0].get("reasoning") == "Reasoning for turn 1."
assert asst_msgs[0].get("content") == "Answer for turn 1."
assert asst_msgs[1].get("reasoning") == "Reasoning for turn 2."
assert asst_msgs[1].get("content") == "Answer for turn 2."
def test_redacted_thinking_block_is_accepted(self):
"""Anthropic clients may echo back redacted thinking blocks.
vLLM should accept these blocks (to avoid 400 validation errors)
and ignore them when constructing the OpenAI-format prompt.
"""
request = _make_request(
[
{"role": "user", "content": "Hello"},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "Thinking...",
"signature": "sig_think",
},
{
"type": "redacted_thinking",
"data": "BASE64_OR_OTHER_OPAQUE_DATA",
},
{"type": "text", "text": "Hi!"},
],
},
{"role": "user", "content": "Continue"},
]
)
result = _convert(request)
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
assert len(asst_msgs) == 1
asst = asst_msgs[0]
# Redacted thinking is ignored, normal thinking still becomes reasoning.
assert asst.get("reasoning") == "Thinking..."
assert asst.get("content") == "Hi!"
+1 -2
View File
@@ -4,12 +4,11 @@ import weakref
import pytest
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.sampling_params import SamplingParams
from ..openai.test_vision import TEST_IMAGE_ASSETS
@pytest.fixture(scope="function")
def text_llm():
+1 -2
View File
@@ -6,13 +6,12 @@ import logging
import pytest
import regex as re
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
from vllm import LLM
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
from vllm.v1.metrics import loggers as stat_loggers
from vllm.v1.metrics.reader import Counter, Metric
from ..openai.test_vision import TEST_IMAGE_ASSETS
def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
return [
@@ -7,11 +7,10 @@ import openai
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from vllm.assets.audio import AudioAsset
from vllm.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio
from ...utils import RemoteOpenAIServer
MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
TEST_AUDIO_URLS = [
AudioAsset("winning_call").url,
@@ -0,0 +1,175 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import base64
import json
import openai
import pytest
import pytest_asyncio
from tests.conftest import VideoTestAssets
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
@pytest.fixture
def server():
args = [
"--max-model-len",
"16384",
"--enforce-eager",
"--limit-mm-per-prompt",
json.dumps({"audio": 3, "video": 3}),
]
with RemoteOpenAIServer(
MODEL_NAME,
args,
) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server):
async with server.get_async_client() as async_client:
yield async_client
@pytest.mark.core_model
@pytest.mark.asyncio
async def test_online_audio_in_video(
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
):
"""Test video input with `audio_in_video=True`"""
# we don't use video_urls above because they missed audio stream.
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
video_base64 = base64.b64encode(f.read()).decode("utf-8")
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this video?"},
{
"type": "video_url",
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
},
],
}
]
# multi-turn to test mm processor cache as well
for _ in range(2):
chat_completion = await client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
max_tokens=16,
extra_body={
"mm_processor_kwargs": {
"use_audio_in_video": True,
}
},
)
assert len(chat_completion.choices) == 1
choice = chat_completion.choices[0]
assert choice.finish_reason == "length"
@pytest.mark.core_model
@pytest.mark.asyncio
async def test_online_audio_in_video_multi_videos(
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
):
"""Test multi-video input with `audio_in_video=True`"""
# we don't use video_urls above because they missed audio stream.
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
video_base64 = base64.b64encode(f.read()).decode("utf-8")
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in these two videos?"},
{
"type": "video_url",
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
},
{
"type": "video_url",
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
},
],
}
]
# multi-turn to test mm processor cache as well
for _ in range(2):
chat_completion = await client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
max_tokens=16,
extra_body={
"mm_processor_kwargs": {
"use_audio_in_video": True,
}
},
)
assert len(chat_completion.choices) == 1
choice = chat_completion.choices[0]
assert choice.finish_reason == "length"
@pytest.mark.core_model
@pytest.mark.asyncio
async def test_online_audio_in_video_interleaved(
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
):
"""Test interleaved video/audio input with `audio_in_video=True`"""
# we don't use video_urls above because they missed audio stream.
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
video_base64 = base64.b64encode(f.read()).decode("utf-8")
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in these two videos?"},
{
"type": "video_url",
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
},
{
"type": "audio_url",
"audio_url": {"url": f"data:audio/mp4;base64,{video_base64}"},
},
{
"type": "video_url",
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
},
],
}
]
with pytest.raises(
openai.BadRequestError,
match="use_audio_in_video requires equal number of audio and video items",
):
await client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
max_tokens=16,
extra_body={
"mm_processor_kwargs": {
"use_audio_in_video": True,
}
},
)
@@ -8,8 +8,8 @@ import pytest
import pytest_asyncio
from huggingface_hub import snapshot_download
from ...conftest import AudioTestAssets
from ...utils import RemoteOpenAIServer
from tests.conftest import AudioTestAssets
from tests.utils import RemoteOpenAIServer
# NOTE - the tests in this module are currently analogous to test_chat, but are
# separated to avoid OOM killing due to module-scoped servers, since we
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from ...utils import VLLM_PATH, RemoteOpenAIServer
from tests.utils import VLLM_PATH, RemoteOpenAIServer
chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
assert chatml_jinja_path.exists()
@@ -8,7 +8,7 @@ from typing import Any, NamedTuple
import openai # use the official client for correctness check
import pytest
from ...utils import RemoteOpenAIServer
from tests.utils import RemoteOpenAIServer
# # any model with a chat template should work here
MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
@@ -7,11 +7,10 @@ import openai
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from vllm.multimodal.utils import encode_video_url, fetch_video
from vllm.platforms import current_platform
from ...utils import RemoteOpenAIServer
MODEL_NAME = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
MAXIMUM_VIDEOS = 3
@@ -8,12 +8,11 @@ import pytest
import pytest_asyncio
from transformers import AutoProcessor
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
from vllm.multimodal.media import MediaWithBytes
from vllm.multimodal.utils import encode_image_url, fetch_image
from vllm.platforms import current_platform
from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
MAXIMUM_IMAGES = 2
@@ -8,10 +8,9 @@ import pytest
import requests
import torch
from tests.utils import RemoteOpenAIServer
from vllm.utils.serial_utils import tensor2base64
from ...utils import RemoteOpenAIServer
@pytest.mark.parametrize(
"model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
@@ -14,7 +14,7 @@ import torch
from openai import BadRequestError
from transformers import AutoConfig
from ...utils import RemoteOpenAIServer
from tests.utils import RemoteOpenAIServer
# any model with a chat template should work here
MODEL_NAME = "facebook/opt-125m"
@@ -11,11 +11,10 @@ import pytest
import regex as re
import torch
from tests.utils import RemoteOpenAIServer
from vllm.config import ModelConfig
from vllm.renderers.embed_utils import safe_load_prompt_embeds
from ...utils import RemoteOpenAIServer
@pytest.mark.asyncio
async def test_empty_prompt():
@@ -9,6 +9,7 @@ import pytest
import pytest_asyncio
import torch.cuda
from tests.utils import RemoteOpenAIServer
from vllm.engine.arg_utils import EngineArgs
from vllm.model_executor.model_loader.tensorizer import (
TensorizerConfig,
@@ -17,8 +18,6 @@ from vllm.model_executor.model_loader.tensorizer import (
)
from vllm.platforms import current_platform
from ...utils import RemoteOpenAIServer
MODEL_NAME = "unsloth/llama-3.2-1b-Instruct"
LORA_PATH = "davzoku/finqa_adapter_1b"
@@ -6,11 +6,10 @@ import tempfile
import pytest
from tests.utils import RemoteOpenAIServer
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
from vllm.tokenizers import get_tokenizer
from ...utils import RemoteOpenAIServer
MODEL_NAME = "Qwen/Qwen3-0.6B"
MODEL_PATH = os.path.join(tempfile.gettempdir(), "qwen3_06b")
@@ -137,6 +137,59 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_streaming_logprobs(client: OpenAI, model_name: str):
"""Test that streaming with logprobs returns valid logprob data on
output_text.delta events and that top_logprobs has the requested count."""
response = await client.responses.create(
model=model_name,
input="Say hello.",
stream=True,
top_logprobs=3,
include=["message.output_text.logprobs"],
)
events = []
async for event in response:
events.append(event)
assert len(events) > 0
# Collect all output_text.delta events that carry logprobs
text_delta_events = [e for e in events if e.type == "response.output_text.delta"]
assert len(text_delta_events) > 0, "Expected at least one text delta event"
for delta_event in text_delta_events:
logprobs = delta_event.logprobs
assert logprobs is not None, "logprobs should be present on text delta events"
assert len(logprobs) > 0, "logprobs list should not be empty"
for lp in logprobs:
# Each logprob entry must have a token and a logprob value
assert lp.token is not None
assert isinstance(lp.logprob, float)
assert lp.logprob <= 0.0, f"logprob should be <= 0, got {lp.logprob}"
# top_logprobs should have up to 3 entries
assert lp.top_logprobs is not None
assert len(lp.top_logprobs) <= 3
for tl in lp.top_logprobs:
assert tl.token is not None
assert isinstance(tl.logprob, float)
# Verify that top_logprobs are actually populated, not always empty
all_top_logprobs = [
tl for e in text_delta_events for lp in e.logprobs for tl in lp.top_logprobs
]
assert len(all_top_logprobs) > 0, (
"Expected at least one top_logprobs entry across all delta events"
)
# Verify the completed event still has valid output
completed = events[-1]
assert completed.type == "response.completed"
assert completed.response.status == "completed"
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
@@ -1,279 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Integration tests for GPT-OSS structural tags functionality (PR #25515)."""
import json
from unittest.mock import Mock
import pytest
from vllm.entrypoints.mcp.tool_server import ToolServer
from vllm.reasoning.gptoss_reasoning_parser import (
GptOssReasoningParser,
)
from vllm.sampling_params import StructuredOutputsParams
class TestGptOssStructuralTagsIntegration:
"""Integration tests for structural tags in GPT-OSS tool calls."""
@pytest.fixture
def mock_tokenizer(self):
"""Create a mock tokenizer."""
tokenizer = Mock()
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
return tokenizer
@pytest.fixture
def gptoss_parser(self, mock_tokenizer):
"""Create a real GptOssReasoningParser instance."""
return GptOssReasoningParser(mock_tokenizer)
@pytest.fixture
def tool_server_with_python(self):
"""Create a tool server with Python tool enabled."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
return tool_server
@pytest.fixture
def tool_server_empty(self):
"""Create a tool server with no tools."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(return_value=False)
return tool_server
def test_end_to_end_no_tools(self, gptoss_parser):
"""Test end-to-end flow when no tools are available."""
# Test the parser directly
result = gptoss_parser.prepare_structured_tag(None, None)
parsed_result = json.loads(result)
# Verify basic structure
assert parsed_result["type"] == "structural_tag"
assert parsed_result["format"]["type"] == "triggered_tags"
assert len(parsed_result["format"]["tags"]) == 1
# Verify only analysis channel is allowed
analysis_tag = parsed_result["format"]["tags"][0]
assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
assert analysis_tag["content"]["type"] == "any_text"
assert analysis_tag["end"] == "<|end|>"
# Verify triggers
assert parsed_result["format"]["triggers"] == ["<|channel|>analysis"]
assert parsed_result["format"]["stop_after_first"] is False
def test_end_to_end_with_python_tool(self, gptoss_parser, tool_server_with_python):
"""Test end-to-end flow with Python tool enabled."""
result = gptoss_parser.prepare_structured_tag(None, tool_server_with_python)
parsed_result = json.loads(result)
# Should have analysis tag + 2 python tags
assert len(parsed_result["format"]["tags"]) == 3
# Verify all expected tags are present
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
expected_begins = [
"<|channel|>analysis<|message|>",
"<|channel|>commentary to=python",
"<|channel|>analysis to=python",
]
for expected in expected_begins:
assert expected in tag_begins
# Verify triggers include commentary
assert "<|channel|>analysis" in parsed_result["format"]["triggers"]
assert "<|channel|>commentary to=" in parsed_result["format"]["triggers"]
def test_structured_outputs_params_integration(
self, gptoss_parser, tool_server_with_python
):
"""Test integration with StructuredOutputsParams."""
# Generate structural tag
structural_tag = gptoss_parser.prepare_structured_tag(
None, tool_server_with_python
)
# Create StructuredOutputsParams
params = StructuredOutputsParams(structural_tag=structural_tag)
# Verify the tag is properly stored and accessible
assert params.structural_tag == structural_tag
# Verify the tag is valid JSON
parsed_tag = json.loads(params.structural_tag)
assert parsed_tag["type"] == "structural_tag"
@pytest.mark.parametrize(
"browser, python, container, expected_tags",
[
# No tools
(False, False, False, 1),
# Single tool
(True, False, False, 3),
# Multiple tools
(True, True, False, 5),
# All tools
(True, True, True, 7),
],
)
def test_tool_server_interaction_flow(
self, gptoss_parser, browser, python, container, expected_tags
):
"""Test the complete tool server interaction flow."""
# Create a mock ToolServer
tool_server = Mock(spec=ToolServer)
# Simulate tool availability based on parameters
tool_server.has_tool = Mock(
side_effect=lambda tool: {
"browser": browser,
"python": python,
"container": container,
}.get(tool, False)
)
# Run the parser and verify results
result = gptoss_parser.prepare_structured_tag(None, tool_server)
parsed_result = json.loads(result)
# Validate number of tags
assert len(parsed_result["format"]["tags"]) == expected_tags
# Verify tool-specific tags exist for enabled tools
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
for tool, enabled in {
"browser": browser,
"python": python,
"container": container,
}.items():
if enabled:
assert f"<|channel|>commentary to={tool}" in tag_begins
assert f"<|channel|>analysis to={tool}" in tag_begins
def test_original_tag_preservation(self, gptoss_parser, tool_server_with_python):
"""Test that original tags are preserved when provided."""
original_tag = '{"type": "custom_tag", "data": "preserved"}'
result = gptoss_parser.prepare_structured_tag(
original_tag, tool_server_with_python
)
# Should return original tag unchanged
assert result == original_tag
@pytest.mark.parametrize(
"tools",
[
[],
["browser"],
["python"],
["container"],
["browser", "python"],
["browser", "container"],
["python", "container"],
["browser", "python", "container"],
],
)
def test_json_validity_comprehensive(self, gptoss_parser, tools):
"""Test JSON validity across all possible tool combinations."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
result = gptoss_parser.prepare_structured_tag(None, tool_server)
# Should be valid JSON
parsed_result = json.loads(result)
# Should have correct structure
assert parsed_result["type"] == "structural_tag"
assert "format" in parsed_result
assert "tags" in parsed_result["format"]
assert "triggers" in parsed_result["format"]
# Tag count should be: 1 (analysis) + 2 * len(tools)
expected_tag_count = 1 + (2 * len(tools))
assert len(parsed_result["format"]["tags"]) == expected_tag_count
def test_error_handling_invalid_tool_server(self, gptoss_parser):
"""Test error handling with invalid tool server."""
# Tool server that raises exceptions
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(side_effect=Exception("Tool server error"))
# Should handle gracefully and still return a valid tag
with pytest.raises(Exception, match="Tool server error"):
gptoss_parser.prepare_structured_tag(None, tool_server)
def test_concurrent_requests_isolation(self, gptoss_parser):
"""Test that concurrent requests don't interfere with each other."""
# Simulate concurrent requests with different tool servers
tool_server_1 = Mock(spec=ToolServer)
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
tool_server_2 = Mock(spec=ToolServer)
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
# Generate tags concurrently
result_1 = gptoss_parser.prepare_structured_tag(None, tool_server_1)
result_2 = gptoss_parser.prepare_structured_tag(None, tool_server_2)
# Parse results
parsed_1 = json.loads(result_1)
parsed_2 = json.loads(result_2)
# Verify they have different tool configurations
tags_1 = [tag["begin"] for tag in parsed_1["format"]["tags"]]
tags_2 = [tag["begin"] for tag in parsed_2["format"]["tags"]]
# Result 1 should have python tags
assert "<|channel|>commentary to=python" in tags_1
assert "<|channel|>commentary to=browser" not in tags_1
# Result 2 should have browser tags
assert "<|channel|>commentary to=browser" in tags_2
assert "<|channel|>commentary to=python" not in tags_2
def test_tag_format_consistency(self, gptoss_parser):
"""Test that all generated tags follow consistent format."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(
side_effect=lambda tool: tool in ["python", "browser"]
)
result = gptoss_parser.prepare_structured_tag(None, tool_server)
parsed_result = json.loads(result)
# Verify all tags have required fields
for tag in parsed_result["format"]["tags"]:
assert "begin" in tag
assert "content" in tag
assert "end" in tag
assert tag["content"]["type"] == "any_text"
assert tag["end"] == "<|end|>"
# Verify begin format
assert tag["begin"].startswith("<|channel|>")
def test_trigger_configuration(self, gptoss_parser):
"""Test trigger configuration for different tool setups."""
# Test with no tools
result_no_tools = gptoss_parser.prepare_structured_tag(None, None)
parsed_no_tools = json.loads(result_no_tools)
assert parsed_no_tools["format"]["triggers"] == ["<|channel|>analysis"]
# Test with tools
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
result_with_tools = gptoss_parser.prepare_structured_tag(None, tool_server)
parsed_with_tools = json.loads(result_with_tools)
expected_triggers = ["<|channel|>analysis", "<|channel|>commentary to="]
assert set(parsed_with_tools["format"]["triggers"]) == set(expected_triggers)
@@ -5,7 +5,7 @@ import json
import pytest
from tests.entrypoints.openai.tool_parsers.utils import (
from tests.tool_parsers.utils import (
run_tool_extraction,
run_tool_extraction_streaming,
)
@@ -0,0 +1,360 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import random
from typing import Any
import openai
import pytest
from transformers import AutoTokenizer
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.engine.protocol import (
DeltaMessage,
)
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
from ....utils import RemoteOpenAIServer
MODEL = "ibm-granite/granite-4.0-h-tiny"
@pytest.fixture(scope="module")
def server():
model = MODEL
args_for_model = [
"--enforce-eager",
"--enable-auto-tool-choice",
"--tool-call-parser",
"granite4",
"--tokenizer",
"ibm-granite/granite-4.0-h-tiny",
"--max-model-len",
"4096",
"--max-num-seqs",
"2",
]
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
yield server
def create_complex_input(create_string_args: bool):
coord_arg: dict | str = {
"coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
"coordinate_type": "latlong",
}
if create_string_args:
# test granite behavior
coord_arg = json.dumps(coord_arg)
return [
{"name": "find_bbox", "arguments": coord_arg},
{
"name": "get_stock_price",
"arguments": {
"symbol": "AAPL",
"start_date": "2021-01-01",
"end_date": "2021-12-31",
},
},
{"name": "find_bbox", "arguments": coord_arg},
]
def random_chunks(s: str, min_len: int, max_len: int):
chunks = []
i = 0
n = len(s)
while i < n:
size = random.randint(min_len, max_len)
chunks.append(s[i : i + size])
i += size
return chunks
@pytest.fixture(scope="module")
def tokenizer():
return AutoTokenizer.from_pretrained(MODEL)
# create a variety of input chunk sizes
@pytest.mark.parametrize(
"min_chunk, max_chunk",
[
(1, 1),
(1, 2),
(5, 7),
(6, 20),
],
)
def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
input_dicts = create_complex_input(True)
formatted_tcs = [
"<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
]
text_messages = [
"Here goes the bbox call: \n",
" Now the stock price call: \n ",
" Now another bbox call: \n ",
" See? I'm a helpful assistant.",
]
test_input = (
text_messages[0]
+ formatted_tcs[0]
+ text_messages[1]
+ formatted_tcs[1]
+ text_messages[2]
+ formatted_tcs[2]
+ text_messages[3]
)
any_chat_request = ChatCompletionRequest(
seed=42,
model=MODEL,
messages=[],
)
parser = Granite4ToolParser(tokenizer=tokenizer)
delta_messages = list[DeltaMessage]()
for text in random_chunks(test_input, min_chunk, max_chunk):
delta = parser.extract_tool_calls_streaming(
previous_text="",
current_text="",
delta_text=text,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[],
request=any_chat_request,
)
if delta is not None:
delta_messages.append(delta)
content = ""
tool_calls = list[dict[str, Any]]()
current_name = "__start__"
current_args = ""
for msg in delta_messages:
if msg.content:
content += msg.content
for tool_call in msg.tool_calls:
if delta_func := tool_call.function:
if delta_func.name is not None:
if current_name == "__start__":
current_name = delta_func.name
if delta_func.name != current_name:
tool_calls.append(
{
"name": current_name,
"arguments": json.loads(current_args),
}
)
current_name = delta_func.name
current_args = ""
if delta_func.arguments:
current_args += delta_func.arguments
if current_name != "__start__":
tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
assert content == "".join(text_messages)
assert tool_calls == create_complex_input(False)
tools = [
{
"type": "function",
"function": {
"name": "get_acme_region_name_for_transaction_id",
"description": "Returns ACME transaction/transaction ID information"
" including ACME regions\n\nArgs:\n start_time "
"(str): Start date and time in datetime format "
'"%Y-%m-%dT%H:%M:%S.%f"\n end_time (str): End '
"date and time in datetime format "
'"%Y-%m-%dT%H:%M:%S.%f"\n size (int, optional): '
"Number of ACME Transaction IDs to return\n "
"order (str, optional): Sort by most run "
"transaction IDs. The value can be 'asc' for "
"ascending or 'desc' for descending\n "
"transaction_id (str, optional): ACME Transaction "
"ID to filter on\n acme_region (str, optional): "
"ACME Region to filter on\nReturns:\n - A "
"dictionary containing a list of ACME transaction "
"ids and the ACME regions they run in:\n {\n"
' "Number of transaction IDs" : int,\n'
' "Total transaction IDs available": int'
',\n "ACME Transaction IDs": [\n '
' {\n "Transaction ID": '
'str,\n "Number of runs": int,\n'
' "ACME Regions": [str],\n '
" },\n ...\n ],"
'\n "Start time" : datetime,\n '
' "End time" : datetime,\n '
' "Order" : str\n }\n '
" - If no ACME region found for transaction id, "
'returns:\n {"Success": "No ACME region '
'found for transaction id."}\n - If an error '
'occurs, returns:\n {"Error": "{exception'
' message}"}',
"parameters": {
"properties": {
"start_time": {},
"end_time": {},
"size": {"default": 500},
"order": {"default": "desc"},
"transaction_id": {"default": None},
"acme_region": {"default": None},
},
"required": ["start_time", "end_time"],
"type": "object",
},
},
}
]
tools2 = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {
"description": "The city and state, e.g. San Francisco, CA",
"type": "string",
}
},
"required": ["location"],
},
},
},
{
"type": "function",
"function": {
"name": "get_stock_price",
"description": "Retrieves the current stock price for a given "
"ticker symbol. The ticker symbol must be a valid "
"symbol for a publicly traded company on a major US"
" stock exchange like NYSE or NASDAQ. The tool will"
" return the latest trade price in USD. It should "
"be used when the user asks about the current or "
"most recent price of a specific stock. It will not"
" provide any other information about the stock or"
" company.",
"parameters": {
"type": "object",
"properties": {
"ticker": {
"description": "The stock ticker symbol, e.g."
" AAPL for Apple Inc.",
"type": "string",
}
},
},
},
},
]
messages = [
{
"content": "\n\nSystem: You are a helpful, precise, and methodical AI"
" assistant that uses tool outputs provided inline.\nAlways"
" assume the current datetime is 2026-01-29T13:59:09.238901"
"+00:00.\n\nIf you receive a ToolMessage with `tool_call_id"
'` equal to "get_time_range" (or "time_range_tool"), you '
"MUST:\n 1. Parse that JSON and use the values `start` and"
" `end` directly when calling other tools.\n 2. Do not "
"re-call or re-compute the time range.\n 3. Pass resolved "
"values (ISO strings) as arguments to any subsequent tool "
"(do not pass function metadata or placeholders).\n 4. If "
"a tool requires datetime objects rather than strings, "
"convert the ISO strings into language-native datetime "
"objects before invoking.\n\nAlways return fully resolved "
"arguments in correct types (e.g., ISO datetime strings or"
" datetime objects) and never include placeholders like "
'"<start>".\n\n',
"role": "system",
},
{
"content": "What are the transaction IDs that ran in the"
" ACME region A9345 over the last two months?",
"role": "user",
},
{
"content": '["2026-01-26T09: 51: 55.467722Z", "2026-01-27T09: 51: 55.467722Z"]',
"role": "tool",
"tool_call_id": "time_range_tool",
},
]
messages2 = [{"role": "user", "content": "What's stock price for IBM?"}]
messages3 = [{"role": "user", "content": "What's the current weather in New York?"}]
def get_args(client: openai.OpenAI, _tools, _messages, _stop):
response = client.chat.completions.create(
model=MODEL,
messages=_messages,
temperature=0,
tools=_tools,
max_tokens=200,
stop=_stop,
tool_choice="auto",
)
return response.choices[0].message.tool_calls[0].function.arguments
async def get_args_streaming(
async_client: openai.AsyncOpenAI, _tools, _messages, _stop
):
stream = await async_client.chat.completions.create(
model=MODEL,
messages=_messages,
temperature=0,
tools=_tools,
max_tokens=200,
stop=_stop,
tool_choice="auto",
stream=True,
)
full_call = []
async for chunk in stream:
tc = chunk.choices[0].delta.tool_calls
if tc and tc[0].function.arguments:
full_call.append(tc[0].function.arguments)
return "".join(full_call)
async def run_scenario(server: RemoteOpenAIServer, _tools, _messages, _stop):
non_streaming = get_args(server.get_client(), _tools, _messages, _stop)
json.loads(non_streaming) # verify that it is json loadable
streaming = await get_args_streaming(
server.get_async_client(), _tools, _messages, _stop
)
json.loads(streaming)
assert non_streaming == streaming, f"{non_streaming=}, {streaming=}"
@pytest.mark.asyncio
async def test_stop_sequence_interference(server: RemoteOpenAIServer):
print("Testing scenario 1")
await run_scenario(server, tools, messages, "veroniqueprattyushveroniqueprattyush")
print("Testing scenario 2")
await run_scenario(
server, tools2, messages2, "veroniqueprattyushveroniqueprattyush"
)
print("Testing scenario 3")
await run_scenario(server, tools2, messages3, "prattyush")
@@ -3,29 +3,22 @@
import json
import openai
import pytest
import pytest_asyncio
from huggingface_hub import snapshot_download
from typing_extensions import TypedDict
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import ToolParser
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
from ....utils import RemoteOpenAIServer
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
SERVER_ARGS = [
"--enforce-eager",
"--enable-auto-tool-choice",
"--tool-call-parser",
"hermes",
"--enable-lora",
"--lora-modules",
f"{LORA_MODEL}={LORA_MODEL}",
"--tokenizer",
f"{LORA_MODEL}",
]
TOOLS = [
{
"type": "function",
@@ -50,6 +43,75 @@ TOOLS = [
}
]
class ServerConfig(TypedDict, total=False):
model: str
arguments: list[str]
model_arg: str
tool_parser: ToolParser
CONFIGS: dict[str, ServerConfig] = {
"llama": {
"model": "meta-llama/Llama-3.2-1B-Instruct",
"arguments": [
"--enforce-eager",
"--enable-auto-tool-choice",
"--tool-call-parser",
"hermes",
"--enable-lora",
"--lora-modules",
f"{LORA_MODEL}={LORA_MODEL}",
"--tokenizer",
f"{LORA_MODEL}",
],
"model_arg": LORA_MODEL,
"tool_parser": Hermes2ProToolParser,
},
"granite4": {
"model": "ibm-granite/granite-4.0-h-tiny",
"arguments": [
"--enforce-eager",
"--enable-auto-tool-choice",
"--tool-call-parser",
"granite4",
"--tokenizer",
"ibm-granite/granite-4.0-h-tiny",
"--max-model-len",
"4096",
"--max-num-seqs",
"2",
],
"model_arg": "ibm-granite/granite-4.0-h-tiny",
"tool_parser": Granite4ToolParser,
},
}
# for each server config, download the model and return the config
@pytest.fixture(scope="session", params=CONFIGS.keys())
def server_config(request):
config = CONFIGS[request.param]
# download model and tokenizer using transformers
snapshot_download(config["model"])
yield CONFIGS[request.param]
@pytest.fixture(scope="module")
def server(request, server_config: ServerConfig):
model = server_config["model"]
args_for_model = server_config["arguments"]
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
yield server
@pytest_asyncio.fixture
async def client(server: RemoteOpenAIServer):
async with server.get_async_client() as async_client:
yield async_client
PRODUCT_TOOLS = [
{
"type": "function",
@@ -87,186 +149,182 @@ PRODUCT_MESSAGES = [
@pytest.mark.asyncio
async def test_non_streaming_tool_call():
async def test_non_streaming_tool_call(
client: openai.AsyncOpenAI, server_config: ServerConfig
):
"""Test tool call in non-streaming mode."""
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
client = server.get_async_client()
response = await client.chat.completions.create(
model=LORA_MODEL,
messages=MESSAGES,
tools=TOOLS,
tool_choice="auto",
temperature=0.0,
)
response = await client.chat.completions.create(
model=server_config["model_arg"],
messages=MESSAGES,
tools=TOOLS,
tool_choice="auto",
temperature=0.0,
)
assert response.choices
choice = response.choices[0]
message = choice.message
assert response.choices
choice = response.choices[0]
message = choice.message
assert choice.finish_reason == "tool_calls"
assert message.tool_calls is not None
assert choice.finish_reason == "tool_calls"
assert message.tool_calls is not None
tool_call = message.tool_calls[0]
assert tool_call.type == "function"
assert tool_call.function.name == "get_current_weather"
tool_call = message.tool_calls[0]
assert tool_call.type == "function"
assert tool_call.function.name == "get_current_weather"
arguments = json.loads(tool_call.function.arguments)
assert "location" in arguments
assert "Boston" in arguments["location"]
print("\n[Non-Streaming Test Passed]")
print(f"Tool Call: {tool_call.function.name}")
print(f"Arguments: {arguments}")
arguments = json.loads(tool_call.function.arguments)
assert "location" in arguments
assert "Boston" in arguments["location"]
print("\n[Non-Streaming Test Passed]")
print(f"Tool Call: {tool_call.function.name}")
print(f"Arguments: {arguments}")
@pytest.mark.asyncio
async def test_streaming_tool_call():
async def test_streaming_tool_call(
client: openai.AsyncOpenAI, server_config: ServerConfig
):
"""Test tool call in streaming mode."""
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
client = server.get_async_client()
stream = await client.chat.completions.create(
model=LORA_MODEL,
messages=MESSAGES,
tools=TOOLS,
tool_choice="auto",
temperature=0.0,
stream=True,
)
stream = await client.chat.completions.create(
model=server_config["model_arg"],
messages=MESSAGES,
tools=TOOLS,
tool_choice="auto",
temperature=0.0,
stream=True,
)
tool_call_chunks = {}
async for chunk in stream:
if not chunk.choices:
continue
tool_call_chunks = {}
async for chunk in stream:
if not chunk.choices:
continue
delta = chunk.choices[0].delta
if not delta or not delta.tool_calls:
continue
delta = chunk.choices[0].delta
if not delta or not delta.tool_calls:
continue
for tool_chunk in delta.tool_calls:
index = tool_chunk.index
if index not in tool_call_chunks:
tool_call_chunks[index] = {"name": "", "arguments": ""}
for tool_chunk in delta.tool_calls:
index = tool_chunk.index
if index not in tool_call_chunks:
tool_call_chunks[index] = {"name": "", "arguments": ""}
if tool_chunk.function.name:
tool_call_chunks[index]["name"] += tool_chunk.function.name
if tool_chunk.function.arguments:
tool_call_chunks[index]["arguments"] += (
tool_chunk.function.arguments
)
if tool_chunk.function.name:
tool_call_chunks[index]["name"] += tool_chunk.function.name
if tool_chunk.function.arguments:
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
assert len(tool_call_chunks) == 1
reconstructed_tool_call = tool_call_chunks[0]
assert len(tool_call_chunks) == 1
reconstructed_tool_call = tool_call_chunks[0]
assert reconstructed_tool_call["name"] == "get_current_weather"
assert reconstructed_tool_call["name"] == "get_current_weather"
arguments = json.loads(reconstructed_tool_call["arguments"])
assert "location" in arguments
assert "Boston" in arguments["location"]
print("\n[Streaming Test Passed]")
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
print(f"Reconstructed Arguments: {arguments}")
arguments = json.loads(reconstructed_tool_call["arguments"])
assert "location" in arguments
assert "Boston" in arguments["location"]
print("\n[Streaming Test Passed]")
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
print(f"Reconstructed Arguments: {arguments}")
@pytest.mark.asyncio
async def test_non_streaming_product_tool_call():
async def test_non_streaming_product_tool_call(
client: openai.AsyncOpenAI, server_config: ServerConfig
):
"""Test tool call integer and boolean parameters in non-streaming mode."""
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
client = server.get_async_client()
response = await client.chat.completions.create(
model=LORA_MODEL,
messages=PRODUCT_MESSAGES,
tools=PRODUCT_TOOLS,
tool_choice="auto",
temperature=0.66,
)
response = await client.chat.completions.create(
model=server_config["model_arg"],
messages=PRODUCT_MESSAGES,
tools=PRODUCT_TOOLS,
tool_choice="auto",
temperature=0.66,
)
assert response.choices
choice = response.choices[0]
message = choice.message
assert response.choices
choice = response.choices[0]
message = choice.message
assert choice.finish_reason == "tool_calls"
assert message.tool_calls is not None
assert choice.finish_reason == "tool_calls"
assert message.tool_calls is not None
tool_call = message.tool_calls[0]
assert tool_call.type == "function"
assert tool_call.function.name == "get_product_info"
tool_call = message.tool_calls[0]
assert tool_call.type == "function"
assert tool_call.function.name == "get_product_info"
arguments = json.loads(tool_call.function.arguments)
assert "product_id" in arguments
assert "inserted" in arguments
arguments = json.loads(tool_call.function.arguments)
assert "product_id" in arguments
assert "inserted" in arguments
product_id = arguments.get("product_id")
inserted = arguments.get("inserted")
product_id = arguments.get("product_id")
inserted = arguments.get("inserted")
assert isinstance(product_id, int)
assert product_id == 7355608
assert isinstance(inserted, bool)
assert inserted is True
assert isinstance(product_id, int)
assert product_id == 7355608
assert isinstance(inserted, bool)
assert inserted is True
print("\n[Non-Streaming Product Test Passed]")
print(f"Tool Call: {tool_call.function.name}")
print(f"Arguments: {arguments}")
print("\n[Non-Streaming Product Test Passed]")
print(f"Tool Call: {tool_call.function.name}")
print(f"Arguments: {arguments}")
@pytest.mark.asyncio
async def test_streaming_product_tool_call():
async def test_streaming_product_tool_call(
client: openai.AsyncOpenAI, server_config: ServerConfig
):
"""Test tool call integer and boolean parameters in streaming mode."""
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
client = server.get_async_client()
stream = await client.chat.completions.create(
model=LORA_MODEL,
messages=PRODUCT_MESSAGES,
tools=PRODUCT_TOOLS,
tool_choice="auto",
temperature=0.66,
stream=True,
)
stream = await client.chat.completions.create(
model=server_config["model_arg"],
messages=PRODUCT_MESSAGES,
tools=PRODUCT_TOOLS,
tool_choice="auto",
temperature=0.66,
stream=True,
)
tool_call_chunks = {}
async for chunk in stream:
if not chunk.choices:
continue
tool_call_chunks = {}
async for chunk in stream:
if not chunk.choices:
continue
delta = chunk.choices[0].delta
if not delta or not delta.tool_calls:
continue
delta = chunk.choices[0].delta
if not delta or not delta.tool_calls:
continue
for tool_chunk in delta.tool_calls:
index = tool_chunk.index
if index not in tool_call_chunks:
tool_call_chunks[index] = {"name": "", "arguments": ""}
for tool_chunk in delta.tool_calls:
index = tool_chunk.index
if index not in tool_call_chunks:
tool_call_chunks[index] = {"name": "", "arguments": ""}
if tool_chunk.function.name:
tool_call_chunks[index]["name"] += tool_chunk.function.name
if tool_chunk.function.arguments:
tool_call_chunks[index]["arguments"] += (
tool_chunk.function.arguments
)
if tool_chunk.function.name:
tool_call_chunks[index]["name"] += tool_chunk.function.name
if tool_chunk.function.arguments:
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
assert len(tool_call_chunks) == 1
reconstructed_tool_call = tool_call_chunks[0]
assert len(tool_call_chunks) == 1
reconstructed_tool_call = tool_call_chunks[0]
assert reconstructed_tool_call["name"] == "get_product_info"
assert reconstructed_tool_call["name"] == "get_product_info"
arguments = json.loads(reconstructed_tool_call["arguments"])
assert "product_id" in arguments
assert "inserted" in arguments
arguments = json.loads(reconstructed_tool_call["arguments"])
assert "product_id" in arguments
assert "inserted" in arguments
# Handle type coercion for streaming test as well
product_id = arguments.get("product_id")
inserted = arguments.get("inserted")
# Handle type coercion for streaming test as well
product_id = arguments.get("product_id")
inserted = arguments.get("inserted")
assert isinstance(product_id, int)
assert product_id == 7355608
assert isinstance(inserted, bool)
assert inserted is True
assert isinstance(product_id, int)
assert product_id == 7355608
assert isinstance(inserted, bool)
assert inserted is True
print("\n[Streaming Product Test Passed]")
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
print(f"Reconstructed Arguments: {arguments}")
print("\n[Streaming Product Test Passed]")
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
print(f"Reconstructed Arguments: {arguments}")
@pytest.fixture
@@ -276,9 +334,10 @@ def qwen_tokenizer() -> TokenizerLike:
return get_tokenizer("Qwen/Qwen3-32B")
@pytest.fixture
def hermes_parser(qwen_tokenizer: TokenizerLike) -> Hermes2ProToolParser:
return Hermes2ProToolParser(qwen_tokenizer)
@pytest.fixture(params=CONFIGS.keys())
def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
config = CONFIGS[request.param]
return config["tool_parser"](qwen_tokenizer)
@pytest.fixture
@@ -292,7 +351,7 @@ def any_chat_request() -> ChatCompletionRequest:
def test_hermes_parser_streaming_just_forward_text(
qwen_tokenizer: TokenizerLike,
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
text = """This is some prior text that has nothing to do with tool calling."""
@@ -324,7 +383,7 @@ def test_hermes_parser_streaming_just_forward_text(
def test_hermes_parser_streaming_failure_case_bug_19056(
qwen_tokenizer: TokenizerLike,
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
text = """<tool_call>
@@ -358,7 +417,7 @@ def test_hermes_parser_streaming_failure_case_bug_19056(
def test_hermes_parser_streaming(
qwen_tokenizer: TokenizerLike,
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
text = '<tool_call>\
@@ -387,16 +446,20 @@ def test_hermes_parser_streaming(
delta_messages.append(delta)
print(delta_messages)
assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
tool_call_args = "".join(
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
)
assert tool_call_args == (
'{"location":"San Francisco, California, United States", "unit": "celsius"}'
# load to normalize whitespace
tool_call_args = json.loads(
"".join(
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
)
)
assert tool_call_args == {
"location": "San Francisco, California, United States",
"unit": "celsius",
}
def test_hermes_parser_non_streaming_no_tool_call(
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
text = """This is not a tool call."""
@@ -410,7 +473,7 @@ def test_hermes_parser_non_streaming_no_tool_call(
def test_hermes_parser_non_streaming_tool_call_between_tags(
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
text = """<tool_call>
@@ -428,9 +491,12 @@ def test_hermes_parser_non_streaming_tool_call_between_tags(
def test_hermes_parser_non_streaming_tool_call_until_eos(
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
if isinstance(hermes_parser, Granite4ToolParser):
pytest.skip(reason="The Granite4 tool parser enforces a complete response")
text = """<tool_call>
{"name": "final_answer", "arguments": {"trigger": true}}"""
tool_call = hermes_parser.extract_tool_calls(
@@ -445,7 +511,7 @@ def test_hermes_parser_non_streaming_tool_call_until_eos(
def test_hermes_parser_non_streaming_tool_call_invalid_json(
hermes_parser: Hermes2ProToolParser,
hermes_parser: ToolParser,
any_chat_request: ChatCompletionRequest,
) -> None:
# Missing closing brace to trigger exception
@@ -7,7 +7,7 @@ from unittest.mock import MagicMock
import pytest
from tests.entrypoints.openai.tool_parsers.utils import (
from tests.tool_parsers.utils import (
run_tool_extraction,
run_tool_extraction_streaming,
)
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
import pytest
from tests.entrypoints.openai.tool_parsers.utils import (
from tests.tool_parsers.utils import (
run_tool_extraction,
run_tool_extraction_streaming,
)
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
import pytest
from tests.entrypoints.openai.tool_parsers.utils import (
from tests.tool_parsers.utils import (
run_tool_extraction,
run_tool_extraction_streaming,
)
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
import pytest
from tests.entrypoints.openai.tool_parsers.utils import (
from tests.tool_parsers.utils import (
run_tool_extraction,
run_tool_extraction_streaming,
)
@@ -0,0 +1,310 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
Validates that the Cohere v2 embed endpoint works correctly with standard
embedding models, covering text embedding, embedding type conversions,
response structure, batching, normalisation, and semantic similarity.
"""
import base64
import struct
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
DTYPE = "bfloat16"
MODELS: list[tuple[str, list[str]]] = [
("intfloat/multilingual-e5-small", []),
(
"Snowflake/snowflake-arctic-embed-m-v1.5",
[
"--trust_remote_code",
"--hf_overrides",
'{"matryoshka_dimensions":[256]}',
],
),
]
@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
def model_config(request):
return request.param
@pytest.fixture(scope="module")
def model_name(model_config):
return model_config[0]
@pytest.fixture(scope="module")
def server(model_config):
name, extra_args = model_config
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512",
"--gpu-memory-utilization",
"0.02",
] + extra_args
with RemoteOpenAIServer(name, args) as remote_server:
yield remote_server
def _cohere_embed(
server: RemoteOpenAIServer,
model_name: str,
texts: list[str] | None = None,
images: list[str] | None = None,
input_type: str | None = None,
embedding_types: list[str] | None = None,
) -> dict:
body: dict = {"model": model_name}
if input_type is not None:
body["input_type"] = input_type
if texts is not None:
body["texts"] = texts
if images is not None:
body["images"] = images
if embedding_types is not None:
body["embedding_types"] = embedding_types
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()
def _openai_embed(
server: RemoteOpenAIServer, model_name: str, texts: list[str]
) -> dict:
body = {"model": model_name, "input": texts, "encoding_format": "float"}
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
resp.raise_for_status()
return resp.json()
def _cosine_sim(a: list[float], b: list[float]) -> float:
va, vb = np.array(a), np.array(b)
return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
# -----------------------------------------------------------
# Text embedding tests
# -----------------------------------------------------------
def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server, model_name, texts=["hello world"], embedding_types=["float"]
)
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
"""An input_type not defined in the model's prompt config should be
rejected with a 400 error."""
body = {
"model": model_name,
"input_type": "nonexistent_type",
"texts": ["hello world"],
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 400
assert "Unsupported input_type" in resp.json()["error"]["message"]
def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
"""Omitting input_type should always work (no prompt prefix applied)."""
body = {
"model": model_name,
"texts": ["hello world"],
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
float embeddings for a generic model."""
texts = ["hello world"]
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
v1 = _openai_embed(server, model_name, texts)
cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["test"],
embedding_types=["float", "binary", "ubinary"],
)
dim = len(r["embeddings"]["float"][0])
assert len(r["embeddings"]["binary"][0]) == dim // 8
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
def test_response_structure(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
assert "id" in r
assert "embeddings" in r
assert "texts" in r
assert r["texts"] == ["test"]
assert "meta" in r
assert r["meta"]["api_version"]["version"] == "2"
assert "billed_units" in r["meta"]
assert r["meta"]["billed_units"]["input_tokens"] > 0
assert r["meta"]["billed_units"]["image_tokens"] == 0
def test_batch(server: RemoteOpenAIServer, model_name: str):
texts = ["apple", "banana", "cherry"]
r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
assert len(r["embeddings"]["float"]) == 3
dim = len(r["embeddings"]["float"][0])
for emb in r["embeddings"]["float"]:
assert len(emb) == dim
def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server, model_name, texts=["hello world"], embedding_types=["float"]
)
emb = np.array(r["embeddings"]["float"][0])
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["machine learning", "deep learning", "chocolate cake recipe"],
embedding_types=["float"],
)
embs = r["embeddings"]["float"]
cos_related = _cosine_sim(embs[0], embs[1])
cos_unrelated = _cosine_sim(embs[0], embs[2])
assert cos_related > cos_unrelated
def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
body = {"model": model_name}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 400
def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["test encoding"],
embedding_types=["float", "base64"],
)
float_emb = r["embeddings"]["float"][0]
b64_str = r["embeddings"]["base64"][0]
decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
# -----------------------------------------------------------
# Truncation tests
# -----------------------------------------------------------
def _cohere_embed_raw(
server: RemoteOpenAIServer,
body: dict,
) -> requests.Response:
return requests.post(server.url_for("/v2/embed"), json=body)
def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
"""truncate=END should silently truncate long input."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "END",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
"""truncate=START should silently truncate long input from the start."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "START",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
"""truncate=NONE should error when input exceeds model context."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "NONE",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 400
def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
"""START and END truncation should produce different embeddings
when the input is long enough to actually be truncated.
We construct input with distinct tokens at the start vs end
so that keeping different halves produces different embeddings.
"""
start_words = " ".join([f"alpha{i}" for i in range(300)])
end_words = " ".join([f"omega{i}" for i in range(300)])
long_text = start_words + " " + end_words
body_end = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "END",
}
body_start = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "START",
}
r_end = _cohere_embed_raw(server, body_end).json()
r_start = _cohere_embed_raw(server, body_start).json()
emb_end = r_end["embeddings"]["float"][0]
emb_start = r_start["embeddings"]["float"][0]
cos = _cosine_sim(emb_end, emb_start)
assert cos < 0.99, (
f"START and END truncation should produce different embeddings "
f"for long input, but cosine similarity was {cos}"
)
@@ -0,0 +1,135 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
Validates image embedding, batching, normalisation, and embedding type
conversions through the /v2/embed endpoint.
"""
import base64
import struct
import zlib
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "google/siglip-so400m-patch14-384"
DTYPE = "bfloat16"
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"64",
"--gpu-memory-utilization",
"0.3",
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
raw = b""
for _ in range(h):
raw += b"\x00" + bytes([r, g, b]) * w
compressed = zlib.compress(raw)
def chunk(ctype: bytes, cdata: bytes) -> bytes:
c = ctype + cdata
return (
struct.pack(">I", len(cdata))
+ c
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
)
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
png = (
b"\x89PNG\r\n\x1a\n"
+ chunk(b"IHDR", ihdr)
+ chunk(b"IDAT", compressed)
+ chunk(b"IEND", b"")
)
return "data:image/png;base64," + base64.b64encode(png).decode()
def _cohere_embed(
server: RemoteOpenAIServer,
texts: list[str] | None = None,
images: list[str] | None = None,
embedding_types: list[str] | None = None,
) -> dict:
body: dict = {"model": MODEL_NAME}
if texts is not None:
body["texts"] = texts
if images is not None:
body["images"] = images
if embedding_types is not None:
body["embedding_types"] = embedding_types
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()
def test_image_embed(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(255, 0, 0)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float"],
)
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0
assert r["meta"]["billed_units"]["image_tokens"] > 0
assert r["meta"]["billed_units"]["input_tokens"] == 0
def test_image_batch(server: RemoteOpenAIServer):
red = _make_tiny_png(255, 0, 0)
blue = _make_tiny_png(0, 0, 255)
r = _cohere_embed(
server,
images=[red, blue],
embedding_types=["float"],
)
assert len(r["embeddings"]["float"]) == 2
def test_image_l2_normalized(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(0, 255, 0)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float"],
)
emb = np.array(r["embeddings"]["float"][0])
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
def test_image_embedding_types(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(128, 128, 128)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float", "binary", "ubinary"],
)
dim = len(r["embeddings"]["float"][0])
assert len(r["embeddings"]["binary"][0]) == dim // 8
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
"""SigLIP also supports text-only embedding via /v2/embed."""
r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0
@@ -0,0 +1,102 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
Verifies that both endpoints produce identical float embeddings when
no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
"""
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "BAAI/bge-base-en-v1.5"
DTYPE = "bfloat16"
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512",
"--gpu-memory-utilization",
"0.02",
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def _cohere_embed(
server: RemoteOpenAIServer,
texts: list[str],
) -> list[list[float]]:
body = {
"model": MODEL_NAME,
"texts": texts,
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()["embeddings"]["float"]
def _openai_embed(
server: RemoteOpenAIServer,
texts: list[str],
) -> list[list[float]]:
body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
resp.raise_for_status()
return [item["embedding"] for item in resp.json()["data"]]
def test_single_text_parity(server: RemoteOpenAIServer):
"""A single text should produce identical embeddings via both APIs."""
texts = ["the quick brown fox jumps over the lazy dog"]
v2 = _cohere_embed(server, texts)
v1 = _openai_embed(server, texts)
np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
def test_batch_parity(server: RemoteOpenAIServer):
"""A batch of texts should produce identical embeddings via both APIs,
in the same order."""
texts = [
"machine learning",
"deep learning",
"natural language processing",
]
v2 = _cohere_embed(server, texts)
v1 = _openai_embed(server, texts)
assert len(v2) == len(v1) == 3
for i in range(3):
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
def test_token_count_parity(server: RemoteOpenAIServer):
"""Both APIs should report the same prompt token count."""
texts = ["hello world"]
v2_resp = requests.post(
server.url_for("/v2/embed"),
json={
"model": MODEL_NAME,
"texts": texts,
"embedding_types": ["float"],
},
)
v1_resp = requests.post(
server.url_for("/v1/embeddings"),
json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
)
v2_resp.raise_for_status()
v1_resp.raise_for_status()
v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
assert v2_tokens == v1_tokens
@@ -0,0 +1,208 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for EmbedIOProcessor."""
import pytest
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedRequest,
)
class TestResolveTruncation:
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
@staticmethod
def _make_request(**kwargs) -> CohereEmbedRequest:
defaults = {
"model": "test",
"input_type": "search_document",
"texts": ["hello"],
}
return CohereEmbedRequest(**(defaults | kwargs))
def test_truncate_end_default(self):
req = self._make_request()
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side is None
def test_truncate_end_explicit(self):
req = self._make_request(truncate="END")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side is None
def test_truncate_end_with_max_tokens(self):
req = self._make_request(truncate="END", max_tokens=128)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == 128
assert side is None
def test_truncate_none(self):
req = self._make_request(truncate="NONE")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens is None
assert side is None
def test_truncate_none_with_max_tokens(self):
"""truncate=NONE should NOT set truncate_prompt_tokens; the
max_tokens limit is enforced separately via _check_max_tokens."""
req = self._make_request(truncate="NONE", max_tokens=10)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens is None
assert side is None
def test_truncate_start(self):
req = self._make_request(truncate="START")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side == "left"
def test_truncate_start_with_max_tokens(self):
req = self._make_request(truncate="START", max_tokens=64)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == 64
assert side == "left"
class TestApplyStPrompt:
"""Unit tests for EmbedIOProcessor._apply_task_instruction."""
@staticmethod
def _make_handler(task_instructions: dict[str, str] | None):
handler = object.__new__(EmbedIOProcessor)
handler.task_instructions = task_instructions
return handler
def test_no_prompts_configured(self):
handler = self._make_handler(None)
texts = ["hello", "world"]
assert handler._apply_task_instruction(texts, "query") is texts
def test_matching_input_type(self):
handler = self._make_handler({"query": "search_query: "})
result = handler._apply_task_instruction(["hello"], "query")
assert result == ["search_query: hello"]
def test_non_matching_input_type(self):
handler = self._make_handler({"query": "search_query: "})
texts = ["hello"]
assert handler._apply_task_instruction(texts, "document") is texts
def test_multiple_texts(self):
handler = self._make_handler(
{"query": "Represent this sentence for searching: "}
)
result = handler._apply_task_instruction(["a", "b", "c"], "query")
assert result == [
"Represent this sentence for searching: a",
"Represent this sentence for searching: b",
"Represent this sentence for searching: c",
]
def test_empty_prefix_returns_unchanged(self):
handler = self._make_handler({"passage": ""})
texts = ["hello"]
assert handler._apply_task_instruction(texts, "passage") is texts
class TestLoadTaskInstructions:
"""Unit tests for EmbedIOProcessor._load_task_instructions."""
def test_no_attribute(self):
class FakeConfig:
pass
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
def test_with_task_instructions(self):
class FakeConfig:
task_instructions = {
"retrieval.query": "Represent the query: ",
"retrieval.passage": "",
}
result = EmbedIOProcessor._load_task_instructions(FakeConfig())
assert result == {
"retrieval.query": "Represent the query: ",
"retrieval.passage": "",
}
def test_empty_dict(self):
class FakeConfig:
task_instructions = {}
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
def test_non_dict(self):
class FakeConfig:
task_instructions = "not a dict"
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
class TestCheckMaxTokens:
"""Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
@staticmethod
def _fake_output(n_tokens: int):
class _Out:
def __init__(self, n: int):
self.prompt_token_ids = list(range(n))
return _Out(n_tokens)
def test_none_check_is_noop(self):
outs = [self._fake_output(100)]
EmbedIOProcessor._check_cohere_max_tokens(outs, None)
def test_within_limit(self):
outs = [self._fake_output(5), self._fake_output(3)]
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
def test_exceeds_limit(self):
outs = [self._fake_output(3), self._fake_output(10)]
with pytest.raises(ValueError, match="exceeds max_tokens=5"):
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
def test_exact_limit(self):
outs = [self._fake_output(5)]
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
class TestValidateInputType:
"""Unit tests for EmbedIOProcessor._validate_input_type."""
@staticmethod
def _make_handler(task_instructions: dict[str, str] | None):
handler = object.__new__(EmbedIOProcessor)
handler.task_instructions = task_instructions
return handler
def test_none_input_type_always_accepted(self):
handler = self._make_handler(None)
handler._validate_input_type(None)
handler_with = self._make_handler({"query": "q: "})
handler_with._validate_input_type(None)
def test_no_prompts_rejects(self):
handler = self._make_handler(None)
with pytest.raises(ValueError, match="does not define any input_type"):
handler._validate_input_type("anything")
def test_known_type_accepted(self):
handler = self._make_handler({"query": "q: ", "document": "d: "})
handler._validate_input_type("query")
handler._validate_input_type("document")
def test_unknown_type_rejected(self):
handler = self._make_handler({"query": "q: ", "document": "d: "})
with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
handler._validate_input_type("other")
def test_error_lists_supported(self):
handler = self._make_handler({"a": "", "b": ""})
with pytest.raises(ValueError, match="Supported values: a, b"):
handler._validate_input_type("z")
@@ -0,0 +1,129 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
underlying packing helpers, plus Cohere-specific serving helpers."""
import base64
import struct
import numpy as np
import pytest
from vllm.entrypoints.pooling.embed.protocol import (
build_typed_embeddings,
)
@pytest.fixture
def sample_embeddings() -> list[list[float]]:
return [
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
[-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
]
class TestBuildTypedEmbeddingsFloat:
def test_float_passthrough(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float"])
assert result.float == sample_embeddings
assert result.binary is None
def test_empty_input(self):
result = build_typed_embeddings([], ["float"])
assert result.float == []
class TestBuildTypedEmbeddingsBinary:
def test_binary_packing(self):
# 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
# signed: 170 - 128 = 42
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
assert result.binary[0] == [42]
def test_ubinary_packing(self):
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
result = build_typed_embeddings(embs, ["ubinary"])
assert result.ubinary is not None
assert result.ubinary[0] == [170] # 0b10101010
def test_binary_all_positive(self):
embs = [[0.1] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
assert result.binary[0] == [127]
def test_binary_all_negative(self):
embs = [[-0.1] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# all bits = 0, signed: 0 - 128 = -128
assert result.binary[0] == [-128]
def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["binary"])
assert result.binary is not None
for orig, packed in zip(sample_embeddings, result.binary):
assert len(packed) == len(orig) // 8
def test_zero_treated_as_positive(self):
embs = [[0.0] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
assert result.binary[0] == [127]
def test_non_multiple_of_8_raises(self):
embs = [[0.1] * 7]
with pytest.raises(ValueError, match="multiple of 8"):
build_typed_embeddings(embs, ["binary"])
def test_ubinary_non_multiple_of_8_raises(self):
embs = [[0.1] * 10]
with pytest.raises(ValueError, match="multiple of 8"):
build_typed_embeddings(embs, ["ubinary"])
class TestBuildTypedEmbeddingsBase64:
def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["base64"])
assert result.base64 is not None
assert len(result.base64) == 2
for orig, b64_str in zip(sample_embeddings, result.base64):
decoded = base64.b64decode(b64_str)
n = len(orig)
values = struct.unpack(f"<{n}f", decoded)
np.testing.assert_allclose(orig, values, rtol=1e-5)
def test_base64_byte_length(self):
embs = [[0.1, 0.2, 0.3]]
result = build_typed_embeddings(embs, ["base64"])
assert result.base64 is not None
raw = base64.b64decode(result.base64[0])
assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
class TestBuildTypedEmbeddingsMultiple:
def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(
sample_embeddings,
["float", "binary", "ubinary", "base64"],
)
assert result.float is not None
assert result.binary is not None
assert result.ubinary is not None
assert result.base64 is not None
def test_subset_types(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
assert result.float is not None
assert result.binary is not None
assert result.ubinary is None
assert result.base64 is None
def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
assert result.float is not None
@@ -0,0 +1,434 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
Tests both contiguous and non-contiguous (cross-layer unified) KV cache
layouts against a pure-PyTorch reference implementation.
"""
import pytest
import torch
from vllm.platforms import current_platform
FP8_DTYPE = current_platform.fp8_dtype()
NUM_BLOCKS = 128
def to_float8(x, dtype=None):
if dtype is None:
dtype = FP8_DTYPE
finfo = torch.finfo(dtype)
min_val, max_val = x.aminmax()
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
scale = finfo.max / amax * 0.1
x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
return x_scl_sat.to(dtype), scale.float().reciprocal()
def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
"""Create a standard contiguous fp8 KV cache (HND layout)."""
raw = torch.randn(
num_blocks,
2,
num_kv_heads,
block_size,
head_size,
dtype=torch.bfloat16,
device="cuda",
)
kv_cache, scale = to_float8(raw)
return kv_cache, scale
def make_cross_layer_kv_cache(
num_blocks,
num_kv_heads,
block_size,
head_size,
num_layers=4,
):
"""
Create a non-contiguous per-layer view mimicking cross-layer allocation.
Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size)
with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
"""
raw = torch.randn(
num_blocks,
2,
num_kv_heads,
num_layers,
block_size,
head_size,
dtype=torch.bfloat16,
device="cuda",
)
fp8_full, scale = to_float8(raw)
layer_view = fp8_full[:, :, :, 0, :, :]
assert not layer_view.is_contiguous(), (
f"Expected non-contiguous view, got strides {layer_view.stride()}"
)
return layer_view, scale
def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
"""Pure PyTorch reference: gather pages and dequantize fp8 -> dequant_dtype."""
batch_size, num_pages_per_seq = block_tables.shape
s = kv_cache.shape
out = torch.zeros(
batch_size * num_pages_per_seq + 1,
s[1],
s[2],
s[3],
s[4],
dtype=dequant_dtype,
device=kv_cache.device,
)
for b in range(batch_size):
for p in range(num_pages_per_seq):
page_idx = block_tables[b, p].item()
if page_idx <= 0:
continue
mock_idx = b * num_pages_per_seq + p + 1
out[mock_idx, 0] = (kv_cache[page_idx, 0].float() * k_scale.item()).to(
dequant_dtype
)
out[mock_idx, 1] = (kv_cache[page_idx, 1].float() * v_scale.item()).to(
dequant_dtype
)
return out
@pytest.mark.parametrize("num_kv_heads", [1, 8])
@pytest.mark.parametrize("head_size", [64, 128])
@pytest.mark.parametrize("block_size", [16, 32])
@pytest.mark.parametrize("batch_size", [1, 4])
@pytest.mark.parametrize("num_pages_per_seq", [3, 8])
@pytest.mark.parametrize("contiguous", [True, False])
@torch.inference_mode()
def test_trtllm_kvfp8_dequant(
num_kv_heads: int,
head_size: int,
block_size: int,
batch_size: int,
num_pages_per_seq: int,
contiguous: bool,
):
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
if contiguous:
kv_cache, scale = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
else:
kv_cache, scale = make_cross_layer_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = scale.clone()
v_scale = scale.clone()
block_tables = torch.randint(
1,
NUM_BLOCKS,
(batch_size, num_pages_per_seq),
dtype=torch.int32,
)
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
expected_bt = torch.arange(
1,
batch_size * num_pages_per_seq + 1,
dtype=torch.int32,
device="cuda",
).reshape(batch_size, num_pages_per_seq)
torch.testing.assert_close(mock_block_table, expected_bt)
# Page 0 is padding (never written), compare only pages 1+
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@torch.inference_mode()
def test_block_tables_with_zero_pages():
"""Pages with index <= 0 must be skipped (early return in kernel)."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 8, 16, 64
kv_cache, scale = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = v_scale = scale.clone()
# Mix of valid pages and zeros (padding)
block_tables = torch.tensor(
[[5, 0, 10], [0, 0, 0], [3, 7, 0]],
dtype=torch.int32,
device="cuda",
)
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
# Only compare pages that were actually written (non-zero page indices)
for b in range(block_tables.shape[0]):
for p in range(block_tables.shape[1]):
if block_tables[b, p].item() > 0:
idx = b * block_tables.shape[1] + p + 1
torch.testing.assert_close(
mock_kv_cache[idx],
ref[idx],
atol=1e-3,
rtol=1e-3,
)
@torch.inference_mode()
def test_all_zero_block_tables():
"""All-zero block_tables: kernel should write nothing."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 4, 16, 64
kv_cache, scale = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = v_scale = scale.clone()
block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda")
# Should not crash even though no pages are valid
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
assert mock_kv_cache.shape[0] == 2 * 4 + 1
assert mock_block_table.shape == (2, 4)
@torch.inference_mode()
def test_different_k_v_scales():
"""Verify K and V are dequantized with independent scales."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 8, 16, 64
kv_cache, _ = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
v_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
block_tables = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@torch.inference_mode()
def test_single_page_per_seq():
"""Minimum grid dim 1 = 1 page per sequence."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 8, 16, 128
kv_cache, scale = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = v_scale = scale.clone()
block_tables = torch.tensor([[5], [10], [20]], dtype=torch.int32, device="cuda")
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@torch.inference_mode()
def test_large_page_indices():
"""Page indices near the top of the buffer stress offset arithmetic."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 8, 16, 128
large_num_blocks = 32768
kv_cache, scale = make_contiguous_kv_cache(
large_num_blocks,
num_kv_heads,
block_size,
head_size,
)
k_scale = v_scale = scale.clone()
# Use page indices near the top of the buffer
block_tables = torch.tensor(
[[large_num_blocks - 1, large_num_blocks - 2, 1]],
dtype=torch.int32,
device="cuda",
)
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@torch.inference_mode()
def test_large_block_size():
"""block_size=64 -> HEAD_STRIDE=8192, large tl.arange per thread block."""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 4, 64, 128
kv_cache, scale = make_contiguous_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
)
k_scale = v_scale = scale.clone()
block_tables = torch.randint(
1,
NUM_BLOCKS,
(2, 4),
dtype=torch.int32,
device="cuda",
)
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@torch.inference_mode()
def test_cross_layer_many_layers():
"""
Non-contiguous with 36 layers -- matches real gpt-oss-120b.
Strides are far from contiguous (factor of 36 in the gaps).
"""
from vllm.v1.attention.backends.flashinfer import (
trtllm_prefill_attn_kvfp8_dequant,
)
torch.set_default_device("cuda")
num_kv_heads, block_size, head_size = 8, 16, 64
num_layers = 36
kv_cache, scale = make_cross_layer_kv_cache(
NUM_BLOCKS,
num_kv_heads,
block_size,
head_size,
num_layers=num_layers,
)
k_scale = v_scale = scale.clone()
block_tables = torch.randint(
1,
NUM_BLOCKS,
(4, 6),
dtype=torch.int32,
device="cuda",
)
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
kv_cache,
block_tables,
k_scale,
v_scale,
torch.bfloat16,
)
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
@@ -280,21 +280,22 @@ def test_rms_norm(
assert torch.allclose(ref_residual, ops_residual)
output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
scales = torch.empty(
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
)
if group_size is None:
scales = torch.empty(
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
)
opcheck(
torch.ops._C.rms_norm_dynamic_per_token_quant,
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
)
else:
# TODO(luka/eliza) opcheck is broken?
# Somehow the cloned args are getting mutated in-place,
# which causes the opcheck to fail.
# https://github.com/vllm-project/vllm/issues/36688
return
assert hidden_size % group_size[1] == 0
num_groups = hidden_size // group_size[1]
scales = torch.empty(
(num_groups, num_tokens),
device=x.device,
dtype=torch.float32,
).transpose(0, 1)
opcheck(
torch.ops._C.rms_norm_per_block_quant,
(
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
)
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import is_deep_gemm_supported
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
from vllm.utils.flashinfer import (
has_flashinfer_cutlass_fused_moe,
has_flashinfer_nvlink_one_sided,
)
from vllm.utils.import_utils import (
has_aiter,
has_deep_ep,
@@ -234,15 +237,15 @@ if has_mori():
)
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
FlashInferA2APrepareAndFinalize,
)
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
FlashInferExperts,
)
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
FlashInferNVLinkTwoSidedPrepareAndFinalize,
)
register_prepare_and_finalize(
FlashInferA2APrepareAndFinalize,
FlashInferNVLinkTwoSidedPrepareAndFinalize,
standard_format,
nvfp4_types + fp8_types,
blocked_quantization_support=True,
@@ -263,6 +266,36 @@ else:
FlashInferCutlassMoEPrepareAndFinalize = None
FlashInferExperts = None
if (
has_flashinfer_nvlink_one_sided()
and has_flashinfer_cutlass_fused_moe()
and current_platform.has_device_capability(100)
):
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
FlashInferNVLinkOneSidedPrepareAndFinalize,
)
register_prepare_and_finalize(
FlashInferNVLinkOneSidedPrepareAndFinalize,
standard_format,
nvfp4_types,
blocked_quantization_support=False,
backend="flashinfer_nvlink_one_sided",
supports_apply_weight_on_input=False,
)
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
TrtLlmNvFp4ExpertsModular,
)
register_experts(
TrtLlmNvFp4ExpertsModular,
standard_format,
nvfp4_types,
blocked_quantization_support=False,
supports_expert_map=True,
)
if has_aiter():
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
@@ -159,6 +159,52 @@ def test_quantize_to_fp4(
torch.testing.assert_close(scale_ans, scale_ref)
@pytest.mark.parametrize(
"shape",
[(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
)
@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
@torch.inference_mode()
def test_python_util_matches_cpp_allocation(
shape: tuple[int, int],
is_sf_swizzled_layout: bool,
) -> None:
"""
Verify that the Python utility (create_fp4_output_tensors) allocates
tensors with the same shapes and dtypes as the C++ functional variant
(scaled_fp4_quant_func).
"""
from vllm._custom_ops import create_fp4_output_tensors
torch.set_default_device("cuda:0")
m, n = shape
input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
# C++ functional variant allocates internally
cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
input_tensor, input_scale, is_sf_swizzled_layout
)
# Python utility
py_out, py_scale = create_fp4_output_tensors(
m, n, torch.device("cuda:0"), is_sf_swizzled_layout
)
assert py_out.shape == cpp_out.shape, (
f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
)
assert py_out.dtype == cpp_out.dtype, (
f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
)
assert py_scale.shape == cpp_scale.shape, (
f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
)
assert py_scale.dtype == cpp_scale.dtype, (
f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
)
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
@torch.inference_mode()
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
@@ -19,6 +19,7 @@ from vllm.transformers_utils.runai_utils import (
def test_is_runai_obj_uri():
assert is_runai_obj_uri("gs://some-gcs-bucket/path")
assert is_runai_obj_uri("s3://some-s3-bucket/path")
assert is_runai_obj_uri("az://some-azure-container/path")
assert not is_runai_obj_uri("nfs://some-nfs-path")
@@ -0,0 +1,361 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for EP weight filtering during model loading."""
import glob
import tempfile
import huggingface_hub.constants
import pytest
import torch
from vllm.model_executor.model_loader.ep_weight_filter import (
compute_local_expert_ids,
parse_expert_id,
should_skip_weight,
)
from vllm.model_executor.model_loader.weight_utils import (
safetensors_weights_iterator,
)
# ---------------------------------------------------------------------------
# Unit tests for parse_expert_id
# ---------------------------------------------------------------------------
class TestParseExpertId:
def test_routed_expert(self):
name = "model.layers.0.mlp.experts.42.gate_proj.weight"
assert parse_expert_id(name) == 42
def test_large_expert_id(self):
name = "model.layers.60.mlp.experts.383.down_proj.weight"
assert parse_expert_id(name) == 383
def test_shared_expert(self):
# Shared experts use a different naming convention in most models
name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
assert parse_expert_id(name) is None
def test_attention_weight(self):
name = "model.layers.0.self_attn.q_proj.weight"
assert parse_expert_id(name) is None
def test_embedding(self):
name = "model.embed_tokens.weight"
assert parse_expert_id(name) is None
def test_layernorm(self):
name = "model.layers.0.input_layernorm.weight"
assert parse_expert_id(name) is None
def test_fused_3d_expert(self):
# 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
# They must NOT be filtered — slicing happens later in weight_loader.
name = "model.layers.0.mlp.experts.gate_proj.weight"
assert parse_expert_id(name) is None
def test_fused_3d_expert_down_proj(self):
name = "model.layers.10.mlp.experts.down_proj.weight"
assert parse_expert_id(name) is None
def test_expert_scale(self):
# NVFP4 quantized models have scale tensors for experts
name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
assert parse_expert_id(name) == 100
def test_expert_zero_id(self):
name = "model.layers.0.mlp.experts.0.up_proj.weight"
assert parse_expert_id(name) == 0
# ---------------------------------------------------------------------------
# Unit tests for compute_local_expert_ids
# ---------------------------------------------------------------------------
class TestComputeLocalExpertIds:
def test_ep_disabled(self):
assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
def test_even_split(self):
# 64 experts, EP=8 → 8 per rank
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
assert ids == set(range(0, 8))
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
assert ids == set(range(56, 64))
def test_uneven_split(self):
# 10 experts, EP=3 → ranks get 4, 3, 3
ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
assert len(ids_0) == 4
assert len(ids_1) == 3
assert len(ids_2) == 3
# All experts covered, no overlap
assert ids_0 | ids_1 | ids_2 == set(range(10))
assert ids_0.isdisjoint(ids_1)
assert ids_1.isdisjoint(ids_2)
def test_384_experts_ep8(self):
# Kimi-K2.5 config: 384 experts, EP=8
for rank in range(8):
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
assert len(ids) == 48
# All experts covered
all_ids = set()
for rank in range(8):
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
all_ids |= ids
assert all_ids == set(range(384))
def test_384_experts_ep16(self):
for rank in range(16):
ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
assert len(ids) == 24
def test_384_experts_ep24(self):
# 384 / 24 = 16 exactly
for rank in range(24):
ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
assert len(ids) == 16
# round_robin placement tests
def test_round_robin_basic(self):
# 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
rr = "round_robin"
ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
assert ids_0 == {0, 2, 4, 6}
assert ids_1 == {1, 3, 5, 7}
def test_round_robin_full_coverage(self):
# 384 experts, EP=8: all experts covered, no overlap
rr = "round_robin"
all_ids: set[int] = set()
for rank in range(8):
ids = compute_local_expert_ids(384, 8, rank, placement=rr)
assert ids is not None and len(ids) == 48
assert all_ids.isdisjoint(ids)
all_ids |= ids
assert all_ids == set(range(384))
def test_round_robin_uneven(self):
# 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
rr = "round_robin"
ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
assert ids_0 == {0, 3, 6, 9}
assert ids_1 == {1, 4, 7}
assert ids_2 == {2, 5, 8}
assert ids_0 | ids_1 | ids_2 == set(range(10))
# ---------------------------------------------------------------------------
# Unit tests for should_skip_weight
# ---------------------------------------------------------------------------
class TestShouldSkipWeight:
def setup_method(self):
# Simulate EP=8, rank=0 → experts 0-47
self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
def test_no_filter(self):
assert not should_skip_weight("anything", None)
def test_dense_not_skipped(self):
assert not should_skip_weight(
"model.layers.0.self_attn.q_proj.weight", self.local_ids
)
def test_local_expert_not_skipped(self):
assert not should_skip_weight(
"model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
)
def test_remote_expert_skipped(self):
assert should_skip_weight(
"model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
)
def test_boundary_expert(self):
# Expert 47 is local (last one), 48 is not
assert not should_skip_weight(
"model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
)
assert should_skip_weight(
"model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
)
def test_shared_expert_not_skipped(self):
assert not should_skip_weight(
"model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
)
def test_embedding_not_skipped(self):
assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
def test_fused_3d_expert_not_skipped(self):
# 3D fused-expert tensors (gpt-oss style) have no numeric id.
# Must not be skipped — weight_loader handles slicing later.
assert not should_skip_weight(
"model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
)
# ---------------------------------------------------------------------------
# Integration test: safetensors_weights_iterator with EP filtering
# ---------------------------------------------------------------------------
class TestSafetensorsWeightsIteratorWithEpFilter:
"""Verify that EP filtering produces a strict subset of unfiltered loading
and that all expected dense + local expert weights are present."""
@pytest.fixture(scope="class")
def gpt2_files(self):
"""Download GPT-2 safetensors to a temp dir (shared across class)."""
with tempfile.TemporaryDirectory() as tmpdir:
huggingface_hub.constants.HF_HUB_OFFLINE = False
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf,
)
download_weights_from_hf(
"openai-community/gpt2",
allow_patterns=["*.safetensors"],
cache_dir=tmpdir,
)
files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
assert len(files) > 0
yield files
def test_no_filter_returns_all(self, gpt2_files):
"""With local_expert_ids=None, all weights are returned (no MoE)."""
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
filtered_weights = dict(
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
)
assert set(all_weights.keys()) == set(filtered_weights.keys())
def test_empty_filter_skips_experts_only(self, gpt2_files):
"""GPT-2 has no expert weights, so even an empty local_expert_ids
set should return all weights (all are dense)."""
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
filtered_weights = dict(
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
)
# GPT-2 has no experts, so nothing should be filtered
assert set(all_weights.keys()) == set(filtered_weights.keys())
class TestEpFilterOnSyntheticMoeWeights:
"""Create synthetic safetensors files with expert-like naming and verify
that the filter correctly skips non-local experts."""
@pytest.fixture
def synthetic_moe_files(self, tmp_path):
"""Create synthetic safetensors with expert-patterned tensor names."""
from safetensors.torch import save_file
tensors = {}
# Dense weights
tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
# Expert weights: 8 experts
for expert_id in range(8):
tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
torch.randn(128, 64)
)
tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
torch.randn(128, 64)
)
tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
torch.randn(64, 128)
)
# Shared expert (should never be filtered)
tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
128, 64
)
filepath = str(tmp_path / "model-00001-of-00001.safetensors")
save_file(tensors, filepath)
return [filepath], tensors
def test_no_filter_returns_all(self, synthetic_moe_files):
files, expected = synthetic_moe_files
loaded = dict(safetensors_weights_iterator(files, False))
assert set(loaded.keys()) == set(expected.keys())
def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
files, expected = synthetic_moe_files
# EP=2, rank=0 → experts 0-3
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
loaded = dict(
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
)
# Should have all dense + shared + experts 0-3 only
for name in loaded:
eid = parse_expert_id(name)
if eid is not None:
assert eid in local_ids, f"Non-local expert {eid} was loaded"
# Check expert count: 4 experts × 3 weights = 12
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
assert len(expert_names) == 4 * 3
# Check all dense weights present
assert "model.embed_tokens.weight" in loaded
assert "model.layers.0.self_attn.q_proj.weight" in loaded
assert "model.layers.0.input_layernorm.weight" in loaded
assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
files, expected = synthetic_moe_files
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
loaded = dict(
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
)
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
assert len(expert_names) == 4 * 3
for name in expert_names:
assert parse_expert_id(name) in local_ids
def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
files, _ = synthetic_moe_files
all_expert_names = set()
for rank in range(8):
local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
loaded = dict(
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
)
expert_names = {n for n in loaded if parse_expert_id(n) is not None}
# 1 expert × 3 weights
assert len(expert_names) == 3
all_expert_names |= expert_names
# All 8 experts × 3 weights covered across ranks
assert len(all_expert_names) == 24
def test_tensor_values_match(self, synthetic_moe_files):
"""Filtered tensors have identical values to unfiltered ones."""
files, _ = synthetic_moe_files
all_weights = dict(safetensors_weights_iterator(files, False))
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
filtered = dict(
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
)
for name, tensor in filtered.items():
assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"
@@ -0,0 +1,68 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for CPU unquantized GEMM dispatch behavior."""
import pytest
import torch
from vllm.model_executor.layers import utils
from vllm.platforms import current_platform
@pytest.fixture(scope="module")
def _mock_zentorch_linear_unary():
"""Register a mock zentorch_linear_unary op when zentorch is not installed.
Allows the dispatch tests to run in CI without a real zentorch build.
Skips registration when zentorch is already available.
"""
if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
yield
return
lib_def = torch.library.Library("zentorch", "DEF")
lib_def.define(
"zentorch_linear_unary("
"Tensor input, "
"Tensor weight, "
"Tensor? bias, "
"bool is_weight_prepacked=False"
") -> Tensor"
)
lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
lib_impl.impl(
"zentorch_linear_unary",
lambda input, weight, bias, is_weight_prepacked=False: (
torch.nn.functional.linear(input, weight, bias)
),
)
yield
lib_impl._destroy()
lib_def._destroy()
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
layer = torch.nn.Linear(16, 8, bias=True)
x = torch.randn(4, 16)
expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
output = layer.cpu_linear(x, layer.weight, layer.bias)
torch.testing.assert_close(output, expected)
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
layer = torch.nn.Linear(16, 8, bias=True)
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
assert layer.weight.numel() == 0
@@ -103,6 +103,10 @@ AITER_MODEL_LIST = [
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
),
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
pytest.param(
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", # hyperclovax
marks=[large_gpu_mark(min_gb=32)],
),
],
)
@pytest.mark.parametrize("max_tokens", [32])
@@ -1,8 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
from dataclasses import asdict
import pytest
import pytest_asyncio
from mistral_common.audio import Audio
from mistral_common.protocol.instruct.chunk import RawAudio
from mistral_common.protocol.transcription.request import (
@@ -17,18 +19,21 @@ from vllm.assets.audio import AudioAsset
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM
from ....utils import ROCM_ENGINE_KWARGS
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
ENGINE_CONFIG = dict(
model=MODEL_NAME,
max_model_len=8192,
max_num_seqs=4,
limit_mm_per_prompt={"audio": 1},
config_format="mistral",
load_format="mistral",
tokenizer_mode="mistral",
enforce_eager=True,
gpu_memory_utilization=0.9,
)
ENGINE_CONFIG = {
"model": MODEL_NAME,
"max_model_len": 8192,
"max_num_seqs": 4,
"limit_mm_per_prompt": {"audio": 1},
"config_format": "mistral",
"load_format": "mistral",
"tokenizer_mode": "mistral",
"enforce_eager": True,
"gpu_memory_utilization": 0.9,
**ROCM_ENGINE_KWARGS,
}
EXPECTED_TEXT = [
@@ -49,6 +54,14 @@ EXPECTED_TEXT = [
]
def _normalize(texts: list[str]) -> list[str]:
# The model occasionally transcribes "OBS" as "a base hit" and
# "oh, my" as "oh my", but both are acoustically valid. Normalise so
# the assertion is stable across runs and hardware.
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
return texts
@pytest.fixture
def audio_assets() -> list[AudioAsset]:
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
@@ -60,15 +73,27 @@ def tokenizer() -> MistralTokenizer:
@pytest.fixture
def engine() -> LLM:
def engine():
engine_args = EngineArgs(**ENGINE_CONFIG)
return LLM(**asdict(engine_args))
llm = LLM(**asdict(engine_args))
try:
yield llm
finally:
with contextlib.suppress(Exception):
llm.llm_engine.engine_core.shutdown()
import torch
torch.accelerator.empty_cache()
@pytest.fixture
def async_engine() -> AsyncLLM:
@pytest_asyncio.fixture
async def async_engine():
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
return AsyncLLM.from_engine_args(engine_args)
llm = AsyncLLM.from_engine_args(engine_args)
try:
yield llm
finally:
llm.shutdown()
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
@@ -108,8 +133,13 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
sampling_params=sampling_params,
)
texts = [out.outputs[0].text for out in outputs]
assert texts == EXPECTED_TEXT
texts = _normalize([out.outputs[0].text for out in outputs])
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
assert got == expected, (
f"Output mismatch at index {i}:\n"
f" got: {got!r}\n"
f" expected: {expected!r}"
)
@pytest.mark.asyncio
@@ -149,9 +179,17 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)
output_tokens_list.append(output_tokens)
texts = [
tokenizer.decode(output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE)
for output_tokens in output_tokens_list
]
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
assert texts == EXPECTED_TEXT
texts = _normalize(
[
tokenizer.decode(
output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE
)
for output_tokens in output_tokens_list
]
)
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
assert got == expected, (
f"Output mismatch at index {i}:\n"
f" got: {got!r}\n"
f" expected: {expected!r}"
)
+6 -2
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from transformers import CLIPModel
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
@@ -50,13 +51,16 @@ def _run_test(
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
).squeeze(0)
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
).squeeze(0)
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
all_outputs.append(pooled_output.tolist())
hf_outputs = all_outputs
@@ -4,6 +4,7 @@
from typing import Any
import pytest
import torch
from transformers import SiglipModel
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
@@ -68,12 +69,15 @@ def _run_test(
if "pixel_values" in inputs:
pooled_output = hf_model.model.get_image_features(
pixel_values=inputs.pixel_values,
).squeeze(0)
)
else:
pooled_output = hf_model.model.get_text_features(
input_ids=inputs.input_ids,
).squeeze(0)
)
if not isinstance(pooled_output, torch.Tensor):
pooled_output = pooled_output.pooler_output
pooled_output = pooled_output.squeeze(0)
all_outputs.append(pooled_output.tolist())
hf_outputs = all_outputs
@@ -34,8 +34,22 @@ MODELS = [
]
def create_mm_data(num_videos: int) -> dict[str, list]:
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
# stays fast even without a GPU.
mm_data = dict[str, list](video=[], audio=[])
for i in range(num_videos):
rng = np.random.RandomState(i)
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
mm_data["video"].append(video)
mm_data["audio"].append((audio, sr))
return mm_data
@pytest.mark.parametrize("model_id", MODELS)
def test_audio_in_video_cache_correctness(model_id: str) -> None:
@pytest.mark.parametrize("num_videos", [1, 2])
def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None:
"""
Regression test for https://github.com/vllm-project/vllm/pull/36800
@@ -47,7 +61,7 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1},
limit_mm_per_prompt={"audio": num_videos, "image": 0, "video": num_videos},
mm_processor_cache_gb=1,
)
@@ -65,17 +79,12 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
video_token_id = baseline_processor.info.get_hf_config().video_token_id
rng = np.random.RandomState(0)
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
# stays fast even without a GPU.
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
mm_data = {"video": [video], "audio": [(audio, sr)]}
mm_data = create_mm_data(num_videos)
hf_processor_mm_kwargs = {"use_audio_in_video": True}
def run(processor):
return processor(
[video_token_id],
[video_token_id] * num_videos,
mm_items=baseline_processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
)["prompt_token_ids"]
+104
View File
@@ -0,0 +1,104 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E tests for online MXFP8 quantization.
Loads a BF16 model with ``--quantization mxfp8`` (online quantization) and
compares log-probabilities against the same model served in BF16 without
quantization. This exercises the full pipeline: config parsing,
``Mxfp8OnlineLinearMethod``, ``Mxfp8OnlineMoEMethod``, weight loading,
online quantization / shuffling, and inference through ``apply_monolithic``.
Layer skipping (``modules_to_not_convert``) is configured in the model's
``config.json`` under ``quantization_config`` and is not tested here.
``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
diverse prompts from ``tests/prompts/example.txt``.
"""
import pytest
from tests.quantization.utils import is_quant_method_supported
from ..utils import check_logprobs_close
# A small MoE model that fits on a single GPU and has both linear + MoE layers.
MOE_MODEL = "Qwen/Qwen3-30B-A3B"
# A small dense model (no MoE) to validate the linear-only path.
DENSE_MODEL = "Qwen/Qwen3-0.6B"
MAX_MODEL_LEN = 1024
MAX_TOKENS = 4
NUM_LOG_PROBS = 8
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
)
@pytest.mark.quant_model
@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
def test_mxfp8_logprobs(
vllm_runner,
example_prompts,
model: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Compare BF16 baseline logprobs against online MXFP8-quantized model.
Runs the same model twice -- once in BF16 (baseline) and once with
online MXFP8 quantization -- then checks that the top log-probabilities
are close. Only 4 tokens are generated to keep the test fast while
still catching numerical divergence.
"""
with monkeypatch.context() as m:
m.setenv("TOKENIZERS_PARALLELISM", "true")
with vllm_runner(
model,
max_model_len=MAX_MODEL_LEN,
enforce_eager=True,
) as vllm_model:
baseline_outputs = vllm_model.generate_greedy_logprobs(
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
)
with vllm_runner(
model,
max_model_len=MAX_MODEL_LEN,
enforce_eager=True,
quantization="mxfp8",
) as vllm_model:
test_outputs = vllm_model.generate_greedy_logprobs(
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
)
check_logprobs_close(
outputs_0_lst=baseline_outputs,
outputs_1_lst=test_outputs,
name_0="bf16",
name_1="mxfp8",
)
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
)
@pytest.mark.quant_model
@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
def test_mxfp8_generation(vllm_runner, model: str) -> None:
"""Smoke test: verify online MXFP8 model generates coherent text."""
prompt = "1 2 3 4 5"
with vllm_runner(
model,
enforce_eager=True,
quantization="mxfp8",
max_model_len=MAX_MODEL_LEN,
) as vllm_model:
output = vllm_model.generate_greedy([prompt], max_tokens=5)
generated = output[0][1]
assert len(generated) > len(prompt), (
f"MXFP8 model produced no new tokens. Output: {generated!r}"
)
+1 -1
View File
@@ -320,7 +320,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
),
"HyperCLOVAXForCausalLM": _HfExamplesInfo(
"naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
trust_remote_code=True,
),
"InternLMForCausalLM": _HfExamplesInfo(
+11
View File
@@ -4,6 +4,7 @@ import base64
from pathlib import Path
from unittest.mock import patch
import librosa
import numpy as np
import pytest
@@ -71,3 +72,13 @@ def test_audio_media_io_encode_base64(dummy_audio):
decoded = base64.b64decode(out)
assert decoded == b"dummy_wav_data"
mock_write.assert_called_once()
def test_audio_media_io_from_video(video_assets):
audio_io = AudioMediaIO()
video_path = video_assets[0].video_path
with open(video_path, "rb") as f:
audio, sr = audio_io.load_bytes(f.read())
audio_ref, sr_ref = librosa.load(video_path, sr=None)
assert sr == sr_ref
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
+139 -1
View File
@@ -1,11 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from unittest.mock import Mock
import pytest
from transformers import AutoTokenizer
from vllm.entrypoints.mcp.tool_server import ToolServer
from vllm.reasoning import ReasoningParser
from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser
from vllm.reasoning.gptoss_reasoning_parser import (
GptOssReasoningParser,
from_builtin_tool_to_tag,
no_func_reasoning_tag,
)
REASONING_MODEL_NAME = "openai/gpt-oss-120b"
@@ -142,3 +150,133 @@ def test_gptoss_is_reasoning_end(
output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output)
actual_is_reasoning_end = parser.is_reasoning_end(output_ids)
assert is_reasoning_end == actual_is_reasoning_end
class TestGptOssStructuralTags:
"""Test cases for GptOssReasoningParser structural tag functionality."""
@pytest.fixture
def mock_tokenizer(self):
"""Create a mock tokenizer for testing."""
tokenizer = Mock()
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
return tokenizer
@pytest.fixture
def reasoning_parser(self, mock_tokenizer):
"""Create a GptOssReasoningParser instance."""
return GptOssReasoningParser(mock_tokenizer)
def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
"""Test prepare_structured_tag with no tool server."""
result = reasoning_parser.prepare_structured_tag(None, None)
expected = json.dumps(no_func_reasoning_tag)
assert result == expected
# Verify the structure is correct
parsed = json.loads(result)
assert parsed["type"] == "structural_tag"
assert parsed["format"]["type"] == "triggered_tags"
assert len(parsed["format"]["tags"]) == 1
assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
"""Test prepare_structured_tag when original_tag is provided."""
original_tag = '{"custom": "tag"}'
result = reasoning_parser.prepare_structured_tag(original_tag, None)
# Should return the original tag unchanged
assert result == original_tag
def test_from_builtin_tool_to_tag(self):
"""Test from_builtin_tool_to_tag function."""
tags = from_builtin_tool_to_tag("python")
assert len(tags) == 2
assert tags[0]["begin"] == "<|channel|>commentary to=python"
assert tags[0]["content"]["type"] == "any_text"
assert tags[0]["end"] == "<|end|>"
assert tags[1]["begin"] == "<|channel|>analysis to=python"
assert tags[1]["content"]["type"] == "any_text"
assert tags[1]["end"] == "<|end|>"
@pytest.mark.parametrize(
"tools",
[
[],
["browser"],
["python"],
["container"],
["browser", "python"],
["browser", "container"],
["python", "container"],
["browser", "python", "container"],
],
)
def test_json_validity_comprehensive(self, reasoning_parser, tools):
"""Test JSON validity across all possible tool combinations."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
result = reasoning_parser.prepare_structured_tag(None, tool_server)
parsed_result = json.loads(result)
assert parsed_result["type"] == "structural_tag"
assert "format" in parsed_result
assert "tags" in parsed_result["format"]
assert "triggers" in parsed_result["format"]
# Tag count should be: 1 (analysis) + 2 * len(tools)
expected_tag_count = 1 + (2 * len(tools))
assert len(parsed_result["format"]["tags"]) == expected_tag_count
# Verify triggers are correctly configured
expected_triggers = ["<|channel|>analysis"]
if tools:
expected_triggers.append("<|channel|>commentary to=")
assert set(parsed_result["format"]["triggers"]) == set(expected_triggers)
def test_no_cross_request_state_pollution(self, reasoning_parser):
"""Test that sequential calls with different tool servers produce
independent results, guarding against shared mutable state
(e.g. missing deepcopy in tag_with_builtin_funcs)."""
tool_server_1 = Mock(spec=ToolServer)
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
tool_server_2 = Mock(spec=ToolServer)
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1)
result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2)
tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]]
tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]]
assert "<|channel|>commentary to=python" in tags_1
assert "<|channel|>commentary to=browser" not in tags_1
assert "<|channel|>commentary to=browser" in tags_2
assert "<|channel|>commentary to=python" not in tags_2
def test_tag_format_consistency(self, reasoning_parser):
"""Test that all generated tags follow consistent format,
catching malformed tags from from_builtin_tool_to_tag."""
tool_server = Mock(spec=ToolServer)
tool_server.has_tool = Mock(
side_effect=lambda tool: tool in ["python", "browser"]
)
result = reasoning_parser.prepare_structured_tag(None, tool_server)
parsed_result = json.loads(result)
for tag in parsed_result["format"]["tags"]:
assert "begin" in tag
assert "content" in tag
assert "end" in tag
assert tag["content"]["type"] == "any_text"
assert tag["end"] == "<|end|>"
assert tag["begin"].startswith("<|channel|>")
+1 -1
View File
@@ -55,7 +55,7 @@ def test_gc():
# The memory allocated for model and KV cache should be released.
# The memory allocated for PyTorch and others should be less than 50MB.
# Usually, it's around 10MB.
allocated = torch.cuda.memory_allocated()
allocated = torch.accelerator.memory_allocated()
assert allocated < 50 * 1024 * 1024
+37
View File
@@ -0,0 +1,37 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import mock_open, patch
from vllm.platforms import _is_amd_zen_cpu
def test_is_amd_zen_cpu_detects_amd_with_avx512():
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
with (
patch("os.path.exists", return_value=True),
patch("builtins.open", mock_open(read_data=cpuinfo)),
):
assert _is_amd_zen_cpu()
def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
with (
patch("os.path.exists", return_value=True),
patch("builtins.open", mock_open(read_data=cpuinfo)),
):
assert not _is_amd_zen_cpu()
def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
with (
patch("os.path.exists", return_value=True),
patch("builtins.open", mock_open(read_data=cpuinfo)),
):
assert not _is_amd_zen_cpu()
def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
with patch("os.path.exists", return_value=False):
assert not _is_amd_zen_cpu()
+378
View File
@@ -0,0 +1,378 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from dataclasses import dataclass, field
from types import NoneType
from typing import Any
import pytest
from tests.tool_parsers.utils import run_tool_extraction
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers import ToolParserManager
@dataclass
class ToolParserTestConfig:
"""Configuration for a tool parser's common tests.
This dataclass contains all the test data and expected results needed
to run the common test suite for a parser. Each parser test file
creates one instance of this config with parser-specific values.
Attributes:
parser_name: Name used with ToolParserManager (e.g., "mistral")
Test data (model outputs):
no_tool_calls_output: Plain text without any tool syntax
single_tool_call_output: One tool call with simple arguments
parallel_tool_calls_output: Multiple tool calls in one response
various_data_types_output: Tool with various data types
empty_arguments_output: Tool call with no parameters
surrounding_text_output: Tool call mixed with regular text
escaped_strings_output: Tool call with escaped chars
malformed_input_outputs: List of invalid inputs
Expected results:
single_tool_call_expected_name: Expected function name
single_tool_call_expected_args: Expected arguments dict
parallel_tool_calls_count: Number of tools in parallel test
parallel_tool_calls_names: Function names in order
single_tool_call_expected_content: Content field when tool called
parallel_tool_calls_expected_content: Content for parallel test
xfail markers:
xfail_streaming: Mapping test name to xfail reason (streaming only)
xfail_nonstreaming: Mapping test name to xfail reason (non-streaming)
Special flags:
allow_empty_or_json_empty_args: True if "" or "{}" both valid for empty args
supports_typed_arguments: True if the parser supports typed function arguments
"""
# Parser identification
parser_name: str
# Test data - model outputs for each common test
no_tool_calls_output: str
single_tool_call_output: str
parallel_tool_calls_output: str
various_data_types_output: str
empty_arguments_output: str
surrounding_text_output: str
escaped_strings_output: str
malformed_input_outputs: list[str]
# Expected results for specific tests (optional overrides)
single_tool_call_expected_name: str = "get_weather"
single_tool_call_expected_args: dict[str, Any] = field(
default_factory=lambda: {"city": "Tokyo"}
)
parallel_tool_calls_count: int = 2
parallel_tool_calls_names: list[str] = field(
default_factory=lambda: ["get_weather", "get_time"]
)
# xfail configuration - maps test name to xfail reason
xfail_streaming: dict[str, str] = field(default_factory=dict)
xfail_nonstreaming: dict[str, str] = field(default_factory=dict)
# Content expectations (some parsers strip content, others don't)
single_tool_call_expected_content: str | None = None
parallel_tool_calls_expected_content: str | None = None
# Special assertions for edge cases
allow_empty_or_json_empty_args: bool = True # "{}" or "" for empty args
supports_typed_arguments: bool = True
class ToolParserTests:
"""Mixin class providing common test suite for tool parsers.
To use this mixin in a parser test file:
1. Create a test_config fixture that returns a ToolParserTestConfig instance
2. Inherit from this class
3. Add parser-specific tests as additional methods
Example:
class TestMistralToolParser(ToolParserTests):
@pytest.fixture
def test_config(self) -> ToolParserTestConfig:
return ToolParserTestConfig(
parser_name="mistral",
no_tool_calls_output="Plain text...",
# ... other config ...
)
# Parser-specific tests
def test_mistral_specific_feature(self, tool_parser):
# Custom test logic
pass
"""
@pytest.fixture
def test_config(self) -> ToolParserTestConfig:
"""Override this to provide parser-specific configuration."""
raise NotImplementedError(
"Subclass must provide test_config fixture returning ToolParserTestConfig"
)
@pytest.fixture
def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike:
"""Override this to provide parser-specific tokenizer."""
return default_tokenizer
@pytest.fixture
def tool_parser(self, test_config: ToolParserTestConfig, tokenizer: TokenizerLike):
return ToolParserManager.get_tool_parser(test_config.parser_name)(tokenizer)
@pytest.fixture(params=[True, False])
def streaming(self, request: pytest.FixtureRequest) -> bool:
return request.param
def test_no_tool_calls(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser handles plain text without tool syntax."""
# Apply xfail markers if configured
test_name = "test_no_tool_calls"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser, test_config.no_tool_calls_output, streaming=streaming
)
assert content == test_config.no_tool_calls_output, (
f"Expected content to match input, got {content}"
)
assert len(tool_calls) == 0, f"Expected no tool calls, got {len(tool_calls)}"
def test_single_tool_call_simple_args(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser extracts one tool with simple arguments."""
# Apply xfail markers if configured
test_name = "test_single_tool_call_simple_args"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser, test_config.single_tool_call_output, streaming=streaming
)
# Content check (some parsers strip it)
if test_config.single_tool_call_expected_content is not None:
assert content == test_config.single_tool_call_expected_content
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
assert tool_calls[0].type == "function"
assert tool_calls[0].function.name == test_config.single_tool_call_expected_name
args = json.loads(tool_calls[0].function.arguments)
for key, value in test_config.single_tool_call_expected_args.items():
assert args.get(key) == value, (
f"Expected {key}={value}, got {args.get(key)}"
)
def test_parallel_tool_calls(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser handles multiple tools in one response."""
# Apply xfail markers if configured
test_name = "test_parallel_tool_calls"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser,
test_config.parallel_tool_calls_output,
streaming=streaming,
)
assert len(tool_calls) == test_config.parallel_tool_calls_count, (
f"Expected {test_config.parallel_tool_calls_count} "
f"tool calls, got {len(tool_calls)}"
)
# Verify tool names match expected
for i, expected_name in enumerate(test_config.parallel_tool_calls_names):
assert tool_calls[i].type == "function"
assert tool_calls[i].function.name == expected_name
# Verify unique IDs
ids = [tc.id for tc in tool_calls]
assert len(ids) == len(set(ids)), "Tool call IDs should be unique"
def test_various_data_types(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser handles all JSON types in arguments."""
# Apply xfail markers if configured
test_name = "test_various_data_types"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser,
test_config.various_data_types_output,
streaming=streaming,
)
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
args = json.loads(tool_calls[0].function.arguments)
# Verify all expected fields present
required_fields_types = {
"string_field": str,
"int_field": int,
"float_field": float,
"bool_field": bool,
"null_field": NoneType,
"array_field": list,
"object_field": dict,
}
for required_field, expected_type in required_fields_types.items():
assert required_field in args, (
f"Expected field '{required_field}' in arguments"
)
if test_config.supports_typed_arguments:
found_type = type(args[required_field])
assert found_type is expected_type, (
f"Expected field '{required_field}' to have type {expected_type}, "
f"got {found_type}"
)
def test_empty_arguments(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser handles parameterless tool calls."""
# Apply xfail markers if configured
test_name = "test_empty_arguments"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser, test_config.empty_arguments_output, streaming=streaming
)
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
args = tool_calls[0].function.arguments
if test_config.allow_empty_or_json_empty_args:
assert args in ["{}", ""], f"Expected empty args, got {args}"
else:
assert args == "{}", f"Expected {{}}, got {args}"
def test_surrounding_text(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser extracts tools from mixed content."""
# Apply xfail markers if configured
test_name = "test_surrounding_text"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser, test_config.surrounding_text_output, streaming=streaming
)
assert len(tool_calls) >= 1, (
f"Expected at least 1 tool call, got {len(tool_calls)}"
)
def test_escaped_strings(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser handles escaped characters in arguments."""
# Apply xfail markers if configured
test_name = "test_escaped_strings"
self.apply_xfail_mark(request, test_config, test_name, streaming)
content, tool_calls = run_tool_extraction(
tool_parser, test_config.escaped_strings_output, streaming=streaming
)
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
args = json.loads(tool_calls[0].function.arguments)
# At minimum, verify we can parse and have expected fields
# Exact escaping behavior varies by parser
assert len(args) > 0, "Expected some arguments with escaped strings"
def test_malformed_input(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
streaming: bool,
):
"""Verify parser gracefully handles invalid syntax."""
# Apply xfail markers if configured
test_name = "test_malformed_input"
self.apply_xfail_mark(request, test_config, test_name, streaming)
for malformed_input in test_config.malformed_input_outputs:
# Should not raise exception
content, tool_calls = run_tool_extraction(
tool_parser, malformed_input, streaming=streaming
)
# Parser should handle gracefully (exact behavior varies)
def test_streaming_reconstruction(
self,
request: pytest.FixtureRequest,
tool_parser: Any,
test_config: ToolParserTestConfig,
):
"""Verify streaming produces same result as non-streaming."""
test_name = "test_streaming_reconstruction"
self.apply_xfail_mark(request, test_config, test_name, True)
test_output = test_config.single_tool_call_output
# Non-streaming result
content_non, tools_non = run_tool_extraction(
tool_parser, test_output, streaming=False
)
# Streaming result
content_stream, tools_stream = run_tool_extraction(
tool_parser, test_output, streaming=True
)
# Compare results
assert content_non == content_stream, "Content should match between modes"
assert len(tools_non) == len(tools_stream), "Tool count should match"
if len(tools_non) > 0:
assert tools_non[0].function.name == tools_stream[0].function.name
assert tools_non[0].function.arguments == tools_stream[0].function.arguments
def apply_xfail_mark(self, request, test_config, test_name, streaming):
reason = None
if streaming and test_name in test_config.xfail_streaming:
reason = test_config.xfail_streaming[test_name]
elif not streaming and test_name in test_config.xfail_nonstreaming:
reason = test_config.xfail_nonstreaming[test_name]
if reason is not None:
mark = pytest.mark.xfail(reason=reason, strict=True)
request.node.add_marker(mark)

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